"""Run all Phase 3 ablations in sequence and print a comparison table.

Each ablation changes one variable vs the baseline. Results are logged to
results/runs.jsonl automatically by the training loop.

Usage:
    python scripts/run_ablations.py                   # all ablations
    python scripts/run_ablations.py --skip-finbert    # skip domain-pretrained run
"""

from __future__ import annotations

import argparse
import json
import logging
import os
import sys
from pathlib import Path

sys.path.insert(0, str(Path(__file__).resolve().parents[1] / "src"))

from rich.console import Console
from rich.table import Table

from finner.config import settings, RESULTS_DIR, CHECKPOINTS_DIR
from finner.model.train import RunConfig, train

logging.basicConfig(level=logging.INFO, format="%(levelname)s %(message)s")
console = Console()

os.environ.setdefault("PYTORCH_MPS_HIGH_WATERMARK_RATIO", "0.0")

# ── Ablation definitions ────────────────────────────────────────────────────
# Each entry: (run_name, RunConfig kwargs overrides on top of BASE_CFG)

BASE = dict(
    model_name="bert-base-uncased",
    learning_rate=2e-5,
    warmup_ratio=0.06,
    num_epochs=3,
    batch_size=4,
    weight_decay=0.01,
    lr_scheduler_type="linear",
    disc_lr_factor=0.0,
    focal_loss_gamma=0.0,
    o_token_weight=1.0,
    gradient_accumulation_steps=4,
    gradient_checkpointing=True,
)

ABLATIONS: list[tuple[str, dict]] = [
    # 1. Discriminative LR: lower encoder layers get 10× smaller LR than head
    ("disc_lr_0.1",   {"disc_lr_factor": 0.1,  "notes": "disc_lr: encoder=lr*0.1"}),
    # 2. Cosine LR schedule instead of linear
    ("cosine_sched",  {"lr_scheduler_type": "cosine", "notes": "cosine_lr_decay"}),
    # 3. Focal loss γ=2 — down-weights easy O-token examples
    ("focal_g2",      {"focal_loss_gamma": 2.0, "notes": "focal_loss_gamma=2"}),
    # 4. Class-weighted CE: O token gets 0.3× weight to reduce its dominance
    ("weighted_ce",   {"o_token_weight": 0.3,   "notes": "o_token_weight=0.3"}),
    # 5. Longer warmup (10% vs 6%)
    ("warmup_10pct",  {"warmup_ratio": 0.10,    "notes": "warmup_ratio=0.10"}),
    # 6. Domain-pretrained: FinBERT (same BERT arch, pretrained on financial text)
    ("finbert",       {"model_name": "ProsusAI/finbert", "notes": "finbert_domain_pretrained"}),
]


def run_ablation(name: str, overrides: dict) -> dict:
    cfg_kwargs = {**BASE, **overrides}
    cfg = RunConfig(**cfg_kwargs)
    console.print(f"\n[bold cyan]▶ Running ablation: {name}[/bold cyan]")
    console.print(f"  Config: {overrides}")
    result = train(cfg, run_name=name)
    console.print(f"  Best val entity-F1: [bold green]{result['best_val_entity_f1']:.4f}[/bold green]  (epoch {result['best_epoch']})")
    return result


def load_baseline_f1() -> float:
    baseline_path = RESULTS_DIR / "baseline.json"
    if baseline_path.exists():
        with open(baseline_path) as f:
            return json.load(f).get("overall_f1", 0.0)
    return 0.0


def print_comparison(results: list[dict], baseline_f1: float) -> None:
    table = Table(title="Phase 3 Ablation Results", show_header=True)
    table.add_column("Run", style="cyan")
    table.add_column("Val entity-F1", justify="right")
    table.add_column("Δ vs baseline", justify="right")
    table.add_column("Best epoch", justify="right")

    # baseline row
    table.add_row("baseline", f"{baseline_f1:.4f}", "—", "—")

    for r in sorted(results, key=lambda x: x["best_val_entity_f1"], reverse=True):
        delta = r["best_val_entity_f1"] - baseline_f1
        delta_str = f"[green]+{delta:.4f}[/green]" if delta >= 0 else f"[red]{delta:.4f}[/red]"
        table.add_row(r["run_name"], f"{r['best_val_entity_f1']:.4f}", delta_str, str(r["best_epoch"]))

    console.print(table)


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--skip-finbert", action="store_true")
    parser.add_argument("--only", default=None, help="Run only this ablation name")
    parser.add_argument("--skip", default=None, help="Comma-separated list of ablation names to skip")
    args = parser.parse_args()

    baseline_f1 = load_baseline_f1()
    console.print(f"\n[bold]Baseline val entity-F1: {baseline_f1:.4f}[/bold]")

    ablations = ABLATIONS
    if args.skip_finbert:
        ablations = [(n, c) for n, c in ablations if n != "finbert"]
    if args.only:
        ablations = [(n, c) for n, c in ablations if n == args.only]
    if args.skip:
        skip_set = {s.strip() for s in args.skip.split(",")}
        ablations = [(n, c) for n, c in ablations if n not in skip_set]

    results = []
    for name, overrides in ablations:
        try:
            r = run_ablation(name, overrides)
            results.append(r)
        except Exception as exc:
            console.print(f"[red]Ablation {name} failed: {exc}[/red]")

    if results:
        print_comparison(results, baseline_f1)

        # Save ablation table
        ablation_table = {
            "baseline_f1": baseline_f1,
            "ablations": [
                {"run": r["run_name"], "val_f1": r["best_val_entity_f1"],
                 "delta": r["best_val_entity_f1"] - baseline_f1,
                 "best_epoch": r["best_epoch"]}
                for r in results
            ]
        }
        out = RESULTS_DIR / "ablation_table.json"
        with open(out, "w") as f:
            json.dump(ablation_table, f, indent=2)
        console.print(f"\n[bold green]✓ Ablation table saved → {out}[/bold green]")


if __name__ == "__main__":
    main()