"""Assemble calibration samples: join deterministic forecast Tx, observed Tx, and the
per-fold Tx90p threshold into (X, y) per lead block.

Because the reforecast is a single control member, the "probability of an EHD in a block"
is learned directly from block-aggregated forecast features (single-member MOS), not from an
ensemble exceedance fraction. This module produces, for each (init, land cell, block):

  features X:
    f_max      block max of forecast daily Tx (degC)
    f_mean     block mean forecast daily Tx
    f_min      block min forecast daily Tx
    f_exceed   fraction of block days with forecast Tx >= threshold
    f_maxanom  f_max minus the block-mean threshold (how far the hottest day clears the bar)
    lead_id    block ordinal 1..6 (lead dependence)
    sin_doy    sin(2*pi*init_doy/365)   (seasonal cycle)
    cos_doy    cos(2*pi*init_doy/365)
    lat, lon   grid location
  labels y:
    y_primary   observed >=1 EHD in block (0/1, NaN if any block day obs missing)
    y_secondary observed heatwave (>=3 consecutive EHD) in block, straddle rule

The threshold array is supplied by the caller and MUST be fit excluding the held-out year
(leakage rule). Observed EHD uses obs Tx vs that same threshold.
"""
from __future__ import annotations
import numpy as np
import pandas as pd
from .targets import block_targets

FEATURE_NAMES = ["f_max", "f_mean", "f_min", "f_exceed", "f_maxanom",
                 "lead_id", "sin_doy", "cos_doy", "lat", "lon"]


def build_obs_index(obs_da):
    """Return (values[time,lat,lon], {yyyymmdd:int -> row}) from a cftime-indexed obs DA."""
    vals = obs_da.values
    lut = {}
    for i, t in enumerate(obs_da["time"].values):
        lut[t.year * 10000 + t.month * 100 + t.day] = i
    return vals, lut


def _doy_noleap(dates: pd.DatetimeIndex) -> np.ndarray:
    doy = np.asarray(dates.dayofyear).copy()
    isleap = np.asarray(dates.is_leap_year)
    doy[isleap & (np.asarray(dates.month) > 2)] -= 1
    return np.clip(doy, 1, 365)


def assemble_year(fc_year, obs_vals, obs_lut, thr_by_doy, blocks, land_iy, land_ix,
                  min_run=3):
    """Assemble all samples for one year's forecast file.

    fc_year   : DataArray (init, lead_day, lat, lon) forecast daily Tx (degC)
    thr_by_doy: DataArray (doy, lat, lon) production threshold min(Tx90p,35), fold-fit
    land_iy/ix: integer indices of land cells (1-D, same length)
    Returns dict block -> dict(X, y_primary, y_secondary, init_idx, cell_id)
    """
    ninit = fc_year.sizes["init"]
    nday = fc_year.sizes["lead_day"]
    ncell = land_iy.size
    fc = fc_year.values                                   # (init, lead, lat, lon)
    thr = thr_by_doy.values                               # (365, lat, lon)
    vdates = pd.to_datetime(fc_year["valid_date"].values.reshape(ninit, nday).ravel())
    vdoy = _doy_noleap(vdates).reshape(ninit, nday)
    vymd = (np.asarray(vdates.year) * 10000 + np.asarray(vdates.month) * 100
            + np.asarray(vdates.day)).reshape(ninit, nday)
    init_doy = _doy_noleap(pd.to_datetime(fc_year["init"].values))

    # gather per (init, lead, cell): forecast Tx, threshold, obs Tx
    fc_c = fc[:, :, land_iy, land_ix]                     # (init, lead, cell)
    thr_c = np.empty((ninit, nday, ncell), dtype="float64")
    obs_c = np.empty((ninit, nday, ncell), dtype="float64")
    for it in range(ninit):
        thr_c[it] = thr[vdoy[it] - 1][:, land_iy, land_ix]
        rows = np.array([obs_lut.get(int(v), -1) for v in vymd[it]])
        ok = rows >= 0
        tmp = np.full((nday, ncell), np.nan)
        if ok.any():
            tmp[ok] = obs_vals[rows[ok]][:, land_iy, land_ix]
        obs_c[it] = tmp

    fc_ehd = (fc_c >= thr_c)                              # forecast EHD days
    obs_ehd = np.where(np.isnan(obs_c), np.nan, (obs_c >= thr_c).astype(float))

    out = {}
    lead_ids = {name: k + 1 for k, name in enumerate(blocks)}
    for name, (d0, d1) in blocks.items():
        sl = slice(d0, d1 + 1)
        fb = fc_c[:, sl, :]                               # (init, dblk, cell)
        tb = thr_c[:, sl, :]
        f_max = np.nanmax(fb, axis=1)
        f_mean = np.nanmean(fb, axis=1)
        f_min = np.nanmin(fb, axis=1)
        f_exceed = np.nanmean(fc_ehd[:, sl, :].astype(float), axis=1)
        f_maxanom = f_max - np.nanmean(tb, axis=1)
        n = ninit * ncell
        lead = np.full(n, lead_ids[name], dtype=float)
        sin_d = np.repeat(np.sin(2 * np.pi * init_doy / 365.0), ncell)
        cos_d = np.repeat(np.cos(2 * np.pi * init_doy / 365.0), ncell)
        latf = np.tile(fc_year["lat"].values[land_iy], ninit)
        lonf = np.tile(fc_year["lon"].values[land_ix], ninit)
        X = np.column_stack([f_max.ravel(), f_mean.ravel(), f_min.ravel(),
                             f_exceed.ravel(), f_maxanom.ravel(), lead,
                             sin_d, cos_d, latf, lonf])
        # labels via block reduction over the daily obs EHD series (per cell)
        yp, ys = block_targets(np.moveaxis(obs_ehd, 0, 0)  # (init, day, cell) -> need (day, ...)
                               .transpose(1, 0, 2), {name: (d0, d1)}, min_run=min_run)
        y_primary = yp[name].ravel()
        y_secondary = ys[name].ravel()
        # severity predictand (v2): block-max observed severity = max_d(Tx_obs - Thr).
        # sev >= 0  <=>  y_primary == 1 (identical event); NaN where the block has any
        # missing obs day, matching y_primary's NaN rule so masks stay aligned.
        ob = obs_c[:, sl, :]
        blk_missing = np.any(np.isnan(ob), axis=1)
        with np.errstate(all="ignore"):
            sev = np.nanmax(ob - tb, axis=1)
        sev = np.where(blk_missing, np.nan, sev).ravel()
        init_idx = np.repeat(np.arange(ninit), ncell)
        cell_id = np.tile(np.arange(ncell), ninit)
        out[name] = dict(X=X, y_primary=y_primary, y_secondary=y_secondary, sev=sev,
                         init_idx=init_idx, cell_id=cell_id)
    return out


