03 — Analysis: Sinervo 2010 h_r mechanism for Iberian Lacertidae

import json
from datetime import date
from pathlib import Path

import numpy as np
import pandas as pd
import xarray as xr

T_B_LACERTIDAE = 35.4              # degC, family mean preferred body T
H_R_THRESHOLD_LACERTIDAE = 3.1     # h, calibrated activity-restriction threshold

CRITICAL_MONTHS = (4, 5)           # April–May (61 days, temperate window)

# SOM Equation S2 coefficients.
SLOPE = 0.74
INTERCEPT = 6.1

# Rare-species flag threshold (DOMAIN.md grid-coupling caveat).
LOW_N_CELL_THRESHOLD = 10

DECADES = {
    "2020s": (2020, 2029),
    "2030s": (2030, 2039),
}

ROOT = Path("..").resolve()
INTERMEDIATE = ROOT / "data" / "intermediate"
RESULTS = ROOT / "results"
RESULTS_TABLES = RESULTS / "tables"
RESULTS_TABLES.mkdir(parents=True, exist_ok=True)
(RESULTS / "logs").mkdir(parents=True, exist_ok=True)

DESTINE_NC = INTERMEDIATE / "destine_iberia_daily_t2m_nside128.nc"
GBIF_PARQUET = INTERMEDIATE / "gbif_lacertidae_presence_nside128.parquet"
GBIF_PARQUET_64 = INTERMEDIATE / "gbif_lacertidae_presence_nside64.parquet"
SUBDAILY_PROBE_LOG = ROOT / "data" / "destine" / "sub_daily" / "probe.log"

EXTINCTION_MASK_NC = INTERMEDIATE / "extinction_mask_nside128.nc"
HR_PER_CELL_NC = INTERMEDIATE / "hr_per_cell_year_nside128.nc"
PER_SPECIES_CSV = RESULTS_TABLES / "local_extinction_per_species.csv"
HEADLINE_JSON = RESULTS / "headline.json"

probe_text = (
    SUBDAILY_PROBE_LOG.read_text().strip()
    if SUBDAILY_PROBE_LOG.exists() else "probe log absent"
)
subdaily_success = probe_text.startswith("SUCCESS")
print(f"Sub-daily probe status: {'SUCCESS' if subdaily_success else 'unavailable'}")
print(f"  message: {probe_text!r}")

print(f"\nLoading {DESTINE_NC}")
ds = xr.open_dataset(DESTINE_NC)
print(f"  dims: {dict(ds.sizes)}")
print(f"  time: {pd.Timestamp(ds.time.values[0]).date()} .. "
      f"{pd.Timestamp(ds.time.values[-1]).date()}")

tmax_daily = ds["tmax_daily"]       # (time, cell), degC
cells = ds["cell"].values.astype(np.int64)
n_cells = cells.size

print(f"\nLoading {GBIF_PARQUET}")
presence = pd.read_parquet(GBIF_PARQUET)
print(f"  rows: {len(presence):,}, species: {presence['species'].nunique()}, "
      f"cells: {presence['cell_id_nside128'].nunique()}")

# Sanity: every presence cell must be in the DestinE cell list.
cells_set = set(cells.tolist())
unknown_cells = set(presence["cell_id_nside128"]) - cells_set
if unknown_cells:
    print(f"  WARNING: {len(unknown_cells)} presence cells absent from DestinE grid "
          f"(will be dropped from extinction calculation)")
    presence = presence[presence["cell_id_nside128"].isin(cells_set)].copy()

print("\n--- Computing h_r per cell x per day (Equation S2) ---")
h_r_daily = xr.apply_ufunc(
    lambda t: np.where(
        t > T_B_LACERTIDAE,
        SLOPE * (t - T_B_LACERTIDAE) + INTERCEPT,
        0.0,
    ),
    tmax_daily,
    dask="allowed",
).rename("h_r_daily")
h_r_daily.attrs.update({
    "units": "hours",
    "long_name": "daily hours of activity restriction (Sinervo SOM Eq. S2)",
    "T_b_used_degC": T_B_LACERTIDAE,
    "formula": (
        f"({SLOPE} * (Tmax - T_b) + {INTERCEPT}) if Tmax > T_b else 0"
    ),
})

