"""
evaluate.py – Evaluation, quantitative metrics, and error analysis.

Usage:
    python src/evaluate.py --model mobilenet_v2 --split test
    python src/evaluate.py --model simple_cnn   --split val
"""

import argparse
import logging
import sys
from pathlib import Path

import matplotlib.pyplot as plt
import numpy as np
import seaborn as sns
import torch
import yaml
from sklearn.metrics import (
    accuracy_score,
    classification_report,
    confusion_matrix,
    f1_score,
    precision_score,
    recall_score,
    roc_auc_score,
    roc_curve,
)
from tqdm import tqdm

sys.path.insert(0, str(Path(__file__).parent))

from dataset import build_dataloaders, IMAGENET_MEAN, IMAGENET_STD, LABEL_NAMES
from model import build_model

logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s | %(levelname)-8s | %(name)s | %(message)s",
    datefmt="%Y-%m-%d %H:%M:%S",
)
logger = logging.getLogger(__name__)


# ─── Inference pass ───────────────────────────────────────────────────────────

@torch.no_grad()
def collect_predictions(model, loader, device):
    """Run the model on a DataLoader and collect all preds + probs + raw images."""
    model.eval()
    all_labels, all_preds, all_probs, all_images = [], [], [], []

    for images, labels in tqdm(loader, desc="Evaluating"):
        images = images.to(device)
        logits = model(images)
        probs  = torch.softmax(logits, dim=1)
        preds  = probs.argmax(dim=1)
        all_images.extend(images.cpu())
        all_preds.extend(preds.cpu().numpy())
        all_probs.extend(probs.cpu().numpy())
        all_labels.extend(labels.numpy())

    return (
        np.array(all_labels),
        np.array(all_preds),
        np.array(all_probs),
        all_images,
    )


# ─── Plot helpers ─────────────────────────────────────────────────────────────

def _denormalize(tensor: torch.Tensor) -> np.ndarray:
    """Reverse ImageNet normalisation for visualisation."""
    mean = np.array(IMAGENET_MEAN)
    std  = np.array(IMAGENET_STD)
    img  = tensor.permute(1, 2, 0).numpy()
    return (img * std + mean).clip(0.0, 1.0)


def plot_confusion_matrix(labels, preds, output_dir: Path, model_name: str) -> None:
    cm = confusion_matrix(labels, preds)
    fig, ax = plt.subplots(figsize=(6, 5))
    sns.heatmap(
        cm, annot=True, fmt="d", cmap="Blues",
        xticklabels=list(LABEL_NAMES.values()),
        yticklabels=list(LABEL_NAMES.values()),
        ax=ax,
    )
    ax.set_xlabel("Predicted", fontsize=12)
    ax.set_ylabel("True", fontsize=12)
    ax.set_title(f"Confusion Matrix – {model_name}", fontsize=13, fontweight="bold")
    plt.tight_layout()
    out = output_dir / f"confusion_matrix_{model_name}.png"
    fig.savefig(out, dpi=150)
    plt.close(fig)
    logger.info(f"Saved confusion matrix → {out}")


def plot_misclassified(
    labels, preds, images, output_dir: Path, model_name: str, n: int = 12
) -> None:
    wrong_idx = np.where(labels != preds)[0]
    if len(wrong_idx) == 0:
        logger.info("No misclassified samples – perfect predictions!")
        return

    wrong_idx = wrong_idx[:n]
    cols = 4
    rows = (len(wrong_idx) + cols - 1) // cols
    fig, axes = plt.subplots(rows, cols, figsize=(cols * 3, rows * 3.2))
    axes = np.array(axes).flatten()

    for i, idx in enumerate(wrong_idx):
        img = _denormalize(images[idx])
        axes[i].imshow(img, cmap="gray" if img.std() < 0.05 else None)
        true_lbl = LABEL_NAMES[labels[idx]]
        pred_lbl = LABEL_NAMES[preds[idx]]
        axes[i].set_title(f"True: {true_lbl}\nPred: {pred_lbl}", fontsize=8, color="crimson")
        axes[i].axis("off")
    for j in range(i + 1, len(axes)):
        axes[j].axis("off")

    fig.suptitle(
        f"Misclassified Samples – {model_name}\n"
        f"({len(wrong_idx)} of {len(labels)} shown)",
        fontsize=12, fontweight="bold", y=1.01,
    )
    plt.tight_layout()
    out = output_dir / f"misclassified_{model_name}.png"
    fig.savefig(out, dpi=150, bbox_inches="tight")
    plt.close(fig)
    logger.info(f"Saved misclassified grid → {out}")


def plot_roc_curve(labels, probs, output_dir: Path, model_name: str) -> float:
    auc = roc_auc_score(labels, probs[:, 1])
    fpr, tpr, _ = roc_curve(labels, probs[:, 1])
    fig, ax = plt.subplots(figsize=(6, 5))
    ax.plot(fpr, tpr, label=f"AUC = {auc:.3f}", color="#1976D2", lw=2)
    ax.fill_between(fpr, tpr, alpha=0.08, color="#1976D2")
    ax.plot([0, 1], [0, 1], "k--", lw=1, label="Random (AUC=0.5)")
    ax.set_xlabel("False Positive Rate", fontsize=11)
    ax.set_ylabel("True Positive Rate", fontsize=11)
    ax.set_title(f"ROC Curve – {model_name}", fontsize=13, fontweight="bold")
    ax.legend(fontsize=11)
    ax.spines[["top", "right"]].set_visible(False)
    plt.tight_layout()
    out = output_dir / f"roc_curve_{model_name}.png"
    fig.savefig(out, dpi=150)
    plt.close(fig)
    logger.info(f"Saved ROC curve → {out}")
    return auc


