import matplotlib.pyplot as plt
from matplotlib import gridspec
import numpy as np
import os
import re
from lmfit import Model, Parameters
from lmfit.models import ConstantModel



rgg_filename = ".\\normal_g2_RGG.dat"
mercury_filename = ".\\normal_g2_mercury.dat"

# skip the parts with deadtime
data_arr_rgg = np.genfromtxt(rgg_filename,skip_header=2).T
data_arr_mercury = np.genfromtxt(mercury_filename,skip_header=3,skip_footer=7).T


def one_gaussian_g2_model(tau, c, A_0, tau_fwhm,tau_0):
	"""
	tau: time_bin_vector (ns)
	c: floor_offset
	A: decay_amplitude,
	tau_0: timing offset
	tau_c: time_constant
	"""
	tau_sigma = tau_fwhm/(2*np.sqrt(2*np.log(2)))
	return c + A_0*(np.exp(-(tau-tau_0)**2/(2*(tau_sigma**2))))

plt.clf()
plt.rcParams['font.size'] = '12'
fig, axs = plt.subplots(2,1,figsize=(8.6,6),constrained_layout=True)

axs = [axs[1],axs[0]]

rgg_file = open(rgg_filename)
line1 = re.split(' |,',rgg_file.readline())
line2 = re.split(' |,',rgg_file.readline())
rgg_scaling_factor = float(line2[4][:-1])

rgg_start_time = -2000 #ns
rgg_end_time = 2000 #ns
theory_time_vector_rgg = np.linspace(rgg_start_time,rgg_end_time,1000)

# rotiing ground glass
axs[0].set_xlabel(r"time difference $\tau$ ($\mu s$)")
axs[0].set_ylabel(r"$g^{(2)}(\tau)$")
axs[0].set_ylim(np.asarray((0.95,2.05)))
axs[0].set_xlim((rgg_start_time*1e-3,rgg_end_time*1e-3))
axs[0].errorbar(data_arr_rgg[0]*1e-3,data_arr_rgg[1]*rgg_scaling_factor,yerr=rgg_scaling_factor*np.sqrt(data_arr_rgg[1]),fmt="r.",zorder=1)

tau_fwhm=  277.460526 #+/- 0.58185803 (0.21%) (init = 124)
tau_0=     314.283128 #+/- 0.25266698 (0.08%) (init = 312)
c=        10139.4397 #+/- 1.85375171 (0.02%) (init = 10725)
A_0=      8706.67753 #+/- 17.8400148 (0.20%) (init = 8996)
axs[0].plot(theory_time_vector_rgg*1e-3,one_gaussian_g2_model(theory_time_vector_rgg, c, A_0, tau_fwhm,0)*rgg_scaling_factor,"k-",zorder=2)
print(rgg_scaling_factor)
# get text location
min_y,max_y = axs[0].get_ylim()
text_y = (max_y - min_y)*0.8 + min_y

min_x,max_x = axs[0].get_xlim()
text_x = (max_x - min_x)*0.01 + min_x

# axs[0].text(text_x,text_y,"RGG")
# axs[0].axvline(-delay,color="k",linestyle="--")
# axs[0].axvline(delay,color="k",linestyle="--")

# plot secondary coincidences axis
axs0_2 = axs[0].twinx()
axs0_2.set_ylim(np.asarray(axs[0].get_ylim())/rgg_scaling_factor)
axs0_2.set_ylabel("coincidences")


# Mercury
mercury_file = open(mercury_filename)
line1 = re.split(' |,',mercury_file.readline())
line2 = re.split(' |,',mercury_file.readline())
mercury_scaling_factor = float(line2[4][:-1])


axs[1].set_ylim(np.asarray((0.95,1.21)))
axs[1].set_xlabel(r"time difference $\tau$ (ns)")
axs[1].set_ylabel(r"$g^{(2)}(\tau)$")
axs[1].errorbar(data_arr_mercury[0],data_arr_mercury[1]*mercury_scaling_factor,yerr=mercury_scaling_factor*np.sqrt(data_arr_mercury[1]),fmt="b.",label="Hg-lamp",zorder=3)

hg_start_time = -2 #ns
hg_end_time = 2 #ns
theory_time_vector_hg = np.linspace(hg_start_time,hg_end_time,1000)

tau_fwhm=  0.34824027 #+/- 0.01964496 (5.64%) (init = 0.3)
tau_0=     207.592286 #+/- 0.00824290 (0.00%) (init = 207.5859)
c=         13278.2958 #+/- 13.6996252 (0.10%) (init = 13260)
A_0=       1883.17796 #+/- 93.7232400 (4.98%) (init = 2043)


axs[1].plot(theory_time_vector_hg,one_gaussian_g2_model(theory_time_vector_hg, c, A_0, tau_fwhm,0)*mercury_scaling_factor,"k-",zorder=4)

axs[1].set_xlim((hg_start_time,hg_end_time))
# plot secondary coincidences axis
axs1_2 = axs[1].twinx()
axs1_2.set_ylim(np.asarray(axs[1].get_ylim())/mercury_scaling_factor)
axs1_2.set_ylabel("coincidences")


# get text location
min_y,max_y = axs[1].get_ylim()
text_y = (max_y - min_y)*0.8 + min_y

min_x,max_x = axs[1].get_xlim()
text_x = (max_x - min_x)*0.01 + min_x
# axs[1].text(text_x,text_y,"Hg-lamp")
print(mercury_scaling_factor)

# 
# plt.show()
plt.savefig("normal_g2_compare.eps",format="eps",dpi=600)