import numpy as np
import time
import matplotlib.pyplot as plt
import fpfind.lib.parse_timestamps as parser
from S15lib.g2lib.g2lib import histogram
from fpfind import TSRES
from scipy.optimize import curve_fit
from lmfit import Model,Parameters
from scipy.stats import poisson as poisson_dis
import pickle

def g2exp(x,peak,cen,wid,base):
    return base+peak*np.exp(-2*np.abs(x-cen)/wid)

def time_to_distance(x):
    #us to m
    return x/2*1e-6*(3*1e8)

def distance_to_time(x):
    #m to us
    return x*2/(3*1e8)*1e6

foldername=''
filenames=['12600_-10200_7mA_','7700_-5950_7mA_','1400_-12750_7mA_']
"""
'yes' or 'dumped', 'yes' if you want to redo histogram
Make sure to change filename_dump to prevent overide
switch to 'dumped' to plot
"""
to_dump='dumped'
filename_dump='20260505_0051_singtel_building_3points_hist_0p25ns'
font_size=20

start_time=time.time()

fig, ax = plt.subplots(figsize=(11.69,8.27))
fig.subplots_adjust(left=0.115,right=0.96,top=0.92,bottom=0.1)
#ax.set_title(r'g$^{(2)}(\tau)$ for 3 different points',fontsize=font_size)
ax.set_xlabel('Time delay (μs)',fontsize=font_size)
ax.set_ylabel(r'g$^{\mathregular{(2)}}$(τ)',fontsize=font_size)
secax = ax.secondary_xaxis('top', functions=(time_to_distance, distance_to_time))
secax.set_xlabel('Distance (m)',fontsize=font_size)
ax.tick_params(axis='both', labelsize=font_size)
secax.tick_params(axis='x', labelsize=font_size)
ax.set_xlim(1,9.5)
ax.set_ylim(1,1.35)

hy_combined=[]
hx_combined=[]

if to_dump=='yes':
    for i in range(len(filenames)):
        filename=filenames[i]

        t,p=parser.read_a1(foldername+filename+'.ts',legacy=True,resolution=TSRES.PS4,fractional=False)
        print("time taken to parse timestamps {}s".format(time.time()-start_time))

        ta=t[(p&1).astype(bool)]
        tb=t[(p&4).astype(bool)]

        s1=len(ta)
        s2=len(tb)
        tacq=(np.max(t)-np.min(t))
        tbin=256/4
        acc=s1/(tacq*4*1e-12)*s2/(tacq*4*1e-12)*(tacq*4*1e-12)*(tbin*4*1e-12)
        print("singles1 {}/s, singles2 {}/s, tacq {}s, bin size {}ns, acc {}/s".format(s1/(tacq*4*1e-12),s2/(tacq*4*1e-12),(tacq*4*1e-12),tbin/256,acc))
        #print(s1,s2,(tacq*4*1e-12),(tbin*4*1e-12),acc)

        hy,hx=histogram(ta,tb,duration=4250*2*256, center=5250*256,resolution=tbin)
        print("time taken to parse timestamps and histogram {}s".format(time.time()-start_time))
        hy_combined.append(hy)
        hx_combined.append(hx)

    with open(foldername+filename_dump+'.pickle',"wb") as f:
        pickle.dump((hx_combined,hy_combined),f)

elif to_dump=='dumped':
    with open(foldername+filename_dump+'.pickle',"rb") as f:
        (hx_combined,hy_combined)=pickle.load(f)

    for i in range(3):
        hx=hx_combined[i]/256
        print(hx)
        hy=hy_combined[i]
        acc=np.mean(hy)
        max_time=hx[np.argmax(hy)]
        gmodel = Model(g2exp)
        params=Parameters()
        params.add('peak',value=0.3*acc,min=0,max=0.7*acc)
        params.add('cen',value=max_time,min=0.9*max_time,max=1.1*max_time)
        params.add('wid',value=4,min=0,max=100)
        params.add('base',value=acc,min=-0.9*acc,max=1.1*acc)

        result = gmodel.fit(hy, params, weights=1/np.sqrt(hy), x=hx)
        print(result.fit_report())
        print(result.summary()['params'][1][1]*10**-9*299792458/2)
        print(result.summary()['params'][1][7]*10**-9*299792458/2)
        #print(result.summary()['params'][0][7])
        #print(result.summary()['params'][0][1]/result.summary()['params'][3][1])
        #print(result.summary()['params'][0][7]/result.summary()['params'][3][1])
        #print(result.summary()['params'][2][1])
        #print(result.summary()['params'][2][7])
        print('\n')

    for i in range(len(hx_combined)):
        if i==0:
            acc=np.mean(hy_combined[i])
            ax.plot(hx_combined[i]/256/1e3,hy_combined[i]/acc,label='(2,12)',c='k')
            plt.text(1.43,np.max(hy_combined[i])/acc+0.015,'(2,12)',horizontalalignment='center',verticalalignment='center',fontsize=font_size,c='k')
        elif i==1:
            acc=np.mean(hy_combined[i])
            ax.plot(hx_combined[i]/256/1e3,hy_combined[i]/acc,label='(9,7)',c='r')
            plt.text(8.6,np.max(hy_combined[i])/acc+0.015,'(9,7)',horizontalalignment='center',verticalalignment='center',fontsize=font_size,c='r')
        elif i==2:
            acc=np.mean(hy_combined[i])
            ax.plot(hx_combined[i]/256/1e3,hy_combined[i]/acc,label='(18,15)',c='g')
            plt.text(8.927,np.max(hy_combined[i])/acc+0.015,'(18,15)',horizontalalignment='center',verticalalignment='center',fontsize=font_size,c='g')

plt.show()