# Author: Simon-Pierre Boucher — simon-pierre.boucher@uqo.ca
"""Standard strata and stratified metric tables with the minimum-cell-size rule (n ≥ 30, else suppressed).
Definitions: docs/evaluation_protocol.md §6."""
from __future__ import annotations

from collections.abc import Sequence
from typing import Any

import numpy as np
import pandas as pd

from ._common import REGION_GROUP_ORDER, derive_price_tercile, derive_region_group
from .metrics import PRIMARY_METRIC, per_record, summarise_records

MIN_CELL_N = 30
STANDARD_STRATA = ["region_group", "urban_class", "property_class", "price_tercile", "year_built_decade", "density_tercile"]
STRATUM_LABELS = {
    "region_group": "Region", "urban_class": "Urban class", "property_class": "Property class",
    "price_tercile": "Price tercile", "year_built_decade": "Construction decade", "density_tercile": "Transaction density",
}
LEVEL_ORDER = {
    "region_group": REGION_GROUP_ORDER,
    "price_tercile": ["low", "mid", "high"],
    "density_tercile": ["thin", "mid", "liquid"],
    "property_class": ["house", "condo", "plex"],
    "urban_class": ["metro", "suburban", "rural"],
}
DEFAULT_STRATA_METRICS = [PRIMARY_METRIC, "mdape", "p95_abs_log_error", "severe_rate", "coverage_90", "bias_log"]


def year_built_decade(year_built: pd.Series) -> pd.Series:
    """Decade bins: <1950, 1950s … 2000s, ≥2010; NaN → NA (row excluded from this stratum)."""
    y = pd.to_numeric(year_built, errors="coerce")
    lab = np.where(y < 1950, "<1950", np.where(y >= 2010, "2010+", (np.floor(y / 10) * 10).astype("float").astype("Int64").astype(str) + "s"))
    lab = pd.Series(lab, index=year_built.index, dtype="string")
    lab[y.isna()] = pd.NA  # missing construction year: excluded from the decade stratum (as for density)
    return lab


def density_tercile(n_sales: pd.Series) -> pd.Series:
    """Terciles of local liquidity (`n_sales_3km_12m`) → thin / mid / liquid."""
    v = pd.to_numeric(n_sales, errors="coerce")
    q1, q2 = np.nanquantile(v, [1 / 3, 2 / 3])
    lab = pd.Series(np.where(v <= q1, "thin", np.where(v <= q2, "mid", "liquid")), index=n_sales.index, dtype="string")
    lab[v.isna()] = pd.NA
    return lab


def add_strata(df: pd.DataFrame) -> pd.DataFrame:
    """Add every standard stratum column that can be derived from the available columns."""
    out = df.copy()
    if "region_group" not in out.columns and "region_name" in out.columns:
        out["region_group"] = derive_region_group(out["region_name"])
    if "price_tercile" not in out.columns and "price" in out.columns:
        out["price_tercile"] = derive_price_tercile(out["price"])
    if "year_built_decade" not in out.columns and "year_built" in out.columns:
        out["year_built_decade"] = year_built_decade(out["year_built"])
    if "density_tercile" not in out.columns and "n_sales_3km_12m" in out.columns:
        out["density_tercile"] = density_tercile(out["n_sales_3km_12m"])
    return out


def _sorted_levels(stratum: str, levels: list) -> list:
    order = LEVEL_ORDER.get(stratum)
    if order:
        return sorted(levels, key=lambda v: (order.index(v) if v in order else len(order), str(v)))
    return sorted(levels, key=str)


def stratified_metrics(
    preds: Any,
    truth: Any,
    strata: Sequence[str] = STANDARD_STRATA,
    metrics: Sequence[str] = DEFAULT_STRATA_METRICS,
    *,
    group_cols: Sequence[str] = ("arm", "committee", "aggregation"),
    min_n: int = MIN_CELL_N,
) -> pd.DataFrame:
    """Long table: group_cols…, stratum, level, n, metric, value, suppressed.

    A cell is suppressed (value = NaN, suppressed = True) when it has fewer than `min_n` records, so that
    unsupported granular estimates never reach the tables. Strata absent from the data are skipped.
    """
    d = add_strata(per_record(preds, truth))
    group_cols = list(group_cols)
    rows = []
    for stratum in strata:
        if stratum not in d.columns:
            continue
        for gkey, g in d.groupby(group_cols, sort=True, dropna=False):
            gkey = gkey if isinstance(gkey, tuple) else (gkey,)
            base = dict(zip(group_cols, gkey))
            levels = _sorted_levels(stratum, [v for v in g[stratum].dropna().unique()])
            for level in levels:
                cell = g[g[stratum] == level]
                n = len(cell)
                summ = summarise_records(cell) if n >= min_n else {}
                for m in metrics:
                    row = dict(base)
                    row.update({"stratum": stratum, "level": str(level), "n": n, "metric": m,
                                "value": float(summ.get(m, np.nan)) if n >= min_n else np.nan, "suppressed": n < min_n})
                    rows.append(row)
    cols = group_cols + ["stratum", "level", "n", "metric", "value", "suppressed"]
    return pd.DataFrame(rows, columns=cols)