# -*- coding: utf-8 -*-
"""
Created on Tue Jan 30 14:07:15 2024

@author: sar1f18
"""

from scipy.special import voigt_profile
from scipy.optimize import curve_fit
import numpy as np
import h5py as h5
import matplotlib.pyplot as plt
import csv


def my_gaus(x,mu,sig,A,C):
    # my_gauss = A/(sig*np.sqrt(2*np.pi))*np.exp(-1/2*((x-mu)/sig)**2)+abs(C)
    my_gauss = abs(A)*sig**2/(sig**2+(x-mu)**2)+abs(C)
    # my_gauss = A/(np.pi*sig*(1+((x-mu)/sig)**2))+C
    return my_gauss

def find_peak_2(line,interval):
    
    find_peak0, peakval0 = find_peak_old(line, interval)

    line_crop = line[interval[0]:interval[1]]
    #qindrange= np.linspace(interval[0],interval[1],num=len(line_crop))
    qindrange = range(len(line_crop))
    mu0 = qindrange[int(len(qindrange)/2)]
    # mu0 = find_peak0-interval[0]
    sig0 = mu0
    A0 = np.max(line_crop)/2
    C0 = np.median(line_crop)/2
    p0s = [mu0, sig0, A0, C0]
    # f = lambda: p
    # popt,pcov = curve_fit(my_gaus,qindrange,line_crop,p0=p0s, bounds = ([interval[0],0,0,0],[interval[1],interval[1]-interval[0],1e-5,1e-5]), maxfev=10000)
    popt,pcov = curve_fit(my_gaus,qindrange,line_crop,p0=p0s, maxfev=10000)
    fitted_dist = [my_gaus(x,popt[0],popt[1],popt[2], popt[3]) for x in qindrange]
    plt.figure()
    plt.plot(qindrange,fitted_dist)
    plt.plot(qindrange,line_crop)        
    check_peak = np.argmax(line_crop)+interval[0]
    find_peak = round(popt[0])+interval[0]
    popt[0] = popt[0]++interval[0]
    peakval = np.max(line_crop)
    return find_peak, peakval, check_peak, popt

def find_peak3(line,interval,q):
    
    find_peak0, peakval0 = find_peak_old(line, interval)

    line_crop = line[interval[0]:interval[1]]
    q_crop = q[interval[0]:interval[1]]
    #qindrange= np.linspace(interval[0],interval[1],num=len(line_crop))
    qindrange = range(len(line_crop))
    mu0 = qindrange[int(len(qindrange)/2)]
    # mu0 = find_peak0-interval[0]
    sig0 = mu0
    A0 = np.max(line_crop)/2
    C0 = np.median(line_crop)/2
    p0s = [mu0, sig0, A0, C0]
    # f = lambda: p
    # popt,pcov = curve_fit(my_gaus,qindrange,line_crop,p0=p0s, bounds = ([interval[0],0,0,0],[interval[1],interval[1]-interval[0],1e-5,1e-5]), maxfev=10000)
    popt,pcov = curve_fit(my_gaus,qindrange,line_crop,p0=p0s, maxfev=10000)
    fitted_dist = [my_gaus(x,popt[0],popt[1],popt[2], popt[3]) for x in qindrange]
    # plt.figure()
    # plt.plot(q_crop,fitted_dist)
    # plt.plot(q_crop,line_crop)        
    check_peak = np.argmax(line_crop)+interval[0]
    find_peak = round(popt[0])+interval[0]
    popt[0] = popt[0]++interval[0]
    peakval = np.max(line_crop)
    return find_peak, peakval, check_peak, popt

def find_peak_old(line,interval):
    line_crop = line[interval[0]:interval[1]]
    find_peak = np.argmax(line_crop)+interval[0]
    peakval = np.max(line_crop)
    return find_peak, peakval

