import os
from astropy.table import Table
import numpy as np
from scipy.stats import gaussian_kde
import matplotlib.pyplot as plt

# Load and process data
data_dir = '../../earthorbitplan/scenarios/farah.h5'
Farah = Table.read(data_dir)[:10000]
Farah.sort('mass1')

# Create subplots
fig, axs = plt.subplots(1, 2)

# increase the font size of the axes
for ax in axs:
    for tick in ax.get_xticklabels() + ax.get_yticklabels():
        tick.set_fontname("Times New Roman")
        tick.set_fontsize(14)

# Mass distribution (log scale)
mass1 = np.log10(Farah['mass1'])
mass2 = np.log10(Farah['mass2'])
xy_mass = np.vstack([mass1, mass2])
z_mass = gaussian_kde(xy_mass)(xy_mass)
idx_mass = z_mass.argsort()
mass1, mass2, z_mass = mass1[idx_mass], mass2[idx_mass], z_mass[idx_mass]
msc = axs[0].scatter(mass1, mass2, c=z_mass, s=5)
axs[0].set_xlabel(r'$\log_{10}(m_1)\ [M_\odot]$',  fontname="Times New Roman", size=16, fontweight="bold")
axs[0].set_ylabel(r'$\log_{10}(m_2)\ [M_\odot]$', fontname="Times New Roman", size=16, fontweight="bold")
cbar1 = fig.colorbar(msc, ax=axs[0])
cbar1.set_label("Event density", fontname="Times New Roman", size=18)


# Spin distribution
spin1z = Farah['spin1z']
spin2z = Farah['spin2z']
xy_spin = np.vstack([spin1z, spin2z])
z_spin = gaussian_kde(xy_spin)(xy_spin)
idx_spin = z_spin.argsort()
spin1z, spin2z, z_spin = spin1z[idx_spin], spin2z[idx_spin], z_spin[idx_spin]
ssc = axs[1].scatter(spin1z, spin2z, c=z_spin, s=5)
axs[1].set_xlabel(r'$\mathrm{spin}_1$', fontname="Times New Roman", size=16, fontweight="bold")
axs[1].set_ylabel(r'$\mathrm{spin}_2$', fontname="Times New Roman", size=16, fontweight="bold")
cbar2 = fig.colorbar(ssc, ax=axs[1])
cbar2.set_label("Event density", fontname="Times New Roman", size=18)

# Adjust layout and figure size
fig.set_size_inches(14, 6)
plt.tight_layout()
plt.show()