print("\n--- Aggregating h_r over April-May per year ---")
months = pd.DatetimeIndex(h_r_daily["time"].values).month
window_mask = np.isin(months, list(CRITICAL_MONTHS))
print(f"  critical-window days kept: {int(window_mask.sum())} "
      f"of {len(months)} total days")

h_r_window = h_r_daily.isel(time=window_mask)
window_year_da = h_r_window["time"].dt.year

h_r_per_year_sum = (
    h_r_window.groupby(window_year_da.rename("year"))
              .sum(dim="time")
              .rename("h_r_sum_critical")
)
h_r_per_year_mean = (
    h_r_window.groupby(window_year_da.rename("year"))
              .mean(dim="time")
              .rename("h_r_mean_critical")
)
print(f"  h_r_sum_critical:  {h_r_per_year_sum.shape}, range "
      f"{float(h_r_per_year_sum.min()):.2f} .. {float(h_r_per_year_sum.max()):.2f} h "
      f"(mean {float(h_r_per_year_sum.mean()):.2f})")
print(f"  h_r_mean_critical: {h_r_per_year_mean.shape}, range "
      f"{float(h_r_per_year_mean.min()):.3f} .. {float(h_r_per_year_mean.max()):.3f} h")

extinction_mask = (h_r_per_year_mean > H_R_THRESHOLD_LACERTIDAE).rename("extinct")
extinction_mask.attrs.update({
    "long_name": "Cell-year flag: h_r_mean_critical > h_r threshold (Lacertidae)",
    "threshold_h": H_R_THRESHOLD_LACERTIDAE,
})

frac_extinct_year = extinction_mask.mean(dim="cell")
print("\nFraction of Iberian cells exceeding threshold, per year:")
print(pd.DataFrame({
    "year": frac_extinct_year["year"].values,
    "frac": frac_extinct_year.values.round(3),
}).to_string(index=False))

ds_out = xr.Dataset(
    data_vars={
        "h_r_sum_critical": h_r_per_year_sum.astype(np.float32),
        "h_r_mean_critical": h_r_per_year_mean.astype(np.float32),
        "extinct": extinction_mask.astype(np.int8),
    },
    coords={
        "year": h_r_per_year_sum["year"].values,
        "cell": ("cell", cells,
                 {"long_name": "HEALPix NESTED pixel index (nside=128)"}),
        "lon": ("cell", ds["lon"].values),
        "lat": ("cell", ds["lat"].values),
    },
    attrs={
        "title": "Per-cell x per-year h_r and extinction mask (Iberian Lacertidae)",
        "history": (
            f"Created {date.today().isoformat()} by notebooks/03_analysis.py"
        ),
        "T_b_used_degC": T_B_LACERTIDAE,
        "h_r_threshold_h": H_R_THRESHOLD_LACERTIDAE,
        "critical_months": list(CRITICAL_MONTHS),
        "subdaily_probe_status": "SUCCESS" if subdaily_success else "unavailable",
        "Conventions": "CF-1.10",
    },
)
ds_out.to_netcdf(
    EXTINCTION_MASK_NC, engine="netcdf4",
    encoding={v: {"zlib": True, "complevel": 4} for v in ds_out.data_vars},
)
print(f"\nSaved {EXTINCTION_MASK_NC}  "
      f"({EXTINCTION_MASK_NC.stat().st_size / 1e6:.2f} MB)")

print("\n--- Per-species local-extinction rate per decade ---")
decade_means = {}
for label, (y_start, y_end) in DECADES.items():
    sel = (extinction_mask["year"] >= y_start) & (extinction_mask["year"] <= y_end)
    # Per-cell decadal frequency of being above threshold.
    decade_means[label] = (
        extinction_mask.isel(year=sel)
                       .mean(dim="year")
                       .astype(np.float32)
    )