def plot_training_curves(history: dict, output_dir: Path, model_name: str) -> None:
    """Optionally called after training to visualise loss/acc curves."""
    if not history:
        return
    epochs = range(1, len(history["train_loss"]) + 1)
    fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(12, 4))

    ax1.plot(epochs, history["train_loss"], label="Train", color="#E53935")
    ax1.plot(epochs, history["val_loss"],   label="Val",   color="#43A047")
    ax1.set_title("Loss"); ax1.set_xlabel("Epoch"); ax1.legend()
    ax1.spines[["top", "right"]].set_visible(False)

    ax2.plot(epochs, history["train_acc"], label="Train", color="#E53935")
    ax2.plot(epochs, history["val_acc"],   label="Val",   color="#43A047")
    ax2.set_title("Accuracy"); ax2.set_xlabel("Epoch"); ax2.legend()
    ax2.spines[["top", "right"]].set_visible(False)

    fig.suptitle(f"Training Curves – {model_name}", fontsize=13, fontweight="bold")
    plt.tight_layout()
    out = output_dir / f"training_curves_{model_name}.png"
    fig.savefig(out, dpi=150)
    plt.close(fig)
    logger.info(f"Saved training curves → {out}")


# ─── Main evaluation ──────────────────────────────────────────────────────────

def evaluate(cfg: dict, model_name_override: str = None, split: str = "test") -> dict:
    """
    Load checkpoint, run evaluation on the chosen split, print metrics, save plots.

    Args:
        cfg:                  Config dict.
        model_name_override:  Override model name from config.
        split:                'val' or 'test'.

    Returns:
        dict of metric values.
    """
    if model_name_override:
        cfg["model"]["name"] = model_name_override

    device     = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    model_name = cfg["model"]["name"]

    ckpt_path = Path(cfg["training"]["checkpoint_dir"]) / f"best_{model_name}.pth"
    if not ckpt_path.exists():
        raise FileNotFoundError(
            f"Checkpoint not found: {ckpt_path}\n"
            "Run `python src/train.py --model {model_name}` first."
        )

    logger.info(f"Loading checkpoint: {ckpt_path}")
    ckpt = torch.load(ckpt_path, map_location=device)
    model = build_model(cfg).to(device)
    model.load_state_dict(ckpt["model_state"])
    logger.info(f"Loaded  epoch={ckpt['epoch']}, best_val_acc={ckpt['val_acc']:.4f}")

    loaders = build_dataloaders(cfg)
    labels, preds, probs, images = collect_predictions(model, loaders[split], device)

    # ── Metrics ──────────────────────────────────────────────────────────────
    acc  = accuracy_score(labels, preds)
    prec = precision_score(labels, preds, average="binary", zero_division=0)
    rec  = recall_score(labels, preds, average="binary", zero_division=0)
    f1   = f1_score(labels, preds, average="binary", zero_division=0)
    auc  = roc_auc_score(labels, probs[:, 1])

    sep = "=" * 58
    print(f"\n{sep}")
    print(f"  Evaluation  |  model={model_name}  |  split={split}")
    print(sep)
    print(f"  Accuracy    : {acc:.4f}")
    print(f"  Precision   : {prec:.4f}  (positive class = PNEUMONIA)")
    print(f"  Recall      : {rec:.4f}  (sensitivity)")
    print(f"  F1-Score    : {f1:.4f}")
    print(f"  AUC-ROC     : {auc:.4f}")
    print(sep)
    print(classification_report(labels, preds, target_names=list(LABEL_NAMES.values())))

    # ── Plots ─────────────────────────────────────────────────────────────────
    output_dir = Path(cfg["evaluation"]["output_dir"])
    output_dir.mkdir(parents=True, exist_ok=True)
    plot_confusion_matrix(labels, preds, output_dir, model_name)
    plot_misclassified(labels, preds, images, output_dir, model_name)
    plot_roc_curve(labels, probs, output_dir, model_name)

    logger.info(f"All evaluation outputs saved to: {output_dir}/")
    return {"acc": acc, "precision": prec, "recall": rec, "f1": f1, "auc": auc}


# ─── CLI ──────────────────────────────────────────────────────────────────────

if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="Evaluate Chest X-Ray classifier")
    parser.add_argument("--config", default="configs/config.yaml")
    parser.add_argument("--model", choices=["simple_cnn", "mobilenet_v2"])
    parser.add_argument("--split", default="test", choices=["val", "test"])
    args = parser.parse_args()

    with open(args.config) as f:
        cfg = yaml.safe_load(f)

    evaluate(cfg, model_name_override=args.model, split=args.split)