"""Tests for the EHD index computation (CLAUDE.md hard constraint) and the leakage
isolation of the climatology (CLAUDE.md leakage rule)."""
import os, sys
import numpy as np
import pandas as pd
import xarray as xr
import cftime

sys.path.insert(0, os.path.join(os.path.dirname(__file__), "..", "src"))
from ehd import climatology as clim


def _synth_obs(years, seed=0):
    """Synthetic gridded obs on a 365_day calendar, 2x2 grid, with a controllable signal."""
    rng = np.random.default_rng(seed)
    times, vals = [], []
    for y in years:
        base = cftime.DatetimeNoLeap(y, 1, 1)
        for d in range(365):
            times.append(base + pd.Timedelta(days=d).to_pytimedelta())
            # smooth seasonal cycle + noise, warmer mid-year
            seasonal = 30 + 5 * np.sin(2 * np.pi * d / 365)
            vals.append(seasonal + rng.normal(0, 1, size=(2, 2)))
    da = xr.DataArray(np.array(vals),
                      coords={"time": times, "lat": [10.0, 11.0], "lon": [105.0, 106.0]},
                      dims=("time", "lat", "lon"), name="tasmax")
    return da


def test_threshold_min_max_rules():
    tx90p = xr.DataArray(np.array([[30.0, 40.0]]), dims=("lat", "lon"))
    tmin = clim.ehd_threshold(tx90p, 35.0, "min").values
    tmax = clim.ehd_threshold(tx90p, 35.0, "max").values
    # min(): percentile when below cap, cap when above
    assert tmin[0, 0] == 30.0 and tmin[0, 1] == 35.0
    # max(): cap when percentile below, percentile when above
    assert tmax[0, 0] == 35.0 and tmax[0, 1] == 40.0


def test_ehd_flag_boundary_is_inclusive():
    # threshold uses >= : a day exactly at threshold IS an EHD
    thr = xr.DataArray(np.full((365, 1, 1), 33.0),
                       coords={"doy": np.arange(1, 366)}, dims=("doy", "lat", "lon"))
    tx = xr.DataArray(np.array([[[32.9]], [[33.0]], [[33.1]]]),
                      dims=("time", "lat", "lon"))
    doy = np.array([1, 1, 1])
    flags = clim.ehd_flags(tx, thr, doy).values
    assert flags[0, 0, 0] == False   # below
    assert flags[1, 0, 0] == True    # exactly at threshold -> EHD (>=)
    assert flags[2, 0, 0] == True    # above


def test_percentile_is_calendar_day_windowed():
    # 90th percentile of a near-normal cell should sit a few sigma above the seasonal mean
    da = _synth_obs(range(2002, 2012))
    tx90p = clim.tx90p_climatology(da, list(range(2002, 2012)), percentile=90, window_days=15)
    assert tx90p.dims == ("doy", "lat", "lon")
    assert tx90p.sizes["doy"] == 365
    # mid-year (hot) climatology must exceed early-year (cool) climatology
    assert float(tx90p.sel(doy=182).mean()) > float(tx90p.sel(doy=1).mean())


def test_climatology_leakage_isolation():
    """Tx90p fit on baseline WITHOUT a contaminated year must be unaffected by that year."""
    clean = _synth_obs(range(2002, 2012), seed=1)
    # a contaminated copy where 2011 is corrupted with absurd heat
    contaminated = clean.copy()
    m = np.array([t.year == 2011 for t in contaminated.time.values])
    contaminated.values[m] += 50.0
    base = list(range(2002, 2011))  # excludes 2011
    a = clim.tx90p_climatology(clean, base)
    b = clim.tx90p_climatology(contaminated, base)
    # identical: the held-out (2011) data never entered the fit
    assert np.allclose(a.values, b.values, equal_nan=True)