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/20230714_EspectrosCristal4iones/Data/')


MOTIONAL_FILES = """000013292-UV_Scan_withcal_optimized_andor
000013293-UV_Scan_withcal_optimized_andor
000013295-UV_Scan_withcal_optimized_andor"""


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()))

print(SeeKeys(MOTIONAL_FILES))

#%%
#carpeta pc nico labo escritorio:
#C:\Users\Usuario\Documents\artiq\artiq_experiments\analisis\plots\20211101_CPT_DosLaseres_v03\Data

CountsRoi1 = []
CountsRoi2 = []
CountsRoi3 = []
CountsRoi4 = []
CountsRoi5 = []

#Amplitudes = []
UV_Freqs    = []
#IR_amps    = []

for i, fname in enumerate(MOTIONAL_FILES.split()):
    print(str(i) + ' - ' + fname)
    data = h5py.File(fname+'.h5', 'r')
    #Amplitudes.append(np.array(data['datasets']['amplitudes']))
    CountsRoi1.append(np.array(data['datasets']['counts_roi1']))
    CountsRoi2.append(np.array(data['datasets']['counts_roi2']))
    CountsRoi3.append(np.array(data['datasets']['counts_roi3']))
    CountsRoi4.append(np.array(data['datasets']['counts_roi4']))
    CountsRoi5.append(np.array(data['datasets']['counts_roi5']))
    UV_Freqs.append(np.array(data['datasets']['UV_Frequencies']))
    #IR_amps.append(np.array(data['datasets']['IR1_measurement_amp']))



#%%

"""
En cristal de 4 iones veo espectros. Primero espectros uv.
La roi1 es la general. Las demas son de cada uno de los 4 iones brillantes del cristal.
"""

i = 0

jvec=[0,2]

step=0.1e8

Desplazamientos = [0, -2.5*step, 1.2*step]

plt.figure()
for j in jvec:
        #plt.errorbar(Amplitudes[j], CountsRoi1[j], yerr=np.sqrt(CountsRoi1[j]), color='red', fmt='-o', capsize=2, markersize=2)
        #plt.plot(Amplitudes[j][1:], CountsRoi1[j][1:], 'o',color='red', markersize=2,label=f'UVamp: {UV_amps[j]}')
    plt.plot([Desplazamientos[j]+f for f in UV_Freqs[j][1:]], CountsRoi2[j][1:], '-o', markersize=2)
    i = i + 1
plt.xlabel('Frecuencia')
plt.ylabel('Cuentas ROI')
#plt.xlim(0.05,0.23)
#plt.ylim(7800,8550)
plt.grid()
plt.legend()

#%%

#mergeo mediciones porque medi variando el piezoB para tener mas rango

Frequencies_vector = []
Counts_vector = []

kk = 2

for counts in [CountsRoi1, CountsRoi2, CountsRoi3, CountsRoi4, CountsRoi5]:
    Frequencies_vector.append([1e-6*2*f for f in [Desplazamientos[kk]+f for f in UV_Freqs[kk][1:]]])
    Counts_vector.append(list(counts[kk][1:]))

ivecs = [1,2,3,4]

#ivecs = [2, 5, 6]

#ivecs = [1]

plt.figure()
for i in range(len(Frequencies_vector)):
    if i in ivecs:
        plt.plot(Frequencies_vector[i], Counts_vector[i],'-o')
plt.grid()
plt.xlabel('Frequency (MHz)')
plt.ylabel('Counts')

#%%
ftrap=22.1
#ahora intento ajustarlos con modelo con micromocion

from scipy.special import jv
from scipy.optimize import curve_fit

def MicromotionSpectra(det, A, beta, x0, gamma, offset):
    ftrap=22.1
    #gamma=30
    P = A*(jv(0, beta)**2)/(((det-x0)**2)+(0.5*gamma)**2)+offset
    i = 1
    #print(P)
    while i <= 1:
        P = P + A*((jv(i, beta))**2)/((((det-x0)+i*ftrap)**2)+(0.5*gamma)**2) + A*((jv(-i, beta))**2)/((((det-x0)-i*ftrap)**2)+(0.5*gamma)**2) 
        i = i + 1
        #print(P)
    return P


popt_vec = []
pcov_vec = []


#uso como refe k=3
jref=3
popt_ref, pcov_ref = curve_fit(MicromotionSpectra, Frequencies_vector[jref], Counts_vector[jref], p0=[1000, 2, 274, 90, 14000], bounds=((0,0,200,20,0),(1e7,100,600,1000,25650)))

freqslong = np.arange(min(Frequencies_vector[jref]), max(Frequencies_vector[jref])+100, (Frequencies_vector[jref][1]-Frequencies_vector[jref][0])*0.01)

print(popt_ref)
plt.figure()
for j in range(1,len(Frequencies_vector)):
    plt.plot(Frequencies_vector[j], Counts_vector[j])
    if j == jref:
        plt.plot(freqslong, MicromotionSpectra(freqslong, *popt_ref))

for i in range(5):
    plt.axvline(popt_ref[2]-i*ftrap, linestyle='dashed', color='black', linewidth=1, zorder=0)
plt.grid()
#%%
for i in range(len(Frequencies_vector)):
    if i != jref:
        popt, pcov = curve_fit(MicromotionSpectra, Frequencies_vector[i], Counts_vector[i], p0=[popt_ref[0], 5, popt_ref[2], 60, popt_ref[4]], bounds=((popt_ref[0]-0.001*popt_ref[0],0,popt_ref[2]-0.001*popt_ref[2],0,popt_ref[4]-0.001*popt_ref[4]),(popt_ref[0]+0.001*popt_ref[0],100,popt_ref[2]+0.001*popt_ref[2],300, popt_ref[4]+0.001*popt_ref[4])))
        popt_vec.append(popt)
        pcov_vec.append(pcov)
    else:
        popt_vec.append(popt_ref)
        pcov_vec.append(pcov_ref)
        
ftrap=22.1

jeval=1

freqslong = np.arange(min(Frequencies_vector[jeval]), max(Frequencies_vector[jeval])+100, (Frequencies_vector[jeval][1]-Frequencies_vector[jeval][0])*0.01)

print(popt_vec[jeval])
plt.figure()
plt.plot(Frequencies_vector[jeval], Counts_vector[jeval])
plt.plot(freqslong, MicromotionSpectra(freqslong, *popt_ref))
plt.axvline(popt_ref[2], linestyle='dashed')
plt.axvline(popt_ref[2]-ftrap, linestyle='dashed')
plt.axvline(popt_ref[2]+ftrap, linestyle='dashed')
plt.axvline(popt_ref[2]-2*ftrap, linestyle='dashed')
plt.axvline(popt_ref[2]+2*ftrap, linestyle='dashed')
plt.axvline(popt_ref[2]-3*ftrap, linestyle='dashed')
plt.axvline(popt_ref[2]+3*ftrap, linestyle='dashed')