#!/usr/bin/env python

import atom_interaction as AI
import matplotlib.pyplot as plt
import numpy as np
import seaborn as sns

sns.set_context("poster")
sns.set_style("white")
sns.set_style("ticks", {"xtick.direction": "in",
                        "ytick.direction": "in"})


""" Useful parameters """
# all the times are expressed in nanoseconds. Frequencies in GHz.
l = 0.03
gamma_0 = 1 / 26.2  # e-9
gamma_p = 1 / 13.5  # e-9
dt_tx = 2
dt_rx = 5
eta = 3.68e-3 * 2 * 2
rx_det_eff = 0.0078
rx_decaying_offset = 3.11e-7
rx_rising_offset = 3.2e-7
rx_decaying_offset_err = 4.67e-9
rx_rising_offset_err = 4.95e-9
rx_scalefactor = gamma_0 * eta * rx_det_eff * dt_rx

"""
data import
"""
infile_tx_decaying = 'non_reversed_with_alternating_09_26_onwards_combined_histo_tx_2ns'
infile_tx_rising = 'reversed_with_alternating_09_26_onwards_combined_histo_tx_2ns'
infile_rx_decaying = 'non_reversed_with_alternating_09_26_onwards_combined_histo_rx_5ns'
infile_rx_rising = 'reversed_with_alternating_09_26_onwards_combined_histo_rx_5ns_4'

raw_tx_decaying = np.genfromtxt(infile_tx_decaying)
raw_rx_decaying = np.genfromtxt(infile_rx_decaying)
raw_tx_rising = np.genfromtxt(infile_tx_rising)
raw_rx_rising = np.genfromtxt(infile_rx_rising)

time_rise = (raw_tx_rising[:, 1] - 878.5)
time_deca = (raw_tx_decaying[:, 1] - 881.5)

# integration limits for accidental noise subtraction
idx = np.arange(np.argmin(abs(time_rise + 150)),
                np.argmin(abs(time_rise - 300)))
idx_r = np.arange(np.argmin(abs(time_rise + 150)),
                  np.argmin(abs(time_rise + 100)))
idx_d = np.arange(np.argmin(abs(time_deca + 100)),
                  np.argmin(abs(time_deca + 50)))

bg_ra = np.mean(raw_tx_rising[idx_r, 4])
bg_rn = np.mean(raw_tx_rising[idx_r, 8])
delta_r = (raw_tx_rising[idx, 8] - bg_rn - raw_tx_rising[idx, 4] + bg_ra) / \
    (sum(raw_tx_rising[idx, 8] - bg_rn)) / dt_tx

bg_da = np.mean(raw_tx_decaying[idx_d, 4])
bg_dn = np.mean(raw_tx_decaying[idx_d, 8])
delta_d = (raw_tx_decaying[idx, 8] - bg_dn -
           raw_tx_decaying[idx, 4] + bg_da) / \
          (sum(raw_tx_decaying[idx, 8] - bg_dn)) / dt_tx

time_r = time_rise[idx]
time_d = time_deca[idx]


P_r = (raw_rx_rising[:, 4] - rx_rising_offset) / rx_scalefactor
P_r_err = (raw_rx_rising[:, 5] + rx_rising_offset_err) / rx_scalefactor
time_pr = (raw_rx_rising[:, 1] - (878.5 + 12))
P_d = (raw_rx_decaying[:, 4] - rx_decaying_offset) / rx_scalefactor
P_d_err = (raw_rx_decaying[:, 5] + rx_decaying_offset_err) / rx_scalefactor
time_pd = (raw_rx_decaying[:, 1] - (881.5 + 12))

""" Using experimental data """
# define functions from the experimental data to use in the ODE


def delta_rise(t):
    return np.interp(t, time_r, delta_r)


def delta_decay(t):
    return np.interp(t, time_r, delta_d)

# Integration of the ODE
y_d = AI.my_int(delta_decay, time_d, dt_tx, gamma_0, l)
y_r = AI.my_int(delta_rise, time_r, dt_tx, gamma_0, l)

""" Plotting and results """
time_v, dt = np.linspace(-400, 400, 1000, retstep=True)

plt.figure()
plt.errorbar(time_pr, P_r, yerr=P_r_err, fmt='o')
plt.plot(time_r, y_r, 'o')
plt.plot(time_v, AI.P_exp_rising(time_v, gamma_0, gamma_p, l))
# plt.xlim(-120, 120)
plt.xlim(time_d[[0, -1]])
plt.ylim(-0.005, 0.030)

plt.figure()
plt.errorbar(time_pd, P_d, yerr=P_d_err, fmt='o')
plt.errorbar(time_d, y_d, fmt='o')
plt.plot(time_v, AI.P_exp_decaying(time_v, gamma_0, gamma_p, l))
# plt.xlim(-40, 200)
plt.xlim(time_d[[0, -1]])
plt.ylim(-0.005, 0.030)


print('Max Pe for decaying: {0:.3f}%\n'
      'Max Pe for rising:  {1:.3f}%\n'
      'Ratio = {2:.3f}\n'
      'Extintion decay: {3:.3f}%\n'
      'Extintion rise: {4:.3f}%\n'.format(np.max(y_d) * 100,
                                          np.max(y_r) * 100,
                                          np.max(y_d) / np.max(y_r),
                                          sum(delta_d) * dt_tx * 100,
                                          sum(delta_r) * dt_tx * 100))


plt.show()