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")

###

data_file=np.load('figure_data/figure_2_data.npz')
freq_O4a=data_file['freq_O4a']
Y_combined=data_file['Y_combined']
sigma_combined=data_file['sigma_combined']
Y_O4a_HL=data_file['Y_O4a_HL']
sigma_O4a_HL=data_file['sigma_O4a_HL']
Y_combined_o123=data_file['Y_combined_o123']
sigma_combined_o123=data_file['sigma_combined_o123']

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

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

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

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

# Labels, grid, etc
ax[0].set_xlabel(r'$f \,[\mathrm{Hz}]$', fontsize=24)
ax[0].set_ylabel(r'$\hat \Omega_{\rm GW}(f)$', fontsize=24)
ax[0].set_xlim(20,160)
ax[0].set_ylim([-1e-5,1e-5])
#ax[0].grid()
ax[0].legend()

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

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

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

# 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].grid()
#ax[1].legend()

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

#plt.show()
exit()