"""Track C performance comparator.

Joins per-estimator timings on the sp side and the R side, emits a
Markdown rollup, a LaTeX table for §6 of the manuscript, and a
log-log scaling figure.
"""
from __future__ import annotations

import json
import shutil
from pathlib import Path

import matplotlib

matplotlib.use("Agg")
import matplotlib.pyplot as plt


HERE = Path(__file__).resolve().parent
ROOT = HERE.parents[1]
RESULTS_DIR = HERE / "results"
FIGURES_DIR = HERE / "figures"
PAPER_TABLES_DIR = ROOT / "Paper-JSS" / "manuscript" / "tables"
PAPER_FIGURES_DIR = ROOT / "Paper-JSS" / "manuscript" / "figures"
PDF_METADATA = {"CreationDate": None}
FIGURES_DIR.mkdir(parents=True, exist_ok=True)
PAPER_TABLES_DIR.mkdir(parents=True, exist_ok=True)
PAPER_FIGURES_DIR.mkdir(parents=True, exist_ok=True)


ESTIMATORS = {
    "01_hdfe":  {"name": "HDFE 2-way FE", "ref": "fixest::feols",
                 "x_label": "N (observations)"},
    "02_csdid": {"name": "CS-DiD simple ATT", "ref": "did::att\\_gt",
                 "x_label": "N (observations)"},
    "03_scm":   {"name": "Classical SCM", "ref": "Synth::synth",
                 "x_label": "n\\_donors"},
    "04_dml":   {"name": "DML PLR (lin. learners)", "ref": "DoubleML::DoubleMLPLR",
                 "x_label": "N (observations)"},
}


def load(estimator: str, side: str) -> list[dict]:
    path = RESULTS_DIR / f"{estimator}_{side}.json"
    if not path.exists():
        return []
    return json.loads(path.read_text())["rows"]


def render_md() -> str:
    lines: list[str] = [
        "# Track C performance report",
        "",
        "Generated by `tests/perf/compare_perf.py`. Hardware:",
        "Apple Silicon arm64, 8 physical cores, 24 GB RAM, macOS 26.",
        "",
        "Each row reports the median wall-clock time across 3 or 5 reps after a warmup. "
        "The ratio column is `R_time / sp_time` (>1 means sp is faster).",
        "",
    ]
    for est, cfg in ESTIMATORS.items():
        py = load(est, "py")
        r  = load(est, "R")
        if not py or not r:
            continue
        lines.append(f"## {est}: {cfg['name']}")
        lines.append("")
        lines.append("| n | sp (s) | R (s) | ratio R/sp |")
        lines.append("|---:|---:|---:|---:|")
        r_by_n = {row["n"]: row for row in r}
        for prow in py:
            rrow = r_by_n.get(prow["n"])
            ratio = rrow["median_time_s"] / prow["median_time_s"] if rrow and prow["median_time_s"] > 0 else None
            ratio_s = f"{ratio:.2f}" if ratio is not None else "—"
            lines.append(f"| {prow['n']} | {prow['median_time_s']:.4f} | "
                          f"{rrow['median_time_s']:.4f} | {ratio_s} |")
        lines.append("")
    return "\n".join(lines) + "\n"


