"""Phase 5: Generate all result artifacts — summary.json + every PNG.

Run after Phase 4 is complete. Idempotent: safe to re-run.

Usage:
    python scripts/run_phase5_report.py
"""

from __future__ import annotations

import logging
import sys
from pathlib import Path

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

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

from finner.eval.report import (
    build_summary,
    plot_training_curve,
    plot_ablation_bar,
    plot_per_class_f1,
    plot_llm_comparison,
)
from finner.eval.confusion import build_confusion_matrix, save_confusion_png, boundary_error_rate
from finner.eval.metrics import compute_seqeval, _reconstruct_sequences
from finner.config import RESULTS_DIR, CHECKPOINTS_DIR, SPLITS_DIR
from finner.data.split import load_jsonl
from finner.eval.run_eval import evaluate_checkpoint
from rich.console import Console
import json

console = Console()


def build_confusion_from_test() -> None:
    """Build confusion matrix from test-set predictions."""
    test_final = RESULTS_DIR / "test_final.json"
    if not test_final.exists():
        console.print("[yellow]Skipping confusion matrix — test_final.json not found.[/yellow]")
        return

    # Re-run predictions on test to get sequence-level data for confusion
    import torch
    from transformers import AutoModelForTokenClassification, AutoTokenizer, DataCollatorForTokenClassification
    from torch.utils.data import DataLoader
    from finner.model.train import NERDataset
    from finner.data.align import batch_align
    from finner.labels import IGNORE_INDEX, ID2LABEL
    from finner.model.build import get_device
    from finner.config import settings

    device = get_device()
    ckpt = CHECKPOINTS_DIR / "best"
    tokenizer = AutoTokenizer.from_pretrained(str(ckpt), use_fast=True)
    model = AutoModelForTokenClassification.from_pretrained(str(ckpt)).to(device)
    model.eval()

    test_examples = load_jsonl(SPLITS_DIR / "test.jsonl")
    aligned = batch_align(tokenizer, test_examples, max_length=settings.max_seq_length)
    collator = DataCollatorForTokenClassification(tokenizer, padding=True, label_pad_token_id=IGNORE_INDEX)
    loader = DataLoader(NERDataset(aligned), batch_size=16, collate_fn=collator)

    flat_labels, flat_preds = [], []
    with torch.no_grad():
        for batch in loader:
            batch = {k: v.to(device) for k, v in batch.items()}
            logits = model(**{k: v for k, v in batch.items() if k != "labels"}).logits
            flat_preds.extend(logits.argmax(-1).cpu().tolist())
            flat_labels.extend(batch["labels"].cpu().tolist())

    true_seqs, pred_seqs = _reconstruct_sequences(flat_labels, flat_preds)
    matrix = build_confusion_matrix(true_seqs, pred_seqs)
    save_confusion_png(matrix)

    boundary = boundary_error_rate(true_seqs, pred_seqs)
    console.print(f"Boundary error rate: {boundary['boundary_error_rate']:.3f}  "
                  f"({boundary['boundary_errors']}/{boundary['total_true_spans']} spans)")

    # Append to test_final.json
    with open(test_final) as f:
        data = json.load(f)
    data["boundary_errors"] = boundary
    with open(test_final, "w") as f:
        json.dump(data, f, indent=2)


def main():
    RESULTS_DIR.mkdir(parents=True, exist_ok=True)

    console.print("[bold cyan]Building confusion matrix...[/bold cyan]")
    build_confusion_from_test()

    console.print("[bold cyan]Building summary.json...[/bold cyan]")
    summary = build_summary()
    console.print(f"  test entity-F1: {summary.get('final_test_entity_f1')}")

    console.print("[bold cyan]Plotting training curve...[/bold cyan]")
    plot_training_curve()

    console.print("[bold cyan]Plotting ablation bar chart...[/bold cyan]")
    plot_ablation_bar()

    console.print("[bold cyan]Plotting per-class F1...[/bold cyan]")
    plot_per_class_f1()

    console.print("[bold cyan]Plotting LLM comparison...[/bold cyan]")
    plot_llm_comparison()

    console.print(f"\n[bold green]✓ Phase 5 complete. All artifacts in results/[/bold green]")
    console.print("\nFiles generated:")
    for f in sorted(RESULTS_DIR.glob("*.json")) + sorted(RESULTS_DIR.glob("*.png")):
        console.print(f"  {f.name}")


if __name__ == "__main__":
    main()