import numpy as np
import matplotlib.pyplot as plt
from matplotlib import gridspec
import cmath
import datetime
import time
from scipy.signal import argrelmax
from scipy import optimize


i=0+1j
pi=np.pi


FSR = 49.96554 # kHz
c = 299792458.
L = c/2/(FSR*1.e3)

R_ETM = 1907.83 # ETMX
R_ITM = 1904.54 # ITMX
#R_ETM = 1900
#R_ITM = 1900

g_ETM = 1.-L/R_ETM
g_ITM = 1.-L/R_ITM

TMS_calc = FSR/2/pi*(2*pi-np.arccos(2*g_ETM*g_ITM-1))
print TMS_calc

data = np.loadtxt('./cavityscan_20181206_6.txt')

t = data[:,0]
trans = data[:,1]
freq = 0.1* (t-97.5) # kHz

trans[np.where(freq>162.)] = 1e-3
df = freq[1]-freq[0]

candidate = argrelmax(trans,order=200)
large = np.where(trans>3)
peaks = np.intersect1d(candidate, large)

trans_p_sort = np.sort(trans[peaks])[::-1]
#print trans_p_sort
freq_p = freq[peaks]
freq_p_sort = freq_p[ np.argsort(trans[peaks])[::-1] ]

n_of_peaks = np.size(freq_p_sort)
#print n_of_peaks


def lor(f,f0,A,pole):
    return A / (1 + ((f-f0)/pole)**2)



peak_list = np.zeros((n_of_peaks,3))

rem = trans

for n in range(0,n_of_peaks):
    use_i = np.where( (freq > freq_p_sort[n] -  20*df) * (freq < freq_p_sort[n] +  20*df)  )
    use_freq = freq[use_i]
    use_trans = rem[use_i] 
    guess = [freq_p_sort[n],trans_p_sort[n],0.015]
    popt,pcov = optimize.curve_fit(lor,use_freq,use_trans,p0=guess)
    f0,A,pole = popt
    peak_list[n,0] = f0
    peak_list[n,1] = A
    peak_list[n,2] = pole
    
    rem = rem - lor(freq,f0,A,pole)


f_list = peak_list[:,0]

#print f_list
peak_list = peak_list[np.argsort(f_list)]
f_list = np.sort(f_list)
#print peak_list
#print f0
#print A
#print pole

#FSR = 50.e3

np.savetxt('./peaklist.txt',peak_list,fmt='%1.4e')


zeroth_i = [1,15,29,43]
first_i = [11,25,39]
second_i = [7,21,35]
third_i = [3,17,31,45]
fourth_i = [13,27,41]

f1u_i = [5,19,33]
f1l_i = [12,26,40]
f2u_i = [6,20,34]
f2l_i = [9,23,37]

f1u_df = -(f_list[f1u_i] - f_list[np.array(f1u_i)-np.array([4,4,4])])
f1l_df = f_list[f1l_i] - f_list[np.array(f1l_i)+np.array([3,3,3])]
f2u_df = -(f_list[f2u_i] - f_list[np.array(f2u_i)-np.array([5,5,5])])
f2l_df = f_list[f2l_i] - f_list[np.array(f2l_i)+np.array([6,6,6])]
f1_df = np.hstack((f1u_df,f1l_df))
f2_df = np.hstack((f2u_df,f2l_df))

f1 = 16.87515e6
f2 = 45.0004e6
FSR_1 = (-f1_df*1.e3 + f1)/338.
FSR_2 = (-f2_df*1.e3 + f2)/901.
FSR_list = np.hstack((FSR_1,FSR_2))
FSR_est = np.average(FSR_list)
FSR_std = np.std(FSR_list)
FSR_err = FSR_std/FSR_est
print 'L = ',c/2./FSR_est, '[m] +/- ', c/2./FSR_est *FSR_err*1000.,' [mm]'
print 'FSR = ',FSR_est*1.e-3, ' +/- ', FSR_std*1.e-3,' [kHz]'

first_f1_i = [8,22,36] 
second_f1_i = [4,10,18,24,32,38,46]
unknown_i = [0,2,14,16,28,30,42,44]


order = np.zeros(n_of_peaks)
order[zeroth_i] = 0
order[first_i] = 1
order[second_i] = 2
order[third_i] = 3
order[fourth_i] = 4

order[f1u_i] = 10
order[f1l_i] = -10
order[f2u_i] = 20
order[f2l_i] = -20

order[first_f1_i] = 11
order[second_f1_i] = 12
order[unknown_i] = 100

#print order

zero = f_list[zeroth_i]

first = f_list[first_i]
first[0] = FSR-zero[1]+first[0]
first[1] = FSR-zero[2]+first[1]
first[2] = FSR-zero[3]+first[2]

second = f_list[second_i]
second[0] = second[0]-zero[0] + FSR
second[1] = -zero[1]+second[1] + FSR
second[2] = -zero[2]+second[2] + FSR

