# G1/G2 for C and S classes, for ZTF.r filter
# (c) SsODNet, B.Carry & J.Berthier, Lagrange/OCA & LTE/Obs. Paris

import rocks
import numpy as np
import matplotlib.pyplot as plt
from scipy.stats import gaussian_kde

# Read the full SsODNet BFT
bft = rocks.load_bft(full=True)

# Only rows with phase functions, then explore arrays
cond = bft["phase_functions.G1.value"].notna()
df = bft[cond].explode(
    [
        "phase_functions.id_filter",
        "phase_functions.G1.value",
        "phase_functions.G2.value",
    ]
)

# Choose only ZTF.r filter
cond_ZTFr = df["phase_functions.id_filter"] == "Palomar/ZTF.r"


fig, ax = plt.subplots()

for cl, color in zip(["C", "S"], ["C0", "C1"]):
    # KDE of (G1,G2) for current class
    cond = cond_ZTFr & (df["taxonomy.class"] == cl)
    kde = gaussian_kde(
        np.vstack(
            (
                df.loc[cond, "phase_functions.G1.value"].to_numpy(float),
                df.loc[cond, "phase_functions.G2.value"].to_numpy(float),
            )
        )
    )

    # Generate 2d probability density grid for contour plotting
    x_grid = np.linspace(0, 1, 200)
    y_grid = np.linspace(0, 1, 200)
    xx, yy = np.meshgrid(x_grid, y_grid)
    density = kde.evaluate(np.vstack((xx.ravel(), yy.ravel()))).reshape(xx.shape)

    # Fill within contours at 10,50, 90% probability
    level = np.max(density) * np.array([0.1, 0.5, 0.9, 1])
    ax.contourf(
        xx,
        yy,
        density,
        levels=level,
        colors=color,
        alpha=[0.1, 0.3, 0.5],
    )

    # Draw limiting contours
    ax.contour(
        xx,
        yy,
        density,
        levels=level[:-1],  # Exclude the last level for contour lines
        colors=color,
        linestyles=["dotted", "dashed", "solid"],
        alpha=[0.3, 0.5, 1.0],
    )

    # Label for classes
    ax.fill_between([], [], color=color, alpha=0.5, label=f"{cl} type")

# Labels for contours
ax.plot([], [], color="black", ls="dotted", alpha=0.3, label=f"10% KDE")
ax.plot([], [], color="black", ls="dashed", alpha=0.5, label=f"50% KDE")
ax.plot([], [], color="black", ls="solid", alpha=1.0, label=f"90% KDE")

# (G1,G2) boundaries from Muinonen et al. 2010
ax.plot([0, 1, 0, 0], [0, 0, 1, 0], linestyle="-.", color="black", label="Muinonen2010")

# Axes
ax.legend(loc="upper right")
ax.set_xlabel("G1")
ax.set_ylabel("G2")

# To disk
fig.tight_layout()
fig.savefig("test-Phase-CS.png", transparent=True, dpi=180)