import matplotlib.pyplot as plt
import numpy as np
from uncertainties import unumpy as unp
from lmfit.models import LinearModel

#IMPORT DATA
data = np.loadtxt('precision_vs_tacq.dat')
acquisition_time = data[:,0] #s
precision = data[:,1] #ns
error = data[:,2]
# GENERATE OUR OWN TIME SERIES FOR MODELS LATER...
acquisition_time_extended = np.arange(0.1,1000,10) # s
# SORT PRECISION IN ASCENDING ORDER
precision = precision[np.argsort(acquisition_time)]
acquisition_time = np.sort(acquisition_time)

log_precision_error = unp.log10(unp.uarray(precision,error)).flatten()[0].std_dev
log_precision_error_sec = unp.log10(unp.uarray(precision,error)*1e9).flatten()[0].std_dev

frac_precision = precision / (acquisition_time * 1e9)
frac_precision_err = error / (acquisition_time * 1e9)
log_frac_precision_error = unp.log10(unp.uarray(frac_precision,frac_precision_err)).flatten()[0].std_dev

############ FRACTIONAL PRECISION ############

# LINEAR FIT
lmodel = LinearModel()
p = lmodel.guess(np.log10(frac_precision), x=np.log10(acquisition_time))
result = lmodel.fit(np.log10(frac_precision), 
          x=np.log10(acquisition_time), 
          weights = 1/log_frac_precision_error,
          params=p)

# PLOT TO CHECK DATA
# plt.figure()
# plt.errorbar(acquisition_time,
#         frac_precision,
#         yerr = frac_precision_err,
#               marker='x',color='black')

# plt.plot(acquisition_time_extended,
#         10**(result.eval(x=np.log10(acquisition_time_extended))),
#               color='black')
# plt.semilogy()
# plt.semilogx()
# plt.show()

# np.savetxt('frac_precision_vs_tacq.dat',
#            np.c_[acquisition_time,
#                  frac_precision,
#                  frac_precision_err],
#            header='t_acq(s)\tfrac_precision\tfrac_precision_error', delimiter='\t')

# np.savetxt('frac_precision_vs_tacq-fit.dat',
#            np.c_[acquisition_time_extended,
#                  10**(result.eval(x=np.log10(acquisition_time_extended)))],
#            header='t_acq(s)\tfrac_precision', delimiter='\t')

############ ABSOLUTE PRECISION ############

# X/Y UNITLESS FOR PAPER TEXT 
# q = lmodel.guess(np.log10(precision), x=np.log10(acquisition_time*1e9)) # precision is in ns, acq is in sec. 1e9 factor makes the model unitless
# linresult = lmodel.fit(np.log10(precision), 
#             x=np.log10(acquisition_time*1e9), 
#             weights = 1/log_precision_error,
#             params=q)

# plt.figure()
# plt.plot(np.log10(acquisition_time), np.log10(precision*1e-9))
# plt.show()

q = lmodel.guess(np.log10(precision*1e-9), x=np.log10(acquisition_time)) # precision is in ns, acq is in sec. 1e9 factor makes the model unitless
q['slope'].set(-0.5,vary=0)
linresult = lmodel.fit(np.log10(precision*1e-9), 
            x=np.log10(acquisition_time), 
            weights = 1/log_precision_error_sec,
            params=q)

plt.figure()
linresult.plot_fit()
plt.ylim(-12,-10)
plt.show()

print('Linear fit, units: sec\n')
print(linresult.fit_report())

#  X/Y NOT UNITLESS, FOR GRAPH
q = lmodel.guess(np.log10(precision), x=np.log10(acquisition_time)) # precision is in ns, acq is in sec. 1e9 factor makes the model unitless
q['slope'].set(-0.5,vary=0)
linresult = lmodel.fit(np.log10(precision), 
            x=np.log10(acquisition_time), 
            weights = 1/log_precision_error,
            params=q)

# PLOT TO CHECK DATA
plt.figure()
plt.errorbar(acquisition_time,
        precision,
        yerr = error,
        marker='x',color='black')

plt.plot(acquisition_time_extended,
        10**(linresult.eval(x=np.log10(acquisition_time_extended))),
              color='black')
# plt.plot(acquisition_time,
#         10**(linresult.eval(x=np.log10(acquisition_time*1e9))),
#               color='black')
plt.semilogy()
plt.semilogx()
plt.show()

# np.savetxt('precision_vs_tacq-fit.dat',
#            np.c_[acquisition_time_extended,
#                  10**(linresult.eval(x=np.log10(acquisition_time_extended)))],
#            header='t_acq(s)\tprecision', delimiter='\t')