# Build per-cell DataFrame for the merge.
per_cell_df = pd.DataFrame({
    "cell_id_nside128": cells,
    **{
        f"extinction_freq_{label}": decade_means[label].values
        for label in DECADES
    },
})

# Merge presence (species x cell) with per-cell decadal frequencies.
joined = presence.merge(per_cell_df, on="cell_id_nside128", how="left")

per_species = (
    joined.groupby("species")
          .agg(
              n_cells=("cell_id_nside128", "nunique"),
              n_records=("n_records", "sum"),
              **{
                  f"local_ext_{label}": (f"extinction_freq_{label}", "mean")
                  for label in DECADES
              },
          )
          .reset_index()
          .sort_values("species")
)
per_species["low_N_warning"] = per_species["n_cells"] < LOW_N_CELL_THRESHOLD
per_species.to_csv(PER_SPECIES_CSV, index=False)
print(f"Saved {PER_SPECIES_CSV}")
print(per_species.round(4).to_string(index=False))

# Cells where ≥1 Lacertidae species was recorded — the "Lacertidae cell
# universe" for Iberia. Local extinction is the fraction of those cells
# above threshold (decadal mean), which matches the Sinervo Table 1
# semantics of "local extinction = fraction of populations lost".
presence_cells = sorted(presence["cell_id_nside128"].unique().tolist())
print(f"\nLacertidae cell universe: {len(presence_cells)} cells "
      f"(of {n_cells} Iberian HEALPix cells)")

family_rates: dict[str, float] = {}
for label in DECADES:
    sel = extinction_mask["cell"].isin(presence_cells)
    rate = float(
        decade_means[label]
        .isel(cell=sel)
        .mean()
        .values
    )
    family_rates[f"local_extinction_{label}"] = rate
    print(f"  Lacertidae local-extinction rate, {label}: {rate:.4f}")

# Reference values from Sinervo Table 1.
sinervo_reference = {
    "lacertidae_local_extinction_2009": 0.034,
    "lacertidae_local_extinction_2050": 0.241,
    "lacertidae_local_extinction_2080": 0.460,
}
print("\nReference (Sinervo et al. 2010 Table 1, Lacertidae):")
for k, v in sinervo_reference.items():
    print(f"  {k}: {v}")
print("Note: 2020s and 2030s horizons are not directly comparable to "
      "Sinervo's 2009/2050/2080 timepoints; reported here for context.")

SUBSTRATE_TB = 31.0          # S3a: Iberolacerta-style cool-adapted T_b
SUBSTRATE_MONTHS = (5, 6)    # S3a: actual Iberian Lacertidae May-Jun breeding


def _per_species_rates(
    tmax_da: xr.DataArray,
    presence_df: pd.DataFrame,
    cell_col: str,
    t_b: float,
    months: tuple[int, ...],
) -> pd.DataFrame:
    """Compute per-species 2020s + 2030s local-extinction rates under a
    (T_b, window) config at the substrate implied by tmax_da's cell
    dimension. Returns one row per species with rate, n_cells, low_N
    flag (n_cells < 10), threshold = H_R_THRESHOLD_LACERTIDAE."""
    hr_d = xr.where(
        tmax_da > t_b,
        SLOPE * (tmax_da - t_b) + INTERCEPT,
        0.0,
    )
    mts = pd.DatetimeIndex(hr_d["time"].values).month
    win = np.isin(mts, list(months))
    hr_w = hr_d.isel(time=win)
    yr = hr_w["time"].dt.year
    hr_mean = hr_w.groupby(yr.rename("year")).mean(dim="time")
    extinct = (hr_mean > H_R_THRESHOLD_LACERTIDAE)
    yrs = extinct["year"].values
    extinct_20 = extinct.isel(year=(yrs >= 2020) & (yrs < 2030)).mean(dim="year")
    extinct_30 = extinct.isel(year=(yrs >= 2030) & (yrs < 2040)).mean(dim="year")
    rows = []
    for sp, sub in presence_df.groupby("species"):
        cells = sub[cell_col].astype(int).unique().tolist()
        mask = np.isin(extinct_20["cell"].values, cells)
        if mask.sum() == 0:
            continue
        r20 = float(extinct_20.isel(cell=mask).mean())
        r30 = float(extinct_30.isel(cell=mask).mean())
        rows.append({
            "species": sp,
            "n_cells": int(mask.sum()),
            "local_ext_2020s": r20,
            "local_ext_2030s": r30,
        })
    df = pd.DataFrame(rows).sort_values("species").reset_index(drop=True)
    df["low_N_warning"] = df["n_cells"] < LOW_N_CELL_THRESHOLD
    return df


