"""
==========
T1 minimum
==========

``RelaxationEvaluation`` is used to get width parameter from a T1 minimum.
As a subclass of ``Relaxation`` it can also be used to calculate Relaxation times.
The basic steps are:

* Determine a T1 minimum with `nmreval.nmr.RelaxationEvaluation.calculate_t1_min`
* Calculate width parameter of a spectral density/coupling constants/... with
  ``RelaxationEvaluation.get_increase``
* Calculate correlation times from these values with ``RelaxationEvaluation.correlation_from_t1``
"""
import numpy as np
import matplotlib.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.4

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

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

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

ax_t1 = plt.figure().add_subplot()
ax_t1.semilogy(inv_temp, t1_values, label='Calculated T1')
ax_t1.semilogy(inv_temp, noisy, 'o', label='Noise')
ax_t1.legend()

plt.show()


# Actual evaluation starts here
# setting necessary parameter
r_eval = RelaxationEvaluation()
r_eval.set_distribution(ColeDavidson)
r_eval.set_coupling(Quadrupolar, (delta, eta))
r_eval.set_data(temperature, noisy)
r_eval.omega = omega

# Find a T1 minumum
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)

ax_min = plt.figure().add_subplot()
ax_min.semilogy(inv_temp, noisy, 'o', label='Data')
ax_min.semilogy(1000/line[0], line[1], '--')
ax_min.semilogy(1000/t1_min_data[0], t1_min_data[1], 'C2X',label='Data minimum')
ax_min.semilogy(1000/t1_min_inter[0], t1_min_inter[1], 'C3P',label='Parabola')
ax_min.set_xlim(4.5, 7)
ax_min.set_ylim(1e-3, 1e-1)
ax_min.legend()

# Vary the first (and for Cole-Davidson, only) parameter of the spectral density
found_gamma, found_height = r_eval.get_increase(t1_min_inter[1], idx=0, mode='distribution')
print(f'Minimum at {found_height} for {found_gamma}; input is {gamma_cd}')
plt.show()

##################################################################################
# Calculation of correlation times uses previously parameter for spectral density
# and prefactor

tau_from_t1, opts = r_eval.correlation_from_t1(mode='mean')
print(f'Used options: {opts}')

ax_tau = plt.figure().add_subplot()
ax_tau.semilogy(inv_temp, tau*gamma_cd, label='Original input')
ax_tau.semilogy(1000/tau_from_t1[:, 0], tau_from_t1[:, 1], 'o', label='Calculated')
ax_tau.legend()
plt.show()