def dif_dir_map(texture,moisture,depth, replicate):
    if texture == 'c' or texture == 'C':
        if moisture =='p' or moisture == 'P':
            if replicate ==1:
                if depth == 1000:
                    location = '//cseg_3/S_Ruiz/DIAD/Jul_2023/mirror_1/processed/k11-23900-diffraction-Diff_CakeAzz.nxs'
                elif depth == 5000:
                    location = '//cseg_3/S_Ruiz/DIAD/Jul_2023/mirror_1/processed/k11-23904-diffraction-Diff_CakeAzz.nxs'
                elif depth == 10000: # actually 8500
                    location = '//cseg_3/S_Ruiz/DIAD/Jul_2023/mirror_1/processed/k11-23907-diffraction-Diff_CakeAzz.nxs'
            elif replicate == 2:
                if depth == 1000:
                    location = '//cseg_3/S_Ruiz/DIAD/Jul_2023/mirror_1/processed/k11-23916-diffraction-Diff_CakeAzz.nxs'
                elif depth == 5000:
                    location = '//cseg_3/S_Ruiz/DIAD/Jul_2023/mirror_1/processed/k11-23919-diffraction-Diff_CakeAzz.nxs'
                elif depth == 10000:
                    location = '//cseg_3/S_Ruiz/DIAD/Jul_2023/mirror_1/processed/k11-23922-diffraction-Diff_CakeAzz.nxs'
            elif replicate == 3:
               if depth == 1000:
                   location = '//cseg_3/S_Ruiz/DIAD/Jul_2023/mirror_1/processed/k11-23928-diffraction-Diff_CakeAzz.nxs'
               elif depth == 5000:
                   location = '//cseg_3/S_Ruiz/DIAD/Jul_2023/mirror_1/processed/k11-23931-diffraction-Diff_CakeAzz.nxs'
               elif depth == 10000: # actually 10500
                   location = '//cseg_3/S_Ruiz/DIAD/Jul_2023/mirror_1/processed/k11-23934-diffraction-Diff_CakeAzz.nxs'
        
        if moisture =='b' or moisture == 'B':
            if replicate ==1:
                if depth == 1000:
                    location = '//cseg_3/S_Ruiz/DIAD/Jul_2023/mirror_1/processed/k11-23986-diffraction-Diff_CakeAzz.nxs'
                elif depth == 3000:
                    location = '//cseg_3/S_Ruiz/DIAD/Jul_2023/mirror_1/processed/k11-23989-diffraction-Diff_CakeAzz.nxs'
                elif depth == 5000:
                    location = '//cseg_3/S_Ruiz/DIAD/Jul_2023/mirror_1/processed/k11-23992-diffraction-Diff_CakeAzz.nxs'
            elif replicate == 2:
                if depth == 1000:
                    location = '//cseg_3/S_Ruiz/DIAD/Jul_2023/mirror_1/processed/k11-23998-diffraction-Diff_CakeAzz.nxs'
                elif depth == 3000:
                    location = '//cseg_3/S_Ruiz/DIAD/Jul_2023/mirror_1/processed/k11-24001-diffraction-Diff_CakeAzz.nxs'
                elif depth == 5000:
                    location = '//cseg_3/S_Ruiz/DIAD/Jul_2023/mirror_1/processed/k11-24004-diffraction-Diff_CakeAzz.nxs'
            elif replicate == 3:
               if depth == 1000:
                   location = '//cseg_3/S_Ruiz/DIAD/Jul_2023/mirror_1/processed/k11-24010-diffraction-Diff_CakeAzz.nxs'
               elif depth == 3000:
                   location = '//cseg_3/S_Ruiz/DIAD/Jul_2023/mirror_1/processed/k11-24013-diffraction-Diff_CakeAzz.nxs'
               elif depth == 5000:
                   location = '//cseg_3/S_Ruiz/DIAD/Jul_2023/mirror_1/processed/k11-24016-diffraction-Diff_CakeAzz.nxs'
                   
        if moisture =='w' or moisture == 'W':
            if replicate ==1:
                if depth == 1000:
                    location = '//cseg_3/S_Ruiz/DIAD/Jul_2023/mirror_1/processed/k11-24023-diffraction-Diff_CakeAzz.nxs'
                elif depth == 5000:
                    location = '//cseg_3/S_Ruiz/DIAD/Jul_2023/mirror_1/processed/k11-24026-diffraction-Diff_CakeAzz.nxs'
                elif depth == 10000:
                    location = '//cseg_3/S_Ruiz/DIAD/Jul_2023/mirror_1/processed/k11-24029-diffraction-Diff_CakeAzz.nxs'
            elif replicate == 2:
                if depth == 1000:
                    location = '//cseg_3/S_Ruiz/DIAD/Jul_2023/mirror_1/processed/k11-24035-diffraction-Diff_CakeAzz.nxs'
                elif depth == 5000:
                    location = '//cseg_3/S_Ruiz/DIAD/Jul_2023/mirror_1/processed/k11-24038-diffraction-Diff_CakeAzz.nxs'
                elif depth == 10000:
                    location = '//cseg_3/S_Ruiz/DIAD/Jul_2023/mirror_1/processed/k11-24041-diffraction-Diff_CakeAzz.nxs'
            elif replicate == 3:
               if depth == 1000:
                   location = '//cseg_3/S_Ruiz/DIAD/Jul_2023/mirror_1/processed/k11-24047-diffraction-Diff_CakeAzz.nxs'
               elif depth == 5000:
                   location = '//cseg_3/S_Ruiz/DIAD/Jul_2023/mirror_1/processed/k11-24050-diffraction-Diff_CakeAzz.nxs'
               elif depth == 10000:
                   location = '//cseg_3/S_Ruiz/DIAD/Jul_2023/mirror_1/processed/k11-24053-diffraction-Diff_CakeAzz.nxs'
    elif texture == 'g' or texture =='G':
         if moisture =='p' or moisture == 'P':
             if replicate ==1:
                 if depth == 1000:
                     location = '//cseg_3/S_Ruiz/DIAD/Jul_2023/mirror_1/processed/k11-23946-diffraction-Diff_CakeAzz.nxs'
                 elif depth == 5000:
                     location = '//cseg_3/S_Ruiz/DIAD/Jul_2023/mirror_1/processed/k11-23949-diffraction-Diff_CakeAzz.nxs'
                 elif depth == 10000: # Actually 9000
                     location = '//cseg_3/S_Ruiz/DIAD/Jul_2023/mirror_1/processed/k11-23952-diffraction-Diff_CakeAzz.nxs'
                 elif depth == 0: #might not work
                    location = '//cseg_3/S_Ruiz/DIAD/Jul_2023/mirror_1/processed/k11-23943-diffraction-Diff_CakeAzz.nxs'         
             elif replicate == 2:
                 if depth == 1000:
                     location = '//cseg_3/S_Ruiz/DIAD/Jul_2023/mirror_1/processed/k11-23958-diffraction-Diff_CakeAzz.nxs'
                 elif depth == 5000:
                     location = '//cseg_3/S_Ruiz/DIAD/Jul_2023/mirror_1/processed/k11-23962-diffraction-Diff_CakeAzz.nxs'
                 elif depth == 10000: #actually 9300
                     location = '//cseg_3/S_Ruiz/DIAD/Jul_2023/mirror_1/processed/k11-23966-diffraction-Diff_CakeAzz.nxs'
                 elif depth == 0: #might not work
                     location = '//cseg_3/S_Ruiz/DIAD/Jul_2023/mirror_1/processed/k11-23954-diffraction-Diff_CakeAzz.nxs'         
             elif replicate == 3:
                if depth == 1000:
                    location = '//cseg_3/S_Ruiz/DIAD/Jul_2023/mirror_1/processed/k11-23974-diffraction-Diff_CakeAzz.nxs'
                elif depth == 5000:
                    location = '//cseg_3/S_Ruiz/DIAD/Jul_2023/mirror_1/processed/k11-23977-diffraction-Diff_CakeAzz.nxs'
                elif depth == 10000: #Actually 8000
                    location = '//cseg_3/S_Ruiz/DIAD/Jul_2023/mirror_1/processed/k11-23980-diffraction-Diff_CakeAzz.nxs'
                elif depth == 0: #This might not work
                    location = '//cseg_3/S_Ruiz/DIAD/Jul_2023/mirror_1/processed/k11-23968-diffraction-Diff_CakeAzz.nxs'         
         if moisture =='w' or moisture == 'W':
             if replicate ==1:
                 if depth == 1000:
                     location = '//cseg_3/S_Ruiz/DIAD/Jul_2023/mirror_1/processed/k11-24068-diffraction-Diff_CakeAzz.nxs' #Need to find data here
                 elif depth == 5000: #Not sure what the robot was doing here
                     location = '//cseg_3/S_Ruiz/DIAD/Jul_2023/mirror_1/processed/k11-24072-diffraction-Diff_CakeAzz.nxs'
                 elif depth == 10000: #Actually 5000
                     location = '//cseg_3/S_Ruiz/DIAD/Jul_2023/mirror_1/processed/k11-24075-diffraction-Diff_CakeAzz.nxs'
                 elif depth == 0:
                        location = '//cseg_3/S_Ruiz/DIAD/Jul_2023/mirror_1/processed/k11-24065-diffraction-Diff_CakeAzz.nxs'
             elif replicate == 2:
                 if depth == 1000:
                     location = '//cseg_3/S_Ruiz/DIAD/Jul_2023/mirror_1/processed/k11-24087-diffraction-Diff_CakeAzz.nxs'
                 elif depth == 5000:
                     location = '//cseg_3/S_Ruiz/DIAD/Jul_2023/mirror_1/processed/k11-24092-diffraction-Diff_CakeAzz.nxs'
                 elif depth == 10000: #actually 7350
                     location = '//cseg_3/S_Ruiz/DIAD/Jul_2023/mirror_1/processed/k11-24095-diffraction-Diff_CakeAzz.nxs'
                 elif depth == 0:
                     location = '//cseg_3/S_Ruiz/DIAD/Jul_2023/mirror_1/processed/k11-24082-diffraction-Diff_CakeAzz.nxs'
             elif replicate == 3:
                if depth == 1000:
                    location = '//cseg_3/S_Ruiz/DIAD/Jul_2023/mirror_1/processed/k11-24105-diffraction-Diff_CakeAzz.nxs'
                elif depth == 5000:
                    location = '//cseg_3/S_Ruiz/DIAD/Jul_2023/mirror_1/processed/k11-24110-diffraction-Diff_CakeAzz.nxs'
                elif depth == 10000:
                    location = '//cseg_3/S_Ruiz/DIAD/Jul_2023/mirror_1/processed/k11-24116-diffraction-Diff_CakeAzz.nxs'
                elif depth == 0:
                    location = '//cseg_3/S_Ruiz/DIAD/Jul_2023/mirror_1/processed/k11-24100-diffraction-Diff_CakeAzz.nxs'       
         if moisture =='b' or moisture == 'B':
             if replicate ==1:
                 if depth == 1000:
                     location = '//cseg_3/S_Ruiz/DIAD/Jul_2023/mirror_1/processed/k11-24126-diffraction-Diff_CakeAzz.nxs'
                 elif depth == 5000:
                     location = '//cseg_3/S_Ruiz/DIAD/Jul_2023/mirror_1/processed/k11-24131-diffraction-Diff_CakeAzz.nxs'
                 elif depth == 10000:
                     location = '//cseg_3/S_Ruiz/DIAD/Jul_2023/mirror_1/processed/k11-24134-diffraction-Diff_CakeAzz.nxs'
                 elif depth == 0:
                         location = '//cseg_3/S_Ruiz/DIAD/Jul_2023/mirror_1/processed/k11-24121-diffraction-Diff_CakeAzz.nxs'    
             elif replicate == 2:
                 if depth == 1000:
                     location = '//cseg_3/S_Ruiz/DIAD/Jul_2023/mirror_1/processed/k11-24144-diffraction-Diff_CakeAzz.nxs'
                 elif depth == 5000:
                     location = '//cseg_3/S_Ruiz/DIAD/Jul_2023/mirror_1/processed/k11-24152-diffraction-Diff_CakeAzz.nxs'
                 elif depth == 10000: # Actually 8000
                     location = '//cseg_3/S_Ruiz/DIAD/Jul_2023/mirror_1/processed/k11-24155-diffraction-Diff_CakeAzz.nxs'
                 elif depth == 0:
                         location = '//cseg_3/S_Ruiz/DIAD/Jul_2023/mirror_1/processed/k11-24139-diffraction-Diff_CakeAzz.nxs'    
             elif replicate == 3:
                if depth == 1000:
                    location = '//cseg_3/S_Ruiz/DIAD/Jul_2023/mirror_1/processed/k11-24165-diffraction-Diff_CakeAzz.nxs'
                elif depth == 5000:
                    location = '//cseg_3/S_Ruiz/DIAD/Jul_2023/mirror_1/processed/k11-24170-diffraction-Diff_CakeAzz.nxs'
                elif depth == 10000: # Actually 8000
                    location = '//cseg_3/S_Ruiz/DIAD/Jul_2023/mirror_1/processed/k11-24174-diffraction-Diff_CakeAzz.nxs'
                elif depth == 0:
                        location = '//cseg_3/S_Ruiz/DIAD/Jul_2023/mirror_1/processed/k11-24160-diffraction-Diff_CakeAzz.nxs'
    return location  