# nside=128 is the native substrate; tmax_daily already has it.
print("\n--- Substrate-sensitivity diagnostic (S3a config) ---")
print(f"  T_b = {SUBSTRATE_TB} °C, window = {SUBSTRATE_MONTHS}")

rates_128 = _per_species_rates(
    tmax_daily, presence, "cell_id_nside128",
    SUBSTRATE_TB, SUBSTRATE_MONTHS,
)
print(f"  nside=128: {len(rates_128)} species with ≥1 presence cell, "
      f"family rate 2020s = {rates_128['local_ext_2020s'].mean():.4f}")

# Downsample nside=128 Tmax to nside=64: average over the 4 NESTED
# children per nside=64 parent. parent = pix >> 2 in NESTED ordering.
parent_n64 = xr.DataArray(
    tmax_daily["cell"].values.astype(np.int64) >> 2,
    dims=["cell"], name="parent_n64",
)
tmax_daily_n64 = (
    tmax_daily.groupby(parent_n64).mean(dim="cell")
              .rename({"parent_n64": "cell"})
)
print(f"  nside=64 Tmax field: {tmax_daily_n64.sizes['cell']} parent cells "
      f"(from {tmax_daily.sizes['cell']} nside=128 children)")

presence_n64 = pd.read_parquet(GBIF_PARQUET_64)
rates_64 = _per_species_rates(
    tmax_daily_n64, presence_n64, "cell_id_nside64",
    SUBSTRATE_TB, SUBSTRATE_MONTHS,
)
print(f"  nside=64:  {len(rates_64)} species with ≥1 presence cell, "
      f"family rate 2020s = {rates_64['local_ext_2020s'].mean():.4f}")

# Compare per-species rankings between the two substrates.
joined = rates_128.set_index("species").join(
    rates_64.set_index("species"),
    how="inner",
    lsuffix="_n128",
    rsuffix="_n64",
)
print(f"  species joinable across substrates: {len(joined)}")

from scipy.stats import spearmanr

def _spear(group):
    """Spearman rho on local_ext_2020s_n128 vs local_ext_2020s_n64."""
    n128 = group["local_ext_2020s_n128"].values
    n64 = group["local_ext_2020s_n64"].values
    if len(group) < 3:
        return float("nan"), len(group)
    if n128.std() == 0 and n64.std() == 0:
        return float("nan"), len(group)
    rho, _ = spearmanr(n128, n64)
    return float(rho), len(group)


# Low-N at nside=128 — the population the rare-species caveat applies to.
non_low_n_128 = joined[joined["n_cells_n128"] >= LOW_N_CELL_THRESHOLD]
low_n_128 = joined[joined["n_cells_n128"] < LOW_N_CELL_THRESHOLD]

rho_all, n_all = _spear(joined)
rho_nonlow, n_nonlow = _spear(non_low_n_128)
rho_low, n_low = _spear(low_n_128)

