"""Unit tests for the LOCKED target/event definitions (Phase 0).

Covers every case named in the prompt's <software> section plus the
known-failure-modes it says fail quietly. Runnable with pytest or directly
(`python tests/test_targets.py`); needs only numpy.
"""
import os
import sys
import numpy as np

sys.path.insert(0, os.path.join(os.path.dirname(__file__), "..", "src"))

from targets.climatology import (  # noqa: E402
    day_of_year_noleap, p90_climatology, ehd_threshold, threshold_for_dates,
)
from targets.events import (  # noqa: E402
    daily_ehd, exceedance, weekly_aggregate, detect_episodes,
    analyze_forecast_sequence, DEFAULT_WEEKS,
)


def _dates(start, n):
    return np.arange(np.datetime64(start), np.datetime64(start) + n)


# ---------- calendar-day index / leap handling ----------

def test_doy_basic_and_wrap():
    assert day_of_year_noleap(np.datetime64("2003-01-01"))[()] == 1
    assert day_of_year_noleap(np.datetime64("2003-12-31"))[()] == 365
    assert day_of_year_noleap(np.datetime64("2003-02-28"))[()] == 59
    # 1 March is index 60 in both leap and non-leap years
    assert day_of_year_noleap(np.datetime64("2003-03-01"))[()] == 60
    assert day_of_year_noleap(np.datetime64("2004-03-01"))[()] == 60


def test_feb29_pooled_with_feb28():
    # 29 Feb maps to 59 (same as 28 Feb), not to 60 (1 Mar)
    assert day_of_year_noleap(np.datetime64("2004-02-29"))[()] == 59
    assert day_of_year_noleap(np.datetime64("2004-02-28"))[()] == 59


# ---------- P90 window inclusion, cyclic wrap, leap ----------

def test_p90_window_and_wrap():
    # Build 3 years of a 1-cell series where Tx = calendar-day index, so the
    # P90 of a +/-15 window centred on d is deterministic.
    dates = _dates("2003-01-01", 366 * 3)  # covers a leap year (2004)
    doy = day_of_year_noleap(dates).astype("float64")
    tx = doy.reshape(-1, 1)  # (T, 1)
    p90 = p90_climatology(tx, dates, ref_years=(2003, 2005), window=15, q=90)
    assert p90.shape == (365, 1)
    # Day 1 window wraps to include Dec (351..365) and Jan (1..16): values
    # present, so P90 is finite and near the top of that set.
    assert np.isfinite(p90[0, 0])
    # Mid-year day 180 window is 165..195; P90 (90th pct) should sit high in it.
    assert 190 <= p90[179, 0] <= 195


# ---------- threshold cap uses min, never max ----------

def test_ehd_threshold_is_min():
    p90 = np.array([30.0, 35.0, 40.0])
    thr = ehd_threshold(p90, cap=35.0)
    assert np.allclose(thr, [30.0, 35.0, 35.0])


# ---------- valid-date alignment ----------

def test_threshold_for_dates_alignment():
    clim = np.arange(365, dtype="float64").reshape(365, 1)  # index i -> value i
    d = np.array([np.datetime64("2003-01-01"), np.datetime64("2003-12-31"),
                  np.datetime64("2004-02-29")])
    got = threshold_for_dates(clim, d)[:, 0]
    assert np.allclose(got, [0.0, 364.0, 58.0])  # doy 1,365,59 -> idx 0,364,58


# ---------- EHD_COUNT -> EHE / HW2 truth table ----------

def test_count_to_ehe_hw2():
    lead = np.arange(0, 7)
    def wk(counts):
        ehd = np.array(counts, dtype=np.int8)
        exc = ehd.astype("float64")  # 1 degC exceedance on event days
        return weekly_aggregate(lead, ehd, exc)["W1"]
    z = wk([0, 0, 0, 0, 0, 0, 0])
    assert (z["EHD_COUNT"], z["EHE_WEEK"], z["HW2_WEEK"]) == (0, 0, 0)
    one = wk([1, 0, 0, 0, 0, 0, 0])
    assert (one["EHD_COUNT"], one["EHE_WEEK"], one["HW2_WEEK"]) == (1, 1, 0)
    two = wk([1, 0, 1, 0, 0, 0, 0])  # 2 EHDs, non-consecutive
    assert (two["EHD_COUNT"], two["EHE_WEEK"], two["HW2_WEEK"]) == (2, 1, 1)


def test_hw2_is_not_episode():
    # 2 non-consecutive EHDs -> HW2_WEEK=1 but NO episode (>=2 consecutive)
    lead = np.arange(0, 7)
    ehd = np.array([1, 0, 1, 0, 0, 0, 0], dtype=np.int8)
    exc = ehd.astype("float64")
    assert weekly_aggregate(lead, ehd, exc)["W1"]["HW2_WEEK"] == 1
    assert detect_episodes(ehd, exc) == []


