#!/usr/bin/env python3
"""Compute temperature-scaled ECE per ablation variant on MVSA-Multiple."""

from __future__ import annotations

import argparse
import json
from pathlib import Path
from typing import Dict, List, Tuple

import matplotlib

matplotlib.use("Agg")
import matplotlib.pyplot as plt
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_multiple_pipeline import (
    MVSALoader,
    create_dataloaders,
    gather_logits_variant,
    load_checkpoint,
    resolve_device,
)


VARIANT_MAP: List[Tuple[str, str]] = [
    ("Full", "full"),
    ("w/o Verification", "w/o_verification"),
    ("w/o Feedback", "w/o_feedback"),
    ("w/o Co-Attention", "w/o_coattn"),
    ("Text-only", "text_only"),
    ("Vision-only", "vision_only"),
    ("w/o Text", "w/o_text"),
    ("w/o Image", "w/o_image"),
]


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description="Compute ECE for MVSA-Multiple")
    parser.add_argument("--checkpoint", default="outputs/mvsa_multiple/clara_mvsa_multiple.pt")
    parser.add_argument("--data-root", default="data/MVSA-Multiple")
    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_multiple")

    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("--seed", type=int, default=None, help="Override split seed from checkpoint config")
    parser.add_argument(
        "--allow-seed-mismatch",
        action="store_true",
        help="Allow using a split seed different from checkpoint training seed (can cause leakage-like overlap).",
    )
    parser.add_argument("--preprocessing-mode", choices=["paper", "strict"], default=None)
    parser.add_argument("--disable-paper-exact-counts", action="store_true")
    parser.add_argument("--n-bins", type=int, default=15)
    parser.add_argument(
        "--lowest-tol",
        type=float,
        default=1e-3,
        help="Tolerance when deciding whether Full has the lowest ECE.",
    )
    parser.add_argument("--device", default="auto", help="auto|cuda|cpu")
    return parser.parse_args()


def softmax_np(logits: np.ndarray) -> np.ndarray:
    shifted = logits - logits.max(axis=1, keepdims=True)
    exps = np.exp(shifted)
    return exps / exps.sum(axis=1, keepdims=True)


def compute_ece(probs: np.ndarray, y_true: np.ndarray, n_bins: int = 15):
    confidence = probs.max(axis=1)
    y_pred = probs.argmax(axis=1)
    correct = (y_pred == y_true).astype(np.float64)

    bins = np.linspace(0.0, 1.0, n_bins + 1)
    bin_conf = np.zeros(n_bins, dtype=np.float64)
    bin_acc = np.zeros(n_bins, dtype=np.float64)
    bin_count = np.zeros(n_bins, dtype=np.int64)

    for idx in range(n_bins):
        lo, hi = bins[idx], bins[idx + 1]
        if idx == 0:
            mask = (confidence >= lo) & (confidence <= hi)
        else:
            mask = (confidence > lo) & (confidence <= hi)

        count = int(mask.sum())
        if count > 0:
            bin_conf[idx] = float(confidence[mask].mean())
            bin_acc[idx] = float(correct[mask].mean())
            bin_count[idx] = count

    total = max(1, len(y_true))
    ece = float(np.sum((bin_count / total) * np.abs(bin_conf - bin_acc)))
    return ece, bin_conf, bin_acc, bin_count


def nll_with_temperature(logits: np.ndarray, y_true: np.ndarray, temperature: float) -> float:
    scaled = logits / max(1e-6, float(temperature))
    shifted = scaled - scaled.max(axis=1, keepdims=True)
    log_probs = shifted - np.log(np.exp(shifted).sum(axis=1, keepdims=True))
    return float(-log_probs[np.arange(len(y_true)), y_true].mean())


def temperature_scale(logits: np.ndarray, y_true: np.ndarray) -> float:
    log_temp_grid = np.linspace(-3.0, 3.0, 601)
    temp_grid = np.exp(log_temp_grid)

    best_temp = 1.0
    best_nll = float("inf")
    for temp in temp_grid:
        nll = nll_with_temperature(logits, y_true, temp)
        if nll < best_nll:
            best_nll = nll
            best_temp = float(temp)

    return best_temp


