from __future__ import annotations

"""Compare baseline (no text) vs caption-fused runs. Compute Delta R@1 report.

Usage:
    python -m scripts.compare_runs \
        --baseline_report out/caption_test/baseline_no_text/eval_report.json \
        --full_report out/caption_test/balanced/eval_report.json \
        --output out/caption_test/comparison.md
"""

import argparse
import json
from pathlib import Path


_DIRECTIONS = (
    "query_to_gallery",
    "gallery_to_query",
)
_KS = (1, 5, 10)


def _load_metrics(report_path: Path) -> dict[str, float]:
    with report_path.open("r", encoding="utf-8") as f:
        data = json.load(f)
    return data.get("metrics", data)


def _format_row(name: str, baseline: dict[str, float], full: dict[str, float]) -> str:
    cells = [name]
    for k in _KS:
        key = f"r@{k}_{name}"
        b = baseline.get(key, float("nan"))
        f_ = full.get(key, float("nan"))
        delta = f_ - b if (b == b and f_ == f_) else float("nan")
        cells.append(f"{b:.4f} -> {f_:.4f} (D {delta:+.4f})")
    return "| " + " | ".join(cells) + " |"


def _interpret_delta(delta: float) -> str:
    if delta >= 0.03:
        return "PASS -- captions informative (D R@1 >= +3%)"
    if delta >= 0.01:
        return "MARGINAL -- consider VLM refinement (+1% <= D < +3%)"
    if delta >= 0:
        return "WEAK -- captions add little signal (< +1%)"
    return "HARMFUL -- captions confuse model (D < 0)"


def build_comparison_markdown(
    baseline: dict[str, float],
    full: dict[str, float],
) -> str:
    lines: list[str] = ["# Caption Quality Test: Comparison Report", ""]

    primary = "query_to_gallery"
    primary_key = f"r@1_{primary}"
    primary_delta = full.get(primary_key, 0.0) - baseline.get(primary_key, 0.0)
    verdict = _interpret_delta(primary_delta)

    lines.append(f"## Primary metric: D R@1 (query->gallery) = {primary_delta:+.4f}")
    lines.append("")
    lines.append(f"**Verdict:** {verdict}")
    lines.append("")

    # Gate value.
    gate = full.get("gate", None)
    if gate is not None:
        lines.append(f"**Fusion gate:** {gate:.4f} (1.0 = text ignored, 0.0 = image ignored)")
        lines.append("")

    lines.append("## All directions x K")
    lines.append("")
    header = "| Direction | R@1 base -> full | R@5 base -> full | R@10 base -> full |"
    sep = "|---|---|---|---|"
    lines.extend([header, sep])
    for direction in _DIRECTIONS:
        lines.append(_format_row(direction, baseline, full))
    lines.append("")

    lines.append("## Decision rule")
    lines.append("")
    lines.append("- D R@1 >= +3% -> captions pass, proceed to production")
    lines.append("- +1% <= D R@1 < +3% -> add VLM refinement, re-run")
    lines.append("- D R@1 < +1% -> redesign caption pipeline")
    lines.append("- D R@1 < 0 -> critical bug, investigate")
    lines.append("")

    return "\n".join(lines)


def main() -> None:
    parser = argparse.ArgumentParser(description="Compare baseline vs caption runs.")
    parser.add_argument("--baseline_report", type=Path, required=True)
    parser.add_argument("--full_report", type=Path, required=True)
    parser.add_argument("--output", type=Path, required=True)
    args = parser.parse_args()

    baseline = _load_metrics(args.baseline_report)
    full = _load_metrics(args.full_report)

    md = build_comparison_markdown(baseline=baseline, full=full)

    args.output.parent.mkdir(parents=True, exist_ok=True)
    with args.output.open("w", encoding="utf-8") as f:
        f.write(md)

    summary = {
        "baseline_metrics": baseline,
        "full_metrics": full,
        "deltas": {
            f"delta_r@{k}_{d}": (
                full.get(f"r@{k}_{d}", 0.0) - baseline.get(f"r@{k}_{d}", 0.0)
            )
            for d in _DIRECTIONS
            for k in _KS
        },
    }
    summary_path = args.output.with_suffix(".json")
    with summary_path.open("w", encoding="utf-8") as f:
        json.dump(summary, f, indent=2)

    print(f"Comparison saved: {args.output}")
    print(f"Summary saved:    {summary_path}")


if __name__ == "__main__":
    main()