def pull_dat(data_file_dir):
    g=h5.File(data_file_dir,'r')
#f=h5.File('27172809_Clay_processed_220825_115123.nxs','r')

    x_cord=np.array(g['entry/diffraction/kb_cs_x'])

    y_cord=np.array(g['entry/diffraction/kb_cs_y'])

    counts =  np.array(g['processed/result/data'])
#gdatax = np.array(g['processed/result/q'])
    twotht_cord = np.array(g['processed/result/2theta'])
    q_cord = np.array(g['processed/intermediate/3-Moving Beam Azimuthal Integration/q'])
    return g,x_cord,y_cord,counts,twotht_cord, q_cord

def ensemble(Amp):
    Amp_sum = np.zeros(np.shape(Amp)[2])
    var_sum = np.zeros(np.shape(Amp)[2])
    for i in range(np.shape(Amp)[0]):
        for j in range(np.shape(Amp)[1]):
            Amp_sum = Amp_sum+ Amp[i,j,:]/np.nansum(Amp[i,j,:])
            
    N = (np.shape(Amp)[0]*np.shape(Amp)[2])
    Amp_mean = Amp_sum/N
    
    for i in range(np.shape(Amp)[0]):
        for j in range(np.shape(Amp)[1]):
            var_sum = var_sum+ (Amp[i,j,:]/np.nansum(Amp[i,j,:])-Amp_mean)**2
    
    var = var_sum/N
    stdev = np.sqrt(var)
    return Amp_sum, Amp_mean, stdev


