#!/usr/bin/env python3
"""Evaluate MVSA-Single checkpoint with notebook-style variants and ablation."""

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.mvsa_single_pipeline import (
    MVSASingleLoader,
    apply_bias_temp_neutral,
    build_classification_outputs,
    create_dataloaders,
    evaluate_ablation,
    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 MVSA-Single checkpoint")
    parser.add_argument("--checkpoint", default="outputs/mvsa_single/clara_mvsa_single.pt")
    parser.add_argument("--data-root", default="data/MVSA-Single")
    parser.add_argument("--text-dir", default=None, help="Default: <data-root>/data")
    parser.add_argument("--label-file", default=None, help="Default: <data-root>/labelResultAll.txt")
    parser.add_argument("--output-dir", default="results/mvsa_single")

    parser.add_argument("--batch-size", type=int, default=None)
    parser.add_argument("--max-length", type=int, default=None)
    parser.add_argument("--num-workers", type=int, default=None)
    parser.add_argument("--train-ratio", type=float, default=None)
    parser.add_argument("--val-ratio", type=float, default=None)
    parser.add_argument("--preprocessing-mode", choices=["paper", "strict"], default=None)

    parser.add_argument("--top2-eps", type=float, default=0.03)
    parser.add_argument(
        "--classification-mode",
        choices=["raw", "bias_temp"],
        default="raw",
        help="Prediction source 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) / "data")
    label_file = args.label_file or str(Path(args.data_root) / "labelResultAll.txt")

    device = resolve_device(args.device)
    model, cfg, ckpt_meta = load_checkpoint(args.checkpoint, device)

    cfg["text_dir"] = text_dir
    cfg["label_file"] = label_file

    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
    if args.train_ratio is not None:
        cfg["train_ratio"] = args.train_ratio
    if args.val_ratio is not None:
        cfg["val_ratio"] = args.val_ratio
    if args.preprocessing_mode is not None:
        cfg["preprocessing_mode"] = args.preprocessing_mode

    loader = MVSASingleLoader(cfg["text_dir"], cfg["label_file"])
    loader.load(
        preprocessing_mode=str(cfg.get("preprocessing_mode", "paper")),
        require_unanimous=bool(cfg.get("require_unanimous", True)),
        require_cross_agree=bool(cfg.get("require_cross_agree", True)),
    )
    train_samples, val_samples, test_samples = loader.split(
        train_ratio=float(cfg.get("train_ratio", 0.7)),
        val_ratio=float(cfg.get("val_ratio", 0.15)),
        seed=int(cfg.get("seed", 42)),
        paper_811=bool(str(cfg.get("preprocessing_mode", "paper")).lower() == "paper"),
    )

    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"])

    pin_memory = bool(cfg.get("pin_memory", True) 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=int(cfg["batch_size"]),
        max_length=int(cfg["max_length"]),
        num_workers=int(cfg["num_workers"]),
        pin_memory=pin_memory,
        persistent_workers=bool(cfg.get("persistent_workers", True)),
        prefetch_factor=int(cfg.get("prefetch_factor", 2)),
        use_mixup_negative=False,
        mixup_alpha=float(cfg.get("mixup_alpha", 0.4)),
        negative_class_boost=float(cfg.get("negative_class_boost", 12.0)),
        min_ratio_negative=float(cfg.get("min_ratio_negative", 0.30)),
        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"])
        adjusted = apply_bias_temp_neutral(logits_test, bias=bias, tau=tau)
        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_weighted"] = float(bias_temp_row["f1_weighted"])

    full_row = next((row for row in ablation["rows"] if row["variant"] == "Full"), None)
    ablation["full_is_highest"] = (
        all(full_row["f1_weighted"] >= row["f1_weighted"] 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_weighted": ckpt_meta.get("best_val_f1_weighted"),
        "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 idx, row in enumerate(cm):
            f.write(names[idx] + "," + ",".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_weighted\n")
        for row in ablation["rows"]:
            f.write(f"{row['variant']},{row['accuracy']:.6f},{row['f1_weighted']:.6f}\n")

    print("\nEvaluation summary:")
    for row in payload["summary"]:
        print(
            f"- {row['variant']:<14} | Acc={row['accuracy']:.4f} | "
            f"F1-Weighted={row['f1_weighted']:.4f}"
        )
    print(
        f"Best variant: {payload['best_variant']} "
        f"(F1-Weighted={payload['best_f1_weighted']:.4f})"
    )

    print(f"\nClassification report ({args.classification_mode}):")
    print(cls_outputs["classification_report_text"])

    print("Ablation summary (test):")
    for row in ablation["rows"]:
        print(
            f"- {row['variant']:<18} | Acc={row['accuracy']:.4f} | "
            f"F1-Weighted={row['f1_weighted']:.4f}"
        )
    print(f"Full highest by F1-Weighted: {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()