import h5py
import matplotlib.pyplot as plt
import numpy as np
import sys
import re
import ast
from scipy.optimize import curve_fit
import os
from scipy import interpolate

#Mediciones barriendo angulo del TISA y viendo kicking de resonancias oscuras

#C:\Users\Usuario\Documents\artiq\artiq_experiments\analisis\plots\20220106_CPT_DosLaseres_v08_TISA_DR\Data

os.chdir('/home/nico/Documents/artiq_experiments/analisis/plots/20230817_RotationalDopplerShift_v5/Data')


"""
en este codigo ploteo espectros CPT de resonancias D-D para configuracion +2/+2 y +2/-2 (usando pentaprisma)
"""

def find_nearest(array, value):
    array = np.asarray(array)
    idx = (np.abs(array - value)).argmin()
    return idx


def Split(array,n):
    length=len(array)/n
    splitlist = []
    jj = 0
    while jj<length:
        partial = []
        ii = 0
        while ii < n:
            partial.append(array[jj*n+ii])
            ii = ii + 1
        splitlist.append(partial)
        jj = jj + 1
    return splitlist


def SeeKeys(files):
    for i, fname in enumerate(files.split()):
        data = h5py.File(fname+'.h5', 'r') # Leo el h5: Recordar que nuestros datos estan en 'datasets'
        print(fname)
        print(list(data['datasets'].keys()))

def Lorentzian( x, A, B, x0, gam ):
    return A * gam**2 / ( gam**2 + ( x - x0 )**2) + B



Piezo1Counts = []
Piezo1Frequencies = []

PIEZOS1_FILES = [32,33,34,35,37,38,39,40,41,42,43,44,45,46]

for i in PIEZOS1_FILES:
    #print(str(i) + ' - ' + fname)
    data = h5py.File(f'VaryingBeamsize/Size1/0000150{i}-IR_Scan_withcal_optimized'+'.h5', 'r')
    Piezo1Counts.append(np.array(data['datasets']['counts_spectrum']))
    Piezo1Frequencies.append(np.array(data['datasets']['IR1_Frequencies']))



#%%
"""
Ahora voy a intentar ajustarlas con una lorentziana que es mejor
"""
import seaborn as sns
"""
Resonancias DD configuracion +2/-2 colineal variando la ubicacion del ion en los haces

Moviendo verticalmente el haz
"""



palette = sns.color_palette("tab10")

pmlocmedvec = np.arange(0,len(PIEZOS1_FILES),1)


#pmlocmedvec = [26]


plt.figure()

#bkg = np.min(PiezoS1Counts[5])
bkg=120


pmdepthsdrver=[]
errorpmdepthsdrver=[]

Intensityver = []
errorIntensityver = []

jj=0
for med in pmlocmedvec:

    Freqs = [2*f*1e-6 for f in Piezo1Frequencies[med][1:]]
    Counts = [c for c in Piezo1Counts[med][1:]]

    popt, pcov = curve_fit(Lorentzian, Freqs, Counts, p0=(-200,2100,435.8,0.05), bounds=((-10000,0,435.5,0),(0,1e4, 436.1, 1)))

    pmdepthsdrver.append(1-(np.min(Lorentzian(Freqs,*popt))-bkg)/(popt[1]-bkg))
    #errorpmdepthsdrver.append(ErrorDRdepth(np.min(Lorentzian(Freqs,*popt)),popt[1], bkg))
        
    Intens = popt[1]
    
    Intensityver.append(Intens)
    # errorIntensity.append(2*np.sqrt(np.mean(Piezo1Counts[med][1:][0:20]))+np.sqrt(bkg))
        
    
    if med not in [800]:
        plt.plot([2*f*1e-6 for f in Piezo1Frequencies[med][1:]], [c for c in Piezo1Counts[med][1:]], '-o', markersize=2, alpha=0.7)
        plt.plot(Freqs,Lorentzian(Freqs,*popt))
    
    jj=jj+1
# plt.xlabel('Frecuencia (MHz)')
# plt.ylabel('Counts')
#plt.xlim(435.2, 436.5)
plt.grid()
# plt.legend()
# #plt.title('Espectros para distintas geometrías')


plt.figure()
plt.plot(np.arange(0,len(Intensityver),1), [i/np.max(Intensityver) for i in Intensityver], '-o',markersize=8)
plt.plot(np.arange(0,len(Intensityver),1), [p for p in pmdepthsdrver], 'o',markersize=8)
plt.xlabel('Ion position')
plt.ylabel('Intensity / DR Relative depth')
#plt.xticks([1,2,3,4,5])
#plt.xlim(200,3200)
plt.ylim(-0.1,1.1)
plt.grid()
#plt.axvline(3, color='salmon')
plt.legend()

#%%
"""
Ahora voy a intentar ajustarlas con una lorentziana que es mejor
"""
import seaborn as sns
"""
Resonancias DD configuracion +2/-2 colineal variando la ubicacion del ion en los haces

Moviendo diagonalmente el haz
"""



palette = sns.color_palette("tab10")

pmlocmedvec = np.arange(0,len(PIEZODIAG_FILES),1)


#pmlocmedvec = [0,1]


plt.figure()

bkg = np.min(PiezoDiagCounts[5])