def plot_reliability(
    results: Dict[str, Dict[str, np.ndarray]],
    variant_map: List[Tuple[str, str]],
    n_bins: int,
    out_path: Path,
) -> None:
    colors = {
        "Full": "#2c6e9e",
        "w/o Verification": "#e07b39",
        "w/o Feedback": "#3aaa5e",
        "w/o Co-Attention": "#9b59b6",
        "Text-only": "#c0392b",
        "Vision-only": "#7f8c8d",
        "w/o Text": "#8e44ad",
        "w/o Image": "#16a085",
    }

    n_cols = 3
    n_rows = int(np.ceil(len(variant_map) / n_cols))
    fig, axes = plt.subplots(n_rows, n_cols, figsize=(12.2, 3.8 * n_rows))
    axes = np.atleast_1d(axes).flatten()

    bin_edges = np.linspace(0.0, 1.0, n_bins + 1)
    bin_centers = 0.5 * (bin_edges[:-1] + bin_edges[1:])
    width = bin_edges[1] - bin_edges[0]

    for ax, (display, _) in zip(axes, variant_map):
        result = results[display]
        color = colors.get(display, "#1f77b4")

        ax.plot([0, 1], [0, 1], "k--", lw=1.1, label="Perfect calibration")

        for idx in range(n_bins):
            if result["bin_count"][idx] > 0:
                low = min(result["bin_conf"][idx], result["bin_acc"][idx])
                high = max(result["bin_conf"][idx], result["bin_acc"][idx])
                ax.bar(
                    bin_centers[idx],
                    high - low,
                    bottom=low,
                    width=width * 0.90,
                    color="tomato",
                    alpha=0.35,
                )

        mask = result["bin_count"] > 0
        ax.bar(
            bin_centers[mask],
            result["bin_acc"][mask],
            width=width * 0.90,
            color=color,
            alpha=0.85,
            label="Accuracy",
        )

        ax.set_xlim(0.0, 1.0)
        ax.set_ylim(0.0, 1.0)
        ax.set_xticks([0.0, 0.25, 0.50, 0.75, 1.00])
        ax.set_yticks([0.0, 0.25, 0.50, 0.75, 1.00])
        ax.set_xlabel("Confidence", fontsize=9)
        ax.set_ylabel("Accuracy", fontsize=9)
        ax.set_title(
            f"{display}\nECE = {result['ece']:.4f}  (T={result['temperature']:.3f})",
            fontsize=10,
            fontweight="bold",
        )
        ax.legend(fontsize=7, loc="upper left")

    for ax in axes[len(variant_map) :]:
        ax.axis("off")

    fig.suptitle("Reliability Diagrams after Temperature Scaling — MVSA-Multiple Ablation", fontsize=13)
    fig.tight_layout()
    fig.savefig(out_path, dpi=150, bbox_inches="tight")
    plt.close(fig)


