import matplotlib.pyplot as plt
import numpy as np
import pandas as pd

###

import matplotlib as mpl
mpl.rcParams['figure.figsize'] = (8,6)
mpl.rcParams['xtick.labelsize'] = 22
mpl.rcParams['ytick.labelsize'] = 22
mpl.rcParams['axes.grid'] = False
mpl.rcParams['grid.linestyle'] = ':'
mpl.rcParams['grid.color'] = 'grey'
mpl.rcParams['lines.linewidth'] = 2
mpl.rcParams['axes.labelsize'] = 24
mpl.rcParams['legend.handlelength'] = 2
mpl.rcParams['legend.fontsize'] = 22

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

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

###


budget_file = pd.read_csv("./Magnetic_budget_O4a.dat")
frequencies = budget_file["frequencies [Hz]"]
magn_budget = budget_file["Magnetic_budget"]
one_sigma = budget_file["1sigma uncertainty"]
frequencies_pi_curve = budget_file["frequencies pi_curve [Hz]"].dropna()
#pi_curve = budget_file["pi curve"].dropna()

pi_curve = np.loadtxt("../../PI_curves/PI_O1_to_O4a_HLV.txt").T

fig, ax = plt.subplots(1,1)
ax.fill_between(frequencies, np.abs(magn_budget) + one_sigma, np.abs(magn_budget), label = r"1$\sigma$", color = sea[0])
ax.fill_between(frequencies, np.abs(magn_budget) + 2*one_sigma, np.abs(magn_budget), label = r"2$\sigma$", color = sea[0], alpha=0.5)
ax.fill_between(frequencies, np.abs(magn_budget) + 3*one_sigma, np.abs(magn_budget), label = r"3$\sigma$", color = sea[0], alpha=0.25)
ax.plot(frequencies, np.abs(magn_budget), label = "O4a budget", color=sea[1])
ax.plot(pi_curve[0], 2*pi_curve[1], label = "O1-O4a ($2\sigma$)", color='k')
plt.xlabel('Frequency [Hz]')
plt.ylabel(r"$\Omega_{\rm GW}(f)$")
ax.set_xlim(20,1000)
ax.set_ylim(1e-12,3e-3)
ax.set_xscale('log')
ax.set_yscale('log')
ax.minorticks_on()
plt.grid(True, which="both", ls="-")
plt.legend(loc='upper left')
fig.tight_layout()
plt.savefig("../../paper_figures/O4a_budget.pdf", dpi = 500, bbox_inches = "tight")
plt.savefig("../../paper_figures/O4a_budget.png", dpi = 500, bbox_inches = "tight")
plt.show()