# Author: Simon-Pierre Boucher — simon-pierre.boucher@uqo.ca
"""LaTeX fragment writers (booktabs, siunitx `\\num{}` cells). Each writer produces `<stem>.tex` and `<stem>.csv`.

Conventions
* A fragment is a complete `tabular` environment without the `table` float or caption (captions live in the paper).
* The last line of the file is `\\end{tabular}` — never a row ending in `\\\\` — so `\\input{}` can be followed by
  a `\\caption` or `\\label` without a spurious blank row.
* Numbers go through `\\num{...}` (siunitx); NaN / suppressed cells are printed as `--` inside braces `{--}`.
* Text cells are escaped (`_ % & # $`), arm names are prettified (`G_debate_r2` → `G (debate r2)`).
"""
from __future__ import annotations

import re
from collections.abc import Sequence
from pathlib import Path

import numpy as np
import pandas as pd

from .metrics import METRIC_LABELS, PRIMARY_METRIC
from .stratify import STRATUM_LABELS

MAIN_METRICS = [PRIMARY_METRIC, "mae_cad", "mdape", "rmse_log", "p90_abs_log_error", "p95_abs_log_error", "bias_log", "severe_rate", "n"]
COVERAGE_METRICS = ["coverage_50", "coverage_80", "coverage_90", "coverage_95", "mean_log_width_90", "pinball_log", "crps_log", "n_with_interval"]
COST_METRICS = ["cost_usd_mean", "latency_s_mean", "tokens_mean", "cost_usd_total", PRIMARY_METRIC, "n"]
DIGITS = {"n": 0, "n_with_interval": 0, "mae_cad": 0, "tokens_mean": 0, "cost_usd_total": 2, "cost_usd_mean": 4, "latency_s_mean": 1,
          "p_boot": 3, "p_holm": 3, "dm_p": 3, "dm_stat": 2}
_ESC = {"&": r"\&", "%": r"\%", "$": r"\$", "#": r"\#", "_": r"\_", "{": r"\{", "}": r"\}", "~": r"\textasciitilde{}", "^": r"\textasciicircum{}"}


def escape(text: object) -> str:
    s = str(text)
    return "".join(_ESC.get(ch, ch) for ch in s)


def arm_pretty(arm: str) -> str:
    """`G_debate_r2` → `G (debate r2)`; `A_hedonic` → `A (hedonic)`; other strings are escaped as-is."""
    m = re.match(r"^([A-Z])_(.+)$", str(arm))
    if not m:
        return escape(arm)
    rest = m.group(2).split("_")
    return f"{m.group(1)} ({escape(' '.join(rest))})"


def num(x: object, digits: int = 3) -> str:
    try:
        v = float(x)  # type: ignore[arg-type]
    except (TypeError, ValueError):
        return "{--}"
    if not np.isfinite(v):
        return "{--}"
    return rf"\num{{{v:.{digits}f}}}"


def tabular(header: Sequence[str], rows: Sequence[Sequence[str]], colspec: str, *, midrules_after: Sequence[int] = ()) -> str:
    """Assemble a booktabs tabular; the returned string ends with `\\end{tabular}` (no trailing newline, no `\\\\`)."""
    lines = [rf"\begin{{tabular}}{{{colspec}}}", r"\toprule", " & ".join(header) + r" \\", r"\midrule"]
    for i, r in enumerate(rows):
        lines.append(" & ".join(r) + r" \\")
        if i in midrules_after and i < len(rows) - 1:
            lines.append(r"\midrule")
    lines += [r"\bottomrule", r"\end{tabular}"]
    return "\n".join(lines)


def write_fragment(tex: str, path: Path, csv_df: pd.DataFrame | None = None) -> Path:
    path = Path(path)
    path.parent.mkdir(parents=True, exist_ok=True)
    body = tex.rstrip()
    assert not body.endswith(r"\\"), "fragment must not end with \\\\"
    path.write_text("% Author: Simon-Pierre Boucher — simon-pierre.boucher@uqo.ca\n% generated by scripts/20_evaluate.py — do not edit\n" + body + "\n", encoding="utf-8")
    if csv_df is not None:
        csv_df.to_csv(path.with_suffix(".csv"), index=False)
    return path


def _label_col(df: pd.DataFrame) -> pd.Series:
    if "label" in df.columns:
        return df["label"].map(escape)
    return df["arm"].map(arm_pretty)


def _metrics_table(df: pd.DataFrame, metrics: Sequence[str], path: Path, *, family_rule: bool = True) -> str:
    metrics = [m for m in metrics if m in df.columns]
    header = ["Arm"] + ["{" + METRIC_LABELS.get(m, escape(m)) + "}" for m in metrics]
    rows, mids = [], []
    labels = _label_col(df)
    fam = df["arm"].astype(str).str[0]
    for i, (_, r) in enumerate(df.iterrows()):
        rows.append([labels.iloc[i]] + [num(r[m], DIGITS.get(m, 3)) for m in metrics])
        if family_rule and i + 1 < len(df) and fam.iloc[i] <= "D" < fam.iloc[i + 1]:
            mids.append(i)
    tex = tabular(header, rows, "l" + "r" * len(metrics), midrules_after=mids)
    write_fragment(tex, path, df[["arm", "committee", "aggregation"] + metrics] if "committee" in df.columns else df[["arm"] + metrics])
    return tex


