"""
This script creates the spectroscopy plots in Figure 1.
"""

import numpy as np
import matplotlib.pyplot as plt
from time import strftime
from matplotlib import colors, ticker
from scipy.io import loadmat
from cmcrameri import cm

datestr = strftime('%Y-%m-%d')
plt.style.use("paperstyle.mplstyle")

pt = 1 / 72  # inches per pt

# %%    Load data for 2D maps

# sigma minus - 20240721 scan 1-25
mat_contents = loadmat('data_FIG1_spectroscopy_sigma_minus.mat')
wf198_freq_sigma_m = mat_contents['wf198_freq'][0]
dProbe_sigma_m = mat_contents['dProbe'][0][0]
dProbeRampStepSize_sigma_m = mat_contents['dProbeRampStepSize'][0][0]
dProbeRampTo_sigma_m = mat_contents['dProbeRampTo'][0][0]
transmission_sigma_m = mat_contents['transmission']

# sigma plus - 20240720 scan 27-58
mat_contents = loadmat('data_FIG1_spectroscopy_sigma_plus.mat')
wf198_freq_sigma_p = mat_contents['wf198_freq'][0]
dProbe_sigma_p = mat_contents['dProbe'][0][0]
dProbeRampStepSize_sigma_p = mat_contents['dProbeRampStepSize'][0][0]
dProbeRampTo_sigma_p = mat_contents['dProbeRampTo'][0][0]
transmission_sigma_p = mat_contents['transmission']

wf198_freq_sigma_m_stepsize = abs(wf198_freq_sigma_m[1] - wf198_freq_sigma_m[0])
wf198_freq_sigma_p_stepsize = abs(wf198_freq_sigma_p[1] - wf198_freq_sigma_p[0])

offset = 194

# %%    Plot

cmap = cm.oslo

fig, ax = plt.subplots(2, 1, figsize=(155 * pt, 172 * pt),
                       sharex='col', sharey='col', dpi=800, constrained_layout=True)

divnorm = colors.TwoSlopeNorm(vmin=0, vcenter=0.5, vmax=1)
im = ax[0].imshow(transmission_sigma_m, aspect='auto', origin='lower',
                  extent=(np.min(wf198_freq_sigma_m) - wf198_freq_sigma_m_stepsize * 0.25 - offset,
                          np.max(wf198_freq_sigma_m) + wf198_freq_sigma_m_stepsize * 0.25 - offset,
                          dProbe_sigma_m - dProbeRampStepSize_sigma_m * 0.5,
                          dProbeRampTo_sigma_m + dProbeRampStepSize_sigma_m * 0.5),
                  interpolation='none', cmap=cmap, norm=divnorm, rasterized=True)
ax[1].imshow(transmission_sigma_p, aspect='auto', origin='lower',
             extent=(np.min(wf198_freq_sigma_p) - wf198_freq_sigma_p_stepsize * 0.25 - offset,
                     np.max(wf198_freq_sigma_p) + wf198_freq_sigma_p_stepsize * 0.25 - offset,
                     dProbe_sigma_p - dProbeRampStepSize_sigma_p * 0.5,
                     dProbeRampTo_sigma_p + dProbeRampStepSize_sigma_p * 0.5),
             interpolation='none', cmap=cmap, norm=divnorm, rasterized=True)
ax[0].hlines(41, np.min(wf198_freq_sigma_p) - wf198_freq_sigma_p_stepsize * 0.25 - offset,
             np.max(wf198_freq_sigma_p) + wf198_freq_sigma_p_stepsize * 0.25 - offset,
             colors='w', alpha=0.7)
ax[1].hlines(41, np.min(wf198_freq_sigma_p) - wf198_freq_sigma_p_stepsize * 0.25 - offset,
             np.max(wf198_freq_sigma_p) + wf198_freq_sigma_p_stepsize * 0.25 - offset,
             colors='w', alpha=0.7)
ax[0].xaxis.set_major_locator(ticker.MultipleLocator(50))
ax[0].yaxis.set_major_locator(ticker.MultipleLocator(5))
cb = fig.colorbar(im, ax=ax[:], shrink=0.9, location='right', pad=0.1)
cb.set_label(r'Probe transmission', labelpad=3)
ax[0].set_ylabel(r'$\delta_\mathrm{p}/(2\pi\cross\mathrm{MHz})$', labelpad=3)
ax[1].set_ylabel(r'$\delta_\mathrm{p}/(2\pi\cross\mathrm{MHz})$', labelpad=3)
ax[1].set_xlabel(r'Trap detuning $\Delta/(2\pi\cross\mathrm{MHz})$', labelpad=3)

plt.savefig(datestr.replace('-', '') + 'FIG1-spectroscopy-data.pdf')
plt.show()