#!/usr/bin/env python3
"""
Composite publication figure aligning:
1) Model-derived event-rate envelope
2) Surface adjustment rate at sites
3) Proximity to equilibrium margin
4) Spatial stability of site environments

Inputs:
- site_sampling_timeseries.csv
- model_event_rate_envelope.csv

Output:
- fig_composite_return_diagnostics.png
"""

import pandas as pd
import matplotlib.pyplot as plt

# ============================================================
# LOAD DATA
# ============================================================

df_sites = pd.read_csv("site_sampling_timeseries.csv")
df_event = pd.read_csv("model_event_rate_envelope.csv")

df_homo = df_sites[df_sites["group"] == "Homo"]
df_civ  = df_sites[df_sites["group"] == "Civilization"]

# ============================================================
# FIGURE SETUP
# ============================================================

fig, axes = plt.subplots(
    4, 1,
    figsize=(7.5, 10.5),
    sharex=True,
    constrained_layout=True
)

# ============================================================
# PANEL 1 — EVENT-RATE ENVELOPE
# ============================================================

axes[0].plot(
    df_event["time_ka"],
    df_event["event_rate_proxy"],
    color="black",
    linewidth=2
)

axes[0].invert_xaxis()
axes[0].set_ylabel("Normalized\n event-rate proxy")
axes[0].set_title(
    "Model-derived global reorganization envelope",
    fontsize=11
)
axes[0].grid(alpha=0.3)

# ============================================================
# PANEL 2 — SURFACE ADJUSTMENT RATE AT SITES
# ============================================================

axes[1].plot(
    df_homo["time_ka"],
    df_homo["mean_abs_dzdt"],
    label="Early Homo",
    linewidth=2
)

axes[1].plot(
    df_civ["time_ka"],
    df_civ["mean_abs_dzdt"],
    label="Early Civilizations",
    linewidth=2
)

axes[1].set_ylabel("Mean |ΔZeff/Δt|\n(m yr⁻¹)")
axes[1].set_title(
    "Surface adjustment rate at occupied sites",
    fontsize=11
)
axes[1].grid(alpha=0.3)
axes[1].legend(frameon=False, loc="upper right")

# ============================================================
# PANEL 3 — PROXIMITY TO EQUILIBRIUM MARGIN
# ============================================================

axes[2].plot(
    df_homo["time_ka"],
    df_homo["mean_zero_contour_distance_deg"],
    linewidth=2
)

axes[2].plot(
    df_civ["time_ka"],
    df_civ["mean_zero_contour_distance_deg"],
    linewidth=2
)

axes[2].set_ylabel("Mean distance to\n equilibrium margin (deg)")
axes[2].set_title(
    "Proximity of sites to equilibrium margin",
    fontsize=11
)
axes[2].grid(alpha=0.3)

# ============================================================
# PANEL 4 — SPATIAL STABILITY (VARIANCE)
# ============================================================

axes[3].plot(
    df_homo["time_ka"],
    df_homo["var_zero_contour_distance_deg2"],
    linewidth=2
)

axes[3].plot(
    df_civ["time_ka"],
    df_civ["var_zero_contour_distance_deg2"],
    linewidth=2
)

axes[3].set_ylabel("Variance of distance\n(deg²)")
axes[3].set_title(
    "Spatial stability of site environments",
    fontsize=11
)
axes[3].set_xlabel("ka BP")
axes[3].grid(alpha=0.3)

# ============================================================
# FINAL OUTPUT
# ============================================================

plt.savefig(
    "fig_composite_return_diagnostics.png",
    dpi=300,
    bbox_inches="tight"
)
plt.close()

print("✓ Composite diagnostic figure generated")