substrate_sensitivity_summary = {
    "config_used": {
        "T_b_degC": SUBSTRATE_TB,
        "critical_months": list(SUBSTRATE_MONTHS),
        "note": (
            "S3a (Iberolacerta T_b + May-Jun). Substrate diagnostic at "
            "baseline config is degenerate because all per-species rates "
            "are zero at both substrates."
        ),
    },
    "n_species_joined": int(len(joined)),
    "family_rate_nside128_2020s": float(rates_128["local_ext_2020s"].mean()),
    "family_rate_nside64_2020s": float(rates_64["local_ext_2020s"].mean()),
    "family_rate_nside128_2030s": float(rates_128["local_ext_2030s"].mean()),
    "family_rate_nside64_2030s": float(rates_64["local_ext_2030s"].mean()),
    "spearman_rho_all_species": rho_all,
    "spearman_rho_n_cells_ge_10": rho_nonlow,
    "spearman_rho_n_cells_lt_10": rho_low,
    "n_species_n_cells_ge_10": int(n_nonlow),
    "n_species_n_cells_lt_10": int(n_low),
    "interpretation": (
        "Per Bombus chain-3 finding (CHAIN_SUMMARY): mechanism is "
        "substrate-robust if rho is high across all species; ranking is "
        "grid-coupled for the n_cells < 10 subset if its rho is "
        "materially lower than the n_cells >= 10 subset's rho."
    ),
}
print(f"\nSubstrate-sensitivity Spearman rho (2020s rates):")
print(f"  all species (n={n_all}):           rho = {rho_all:.3f}")
print(f"  n_cells >= 10 (n={n_nonlow}):      rho = {rho_nonlow:.3f}")
print(f"  n_cells < 10  (n={n_low}):         rho = {rho_low:.3f}")

# Persist the per-species substrate comparison table.
SUBSTRATE_CSV = RESULTS_TABLES / "substrate_sensitivity_per_species.csv"
joined_out = joined.reset_index()[[
    "species",
    "n_cells_n128", "local_ext_2020s_n128", "local_ext_2030s_n128",
    "n_cells_n64", "local_ext_2020s_n64", "local_ext_2030s_n64",
    "low_N_warning_n128",
]]
joined_out.to_csv(SUBSTRATE_CSV, index=False)
print(f"  wrote {SUBSTRATE_CSV}")

print("\n--- Sensitivity matrix: prior-conditional mechanism applicability ---")


def _h_r_extinction_rates(t_b: float, months: tuple[int, ...]) -> dict:
    """Compute Lacertidae local-extinction rates under (T_b, window) config.

    Returns decadal 2020s / 2030s rates restricted to Lacertidae presence
    cells, plus the worst-cell-year h_r_mean value reached.
    """
    hr_daily_c = xr.where(
        tmax_daily > t_b,
        SLOPE * (tmax_daily - t_b) + INTERCEPT,
        0.0,
    )
    months_arr = pd.DatetimeIndex(hr_daily_c["time"].values).month
    window_c = np.isin(months_arr, list(months))
    hr_window_c = hr_daily_c.isel(time=window_c)
    yr_da = hr_window_c["time"].dt.year
    hr_mean_c = (
        hr_window_c.groupby(yr_da.rename("year")).mean(dim="time")
    )
    extinct_c = hr_mean_c > H_R_THRESHOLD_LACERTIDAE
    cell_mask = np.isin(hr_mean_c["cell"].values, presence_cells)
    extinct_lac = extinct_c.isel(cell=cell_mask)
    hr_mean_lac = hr_mean_c.isel(cell=cell_mask)
    yrs_arr = extinct_lac["year"].values
    d20 = float(
        extinct_lac.isel(year=(yrs_arr >= 2020) & (yrs_arr < 2030))
        .mean(dim=("year", "cell"))
    )
    d30 = float(
        extinct_lac.isel(year=(yrs_arr >= 2030) & (yrs_arr < 2040))
        .mean(dim=("year", "cell"))
    )
    return {
        "t_b_degC": t_b,
        "months": list(months),
        "local_extinction_2020s": d20,
        "local_extinction_2030s": d30,
        "max_h_r_mean_h": float(hr_mean_lac.max()),
        "n_cell_years_above_threshold": int((hr_mean_lac > H_R_THRESHOLD_LACERTIDAE).sum()),
    }