def write_main_results(summary: pd.DataFrame, path: Path, metrics: Sequence[str] = MAIN_METRICS) -> str:
    """Arms × accuracy metrics (one row per arm; pass a pre-filtered summary with one row per arm)."""
    return _metrics_table(summary.sort_values("arm"), metrics, path)


def write_coverage(summary: pd.DataFrame, path: Path, metrics: Sequence[str] = COVERAGE_METRICS) -> str:
    """Arms × interval metrics; arms without intervals print `--`."""
    return _metrics_table(summary.sort_values("arm"), metrics, path)


def write_cost(summary: pd.DataFrame, path: Path, metrics: Sequence[str] = COST_METRICS) -> str:
    return _metrics_table(summary.sort_values("arm"), metrics, path)


def write_contrasts(contrasts: pd.DataFrame, path: Path) -> str:
    """Paired contrasts: Δ = metric(B) − metric(A) with bootstrap percentile CI, bootstrap p, Holm p and DM p."""
    header = ["Contrast (B vs A)", "{Metric}", "{$N$}", "{A}", "{B}", r"{$\Delta$}", "{95\\% CI}", "{$p$ (boot)}", "{$p$ (Holm)}", "{$p$ (DM)}"]
    rows = []
    for _, r in contrasts.iterrows():
        ci = rf"[{float(r['ci_lo']):.3f}, {float(r['ci_hi']):.3f}]" if np.isfinite(r["ci_lo"]) else "{--}"
        rows.append([f"{arm_pretty(r['arm_b'])} vs {arm_pretty(r['arm_a'])}", "{" + METRIC_LABELS.get(r["metric"], escape(r["metric"])) + "}",
                     num(r["n"], 0), num(r["value_a"], DIGITS.get(r["metric"], 3)), num(r["value_b"], DIGITS.get(r["metric"], 3)),
                     num(r["diff"], DIGITS.get(r["metric"], 3)), "{" + ci + "}", num(r["p_boot"], 3), num(r.get("p_holm", np.nan), 3), num(r.get("dm_p", np.nan), 3)])
    tex = tabular(header, rows, "llrrrrcrrr")
    write_fragment(tex, path, contrasts)
    return tex


def write_strata(strata: pd.DataFrame, path: Path, metric: str = PRIMARY_METRIC, arms: Sequence[str] | None = None) -> str:
    """Stratified table: rows = stratum × level (with n), columns = arms; suppressed cells print `--`."""
    s = strata[strata["metric"] == metric].copy()
    if "label" not in s.columns:
        s["label"] = s["arm"]
    if arms is not None:
        s = s[s["label"].isin(list(arms))]
    arm_list = list(arms) if arms is not None else sorted(s["label"].unique())
    if s.empty:
        tex = tabular(["Stratum", "{$n$}"], [[escape("no strata available"), "{--}"]], "lr")
        write_fragment(tex, path, s)
        return tex
    piv = s.pivot_table(index=["stratum", "level"], columns="label", values="value", aggfunc="first", sort=False)
    ns = s.groupby(["stratum", "level"], sort=False)["n"].max()
    header = ["Stratum", "{$n$}"] + ["{" + arm_pretty(a) + "}" for a in arm_list]
    rows, mids = [], []
    order = [st for st in STRATUM_LABELS if st in s["stratum"].unique()]
    i = -1
    for st in order:
        block = s[s["stratum"] == st].drop_duplicates(["level"])
        for _, lv in block.iterrows():
            i += 1
            key = (st, lv["level"])
            vals = [num(piv.loc[key, a] if a in piv.columns else np.nan, DIGITS.get(metric, 3)) for a in arm_list]
            rows.append([f"{escape(STRATUM_LABELS.get(st, st))}: {escape(lv['level'])}", num(ns.loc[key], 0)] + vals)
        mids.append(i)
    tex = tabular(header, rows, "lr" + "r" * len(arm_list), midrules_after=mids)
    write_fragment(tex, path, s)
    return tex


FAITH_COLS = [("cited_share_valid", "Valid ids"), ("share_turns_adj_ok", "Adj. $\\pm$40\\%"), ("point_in_hull", "In hull"),
              ("quantile_sane", "Quantiles ok"), ("summary_words", "Words"), ("iqr_log_round0", "IQR$_0$"),
              ("iqr_log_final", "IQR$_R$"), ("anchoring_index", "Anchoring"), ("share_toward_median", "To median")]


def write_faithfulness(faith: pd.DataFrame, path: Path) -> str:
    cols = [(c, l) for c, l in FAITH_COLS if c in faith.columns]
    header = ["Arm", "{Committee}", "{Turns}"] + ["{" + l + "}" for _, l in cols]
    rows = []
    for _, r in faith.iterrows():
        rows.append([arm_pretty(r["arm"]), escape(r.get("committee", "")), num(r.get("n_turns", np.nan), 0)]
                    + [num(r[c], 0 if c == "summary_words" else 3) for c, _ in cols])
    if not rows:
        rows = [["{--}"] * len(header)]
    tex = tabular(header, rows, "ll" + "r" * (len(header) - 2))
    write_fragment(tex, path, faith)
    return tex