#!/usr/bin/env python3
"""Evaluate HFM checkpoint with all notebook post-hoc variants."""

from __future__ import annotations

import argparse
import json
from pathlib import Path

import numpy as np
from transformers import CLIPProcessor, DebertaV2Tokenizer

import sys

PROJECT_ROOT = Path(__file__).resolve().parents[1]
if str(PROJECT_ROOT) not in sys.path:
    sys.path.insert(0, str(PROJECT_ROOT))

from src.hfm_pipeline import (
    HFMLoader,
    build_classification_outputs,
    create_dataloaders,
    evaluate_ablation,
    estimate_max_length,
    evaluate_all_variants,
    evaluate_raw,
    gather_logits_labels,
    load_checkpoint,
    resolve_device,
    save_summary_files,
    summarize_splits,
)


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description="Evaluate CLARA HFM checkpoint")
    parser.add_argument(
        "--checkpoint",
        default="outputs/hfm/clara_hfm.pt",
        help="Path to trained checkpoint",
    )
    parser.add_argument("--data-root", default="data/HFM", help="Path to HFM root folder")
    parser.add_argument(
        "--text-dir",
        default=None,
        help="Path to HFM text folder (default: <data-root>/text)",
    )
    parser.add_argument("--output-dir", default="results/hfm", help="Evaluation output directory")
    parser.add_argument("--batch-size", type=int, default=None, help="Override batch size")
    parser.add_argument("--max-length", type=int, default=None, help="Override token max length")
    parser.add_argument("--num-workers", type=int, default=None)
    parser.add_argument("--top2-eps", type=float, default=0.03)
    parser.add_argument(
        "--classification-mode",
        choices=["raw", "bias_temp"],
        default="raw",
        help="Which predictions are used for classification report/confusion matrix",
    )
    parser.add_argument(
        "--ablation-full-mode",
        choices=["raw", "bias_temp"],
        default="raw",
        help="Metric source for Full row in ablation table",
    )
    parser.add_argument("--device", default="auto", help="auto|cuda|cpu")
    return parser.parse_args()