def quick_diff_plot(texture,moisture,rep,initial,final):
    TextMoistRep_init = dif_dir_map(texture,moisture,initial,rep)
    TextMoistRep_final= dif_dir_map(texture,moisture,final,rep)

    g0,x0,y0,Amp0,twotht0, q0 = pull_dat(TextMoistRep_init)
    gn,xn,yn,Ampn,twothtn, qn = pull_dat(TextMoistRep_final)

    Amp0sm,Amp0mn, stdev0 = ensemble(Amp0)
    Ampnsm, Ampnmn, stdevn = ensemble(Ampn)

    plt.figure()
    plt.plot(q0,Amp0mn, zorder = 2)
    plt.plot(qn,Ampnmn, zorder = 1)
    
def quick_dat_import(texture,moisture,rep,position):
    data_dir = dif_dir_map(texture,moisture,position,rep)
    g,x,y,Amp,twotht, q = pull_dat(data_dir)
    return g,x,y,Amp,twotht, q
    
def quick_diff_sub_plot(texture,moisture,rep,initial,final,ax1,ax2):
    TextMoistRep_init = dif_dir_map(texture,moisture,initial,rep)
    TextMoistRep_final= dif_dir_map(texture,moisture,final,rep)

    g0,x0,y0,Amp0,twotht0, q0 = pull_dat(TextMoistRep_init)
    gn,xn,yn,Ampn,twothtn, qn = pull_dat(TextMoistRep_final)

    Amp0sm,Amp0mn, stdev0 = ensemble(Amp0)
    Ampnsm, Ampnmn, stdevn = ensemble(Ampn)

    
