# Author: Simon-Pierre Boucher — simon-pierre.boucher@uqo.ca
"""Shared helpers: frame coercion, truth join, derived strata, arm selection.

Prediction record contract (docs/architecture.md §1):
    tx_id, split, arm, committee, aggregation, point, q05, q25, q50, q75, q95, cost_usd, latency_s, tokens, run_id
Truth contract (data/processed/analytic.parquet):
    tx_id, price, date, region_group, region_name, mun_code, property_class, subtype, urban_class, price_tercile,
    year_built, floor_area_m2 (+ optional n_sales_3km_12m)
"""
from __future__ import annotations

from typing import Any

import numpy as np
import pandas as pd

try:  # polars is a hard dependency of the package, but the evaluation code only needs it for coercion
    import polars as pl
except ImportError:  # pragma: no cover
    pl = None  # type: ignore[assignment]

PRED_COLUMNS = ["tx_id", "split", "arm", "committee", "aggregation", "point", "q05", "q25", "q50", "q75", "q95",
                "cost_usd", "latency_s", "tokens", "run_id"]
QUANTILE_COLUMNS = ["q05", "q25", "q50", "q75", "q95"]
TAUS = np.array([0.05, 0.25, 0.50, 0.75, 0.95])
ARM_KEY = ["arm", "committee", "aggregation"]
TRUTH_REQUIRED = ["tx_id", "price"]
TRUTH_OPTIONAL = ["date", "region_group", "region_name", "mun_code", "property_class", "subtype", "urban_class",
                  "price_tercile", "year_built", "floor_area_m2", "n_sales_3km_12m", "split"]

# Region groups used in the paper (preregistration §Sample). Any region not listed maps to "other".
REGION_GROUP_MAP = {
    "Montréal": "Montréal", "Montreal": "Montréal", "Laval": "Montréal",
    "Capitale-Nationale": "Québec", "Québec": "Québec", "Quebec": "Québec",
    "Outaouais": "Outaouais",
}
REGION_GROUP_ORDER = ["Montréal", "Québec", "Outaouais", "other"]


def to_pandas(df: Any) -> pd.DataFrame:
    """Accept a polars or pandas frame and return a pandas copy."""
    if pl is not None and isinstance(df, pl.DataFrame):
        return df.to_pandas()
    if isinstance(df, pd.DataFrame):
        return df.copy()
    raise TypeError(f"expected a polars or pandas DataFrame, got {type(df)!r}")


def validate_predictions(preds: pd.DataFrame) -> pd.DataFrame:
    """Check the prediction contract; fill optional cost columns with NaN when absent."""
    missing = [c for c in ["tx_id", "arm", "point"] if c not in preds.columns]
    if missing:
        raise ValueError(f"prediction table is missing required columns {missing}")
    out = preds.copy()
    for c in ["committee", "aggregation", "split", "run_id"]:
        if c not in out.columns:
            out[c] = "none" if c in {"committee"} else ("model" if c == "aggregation" else "unknown")
    for c in QUANTILE_COLUMNS + ["cost_usd", "latency_s", "tokens"]:
        if c not in out.columns:
            out[c] = np.nan
    out["tx_id"] = out["tx_id"].astype(str)
    out["point"] = out["point"].astype(float)
    for c in QUANTILE_COLUMNS + ["cost_usd", "latency_s"]:
        out[c] = pd.to_numeric(out[c], errors="coerce").astype(float)
    out["tokens"] = pd.to_numeric(out["tokens"], errors="coerce")
    if (out["point"] <= 0).any():
        raise ValueError("point estimates must be strictly positive (log error undefined)")
    return out


def derive_region_group(region_name: pd.Series) -> pd.Series:
    return region_name.map(lambda r: REGION_GROUP_MAP.get(str(r), "other")).astype("string")


def derive_price_tercile(price: pd.Series) -> pd.Series:
    """Tercile of the sale price over the supplied truth table (labels low/mid/high)."""
    q1, q2 = np.nanquantile(price.astype(float), [1 / 3, 2 / 3])
    return pd.Series(np.where(price <= q1, "low", np.where(price <= q2, "mid", "high")), index=price.index, dtype="string")


def prepare_truth(truth: Any) -> pd.DataFrame:
    """Coerce the truth table, derive `region_group` and `price_tercile` when absent, parse `date`."""
    t = to_pandas(truth)
    missing = [c for c in TRUTH_REQUIRED if c not in t.columns]
    if missing:
        raise ValueError(f"truth table is missing required columns {missing}")
    t["tx_id"] = t["tx_id"].astype(str)
    t["price"] = t["price"].astype(float)
    if (t["price"] <= 0).any():
        raise ValueError("truth prices must be strictly positive")
    if "region_group" not in t.columns and "region_name" in t.columns:
        t["region_group"] = derive_region_group(t["region_name"])
    if "price_tercile" not in t.columns:
        t["price_tercile"] = derive_price_tercile(t["price"])
    if "date" in t.columns:
        t["date"] = pd.to_datetime(t["date"])
    if t["tx_id"].duplicated().any():
        raise ValueError("truth table has duplicated tx_id")
    return t


def join_truth(preds: Any, truth: Any, *, how: str = "inner") -> pd.DataFrame:
    """Join predictions to truth on tx_id (truth columns suffixed only on collision: `split` → `split_truth`)."""
    p = validate_predictions(to_pandas(preds))
    t = prepare_truth(truth)
    keep = [c for c in TRUTH_REQUIRED + TRUTH_OPTIONAL if c in t.columns]
    t = t[keep].rename(columns={"split": "split_truth"})
    joined = p.merge(t, on="tx_id", how=how, validate="many_to_one")
    if how == "inner" and len(joined) < len(p):
        joined.attrs["n_unmatched"] = len(p) - len(joined)
    return joined


def arm_filter(preds: pd.DataFrame, spec: str | dict[str, Any]) -> pd.DataFrame:
    """Select the rows of one arm. `spec` is an arm name or a dict with keys among arm/committee/aggregation/run_id."""
    if isinstance(spec, str):
        spec = {"arm": spec}
    mask = np.ones(len(preds), dtype=bool)
    for k, v in spec.items():
        mask &= (preds[k].astype(str) == str(v)).to_numpy()
    sub = preds.loc[mask]
    if sub.empty:
        raise ValueError(f"no predictions match {spec}")
    if sub["tx_id"].duplicated().any():
        combos = sub[ARM_KEY + ["run_id"]].drop_duplicates().to_dict("records")
        raise ValueError(f"arm spec {spec} is ambiguous (several rows per tx_id); disambiguate with committee/aggregation/run_id: {combos[:6]}")
    return sub


def arm_label(spec: str | dict[str, Any]) -> str:
    if isinstance(spec, str):
        return spec
    return "/".join(str(spec[k]) for k in ["arm", "committee", "aggregation", "run_id"] if k in spec)


def block_ids(df: pd.DataFrame) -> pd.Series:
    """Block = mun_code × calendar month when both are present, else mun_code, else tx_id (i.i.d. fallback)."""
    if "mun_code" in df.columns and df["mun_code"].notna().any():
        mun = df["mun_code"].astype(str)
        if "date" in df.columns and df["date"].notna().any():
            month = pd.to_datetime(df["date"]).dt.strftime("%Y-%m")
            return (mun + "|" + month).rename("block")
        return mun.rename("block")
    return df["tx_id"].astype(str).rename("block")