"""
=======================
Spin-lattice relaxation
=======================

Example for
"""
import numpy as np
from matplotlib import pyplot as plt

from nmreval.distributions import ColeDavidson
from nmreval.nmr import Relaxation, RelaxationEvaluation
from nmreval.nmr.coupling import Quadrupolar
from nmreval.utils.constants import kB

# Define temperature range
inv_temp = np.linspace(3, 9, num=30)
temperature = 1000/inv_temp

# spectral density parameter
ea = 0.45
tau = 1e-21 * np.exp(ea / kB / temperature)
gamma_cd = 0.1

# interaction parameter
omega = 2*np.pi*46e6
delta = 120e3
eta = 0

r = Relaxation()
r.set_distribution(ColeDavidson)  # the only parameter that has to be set beforehand

t1_values = r.t1(omega, tau, gamma_cd, mode='bpp',
                 prefactor=Quadrupolar.relax(delta, eta))
# add noise
rng = np.random.default_rng(123456789)
noisy = (rng.random(t1_values.size)-0.5) * 0.5 * t1_values + t1_values

# set parameter and data
r_eval = RelaxationEvaluation()
r_eval.set_distribution(ColeDavidson)
r_eval.set_coupling(Quadrupolar, (delta, eta))
r_eval.data(temperature, noisy)
r_eval.omega = omega

t1_min_data, _ = r_eval.calculate_t1_min()  # second argument is None
t1_min_inter, line = r_eval.calculate_t1_min(interpolate=1, trange=(160, 195), use_log=True)

fig, ax = plt.subplots()
ax.semilogy(1000/t1_min_data[0], t1_min_data[1], 'rx', label='Data minimum')
ax.semilogy(1000/t1_min_inter[0], t1_min_inter[1], 'r+', label='Parabola')
ax.semilogy(1000/line[0], line[1])

found_gamma, found_height = r_eval.get_increase(t1_min_inter[1], idx=0, mode='distribution')
print(found_gamma)

plt.axhline(found_height)
plt.show()

#%%
# Now we found temperature and height of the minimum we can calculate the correlation time

plt.semilogy(1000/temperature, tau)
tau_from_t1, opts = r_eval.correlation_from_t1()
print(opts)
plt.semilogy(1000/tau_from_t1[:, 0], tau_from_t1[:, 1], 'o')
plt.show()