def disp_all_diff(texture,zoom_on,xwin,ymax,axs):


    quick_diff_sub_plot(texture,'P',1,1000,10000,0,0)

    quick_diff_sub_plot(texture,'P',2,1000,10000, 0,1)

    quick_diff_sub_plot(texture,'P',3,1000,10000, 0,2)

    if texture == 'c' or texture == 'C':
        quick_diff_sub_plot(texture,'B',1,1000,3000,1,0)
        quick_diff_sub_plot(texture,'B',2,1000,3000,1,1)
        quick_diff_sub_plot(texture,'B',3,1000,3000,1,2)
        
        quick_diff_sub_plot(texture,'W',1,1000,10000,2,0)
        quick_diff_sub_plot(texture,'W',2,1000,10000,2,1)
        quick_diff_sub_plot(texture,'W',3,1000,10000,2,2)
        
    elif texture =='g' or texture=='G':
        quick_diff_sub_plot(texture,'B',1,0,10000,1,0)
        quick_diff_sub_plot(texture,'B',2,0,10000,1,1)
        quick_diff_sub_plot(texture,'B',3,0,10000,1,2)
        
        quick_diff_sub_plot(texture,'W',1,5000,10000,2,0)
        quick_diff_sub_plot(texture,'W',2,0,10000,2,1)
        quick_diff_sub_plot(texture,'W',3,0,10000,2,2)
    
    for ax in axs.flat:
        ax.set(xlabel = 'q A-1', ylabel = 'counts')
        
        if zoom_on==1:
            ax.set_xlim(left = xwin[0], right = xwin[1])
            ax.set_ylim(bottom = 0, top = ymax)
        

        

