"""
=========
Cole-Cole
=========

Example for Cole-Cole distributions
"""
import matplotlib.pyplot as plt
import numpy as np

from nmreval.distributions import ColeCole

x = np.logspace(-5, 5, num=101)

cc = ColeCole

alpha_CC = [0.3, 0.5, 0.7]

fig, axes = plt.subplots(2, 3, constrained_layout=True)

lines = []
for a in alpha_CC:
    axes[0, 0].plot(np.log10(x), cc.correlation(x, 1, a))
    axes[1, 0].plot(np.log10(x), np.log10(cc.specdens(x, 1, a)))
    axes[0, 1].plot(np.log10(x), np.log10(cc.susceptibility(x, 1, a).real))
    axes[1, 1].plot(np.log10(x), np.log10(cc.susceptibility(x, 1, a).imag))
    l, = axes[0, 2].plot(np.log10(x), cc.distribution(x, 1, a),
                         label=rf'$\alpha={a}$')
    lines.append(l)

fig_titles = ('Correlation function', 'Susceptibility (real)', 'Distribution',
              'Spectral density', 'Susceptibility (imag)')
fig_xlabel = (r'$\log(t/\tau_\mathrm{HN})$', r'$\log(\omega\tau_\mathrm{HN})$',
              r'$\log(\tau/\tau_\mathrm{HN})$', r'$\log(\omega\tau_\mathrm{HN})$',
              r'$\log(\omega\tau_\mathrm{HN})$')
fig_ylabel = (r'$C(t)$', r"$\log(\chi'(\omega))$", r'$G(\ln\tau)$',
              r'$\log(J(\omega))$', r"$\log(\chi''(\omega))$")

for title, xlabel, ylabel, ax in zip(fig_titles, fig_xlabel, fig_ylabel, axes.ravel()):
    ax.set_title(title)
    ax.set_xlabel(xlabel)
    ax.set_ylabel(ylabel)

labels = [l.get_label() for l in lines]
leg = fig.legend(lines, labels, loc='center left', bbox_to_anchor=(1.05, 0.50),
                 bbox_transform=axes[1, 1].transAxes)

fig.delaxes(axes[1, 2])

plt.show()