sensitivity_configs = [
    ("baseline_family_aprmay",  35.4, (4, 5)),
    ("s1a_family_mayjun",       35.4, (5, 6)),
    ("s1b_family_aprjun",       35.4, (4, 5, 6)),
    ("s2_iberolacerta_aprmay",  31.0, (4, 5)),
    ("s3a_iberolacerta_mayjun", 31.0, (5, 6)),
    ("s3b_iberolacerta_aprjun", 31.0, (4, 5, 6)),
]
sensitivity_matrix = {
    name: _h_r_extinction_rates(t_b, months)
    for name, t_b, months in sensitivity_configs
}
print(f"{'CONFIG':30s}  2020s    2030s    max_h_r   n_above")
print("-" * 80)
for name, row in sensitivity_matrix.items():
    print(
        f"{name:30s}  "
        f"{row['local_extinction_2020s']*100:5.2f}%   "
        f"{row['local_extinction_2030s']*100:5.2f}%   "
        f"{row['max_h_r_mean_h']:5.2f}h   "
        f"{row['n_cell_years_above_threshold']}"
    )

headline = {
    "written_on": date.today().isoformat(),
    "config": {
        "T_b_degC": T_B_LACERTIDAE,
        "h_r_threshold_h": H_R_THRESHOLD_LACERTIDAE,
        "critical_months": list(CRITICAL_MONTHS),
        "subdaily_probe_status": "SUCCESS" if subdaily_success else "unavailable",
        "n_iberian_cells_nside128": int(n_cells),
        "n_lacertidae_presence_cells": int(len(presence_cells)),
        "n_species": int(presence["species"].nunique()),
    },
    "lacertidae_family_rates_destine": family_rates,
    "sensitivity_matrix": sensitivity_matrix,
    "substrate_sensitivity": substrate_sensitivity_summary,
    "sinervo_reference_table1": sinervo_reference,
    "rare_species_flag": {
        "low_N_cell_threshold": LOW_N_CELL_THRESHOLD,
        "n_low_N_species": int(per_species["low_N_warning"].sum()),
        "low_N_species_list": per_species.loc[
            per_species["low_N_warning"], "species"
        ].tolist(),
        "note": (
            "Rare-species per-species rates (n_cells < 10) are grid-coupled "
            "(Lobo et al. 2007, Hurlbert & Jetz 2007); interpret with caution."
        ),
    },
    "deviations": {
        "tb_source": (
            "Lacertidae family-level Table 1 (T_b=35.4 degC). "
            "Iberian-specific T_b records deferred to v0.2.0."
        ),
        "critical_window": (
            "April-May (61 days, Iberian Podarcis/Iberolacerta/Timon "
            "breeding chronology). Sinervo SOM Sceloporus calibration "
            "uses March-April; SOM says '2 months in temperate zone'."
        ),
        "climate_forcing": (
            "DestinE Climate DT SSP3-7.0 IFS-NEMO at HEALPix nside=128, "
            "2020-2039 (vs Sinervo WorldClim 1.4 + CCCMA/HADCM3/CSIRO a2a, "
            "1975/2020/2050/2080)."
        ),
        "diurnal_reconstruction": (
            "Not applied. Sinervo SOM Eq. S2 is parameterised in daily Tmax; "
            "DestinE 6-hourly samples are aggregated to daily max in "
            "02_data_clean. Sub-daily probe outcome: "
            f"{'SUCCESS' if subdaily_success else 'unavailable - sinusoidal fallback unused'}."
        ),
    },
}
with open(HEADLINE_JSON, "w") as f:
    json.dump(headline, f, indent=2)
print(f"\nSaved {HEADLINE_JSON}")
print(json.dumps(headline, indent=2))