def map_diff(texture,moisture,rep,position,qrange):
    g,x,y,Amp,twotht, q = quick_dat_import(texture,moisture,rep,position)
    qlb = [x for x in q if x<qrange[0]]
    ilb = np.shape(qlb)[0]
    qub = [x for x in q if x<qrange[1]]
    iub = np.shape(qub)[0]
    i_interval =[ilb,iub]
    diff_peak_map = np.zeros(np.shape(Amp)[0:2])
    diff_count_map = np.zeros(np.shape(Amp)[0:2])
    for i in range(np.shape(Amp)[0]):
        for j in range(np.shape(Amp)[1]):
                        
            find_peak,peakval=find_peak_old(Amp[i,j,:],i_interval)
            # find_peak,peakval,check_peak,popt=find_peak_2(Amp[i,j,:],i_interval)
            # find_peak,peakval,check_peak,popt=find_peak3(Amp[i,j,:],i_interval,q)
            
            diff_peak_map[i,j] = q[find_peak]
            diff_count_map[i,j] = peakval
    
    
    return diff_peak_map, diff_count_map



def plot_map_diff(texture,moisture,qrange,rep,position):
    
    diff_peak_map, diff_count_map=map_diff(texture,moisture,rep,position,qrange)
    plt.figure()
    plt.imshow(diff_peak_map, cmap = 'coolwarm')
    plt.colorbar(extend='both')
    
def plot_map_diff_shift(texture,moisture,qrange,rep,E):
    pos0 = 0000
    posn = 5000
    diff_peak_map0, diff_count_map0=map_diff(texture,moisture,rep,pos0,qrange)
    diff_peak_mapn, diff_count_mapn=map_diff(texture,moisture,rep,posn,qrange)
    q0 = np.mean(qrange)
    peak_shift_map = diff_peak_mapn-diff_peak_map0
    
    dq = peak_shift_map
    
    qn = q0+dq

    eps = np.reciprocal(qn)*dq  
    sig = E*eps
    
    plt.figure()
    plt.imshow(peak_shift_map, cmap = 'coolwarm',vmin=0,vmax=.06)
    
    plt.colorbar(extend='both')
    
    return peak_shift_map

def plot_point_diffraction(texture,moisture,qrange,rep,position,ii,jj):
    g,x,y,Amp,twotht, q = quick_dat_import(texture,moisture,rep,position)
    
    i_0 = np.min(np.argwhere(q>=qrange[0]))
    i_n = np.max(np.argwhere(q<=qrange[1]))
    xi = x[ii]
    yj = y[jj]
    
    Amp_ij = Amp[ii,jj,:]
    qwin = q[i_0:i_n]
    Awin = Amp_ij[i_0:i_n]
    i_win = range(len(Awin))
    
    plt.figure()
    plt.plot(qwin,Awin)
    
    print('data from pixel coordinate (x =', xi,' , y = ',yj,')')    
    
    
    return Amp_ij, q
    


Egyp = 1.75e9; # Gypsun elastic modulus


xwin = [1.3,1.8]
ymax = 0.14*0+.6e-5

zoom_on=1
texture = 'G'
qrange = [1.55,1.7] # peak regions

q0 = np.mean(qrange)
dq = plot_map_diff_shift(texture,'b',qrange,1,Egyp)