def render_tex() -> str:
    rows: list[str] = []
    ref_tex = {
        "fixest::feols": "\\pkg{fixest}",
        "did::att\\_gt": "\\pkg{did}",
        "Synth::synth": "\\pkg{Synth}",
        "DoubleML::DoubleMLPLR": "\\pkg{DoubleML}",
    }
    for est, cfg in ESTIMATORS.items():
        py = load(est, "py")
        r = load(est, "R")
        if not py or not r:
            continue
        r_by_n = {row["n"]: row for row in r}
        ratios = []
        for prow in py:
            rrow = r_by_n.get(prow["n"])
            if rrow and prow["median_time_s"] > 0:
                ratios.append(rrow["median_time_s"] / prow["median_time_s"])
        if not ratios:
            continue
        max_n = max(prow["n"] for prow in py)
        last = next(p for p in py if p["n"] == max_n)
        last_r = r_by_n.get(max_n, {}).get("median_time_s", float("nan"))
        ratio_max = last_r / last["median_time_s"] if last["median_time_s"] > 0 else None
        ratio_str = (
            f"sp $\\mathbf{{{ratio_max:.1f}\\times}}$"
            if ratio_max and ratio_max > 1.5 else
            f"R $\\mathbf{{{1/ratio_max:.1f}\\times}}$"
            if ratio_max and ratio_max < 0.67 else
            f"tie ${ratio_max:.2f}\\times$"
        )
        est_tex = est.replace('_', '\\_')
        max_n_tex = f"{max_n:,}".replace(",", "{,}")
        rows.append(
            f"\\code{{{est_tex}}} & {cfg['name']} & "
            f"{ref_tex.get(cfg['ref'], cfg['ref'])} & {max_n_tex} & "
            f"{last['median_time_s']:.3f} & {last_r:.3f} & {ratio_str} \\\\"
        )
    body = "\n".join(rows)
    return (
        "% AUTO-GENERATED by tests/perf/compare_perf.py\n"
        "\\begingroup\n"
        "\\scriptsize\n"
        "\\setlength{\\tabcolsep}{2pt}\n"
        "\\begin{longtable}{@{}p{0.09\\linewidth}p{0.18\\linewidth}p{0.12\\linewidth}r r r p{0.13\\linewidth}@{}}\n"
        "\\caption{Track C measured performance on Apple Silicon arm64 / 8 cores / 24~GB. "
        "Each row reports median wall-clock time, in seconds, at the largest sample size in the sweep, "
        "across 3 or 5 reps with one warmup, on identical seed-fixed DGPs. "
        "The Winner column reports the faster implementation and multiplicative gap; "
        "bold marks decisive $\\ge 1.5\\times$ wins. Two rows carry caveats discussed in the text "
        "and should not be read as like-for-like algorithmic wins: the DML row reflects "
        "\\pkg{R6} method-dispatch wrapper overhead on the \\proglang{R} side "
        "(cheap \\code{lm} nuisance), and the Callaway--Sant'Anna gap is measured on a "
        "homogeneous-effect DGP, the cleanest case for the vectorised group-time loop.}\n"
        "\\label{tab:track-c-perf}\\\\\n"
        "\\toprule\n"
        "Module & Estimator & R ref. & Max. $N$ & sp & R & Winner \\\\\n"
        "\\midrule\n"
        "\\endfirsthead\n"
        "\\multicolumn{7}{c}{\\textit{(continued)}}\\\\\n"
        "\\toprule\n"
        "Module & Estimator & R ref. & Max. $N$ & sp & R & Winner \\\\\n"
        "\\midrule\n"
        "\\endhead\n"
        "\\bottomrule\n"
        "\\endlastfoot\n"
        f"{body}\n"
        "\\end{longtable}\n"
        "\\endgroup\n"
    )


def render_figure() -> Path:
    fig, axes = plt.subplots(2, 2, figsize=(11, 8))
    for ax, (est, cfg) in zip(axes.flat, ESTIMATORS.items()):
        py = load(est, "py")
        r = load(est, "R")
        if not py or not r:
            ax.text(0.5, 0.5, "no data", ha="center", va="center")
            continue
        py_n = [row["n"] for row in py]
        py_t = [row["median_time_s"] for row in py]
        r_n = [row["n"] for row in r]
        r_t = [row["median_time_s"] for row in r]
        ax.loglog(py_n, py_t, "o-", label="StatsPAI", linewidth=2, markersize=8)
        ax.loglog(r_n, r_t, "s--", label=cfg["ref"], linewidth=2, markersize=7)
        ax.set_xlabel(cfg["x_label"])
        ax.set_ylabel("median wall-clock (s)")
        ax.set_title(f"{est}: {cfg['name']}")
        ax.grid(True, which="both", alpha=0.3)
        ax.legend()
    fig.suptitle(
        "Track C: log-log scaling, StatsPAI vs canonical R reference\n"
        "(Apple Silicon arm64 / 8 cores / 24 GB / macOS 26)",
        y=1.0,
    )
    fig.tight_layout()
    out = FIGURES_DIR / "track_c_loglog.pdf"
    fig.savefig(out, bbox_inches="tight", metadata=PDF_METADATA)
    fig.savefig(FIGURES_DIR / "track_c_loglog.png", bbox_inches="tight", dpi=150)
    plt.close(fig)
    shutil.copyfile(out, PAPER_FIGURES_DIR / "track_c_loglog.pdf")
    shutil.copyfile(FIGURES_DIR / "track_c_loglog.png", PAPER_FIGURES_DIR / "track_c_loglog.png")
    return out


def main() -> None:
    md = render_md()
    tex = render_tex()
    (RESULTS_DIR / "perf_table.md").write_text(md, encoding="utf-8")
    (RESULTS_DIR / "perf_table.tex").write_text(tex, encoding="utf-8")
    (PAPER_TABLES_DIR / "track_c_perf.tex").write_text(tex, encoding="utf-8")
    fig = render_figure()
    print(
        "OK -- wrote perf_table.{md,tex}, "
        "Paper-JSS/manuscript/tables/track_c_perf.tex, and "
        f"{fig.name}"
    )


if __name__ == "__main__":
    main()
