# Author: Simon-Pierre Boucher — simon-pierre.boucher@uqo.ca
"""Shared helpers: feature selection from the feature matrix, prediction records, interval calibration."""
from __future__ import annotations

from dataclasses import dataclass

import numpy as np
import polars as pl

ID_COLS = {"tx_id", "split", "log_price", "price", "date", "parcel_id", "stratum"}
CATEGORICAL = ["region_name", "region_group", "mun_code", "property_class", "subtype", "grid1k", "grid5k", "urban_class"]
QUANTILES = (0.05, 0.25, 0.50, 0.75, 0.95)
QCOLS = ("q05", "q25", "q50", "q75", "q95")


def numeric_features(df: pl.DataFrame) -> list[str]:
    out = []
    for c, dt in df.schema.items():
        if c in ID_COLS or c in CATEGORICAL:
            continue
        if dt.is_numeric() or dt == pl.Boolean:
            out.append(c)
    return out


def categorical_features(df: pl.DataFrame) -> list[str]:
    return [c for c in CATEGORICAL if c in df.columns]


@dataclass
class Prediction:
    tx_id: np.ndarray
    point: np.ndarray            # CAD
    quantiles: np.ndarray | None  # shape (n, 5) CAD, in QUANTILES order

    def to_records(self, split: str, arm: str, run_id: str, latency_s: float = 0.0) -> pl.DataFrame:
        n = len(self.tx_id)
        q = self.quantiles if self.quantiles is not None else np.full((n, 5), np.nan)
        return pl.DataFrame(
            {
                "tx_id": self.tx_id, "split": [split] * n, "arm": [arm] * n, "committee": ["none"] * n, "aggregation": ["model"] * n,
                "point": self.point.astype(float), "q05": q[:, 0], "q25": q[:, 1], "q50": q[:, 2], "q75": q[:, 3], "q95": q[:, 4],
                "cost_usd": [0.0] * n, "latency_s": [latency_s / max(n, 1)] * n, "tokens": [0] * n, "run_id": [run_id] * n,
            }
        )


def residual_quantile_intervals(point_log: np.ndarray, resid_valid: np.ndarray, scale: np.ndarray | None = None) -> np.ndarray:
    """Intervals from empirical validation residual quantiles (optionally scaled per observation)."""
    qs = np.quantile(resid_valid, QUANTILES)
    s = np.ones_like(point_log) if scale is None else scale
    return np.exp(point_log[:, None] + s[:, None] * qs[None, :])


def conformal_scale(valid_point_log: np.ndarray, valid_q_log: np.ndarray, y_valid_log: np.ndarray, nominal: float = 0.90) -> float:
    """Multiplicative factor on the half-width of the [q05,q95] log interval achieving `nominal` coverage on validation."""
    lo, hi = valid_q_log[:, 0], valid_q_log[:, 4]
    half = np.maximum((hi - lo) / 2.0, 1e-6)
    centre = (hi + lo) / 2.0
    score = np.abs(y_valid_log - centre) / half
    k = int(np.ceil((len(score) + 1) * nominal)) - 1
    k = min(max(k, 0), len(score) - 1)
    return float(np.sort(score)[k])


def apply_scale(q_log: np.ndarray, factor: float) -> np.ndarray:
    centre = (q_log[:, 4] + q_log[:, 0]) / 2.0
    return centre[:, None] + factor * (q_log - centre[:, None])