# ---------- episode across LT6/LT7 is ONE 2-day episode ----------

def test_episode_crosses_week_boundary_single_event():
    lead = np.arange(0, 14)  # W1: 0..6, W2: 7..13
    ehd = np.zeros(14, dtype=np.int8)
    ehd[6] = 1  # LT6 (last day of W1)
    ehd[7] = 1  # LT7 (first day of W2)
    exc = ehd.astype("float64") * 2.0
    eps = detect_episodes(ehd, exc)
    assert len(eps) == 1
    assert eps[0]["duration"] == 2

    episodes, weekly = analyze_forecast_sequence(lead, ehd, exc)
    assert len(episodes) == 1
    e = episodes[0]
    # one onset, in the week it starts (W1)
    assert e["onset_week"] == "W1"
    assert weekly["W1"]["HW_ONSET_COUNT"] == 1
    assert weekly["W2"]["HW_ONSET_COUNT"] == 0
    # active in both weeks
    assert weekly["W1"]["HW_ACTIVE"] == 1 and weekly["W2"]["HW_ACTIVE"] == 1
    # HW_DAYS split by day: 1 day in each week
    assert weekly["W1"]["HW_DAYS"] == 1 and weekly["W2"]["HW_DAYS"] == 1
    # severity split by the days that fall in each week (2.0 each)
    assert np.isclose(weekly["W1"]["HW_SEV_ACTIVE"], 2.0)
    assert np.isclose(weekly["W2"]["HW_SEV_ACTIVE"], 2.0)


# ---------- censoring ----------

def test_left_censoring_not_onset():
    # episode active on the earliest available (buffer) day -> left-censored,
    # contributes to HW_ACTIVE/HW_DAYS but is not an onset.
    lead = np.array([-2, -1, 0, 1, 2])
    ehd = np.array([1, 1, 1, 0, 0], dtype=np.int8)
    exc = ehd.astype("float64")
    episodes, weekly = analyze_forecast_sequence(lead, ehd, exc)
    assert len(episodes) == 1
    e = episodes[0]
    assert e["left_censored"] is True
    assert e["onset_week"] is None
    assert weekly["W1"]["HW_ONSET_COUNT"] == 0
    assert weekly["W1"]["HW_ACTIVE"] == 1
    # only the in-horizon day (lead 0) counts toward HW_DAYS
    assert weekly["W1"]["HW_DAYS"] == 1


def test_right_censoring_no_full_stats():
    # a preceding non-event day (LT43) confirms the onset, so only the RIGHT
    # edge is censored (46,47 are post-LT45 buffer with the event still active).
    lead = np.array([43, 44, 45, 46, 47])
    ehd = np.array([0, 1, 1, 1, 1], dtype=np.int8)
    exc = ehd.astype("float64")
    episodes, weekly = analyze_forecast_sequence(lead, ehd, exc)
    e = episodes[0]
    assert e["left_censored"] is False
    assert e["right_censored"] is True
    assert e["onset_week"] == "W7"          # onset still counts (starts at LT44)
    assert e["full_stats_valid"] is False   # full duration/severity NOT valid
    # active-day accounting only inside horizon (LT44, LT45)
    assert weekly["W7"]["HW_DAYS"] == 2


def test_no_full_duration_duplication_across_weeks():
    # a long episode spanning W1..W2 must not put its full duration in both
    lead = np.arange(0, 14)
    ehd = np.zeros(14, dtype=np.int8)
    ehd[4:10] = 1  # LT4..LT9 : 6-day episode, onset in W1
    exc = ehd.astype("float64")
    episodes, weekly = analyze_forecast_sequence(lead, ehd, exc)
    e = episodes[0]
    assert e["duration"] == 6 and e["onset_week"] == "W1"
    # HW_DAYS: LT4,5,6 in W1 (3) ; LT7,8,9 in W2 (3) ; sum == full duration once
    assert weekly["W1"]["HW_DAYS"] == 3 and weekly["W2"]["HW_DAYS"] == 3
    assert weekly["W1"]["HW_DAYS"] + weekly["W2"]["HW_DAYS"] == e["duration"]


def _run_all():
    fns = [v for k, v in sorted(globals().items()) if k.startswith("test_")]
    passed = 0
    for fn in fns:
        fn()
        print("PASS", fn.__name__)
        passed += 1
    print(f"\n{passed}/{len(fns)} target/event unit tests passed")


if __name__ == "__main__":
    _run_all()