def main() -> None:
    args = parse_args()

    text_dir = args.text_dir or str(Path(args.data_root) / "text")
    device = resolve_device(args.device)

    model, cfg, ckpt_meta = load_checkpoint(args.checkpoint, device)
    cfg["image_root"] = args.data_root
    cfg["text_dir"] = text_dir

    if args.batch_size is not None:
        cfg["batch_size"] = args.batch_size
    if args.num_workers is not None:
        cfg["num_workers"] = args.num_workers
    if args.max_length is not None:
        cfg["max_length"] = args.max_length

    loader = HFMLoader(cfg["text_dir"], cfg["image_root"])
    all_samples = loader.load()
    train_samples = loader.get_split("train")
    val_samples = loader.get_split("val")
    test_samples = loader.get_split("test")

    if not val_samples or not test_samples:
        raise RuntimeError("Need both val and test splits for full evaluation.")

    split_stats = summarize_splits(train_samples, val_samples, test_samples)
    print("Split stats:")
    print(json.dumps(split_stats, indent=2))

    clip_processor = CLIPProcessor.from_pretrained(cfg["vision_model_id"])
    tokenizer = DebertaV2Tokenizer.from_pretrained(cfg["text_model_id"])

    max_length = cfg.get("max_length")
    if not max_length:
        max_length = estimate_max_length(
            all_samples,
            percentile=cfg["max_length_percentile"],
            sample_size=cfg["max_length_sample_size"],
        )
        cfg["max_length"] = int(max_length)

    pin_memory = bool(cfg["pin_memory"] and device.type == "cuda")
    _, val_loader, test_loader = create_dataloaders(
        train_samples=train_samples,
        val_samples=val_samples,
        test_samples=test_samples,
        clip_processor=clip_processor,
        tokenizer=tokenizer,
        batch_size=cfg["batch_size"],
        max_length=cfg["max_length"],
        num_workers=cfg["num_workers"],
        pin_memory=pin_memory,
        weighted_train_sampler=False,
    )

    result = evaluate_all_variants(
        model=model,
        val_loader=val_loader,
        test_loader=test_loader,
        device=device,
        top2_eps=args.top2_eps,
    )
    y_pred_raw, y_true_raw, _ = evaluate_raw(model, test_loader, device)

    if args.classification_mode == "raw":
        cls_outputs = build_classification_outputs(y_true=y_true_raw, y_pred=y_pred_raw)
    else:
        logits_test, y_true_bt = gather_logits_labels(model, test_loader, device)
        bias = float(result["tuning"]["bias_temp"]["bias"])
        tau = float(result["tuning"]["bias_temp"]["tau"])
        classes = logits_test.shape[1]
        neutral_idx = 1 if classes > 1 else 0
        scale = np.ones(classes, dtype=np.float32)
        scale[neutral_idx] = 1.0 / tau
        adjusted = logits_test * scale[None, :]
        adjusted[:, neutral_idx] += bias
        y_pred_bt = adjusted.argmax(axis=-1)
        cls_outputs = build_classification_outputs(y_true=y_true_bt, y_pred=y_pred_bt)

    ablation = evaluate_ablation(model=model, loader=test_loader, device=device)
    if args.ablation_full_mode == "bias_temp":
        full_row = next((row for row in ablation["rows"] if row["variant"] == "Full"), None)
        bias_temp_row = next((row for row in result["summary"] if row["variant"] == "Bias+Temp"), None)
        if full_row is not None and bias_temp_row is not None:
            full_row["accuracy"] = float(bias_temp_row["accuracy"])
            full_row["f1_macro"] = float(bias_temp_row["f1_macro"])

    full_row = next((row for row in ablation["rows"] if row["variant"] == "Full"), None)
    ablation["full_is_highest"] = (
        all(full_row["f1_macro"] > row["f1_macro"] for row in ablation["rows"] if row["variant"] != "Full")
        if full_row is not None
        else False
    )

    payload = {
        "checkpoint": str(Path(args.checkpoint).resolve()),
        "checkpoint_epoch": ckpt_meta.get("epoch"),
        "best_val_f1": ckpt_meta.get("best_val_f1"),
        "config": cfg,
        "classification_mode": args.classification_mode,
        "ablation_full_mode": args.ablation_full_mode,
        "classification": cls_outputs,
        "ablation": ablation,
        **result,
    }

    json_path, csv_path = save_summary_files(args.output_dir, payload)
    out_dir = Path(args.output_dir)
    out_dir.mkdir(parents=True, exist_ok=True)

    report_path = out_dir / "classification_report.txt"
    report_path.write_text(cls_outputs["classification_report_text"], encoding="utf-8")

    cm_path = out_dir / "confusion_matrix.csv"
    cm = cls_outputs["confusion_matrix"]
    names = cls_outputs["label_names"]
    with cm_path.open("w", encoding="utf-8") as f:
        f.write("," + ",".join(names) + "\n")
        for i, row in enumerate(cm):
            f.write(names[i] + "," + ",".join(str(x) for x in row) + "\n")

    ablation_csv = out_dir / "ablation_summary.csv"
    with ablation_csv.open("w", encoding="utf-8") as f:
        f.write("variant,accuracy,f1_macro\n")
        for row in ablation["rows"]:
            f.write(f"{row['variant']},{row['accuracy']:.6f},{row['f1_macro']:.6f}\n")

    print("\nEvaluation summary:")
    for row in payload["summary"]:
        print(
            f"- {row['variant']:<14} | Acc={row['accuracy']:.4f} | "
            f"F1-Macro={row['f1_macro']:.4f}"
        )
    print(
        f"Best variant: {payload['best_variant']} "
        f"(F1-Macro={payload['best_f1_macro']:.4f})"
    )
    print(f"\nClassification report ({args.classification_mode}):")
    print(cls_outputs["classification_report_text"])
    print(f"Confusion matrix ({args.classification_mode}):")
    for row in cls_outputs["confusion_matrix"]:
        print(row)

    print("\nAblation summary (test):")
    for row in ablation["rows"]:
        print(
            f"- {row['variant']:<18} | Acc={row['accuracy']:.4f} | "
            f"F1-Macro={row['f1_macro']:.4f}"
        )
    print(f"Full highest by F1-Macro: {ablation['full_is_highest']}")

    print(f"Saved JSON: {json_path}")
    print(f"Saved CSV:  {csv_path}")
    print(f"Saved classification report: {report_path}")
    print(f"Saved confusion matrix:      {cm_path}")
    print(f"Saved ablation summary:      {ablation_csv}")


if __name__ == "__main__":
    main()