third = f_list[third_i]
third[0] = third[0]-zero[0] +2.*FSR
third[1] = -zero[1]+third[1] +2.*FSR
third[2] = -zero[2]+third[2] +2.*FSR
third[3] = -zero[3]+third[3] +2.*FSR

fourth = f_list[fourth_i]
fourth[0] = 3.*FSR-zero[1]+fourth[0]
fourth[1] = 3.*FSR-zero[2]+fourth[1]
fourth[2] = 3.*FSR-zero[3]+fourth[2]


zero[3] = zero[3]-zero[2]-FSR
zero[2] = zero[2]-zero[1]-FSR
zero[1] = zero[1]-zero[0]-FSR
zero[0] = 0


order_hom = [0,0,0,0,1,1,1,2,2,2,3,3,3,3,4,4,4]
f_hom_folded = np.hstack((zero,first,second,third,fourth))
#print f_hom_folded

z,cov = np.polyfit(order_hom,f_hom_folded,1,cov=True)
#print z

#TMS = 34.8 # kHz
TMS = z[0]
print 'TMS = ',TMS, '+/-',np.sqrt(cov[0,0]/(np.size(order_hom)-1))





plt.figure(figsize=(12,7))

plt.semilogy(freq,trans)
#plt.semilogy(freq,rem)
plt.semilogy(freq_p_sort,trans_p_sort,'.')

#plt.semilogy(peak_list[:,0],peak_list[:,1],'o')
for n in range(0,n_of_peaks):
    plt.semilogy(freq,lor(freq,peak_list[n,0],peak_list[n,1],peak_list[n,2]),linewidth=.3,color='k')
    if (order[n]<10) and (order[n]>-1) :
        text = 'Order: '+str(int(order[n]))
    elif order[n]==10:
        text = 'f1u'
    elif order[n]==20:
        text = 'f2u'
    elif order[n]==-10:
        text = 'f1l'
    elif order[n]==-20:
        text = 'f2l'
    elif order[n]==11:
        text = '1st mode\n f1sb'
    elif order[n]==12:
        text = '2nd mode\n f1sb'
    else:
        text = 'Unknown'
    plt.text(peak_list[n,0],peak_list[n,1],text,fontsize=10,rotation=90)

#plt.semilogy(use_freq,lor(use_freq,f0,A,pole))


plt.xlabel('Frequency [kHz]',fontsize=18)
plt.ylabel('Transmission [a.u.]',fontsize=18)

plt.xticks(fontsize=15)
plt.yticks(fontsize=15)
plt.xlim((-11,161) )
plt.ylim((1,8000) )


plt.grid(True,which="Major",ls="-")
plt.grid(True,which="Minor",ls="--")



plt.tight_layout()
plt.savefig("CavityScan.png")
plt.show()

plt.figure(figsize=(10,8))
gs = gridspec.GridSpec(2, 1, height_ratios=[2, 1 ])
plt.subplot(gs[0])
plt.title('Transverse mode spacing (X arm)',fontsize=18)


plt.plot(order_hom,f_hom_folded,'o',color='b')

plt.plot(order_hom,np.array(order_hom)*z[0]+z[1],'-',color='k')

plt.text(2.5,40,'TMS = '+str(round(TMS,2))+' [kHz]',fontsize=20,bbox={'facecolor':'w', 'alpha':1, 'pad':10})

#plt.plot(order,f_list,'o')
#plt.plot([0,0,0,0],zero,'o',color='b')
#plt.plot([1,1,1],first,'o',color='b')
#plt.plot([2,2,2],second,'o',color='b')
#plt.plot([3,3,3,3],third,'o',color='b')
#plt.plot([4,4,4],fourth,'o',color='b')

plt.xlim( (-0.03,4.03) )
plt.xticks((0,1,2,3,4))
plt.xlabel('Order of modes',fontsize=15)
plt.ylabel('Distance from TEM00 [kHz]',fontsize=15)



plt.grid(True,which="Major",ls="-")
plt.grid(True,which="Minor",ls="--")




plt.subplot(gs[1])

plt.plot(order_hom,f_hom_folded-np.array(order_hom)*TMS-z[1],'o',color='b')
plt.plot([-1,5],[0,0],'-',color='k')



plt.xlim( (-0.03,4.03) )
plt.xticks((0,1,2,3,4))

plt.ylim((-0.1,0.1))
plt.grid(True,which="Major",ls="-")
plt.grid(True,which="Minor",ls="--")

plt.tick_params(
    axis='x',          # changes apply to the x-axis
    which='both',      # both major and minor ticks are affected
    bottom='off',      # ticks along the bottom edge are off
    top='on',
    labelbottom='off',
    labeltop='off')  

plt.ylabel('Error [kHz]',fontsize=15)

plt.tight_layout()
plt.savefig("HOMs.png")
plt.show()