from __future__ import division
import numpy as np
from lmfit import Model
import matplotlib.pyplot as plt

R_atom = np.genfromtxt("R_atom")
R_woatom = np.genfromtxt("R_woatom")
print(R_atom)
# Define the fitting function
def lorentzian(x, amp, width, mean,theta):
    top =  (width**2 / 4)
    bottom = (x - mean)**2 + (width/2)**2
    re= amp-theta*amp*top/bottom
    return re
def normal_mode(x, amp, mean, gamma, coupling, kappa, shift,theta):
    atom_freq = x - mean
    cav_freq = atom_freq - shift
    top = kappa**2 * (atom_freq**2 + gamma**2)
    bottom = (cav_freq**2 + kappa**2) * (atom_freq**2 + gamma**2) + 2 * coupling**2 * (atom_freq*cav_freq + kappa*gamma) + coupling**4
    g=coupling
    #theta=np.asarray(Theta,dtype='object')
    del_c=cav_freq
    del_a=atom_freq
    kappa_l=kappa/2
    A=g**2-del_c*del_a+kappa*gamma-2*kappa_l*gamma
    B=del_c*gamma+del_a*kappa-2*kappa_l*del_a
    C=(g**2-del_c*del_a+kappa*gamma)
    D=del_c*gamma+del_a*kappa
    #re=((A*Cnnnn+B*D)/(C**2+D**2))**2+((B*C-A*D)/(C**2+D**2))**2)
    re=amp-theta*amp*((B*C-A*D)**2/(D**2+C**2)**2+(A*C+B*D)**2/(C**2+D**2)**2)
    #re=(1-theta)+re
    
    #re=top/bottom
    return re

def lorentz_fitting(data, init, no_trials = 1):
    # Declare model and initializing data
    lzmod = Model(lorentzian)
    xdata = data[:,0]
    ydata = data[:,1]
    ystderr = data[:,2] / np.sqrt(no_trials)

    # Set parameter hints
    lzmod.set_param_hint('amp', value = init[0], min=0)
    lzmod.set_param_hint('width', value = init[1], min=0)
    lzmod.set_param_hint('mean', value = init[2])
    lzmod.set_param_hint('theta', value = init[3])
    pars = lzmod.make_params()

    fit = lzmod.fit(ydata, pars, x=xdata, weights=1/ystderr, verbose=False)
    print fit.fit_report()

    amp_est = fit.params['amp'].value
    amp_std = fit.params['amp'].stderr
    width_est = fit.params['width'].value
    width_std = fit.params['width'].stderr
    mean_est = fit.params['mean'].value
    mean_std = fit.params['mean'].stderr
    theta_est = fit.params['theta'].value
    theta_std = fit.params['theta'].stderr
    redchi = fit.redchi

    # Return back the result list with the following order:
    result_list = [amp_est, amp_std, width_est, width_std, mean_est, mean_std, theta_est,theta_std,redchi]

    return result_list

def normal_mode_fitting(data, init, no_trials = 1):
    # Declare model and initializing data
    lzmod = Model(normal_mode)
    xdata = data[:,0]
    ydata = data[:,1]
    ystderr = data[:,2] / np.sqrt(no_trials)

    # Set parameter hints
    lzmod.set_param_hint('amp', value = init[0], min=0)
    lzmod.set_param_hint('mean', value = init[1])
    lzmod.set_param_hint('gamma', value = init[2], vary=False)
    lzmod.set_param_hint('coupling', value = init[3], min=0)
    lzmod.set_param_hint('kappa', value = init[4], min=0)
    lzmod.set_param_hint('shift', value = init[5])
    lzmod.set_param_hint('theta', value = init[6])

    pars = lzmod.make_params()

    fit = lzmod.fit(ydata, pars, x=xdata, weights=1/ystderr, verbose=False)
    print fit.fit_report()

    amp_est = fit.params['amp'].value
    amp_std = fit.params['amp'].stderr
    mean_est = fit.params['mean'].value
    mean_std = fit.params['mean'].stderr
    gamma_est = fit.params['gamma'].value
    gamma_std = fit.params['gamma'].stderr
    coupling_est = fit.params['coupling'].value
    coupling_std = fit.params['coupling'].stderr
    kappa_est = fit.params['kappa'].value
    kappa_std = fit.params['kappa'].stderr
    shift_est = fit.params['shift'].value
    shift_std = fit.params['shift'].stderr
    theta_est=fit.params['theta'].value
    theta_std=fit.params['theta'].stderr
    redchi = fit.redchi

    # Return back the result list with the following order:
    result_list = [amp_est, amp_std, mean_est, mean_std, gamma_est, gamma_std, redchi, coupling_est, coupling_std, kappa_est, kappa_std, shift_est, shift_std,theta_est,theta_std]

    coop_est = coupling_est**2 / (2 * kappa_est * gamma_est)
    coop_std = coop_est * np.sqrt((2*coupling_std/coupling_est)**2 + (kappa_std/kappa_est)**2 + (gamma_std/gamma_est)**2)
    print "Cooperativity :", coop_est, "+/-", coop_std


    return result_list

R_woatom_init = np.array([30, 80, 40,0.4])
res_woatom = lorentz_fitting(R_woatom, R_woatom_init)
R_atom_init = [32.7, 30, 6.0659/2, 10, 100, -5,40]
res_atom = normal_mode_fitting(R_atom[0:], R_atom_init)
data_x = np.linspace(-96, 144, 1000)

plt.plot(R_woatom[:,0], R_woatom[:,1], marker='o')
plt.plot(data_x[:], lorentzian(data_x[:], res_woatom[0], res_woatom[2], res_woatom[4],res_woatom[-3]))

plt.plot(R_atom[:,0], R_atom[:,1], marker='o')
plt.plot(data_x[:], normal_mode(data_x[:], res_atom[0], res_atom[2], res_atom[4], res_atom[7], res_atom[9], res_atom[11],res_atom[-2]))


plt.show()
np.savetxt("fit_R_woatom", np.stack((data_x[:], lorentzian(data_x[:], res_woatom[0], res_woatom[2], res_woatom[4],res_woatom[-3])), axis=1))
np.savetxt("fit_R_atom", np.stack((data_x[:], normal_mode(data_x[:], res_atom[0],res_atom[2], res_atom[4], res_atom[7], res_atom[9], res_atom[11],res_atom[-2])), axis=1))