"""Observed (forecast-INDEPENDENT) target products from VnGAD Tx.

Uses the LOCKED, unit-tested definitions (P90 threshold, daily EHD, exceedance,
continuous heat-wave episodes) to build the observed truth the model is trained
against and verified on, and the observed climatology (B0 baseline + paper
'observed climatology' figures). Nothing here depends on the ECMWF GRIB.

Episodes are detected on the CONTINUOUS daily observed series with a small buffer
beyond the counting window so boundary episodes are handled correctly; onsets are
counted only inside the official window.
"""
import numpy as np

from .climatology import threshold_for_dates
from .events import daily_ehd, exceedance, detect_episodes


def build_daily(tx, dates, thr_clim):
    """Return (ehd int8, exceedance float32) for Tx over the given dates."""
    thr = threshold_for_dates(thr_clim, dates)
    return daily_ehd(tx, thr), exceedance(tx, thr).astype("float32")


def episode_catalog(ehd, exc, dates, valid_mask, region_field, window):
    """Per-cell observed heat-wave episode catalog over a continuous series.

    ``window`` = (start_date, end_date) inclusive; an episode is an ONSET in the
    window only if it starts in-window and is not left-censored (touching the
    earliest available day). Duration/severity are 'complete' only if the episode
    is neither left- nor right-censored.
    Returns a list of dict rows.
    """
    T = ehd.shape[0]
    lo = np.datetime64(str(window[0])[:10])
    hi = np.datetime64(str(window[1])[:10])
    dts = np.asarray(dates).astype("datetime64[D]")
    rows = []
    ys, xs = np.where(valid_mask)
    for i, j in zip(ys.tolist(), xs.tolist()):
        eps = detect_episodes(ehd[:, i, j], exc[:, i, j])
        for ep in eps:
            onset = dts[ep["start"]]
            left_c = ep["start"] == 0
            right_c = ep["end"] == T - 1
            rows.append({
                "lat_i": i, "lon_j": j, "region": int(region_field[i, j]),
                "onset": str(onset), "end": str(dts[ep["end"]]),
                "duration": ep["duration"], "i_mean": round(ep["i_mean"], 4),
                "i_max": round(ep["i_max"], 4), "severity": round(ep["sev"], 4),
                "left_censored": bool(left_c), "right_censored": bool(right_c),
                "onset_in_window": bool((onset >= lo) and (onset <= hi) and not left_c),
                "complete": bool(not left_c and not right_c),
            })
    return rows


def cell_climatology(ehd_window, valid_mask, catalog, shape, n_years):
    """Per-cell climatology maps from the in-window daily EHD and episode catalog.

    ``ehd_window`` is the daily EHD restricted to the counting window (T, Y, X).
    Returns a dict of (Y, X) float arrays (NaN off the valid mask).
    """
    Y, X = shape
    nan = lambda: np.full((Y, X), np.nan, dtype="float64")
    ehd_freq = nan()
    ehd_days_per_year = nan()
    ehd_freq[valid_mask] = ehd_window[:, valid_mask].mean(0)
    ehd_days_per_year[valid_mask] = ehd_window[:, valid_mask].sum(0) / n_years

    # per-cell episode stats (onsets in window)
    from collections import defaultdict
    dur = defaultdict(list); sev = defaultdict(list); imax = defaultdict(list); nep = defaultdict(int)
    for r in catalog:
        if not r["onset_in_window"]:
            continue
        key = (r["lat_i"], r["lon_j"])
        nep[key] += 1
        if r["complete"]:
            dur[key].append(r["duration"])
            sev[key].append(r["severity"])
            imax[key].append(r["i_max"])

    eps_per_year = nan(); mean_dur = nan(); p95_dur = nan()
    mean_sev = nan(); max_sev = nan(); mean_imax = nan()
    for (i, j), n in nep.items():
        eps_per_year[i, j] = n / n_years
        if dur[(i, j)]:
            d = np.array(dur[(i, j)]); s = np.array(sev[(i, j)]); im = np.array(imax[(i, j)])
            mean_dur[i, j] = d.mean(); p95_dur[i, j] = np.percentile(d, 95)
            mean_sev[i, j] = s.mean(); max_sev[i, j] = s.max(); mean_imax[i, j] = im.mean()
    return {"ehd_freq": ehd_freq, "ehd_days_per_year": ehd_days_per_year,
            "episodes_per_year": eps_per_year, "mean_duration": mean_dur,
            "p95_duration": p95_dur, "mean_severity": mean_sev,
            "max_severity": max_sev, "mean_intensity_max": mean_imax}