#!/usr/bin/env python3
# Author: Simon-Pierre Boucher — simon-pierre.boucher@uqo.ca
"""Phase 5 — estimate the aggregation weights (inverse-variance, stacking, robust-Bayes scale) on the VALIDATION
transcripts of the reference debate run, and the conformal scales of the LLM arms' median-of-quantiles intervals.
Both are frozen inputs of the test evaluation (never re-estimated on test).

Usage: python3 scripts/34_fit_weights.py [--run validation_G_debate_r2_mixed_random_mid]
Outputs: results/aggregation_weights.json, results/evaluation/llm_conformal_scales.json
"""
from __future__ import annotations

import argparse
import json
import sys
from pathlib import Path

import pandas as pd
import polars as pl

sys.path.insert(0, str(Path(__file__).resolve().parents[1] / "src"))
from agentic_valuation.config import PROCESSED, RESULTS  # noqa: E402
from agentic_valuation.debate.aggregation import fit_weights, save_weights  # noqa: E402
from agentic_valuation.evaluation.calibration import conformalise  # noqa: E402
from agentic_valuation.schemas import DebateTranscript  # noqa: E402


def load_transcripts(run_dir: Path) -> list[DebateTranscript]:
    return [DebateTranscript.model_validate_json(p.read_text()) for p in sorted((run_dir / "transcripts").glob("*.json"))]


def main(argv: list[str] | None = None) -> None:
    ap = argparse.ArgumentParser()
    ap.add_argument("--run", default="validation_G_debate_r2_mixed_random_mid")
    a = ap.parse_args(argv)
    truth_df = pl.read_parquet(PROCESSED / "analytic.parquet", columns=["tx_id", "price", "mun_code", "date"])
    truth = dict(zip(truth_df["tx_id"], truth_df["price"]))
    run_dir = RESULTS / "runs" / a.run
    trs = load_transcripts(run_dir)
    w = fit_weights(trs, truth)
    w["source_run"] = a.run
    path = save_weights(w)
    print(f"weights from {len(trs)} validation transcripts -> {path}")
    print({k: round(v, 3) for k, v in w["inverse_variance"].items()})
    print("stacking", {k: round(v, 3) for k, v in w["stacking"].items() if v > 0.001})
    # conformal scales per LLM arm (validation runs only, median aggregation / median_of_quantiles interval)
    scales: dict[str, dict] = {}
    tp = truth_df.to_pandas().rename(columns={"price": "price"})
    for d in sorted((RESULTS / "runs").glob("validation_*")):
        f = d / "predictions.parquet"
        if not f.exists():
            continue
        p = pl.read_parquet(f).filter(pl.col("aggregation") == "median").to_pandas()
        if p.empty or p["q05"].isna().all():
            continue
        try:
            cs = conformalise(p, tp)
            scales[d.name] = {"arm": str(p["arm"].iloc[0]), "committee": str(p["committee"].iloc[0]), "n": int(len(p)), "scales": cs.__dict__ if hasattr(cs, "__dict__") else str(cs)}
        except Exception as e:  # noqa: BLE001
            scales[d.name] = {"error": str(e)[:200]}
    out = RESULTS / "evaluation" / "llm_conformal_scales.json"
    out.parent.mkdir(parents=True, exist_ok=True)
    out.write_text(json.dumps(scales, indent=1, default=float))
    print(f"conformal scales for {len(scales)} validation runs -> {out}")


if __name__ == "__main__":
    main()