# Author: Simon-Pierre Boucher — simon-pierre.boucher@uqo.ca
"""Evaluate prediction runs against the resale-implied value (docs/fundamental_value.md §5).

Reading. For a sale (i, t1) with resale target V^R (repeat_sales.resale_targets), an arm's estimate ŷ has two errors:
    e^P = log ŷ − log P1          (price-based, the paper's primary metric)
    e^R = log ŷ − log V^R         (resale-based)
and the price itself has e^R_P = log P1 − log V^R = −(ε_{t2} − ε_{t1}) under the repeat-sales assumptions. The
diagnostic contrast is Δ = mean|e^R(arm)| − mean|e^R(price)|: an arm with Δ < 0 is closer to the market-adjusted
resale than the first transaction price was, which is evidence that it captured a component of value obscured by
transaction noise. Δ is estimated with the paired block bootstrap of `evaluation.inference` (blocks = municipality ×
month), the price playing the role of a pseudo-arm `P_price`. Arm-vs-arm contrasts against V^R use the same machinery.
Selection: by default, pairs held < 12 months and pairs whose later leg carries a data-quality flag are excluded.
"""
from __future__ import annotations

from pathlib import Path
from typing import Any

import numpy as np
import pandas as pd

from ..evaluation._common import to_pandas
from ..evaluation.inference import compare_arms
from ..evaluation.metrics import per_record, summarise_records

PRICE_ARM = "P_price"
RESALE_METRICS = ["mean_abs_log_error", "rmse_log", "bias_log"]


def collect_live_runs(runs_dir: Path, *, include_fake: bool = False) -> tuple[pd.DataFrame, list[str]]:
    """All prediction tables under `runs_dir`; run ids starting with `_` (mock / synthetic) are skipped."""
    import polars as pl

    frames, used = [], []
    for d in sorted(p for p in Path(runs_dir).iterdir() if p.is_dir()):
        if d.name.startswith("_") and not include_fake:
            continue
        f = d / "predictions.parquet"
        if not f.exists():
            continue
        df = pl.read_parquet(f).to_pandas()
        if "run_id" not in df.columns:
            df["run_id"] = d.name
        if not include_fake:
            df = df[~df["run_id"].astype(str).str.startswith("_")]
        if df.empty:
            continue
        frames.append(df)
        used.append(d.name)
    if not frames:
        return pd.DataFrame(), used
    return pd.concat(frames, ignore_index=True), used


def select_main_config(preds: pd.DataFrame, *, committee: str = "mixed_random", aggregation: str = "median") -> pd.DataFrame:
    """One (committee, aggregation, run_id) per arm — same preference order as scripts/20_evaluate.py — so that each
    arm has one row per tx_id."""
    out = []
    for arm, g in preds.groupby("arm", sort=True):
        cands = [g[(g["committee"] == committee) & (g["aggregation"] == aggregation)], g[g["aggregation"] == aggregation],
                 g[(g["committee"] == "none") & (g["aggregation"] == "model")], g]
        pick = next(c for c in cands if not c.empty)
        key = pick.sort_values(["committee", "run_id"]).iloc[-1][["committee", "aggregation", "run_id"]]
        sel = g[(g["committee"] == key["committee"]) & (g["aggregation"] == key["aggregation"]) & (g["run_id"] == key["run_id"])]
        out.append(sel.drop_duplicates(subset=["tx_id"], keep="last"))
    return pd.concat(out, ignore_index=True) if out else preds.iloc[0:0]


def resale_truth(truth: Any, targets: pd.DataFrame, *, exclude_short: bool = True, exclude_flagged: bool = True) -> pd.DataFrame:
    """Truth table restricted to sales with a resale target; `price` is replaced by the resale-implied value so that the
    standard metric code measures e^R, and the realised price is kept in `price_realised`."""
    t = to_pandas(truth)
    t["tx_id"] = t["tx_id"].astype(str)
    tg = targets.copy()
    tg["tx_id"] = tg["tx_id"].astype(str)
    if exclude_short:
        tg = tg[~tg["short_hold"].astype(bool)]
    if exclude_flagged:
        tg = tg[~tg["resale_flagged"].astype(bool)]
    tg = tg[np.isfinite(tg["resale_implied_log_value"])]
    cols = ["tx_id", "resale_implied_value", "resale_implied_log_value", "resale_price", "resale_date", "months_held", "index_change"]
    m = t.merge(tg[cols], on="tx_id", how="inner", validate="one_to_one")
    m["price_realised"] = m["price"].astype(float)
    m["price"] = m["resale_implied_value"].astype(float)
    return m


def evaluate_vs_resale(preds: pd.DataFrame, truth: Any, targets: pd.DataFrame, *, exclude_short: bool = True, exclude_flagged: bool = True,
                       B: int = 2000, seed: int = 20261027, arm_pairs: list[tuple[str, str]] | None = None) -> tuple[pd.DataFrame, pd.DataFrame]:
    """Per-arm summary against the realised price and against V^R on the common resale subsample, and the paired
    contrasts (arm vs price-as-estimate; optional arm-vs-arm pairs) against V^R.

    Returns (summary, contrasts). `summary` has one row per arm (plus `P_price`): n, mean|e| / rmse / bias vs price and
    vs resale-implied value. `contrasts` is the tidy output of `compare_arms` with `arm_a = P_price`."""
    tr = resale_truth(truth, targets, exclude_short=exclude_short, exclude_flagged=exclude_flagged)
    p = preds.copy()
    p["tx_id"] = p["tx_id"].astype(str)
    p = p[p["tx_id"].isin(set(tr["tx_id"]))]
    if p.empty:
        return pd.DataFrame(), pd.DataFrame()
    # pseudo-arm: the realised price as an estimate (quantiles absent)
    price_rows = tr[["tx_id", "price_realised"]].rename(columns={"price_realised": "point"})
    price_rows = price_rows[price_rows["tx_id"].isin(set(p["tx_id"]))]
    price_rows = price_rows.assign(arm=PRICE_ARM, committee="none", aggregation="model", run_id="price", split=p["split"].iloc[0] if "split" in p.columns else "n/a")
    all_preds = pd.concat([p, price_rows], ignore_index=True, sort=False)

    truth_price = tr.drop(columns=["price"]).rename(columns={"price_realised": "price"})
    rows = []
    for arm, g in all_preds.groupby("arm", sort=True):
        rp = summarise_records(per_record(g, truth_price))
        rr = summarise_records(per_record(g, tr))
        rows.append({"arm": arm, "n": rp["n"], "mean_abs_log_error_price": rp["mean_abs_log_error"], "rmse_log_price": rp["rmse_log"], "bias_log_price": rp["bias_log"],
                     "mean_abs_log_error_resale": rr["mean_abs_log_error"], "rmse_log_resale": rr["rmse_log"], "bias_log_resale": rr["bias_log"]})
    summary = pd.DataFrame(rows)
    arms = [a for a in all_preds["arm"].unique() if a != PRICE_ARM]
    pairs = [(PRICE_ARM, a) for a in sorted(arms)]
    if arm_pairs:
        pairs += [(a, b) for a, b in arm_pairs if a in arms and b in arms]
    contrasts = compare_arms(all_preds, tr, pairs, metric=RESALE_METRICS, B=B, seed=seed, dm=False, holm_family=True)
    contrasts.insert(0, "target", "resale_implied")
    return summary, contrasts