def forecast_features(fc_year, thr_by_doy, blocks, land_iy, land_ix):
    """Forecast-ONLY feature matrices (no obs, no labels) for the operational producer.

    Builds exactly the FEATURE_NAMES columns that `assemble_year` builds for X (they are all
    functions of the deterministic forecast, the per-doy threshold and grid geometry — no
    observations), so a production model trained via `assemble_year` can score a new init.
    The two paths MUST keep the same column construction; if `assemble_year`'s X changes,
    change this in lockstep.

    fc_year   : DataArray (init, lead_day, lat, lon) forecast daily Tx (degC).
    thr_by_doy: DataArray (doy=1..365, lat, lon) production threshold min(Tx90p,35).
    Returns dict block -> dict(X, init_idx, cell_id).
    """
    ninit = fc_year.sizes["init"]
    nday = fc_year.sizes["lead_day"]
    ncell = land_iy.size
    fc = fc_year.values
    thr = thr_by_doy.values
    vdates = pd.to_datetime(fc_year["valid_date"].values.reshape(ninit, nday).ravel())
    vdoy = _doy_noleap(vdates).reshape(ninit, nday)
    init_doy = _doy_noleap(pd.to_datetime(fc_year["init"].values))

    fc_c = fc[:, :, land_iy, land_ix]                     # (init, lead, cell)
    thr_c = np.empty((ninit, nday, ncell), dtype="float64")
    for it in range(ninit):
        thr_c[it] = thr[vdoy[it] - 1][:, land_iy, land_ix]
    fc_ehd = (fc_c >= thr_c)

    lat_cells = fc_year["lat"].values[land_iy]
    lon_cells = fc_year["lon"].values[land_ix]
    out = {}
    lead_ids = {name: k + 1 for k, name in enumerate(blocks)}
    for name, (d0, d1) in blocks.items():
        sl = slice(d0, d1 + 1)
        fb = fc_c[:, sl, :]
        tb = thr_c[:, sl, :]
        f_max = np.nanmax(fb, axis=1)
        f_mean = np.nanmean(fb, axis=1)
        f_min = np.nanmin(fb, axis=1)
        f_exceed = np.nanmean(fc_ehd[:, sl, :].astype(float), axis=1)
        f_maxanom = f_max - np.nanmean(tb, axis=1)
        n = ninit * ncell
        lead = np.full(n, lead_ids[name], dtype=float)
        sin_d = np.repeat(np.sin(2 * np.pi * init_doy / 365.0), ncell)
        cos_d = np.repeat(np.cos(2 * np.pi * init_doy / 365.0), ncell)
        latf = np.tile(lat_cells, ninit)
        lonf = np.tile(lon_cells, ninit)
        X = np.column_stack([f_max.ravel(), f_mean.ravel(), f_min.ravel(),
                             f_exceed.ravel(), f_maxanom.ravel(), lead,
                             sin_d, cos_d, latf, lonf])
        out[name] = dict(X=X, init_idx=np.repeat(np.arange(ninit), ncell),
                         cell_id=np.tile(np.arange(ncell), ninit))
    return out