#!/usr/bin/env python
import matplotlib

matplotlib.use("agg")
import seaborn as sns

import matplotlib.gridspec as gridspec
import matplotlib.pyplot as plt

matplotlib.rcParams["text.usetex"] = True
matplotlib.rcParams["font.size"] = 9
matplotlib.rcParams["savefig.dpi"] = 300
matplotlib.rcParams["text.latex.preamble"] = [r"\usepackage{amsmath}"]
matplotlib.rcParams["legend.fontsize"] = 9
matplotlib.rcParams["font.family"] = "serif"


cp = sns.color_palette("colorblind", 6)

import h5py


if __name__ == "__main__":

    # Load the h5 data
    fp = h5py.File("GW190425_data.h5", "r")

    # Create the figure itself
    fig, axes = plt.subplots(2, 1, figsize=(3.375, 4.2))
    labels = [
        r"IMRPhenomPv2_NRTidal",
        r"IMRPhenomD_NRTidal",
        r"SEOBNRv4T_surrogate",
        r"TEOBResumS",
    ]
    # Top panel: lambda
    # Plot the high spin prior results
    for i, grp in enumerate(labels):
        d = fp[f"{grp}/HS/lambda_prob"][()]
        axes[0].plot(d[:, 0], d[:, 1], color=cp[i], lw=1)
    # Plot the low spin prior results
    for i, grp in enumerate(labels):
        d = fp[f"{grp}/LS/lambda_prob"][()]
        axes[0].plot(d[:, 0], d[:, 1], color=cp[i], ls="--", lw=1)
    # Plot the high spin prior upper bounds
    for i, grp in enumerate(labels):
        d = fp[f"{grp}/HS/lambda_UL"][()]
        axes[0].axvline(x=[d], color=cp[i], lw=1)
    # Plot the low spin prior upper bounds
    for i, grp in enumerate(labels):
        d = fp[f"{grp}/LS/lambda_UL"][()]
        axes[0].axvline(x=d, color=cp[i], ls="--", lw=1)

    axes[0].set_xlim(0, 1200)
    axes[0].set_ylim(0, 0.004)
    axes[0].set_xlabel(r"$\tilde\Lambda$")
    axes[0].set_ylabel("Posterior Density")

    # Bottom panel: q

    # Plot the high spin prior results
    for i, grp in enumerate(labels):
        d = fp[f"{grp}/HS/q_prob"][()]
        axes[1].plot(d[:, 0], d[:, 1], color=cp[i], label=grp.replace("_", "\_"), lw=1)

    # Plot the low spin prior results
    for i, grp in enumerate(labels):
        d = fp[f"{grp}/LS/q_prob"][()]
        axes[1].plot(d[:, 0], d[:, 1], color=cp[i], ls="--", lw=1)

    axes[1].legend(loc=2, frameon=False)
    axes[1].set_xlim(0.2, 1)
    y_min, y_max = axes[1].get_ylim()
    axes[1].set_ylim(0, y_max)

    axes[1].set_xlabel(r"$q$")
    axes[1].set_ylabel("Posterior Density")

    plt.tight_layout()

    fp.close()
    plt.savefig("GW190425ALL.pdf")