pmdepthsdrdiag=[]
errorpmdepthsdrdiag=[]

Intensitydiag = []
errorIntensitydiag = []

jj=0
for med in pmlocmedvec:

    Freqs = [2*f*1e-6 for f in PiezoDiagFrequencies[med][1:]]
    Counts = [c for c in PiezoDiagCounts[med][1:]]

    if med==2:
        Freqs = Freqs[1:-30]
        Counts = Counts[1:-30]
        popt, pcov = curve_fit(Lorentzian, Freqs, Counts, p0=(-200,2100,435.8,0.05), bounds=((-10000,0,435.5,0),(0,1e4, 436.1, 1)))
    elif med==1:
        Freqs = Freqs[10:-30]
        Counts = Counts[10:-30]
        popt, pcov = curve_fit(Lorentzian, Freqs, Counts, p0=(-200,2100,435.8,0.05), bounds=((-10000,0,435.7,0),(0,1e4, 436.1, 1)))
 
    elif med==5:
        Freqs = Freqs[10:-55]+Freqs[-30:-1]
        Counts = Counts[10:-55]+Counts[-30:-1]
        popt, pcov = curve_fit(Lorentzian, Freqs, Counts, p0=(-200,2100,435.8,0.05), bounds=((-10000,0,435.5,0),(0,1e4, 436.1, 1)))
        
    else:
        popt, pcov = curve_fit(Lorentzian, Freqs, Counts, p0=(-200,2100,435.8,0.05), bounds=((-10000,0,435.5,0),(0,1e4, 436.1, 1)))

    pmdepthsdrdiag.append(1-(np.min(Lorentzian(Freqs,*popt))-bkg)/(popt[1]-bkg))
    errorpmdepthsdrdiag.append(ErrorDRdepth(np.min(Lorentzian(Freqs,*popt)),popt[1], bkg))
        
    Intens = popt[1]
    
    Intensitydiag.append(Intens)
    # errorIntensity.append(2*np.sqrt(np.mean(Piezo1Counts[med][1:][0:20]))+np.sqrt(bkg))
        
    
    if med not in [800]:
        plt.plot([2*f*1e-6 for f in PiezoDiagFrequencies[med][1:]], [c for c in PiezoDiagCounts[med][1:]], '-o', markersize=2, alpha=0.7)
        plt.plot(Freqs,Lorentzian(Freqs,*popt))
    
    jj=jj+1
# plt.xlabel('Frecuencia (MHz)')
# plt.ylabel('Counts')
#plt.xlim(435.2, 436.5)
plt.grid()
# plt.legend()
# #plt.title('Espectros para distintas geometrías')


plt.figure()
plt.plot(np.arange(0,len(Intensitydiag),1), [i/np.max(Intensitydiag) for i in Intensitydiag], '-o',markersize=8)
plt.plot(np.arange(0,len(Intensitydiag),1), [p for p in pmdepthsdrdiag], 'o',markersize=8)
plt.xlabel('Ion position')
plt.ylabel('Intensity / DR Relative depth')
#plt.xticks([1,2,3,4,5])
#plt.xlim(200,3200)
plt.ylim(-0.1,1.1)
plt.grid()
#plt.axvline(3, color='salmon')
plt.legend()


#%%

"""
Ploteo en conjunto. La horizontal sale de RDS_piezo.py
"""
cap=3

plt.figure()
plt.plot(np.arange(0,len(Intensitydiag),1), [i/np.max(Intensitydiag) for i in Intensitydiag], '-o',markersize=8)
#plt.plot(np.arange(0,len(Intensitydiag),1), [p for p in pmdepthsdrdiag], 'o',markersize=8, label='Diagonal')
plt.errorbar(np.arange(0,len(Intensitydiag),1), [p for p in pmdepthsdrdiag], yerr=errorpmdepthsdrdiag, fmt='o',capsize=cap,markersize=8, label='Diagonal')


#plt.plot(np.arange(0,len(Intensityver),1), [i/np.max(Intensityver) for i in Intensityver], '-o',markersize=8)
#plt.plot(np.arange(0,len(Intensityver),1), [p for p in pmdepthsdrver], 'o',markersize=8, label='Vertical')
plt.errorbar(np.arange(0,len(Intensityver),1), [p for p in pmdepthsdrver], yerr=errorpmdepthsdrver, fmt='o', capsize=cap,markersize=8, label='Vertical')


scale = 1.6

#plt.plot([s*scale for s in np.arange(16,len(Intensity),1)-16], [i/np.max(Intensity) for i in Intensity[16:]], '-o',markersize=8)
#plt.plot([s*scale for s in np.arange(16,len(Intensity),1)-16], [p for p in pmdepthsdr[16:]], 'o',markersize=8, label='Horizontal')
plt.errorbar([s*scale for s in np.arange(16,len(Intensity),1)-16], [p for p in pmdepthsdr[16:]], yerr=errorpmdepthsdr[16:], fmt='o', capsize=cap,markersize=8, label='Horizontal')


plt.xlabel('Ion position')
plt.ylabel('Intensity / DR Relative depth')
#plt.xticks([1,2,3,4,5])
plt.xlim(-1,15)
plt.ylim(-0.1,1.1)
plt.grid()
#plt.axvline(3, color='salmon')
plt.legend()