def plot_ece_bar(results: Dict[str, Dict[str, np.ndarray]], variant_map: List[Tuple[str, str]], out_path: Path) -> None:
    names = [name for name, _ in variant_map]
    eces = [float(results[name]["ece"]) for name in names]

    colors = {
        "Full": "#2c6e9e",
        "w/o Verification": "#e07b39",
        "w/o Feedback": "#3aaa5e",
        "w/o Co-Attention": "#9b59b6",
        "Text-only": "#c0392b",
        "Vision-only": "#7f8c8d",
        "w/o Text": "#8e44ad",
        "w/o Image": "#16a085",
    }

    fig, ax = plt.subplots(figsize=(8.4, 4.8))
    bars = ax.bar(names, eces, color=[colors.get(name, "#1f77b4") for name in names], width=0.58)

    for bar, val in zip(bars, eces):
        ax.text(
            bar.get_x() + bar.get_width() / 2,
            bar.get_height() + 0.001,
            f"{val:.4f}",
            ha="center",
            va="bottom",
            fontsize=9.5,
            fontweight="bold",
        )

    ax.set_ylabel("ECE ↓", fontsize=11)
    ax.set_title("ECE after Temperature Scaling — MVSA-Multiple Ablation", fontsize=12)
    ax.set_ylim(0.0, max(eces) * 1.22 if eces else 1.0)
    ax.set_xticks(range(len(names)))
    ax.set_xticklabels(names, rotation=20, ha="right")
    ax.spines["top"].set_visible(False)
    ax.spines["right"].set_visible(False)

    full_ece = results["Full"]["ece"]
    ax.axhline(full_ece, color=colors["Full"], ls="--", lw=1.3, alpha=0.7, label=f"Full ECE = {full_ece:.4f}")
    ax.legend(fontsize=9)

    fig.tight_layout()
    fig.savefig(out_path, dpi=150, bbox_inches="tight")
    plt.close(fig)


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, _ = 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.seed is not None:
        ckpt_seed = cfg.get("seed")
        if ckpt_seed is not None and int(args.seed) != int(ckpt_seed) and not args.allow_seed_mismatch:
            raise ValueError(
                "Seed mismatch detected: "
                f"checkpoint seed={ckpt_seed}, eval seed={args.seed}. "
                "Use --allow-seed-mismatch to override explicitly."
            )
        if ckpt_seed is not None and int(args.seed) != int(ckpt_seed):
            print(
                "WARNING: computing ECE with different split seed "
                f"(checkpoint={ckpt_seed}, eval={args.seed})."
            )
        cfg["seed"] = int(args.seed)
    if args.preprocessing_mode is not None:
        cfg["preprocessing_mode"] = args.preprocessing_mode
    if args.disable_paper_exact_counts:
        cfg["paper_exact_counts"] = False

    loader = MVSALoader(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)),
        paper_exact_counts=bool(cfg.get("paper_exact_counts", False)),
    )
    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"),
    )

    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,
    )

    out_dir = Path(args.output_dir)
    out_dir.mkdir(parents=True, exist_ok=True)

    results: Dict[str, Dict[str, np.ndarray]] = {}
    for display_name, key in VARIANT_MAP:
        print(f"[{key}] calibrating temperature on val...")
        val_logits, val_labels = gather_logits_variant(model, val_loader, device, variant=key)
        temperature = temperature_scale(val_logits, val_labels)

        print(f"[{key}] evaluating ECE on test...")
        test_logits, test_labels = gather_logits_variant(model, test_loader, device, variant=key)
        test_probs = softmax_np(test_logits / max(1e-6, temperature))

        ece, bin_conf, bin_acc, bin_count = compute_ece(test_probs, test_labels, n_bins=args.n_bins)

        results[display_name] = {
            "ece": float(ece),
            "temperature": float(temperature),
            "bin_conf": bin_conf,
            "bin_acc": bin_acc,
            "bin_count": bin_count,
        }
        print(f"  T={temperature:.4f} | ECE={ece:.4f}")

    full_ece = results["Full"]["ece"]
    tol = max(0.0, float(args.lowest_tol))
    full_is_lowest = all(
        full_ece <= (results[name]["ece"] + tol)
        for name, _ in VARIANT_MAP
        if name != "Full"
    )

    csv_path = out_dir / "ece_summary.csv"
    with csv_path.open("w", encoding="utf-8") as f:
        f.write("variant,ece,temperature\n")
        for name, _ in VARIANT_MAP:
            f.write(f"{name},{results[name]['ece']:.6f},{results[name]['temperature']:.6f}\n")

    json_payload = {
        "checkpoint": str(Path(args.checkpoint).resolve()),
        "n_bins": int(args.n_bins),
        "lowest_tolerance": float(tol),
        "full_is_lowest": bool(full_is_lowest),
        "rows": [
            {
                "variant": name,
                "ece": float(results[name]["ece"]),
                "temperature": float(results[name]["temperature"]),
            }
            for name, _ in VARIANT_MAP
        ],
    }
    json_path = out_dir / "ece_summary.json"
    json_path.write_text(json.dumps(json_payload, indent=2), encoding="utf-8")

    diag_path = out_dir / "figure_ece_reliability.png"
    bar_path = out_dir / "figure_ece_bar.png"
    plot_reliability(results, VARIANT_MAP, args.n_bins, diag_path)
    plot_ece_bar(results, VARIANT_MAP, bar_path)

    print("\nECE summary:")
    for name, _ in VARIANT_MAP:
        print(f"- {name:<20} | T={results[name]['temperature']:.4f} | ECE={results[name]['ece']:.4f}")
    print(f"Full lowest ECE: {full_is_lowest}")
    print(f"Saved CSV: {csv_path}")
    print(f"Saved JSON: {json_path}")
    print(f"Saved reliability plot: {diag_path}")
    print(f"Saved ECE bar plot:     {bar_path}")


if __name__ == "__main__":
    main()