import numpy as np
import matplotlib.pyplot as plt
import seaborn as sns
from lmfit import Model, Parameters
# from uncertainties import ufloat


# Function model for the fit
def beat_osc(x, g, r, x0, phi, fr, y1, tau_r, tau0, tau1):
    phi = phi/180*np.pi
    x = x-x0
    f = (x > 0)*g*(2*r*np.exp(-x*(1/(2*tau0)+1/(2*tau1)))*np.cos(fr*x+phi) +
                   np.exp(-x/tau0)+r**2*np.exp(-x/tau1))
    rise = (x <= 0) * g * (2*r*np.cos(phi)+r**2+1) * np.exp(x/tau_r)
    return f+rise+y1


# Fitting
gmod = Model(beat_osc)
pars = Parameters()
pars.add('g', value=30)
pars.add('r', value=.75, min=.1, vary=0)
pars.add('x0', value=0, vary=1)
pars.add('fr', value=2*np.pi*0.266, vary=0)
pars.add('phi', value=180, vary=0)
pars.add('y1', value=4.4, vary=1)
pars.add('tau_r', value=0.3, vary=0)
pars.add('tau0', value=5.65, vary=0)
pars.add('tau1', value=13.1, vary=0)

# file number 1
filename = 'ff.dat'
rawdata = np.genfromtxt(filename)

min_idx = 252
max_idx = 1710
delay = rawdata[min_idx:max_idx, 0]
counts = rawdata[min_idx:max_idx, 1]
counts_err = np.sqrt(counts)

# Remove errors = 0 due to the low statistics
counts_err[counts_err == 0] = 1
idx_min = 0
idx_max = -1
result = gmod.fit(counts[idx_min:idx_max],
                  x=delay[idx_min:idx_max],
                  params=pars,
                  weights=1/counts_err[idx_min:idx_max])
print(result.fit_report())

sns.set_style("white")
plt.errorbar(delay, counts, yerr=counts_err, fmt='o')
plt.plot(delay[idx_min:idx_max], result.best_fit, 'r-')
plt.xlim(-10, 50)
sns.despine(offset=0, trim=True)
# plt.show()

"""
FFT
"""
step = 0.05e-9
plt.figure()
counts = rawdata[265:1250, 1]
plt.plot(counts, 'o')
f_counts = np.fft.rfft(counts)
freq = np.linspace(0, 1/step/2, len(f_counts)+1)
plt.figure()
plt.plot(freq[:-1]/1e6, np.abs(f_counts)**2, 'o-')
plt.xlim(0, 500)
plt.ylim(0, .3e8)
plt.xlabel('Frequency (MHz)')
plt.show()