import numpy as np

###

import matplotlib.pyplot as plt
import matplotlib as mpl
mpl.rcParams['figure.figsize'] = (8,6)
mpl.rcParams['xtick.labelsize'] = 20
mpl.rcParams['ytick.labelsize'] = 20
mpl.rcParams['axes.grid'] = True
mpl.rcParams['grid.linestyle'] = ':'
mpl.rcParams['grid.color'] = 'grey'
mpl.rcParams['lines.linewidth'] = 2
mpl.rcParams['axes.labelsize'] = 20
mpl.rcParams['legend.handlelength'] = 3
mpl.rcParams['legend.fontsize'] = 18

from matplotlib import rc
rc('font', **{'family': 'serif', 'serif': ['Computer Modern']})
rc('text', usetex=True)

import seaborn as sns
sea = sns.color_palette("Set1")
#sead = sns.color_palette("dark")

###
infile = np.load('./combined_spectrum_data.npz')

freq_O4a = infile['freq_O4a']
Y_combined = infile['Y_combined']
sigma_combined = infile['sigma_combined']
sigma_O4a_HL = infile['sigma_O4a_HL']
###

fig, ax = plt.subplots(1, 2, figsize=(18, 6), gridspec_kw={'width_ratios': [2, 1]})

# Plot Data
#plt.plot(freq_O4a, Y_O4a_HL, zorder=-1, alpha=0.9, lw=0.4, color=sea[0], label=r'$\mathrm{O4a}$')
ax[0].plot(freq_O4a, Y_combined, zorder=-1, alpha=0.9, lw=0.5, color=sea[1], label=r'$\mathrm{O1-O4a}$')

# Plot error bars
ax[0].plot(freq_O4a, -sigma_combined, color='black', lw=0.5)
ax[0].plot(freq_O4a, sigma_combined, color='black', lw=0.5, label=r'$1\sigma\,\,\mathrm{CL}\,\, \mathrm{O1-O4a}$')

ax[0].plot(freq_O4a, -sigma_O4a_HL, ':', color='black', lw=0.9, label='$1\sigma\,\,\mathrm{CL}\,\, \mathrm{O4a}$')
ax[0].plot(freq_O4a, sigma_O4a_HL, ':', color='black', lw=0.9,)

# Labels, grid, etc
ax[0].set_xlabel(r'$f \,(\mathrm{Hz})$', fontsize=24)
ax[0].set_ylabel(r'$\hat \Omega_{0}(f)$', fontsize=24)

ax[0].set_xlim(20,160)
ax[0].set_ylim([-1e-5,1e-5])

ax[0].legend()

#plt.plot(freq_O4a, Y_O4a_HL, zorder=-1, alpha=0.9, lw=0.4, color=sea[0], label=r'$\mathrm{O4a}$')
ax[1].plot(freq_O4a, Y_combined, zorder=-1, alpha=0.9, lw=0.5, color=sea[1], label=r'$\mathrm{O1-O4a}$')

# Plot error bars
ax[1].plot(freq_O4a, -sigma_combined, color='black', lw=0.5)
ax[1].plot(freq_O4a, sigma_combined, color='black', lw=0.5, label=r'$1\sigma\,\,\mathrm{CL}\,\, \mathrm{O1-O4a}$')

ax[1].plot(freq_O4a, -sigma_O4a_HL, ':', color='black', lw=0.9, label='$1\sigma\,\,\mathrm{CL}\,\, \mathrm{O4a}$')
ax[1].plot(freq_O4a, sigma_O4a_HL, ':', color='black', lw=0.9,)

# Labels, grid, etc
ax[1].set_xlabel(r'$f \,(\mathrm{Hz})$', fontsize=24)
#ax[1].set_ylabel(r'$\hat \Omega_{\rm GW}(f)$', fontsize=24)

ax[1].set_xlim(20,60)
ax[1].set_ylim([-2e-6,2e-6])

#ax[1].legend()

plt.savefig('../paper_figures/Cross_correlation_spectrum_O1_to_O4a.png', dpi=300, bbox_inches='tight')
plt.savefig('../paper_figures/Cross_correlation_spectrum_O1_to_O4a.pdf', dpi=300, bbox_inches='tight')