from __future__ import annotations

"""Compare baseline vs full-caption runs and compute Delta R@1 report.

Reads eval reports produced by src.eval.evaluate.run_evaluation_from_checkpoint
and produces a markdown + JSON summary.

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 = (
    "drone_to_sat",
    "sat_to_drone",
    "text_to_sat",
    "text_to_drone",
)
_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:
    """Render one markdown row for a direction across R@1, R@5, R@10."""
    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")  # NaN-safe
        cells.append(f"{b:.4f} → {f_:.4f} (Δ {delta:+.4f})")
    return "| " + " | ".join(cells) + " |"


def _interpret_delta(delta: float) -> str:
    """Human-readable caption-quality verdict."""
    if delta >= 0.03:
        return "✅ PASS — captions informative (Δ R@1 ≥ +3%)"
    if delta >= 0.01:
        return "⚠️  MARGINAL — consider VLM refinement (+1% ≤ Δ < +3%)"
    if delta >= 0:
        return "❌ WEAK — captions add little signal (< +1%)"
    return "❌❌ HARMFUL — captions confuse model (Δ < 0)"


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

    # Headline Δ R@1 on primary direction.
    primary = "drone_to_sat"
    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: Δ R@1 ({primary}) = {primary_delta:+.4f}")
    lines.append("")
    lines.append(f"**Verdict:** {verdict}")
    lines.append("")

    # Full table.
    lines.append("## All directions × 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:
        row = _format_row(direction, baseline, full)
        lines.append(row)
    lines.append("")

    # Decision rule recap.
    lines.append("## Decision rule")
    lines.append("")
    lines.append("- Δ R@1 ≥ +3% → captions pass, proceed to World-UAV generation")
    lines.append("- +1% ≤ Δ R@1 < +3% → add VLM refinement, re-run")
    lines.append("- Δ R@1 < +1% → redesign caption pipeline")
    lines.append("- Δ R@1 < 0 → critical bug, investigate caption/image alignment")
    lines.append("")

    return "\n".join(lines)


def main() -> None:
    parser = argparse.ArgumentParser(
        description="Compare baseline vs full-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)

    # Also write machine-readable summary.
    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()