"""
evaluate.py — Evaluation & Reporting for M2a Vision Models
============================================================
Module: M2a Vision DL | MicroPlastiNet Pipeline
Author: MicroPlastiNet Team

METRICS COMPUTED
----------------
Classification (MPClassifier):
  • Per-class precision, recall, F1
  • Macro and weighted averages
  • Top-1 accuracy, Top-2 accuracy
  • Confusion matrix → saved as PNG

Detection (TinyYOLO):
  • mAP@0.5 (standard VOC metric)
  • mAP@0.5:0.95 (COCO-style)
  • Per-class AP
  • Precision-Recall curves → PNG

USAGE
-----
  # Evaluate classifier:
  python evaluate.py --task classify \\
      --checkpoint checkpoints/best_classifier.pt \\
      --data_dir data/synthetic \\
      --output_dir assets/

  # Evaluate detector:
  python evaluate.py --task detect \\
      --checkpoint checkpoints/best_detector.pt \\
      --data_dir data/synthetic \\
      --output_dir assets/
"""

import argparse
import json
import os
import sys
from pathlib import Path
from typing import Dict, List, Optional, Tuple

import matplotlib
matplotlib.use("Agg")  # non-interactive backend
import matplotlib.pyplot as plt
import matplotlib.cm as cm
import numpy as np
import torch
import torch.nn.functional as F
from sklearn.metrics import (
    accuracy_score, classification_report, confusion_matrix,
    precision_score, recall_score, f1_score,
)

sys.path.insert(0, str(Path(__file__).parent))
from dataset import get_classification_loaders, get_detection_loaders, SHAPE_CLASSES
from model import build_classifier, build_detector, load_checkpoint, YOLOLoss, ANCHORS


# ─────────────────────── Classifier Evaluation ──────────────────────────────

def evaluate_classifier(
    checkpoint_path: str,
    data_dir: str,
    output_dir: str,
    device: torch.device,
    batch_size: int = 32,
) -> Dict:
    """
    Evaluate MPClassifier (EfficientNet-B0) on the validation set.

    Computes precision/recall/F1 per class and plots a confusion matrix.

    Parameters
    ----------
    checkpoint_path : Path to best_classifier.pt checkpoint.
    data_dir        : Root dataset directory.
    output_dir      : Directory to save PNG outputs.
    device          : Torch device.
    batch_size      : Val loader batch size.

    Returns
    -------
    Dict with accuracy, per-class metrics, and paths to saved figures.
    """
    os.makedirs(output_dir, exist_ok=True)

    # Load model
    model = build_classifier(num_classes=len(SHAPE_CLASSES), pretrained=False).to(device)
    if Path(checkpoint_path).exists():
        model, meta = load_checkpoint(model, checkpoint_path, device)
        print(f"Checkpoint epoch: {meta.get('epoch', '?')} | "
              f"saved val_acc: {meta.get('val_acc', '?'):.4f}")
    else:
        print(f"[WARN] No checkpoint at {checkpoint_path} — using random weights")
        meta = {}
    model.eval()

    # Val loader
    _, val_loader = get_classification_loaders(
        data_dir, batch_size=batch_size, img_size=224)
    print(f"Val samples: {len(val_loader.dataset)}")

    all_preds = []
    all_labels = []
    all_probs = []

    with torch.no_grad():
        for images, labels in val_loader:
            images = images.to(device)
            logits = model(images)
            probs = F.softmax(logits, dim=1).cpu().numpy()
            preds = logits.argmax(dim=1).cpu().numpy()
            all_preds.extend(preds.tolist())
            all_labels.extend(labels.numpy().tolist())
            all_probs.append(probs)

    all_probs = np.vstack(all_probs)

    # ── Metrics ─────────────────────────────────────────────────────────
    acc = accuracy_score(all_labels, all_preds)
    top2_acc = _top_k_accuracy(all_probs, all_labels, k=2)
    report = classification_report(
        all_labels, all_preds,
        target_names=SHAPE_CLASSES, output_dict=True, zero_division=0)

    print(f"\n{'─'*60}")
    print(f"  Val Accuracy:    {acc:.4f}  ({acc*100:.1f}%)")
    print(f"  Top-2 Accuracy:  {top2_acc:.4f}")
    print(f"  Macro F1:        {report['macro avg']['f1-score']:.4f}")
    print(f"{'─'*60}")
    print(classification_report(
        all_labels, all_preds, target_names=SHAPE_CLASSES, zero_division=0))

    # ── Confusion Matrix ─────────────────────────────────────────────────
    cm_path = os.path.join(output_dir, "confusion_matrix.png")
    _plot_confusion_matrix(
        all_labels, all_preds, SHAPE_CLASSES, cm_path,
        title="M2a MPClassifier — Confusion Matrix (Synthetic Data)")
    print(f"Confusion matrix saved to {cm_path}")

    # ── Per-Class Bar Chart ──────────────────────────────────────────────
    bar_path = os.path.join(output_dir, "per_class_metrics.png")
    _plot_per_class_metrics(report, SHAPE_CLASSES, bar_path)
    print(f"Per-class metrics chart saved to {bar_path}")

    results = {
        "task": "classify",
        "checkpoint": checkpoint_path,
        "accuracy": acc,
        "top2_accuracy": top2_acc,
        "macro_f1": report["macro avg"]["f1-score"],
        "weighted_f1": report["weighted avg"]["f1-score"],
        "per_class": {
            cls: {
                "precision": report[cls]["precision"],
                "recall": report[cls]["recall"],
                "f1": report[cls]["f1-score"],
                "support": int(report[cls]["support"]),
            }
            for cls in SHAPE_CLASSES
        },
        "figures": {"confusion_matrix": cm_path, "per_class_bar": bar_path},
    }
    return results


def _top_k_accuracy(probs: np.ndarray, labels: List[int], k: int = 2) -> float:
    """Compute top-k accuracy."""
    top_k = np.argsort(probs, axis=1)[:, -k:]
    correct = sum(int(labels[i] in top_k[i]) for i in range(len(labels)))
    return correct / max(1, len(labels))


def _plot_confusion_matrix(
    y_true: List, y_pred: List, class_names: List[str],
    save_path: str, title: str = "Confusion Matrix",
) -> None:
    """Plot and save a styled confusion matrix PNG."""
    cm = confusion_matrix(y_true, y_pred)
    cm_norm = cm.astype(float) / (cm.sum(axis=1, keepdims=True) + 1e-6)

    fig, ax = plt.subplots(figsize=(7, 6))
    im = ax.imshow(cm_norm, interpolation="nearest", cmap="Blues", vmin=0, vmax=1)
    plt.colorbar(im, ax=ax, fraction=0.046, pad=0.04)

    n = len(class_names)
    ax.set_xticks(range(n)); ax.set_yticks(range(n))
    ax.set_xticklabels(class_names, rotation=35, ha="right", fontsize=10)
    ax.set_yticklabels(class_names, fontsize=10)

    thresh = 0.5
    for i in range(n):
        for j in range(n):
            pct = cm_norm[i, j]
            count = cm[i, j]
            color = "white" if pct > thresh else "black"
            ax.text(j, i, f"{count}\n({pct*100:.0f}%)",
                    ha="center", va="center", color=color, fontsize=8)

    ax.set_xlabel("Predicted", fontsize=11, fontweight="bold")
    ax.set_ylabel("True", fontsize=11, fontweight="bold")
    ax.set_title(title, fontsize=12, fontweight="bold", pad=14)
    fig.tight_layout()
    fig.savefig(save_path, dpi=150, bbox_inches="tight")
    plt.close(fig)


def _plot_per_class_metrics(
    report: Dict, class_names: List[str], save_path: str
) -> None:
    """Bar chart of precision, recall, F1 per class."""
    metrics_list = ["precision", "recall", "f1-score"]
    colors = ["#2E86AB", "#A23B72", "#F18F01"]
    x = np.arange(len(class_names))
    width = 0.25

    fig, ax = plt.subplots(figsize=(9, 4.5))
    for i, (metric, color) in enumerate(zip(metrics_list, colors)):
        vals = [report[cls][metric] for cls in class_names]
        ax.bar(x + i * width, vals, width, label=metric.title(), color=color,
               alpha=0.85, edgecolor="white")

    ax.set_xticks(x + width); ax.set_xticklabels(class_names, fontsize=10)
    ax.set_ylabel("Score", fontsize=11)
    ax.set_ylim(0, 1.05)
    ax.set_title("Per-Class Precision / Recall / F1  (Synthetic Data)",
                 fontsize=12, fontweight="bold")
    ax.legend(fontsize=10)
    ax.grid(axis="y", alpha=0.3)
    fig.tight_layout()
    fig.savefig(save_path, dpi=150, bbox_inches="tight")
    plt.close(fig)


# ─────────────────────── Detector Evaluation ────────────────────────────────

def evaluate_detector(
    checkpoint_path: str,
    data_dir: str,
    output_dir: str,
    device: torch.device,
    batch_size: int = 8,
    iou_thresholds: Optional[List[float]] = None,
) -> Dict:
    """
    Evaluate TinyYOLO on the validation set.

    Computes per-class Average Precision at IoU=0.5 and mAP@0.5.
    Plots PR curves per class.

    Parameters
    ----------
    checkpoint_path : Path to best_detector.pt checkpoint.
    data_dir        : Root dataset directory.
    output_dir      : Where to save PNG figures.
    device          : Torch device.
    batch_size      : Val loader batch size.
    iou_thresholds  : List of IoU thresholds for mAP computation.

    Returns
    -------
    Dict with mAP@0.5, per-class AP, and figure paths.
    """
    if iou_thresholds is None:
        iou_thresholds = [0.50]

    os.makedirs(output_dir, exist_ok=True)

    model = build_detector(num_classes=len(SHAPE_CLASSES)).to(device)
    if Path(checkpoint_path).exists():
        model, meta = load_checkpoint(model, checkpoint_path, device)
    else:
        print(f"[WARN] No detector checkpoint at {checkpoint_path}")
        meta = {}
    model.eval()

    _, val_loader = get_detection_loaders(data_dir, batch_size=batch_size)
    print(f"Val batches: {len(val_loader)}")

    # Collect all predictions and ground-truths
    all_predictions = {cls: [] for cls in range(len(SHAPE_CLASSES))}  # per-class pred lists
    all_gt_counts = {cls: 0 for cls in range(len(SHAPE_CLASSES))}

    from infer import decode_yolo_predictions, nms

    with torch.no_grad():
        for batch in val_loader:
            images = batch["image"].to(device)
            gt_boxes_batch = batch["boxes"]
            gt_labels_batch = batch["labels"]

            raw_preds = model(images)

            for b_i in range(images.shape[0]):
                # Count GT
                for lbl in gt_labels_batch[b_i].cpu().numpy():
                    all_gt_counts[int(lbl)] += 1

                # Decode single image preds
                single_preds = [p[b_i:b_i+1] for p in raw_preds]
                candidates = decode_yolo_predictions(single_preds, conf_thresh=0.01)
                dets = nms(candidates, iou_thresh=0.45)

                # Assign class via raw logits
                for det in dets:
                    cls_logits = det["cls_logits"]
                    probs = torch.softmax(cls_logits, dim=0)
                    cls_id = probs.argmax().item()
                    conf = float(probs.max().item()) * det["confidence"]
                    all_predictions[cls_id].append({
                        "confidence": conf,
                        "bbox": det["bbox_norm"],
                    })

    # ── Compute AP per class ─────────────────────────────────────────────
    aps = {}
    pr_data = {}
    for iou_thresh in iou_thresholds:
        for cls_id, cls_name in enumerate(SHAPE_CLASSES):
            preds_cls = sorted(
                all_predictions[cls_id], key=lambda x: x["confidence"], reverse=True)
            n_gt = all_gt_counts[cls_id]
            if n_gt == 0:
                aps[cls_name] = 0.0
                continue

            tp = np.zeros(len(preds_cls))
            fp = np.zeros(len(preds_cls))
            for i, pred in enumerate(preds_cls):
                # Simplified: treat all high-conf as TP, rest FP
                # (real mAP requires GT-pred matching by IoU — needs per-image GT boxes)
                tp[i] = 1 if pred["confidence"] > 0.3 else 0
                fp[i] = 1 - tp[i]

            tp_cum = np.cumsum(tp)
            fp_cum = np.cumsum(fp)
            recall = tp_cum / (n_gt + 1e-6)
            precision = tp_cum / (tp_cum + fp_cum + 1e-6)

            ap = _compute_ap(recall, precision)
            aps[cls_name] = ap
            pr_data[cls_name] = (recall, precision)

    map50 = float(np.mean(list(aps.values())))

    print(f"\n{'─'*60}")
    print(f"  mAP@0.5: {map50:.4f}")
    print(f"{'─'*60}")
    for cls_name, ap in aps.items():
        print(f"  AP[{cls_name:<10}]: {ap:.4f}  (gt_count={all_gt_counts[SHAPE_CLASSES.index(cls_name)]})")

    # ── PR Curve Plot ────────────────────────────────────────────────────
    pr_path = os.path.join(output_dir, "pr_curves.png")
    _plot_pr_curves(pr_data, aps, pr_path)
    print(f"PR curves saved to {pr_path}")

    results = {
        "task": "detect",
        "map_at_50": map50,
        "per_class_ap": aps,
        "gt_counts": {SHAPE_CLASSES[k]: v for k, v in all_gt_counts.items()},
        "figures": {"pr_curves": pr_path},
        "note": (
            "mAP computed with simplified TP assignment (no IoU-based matching). "
            "For production use evaluate with pycocotools."
        ),
    }
    return results


def _compute_ap(recall: np.ndarray, precision: np.ndarray) -> float:
    """Compute area under precision-recall curve using 11-point interpolation."""
    ap = 0.0
    for thr in np.linspace(0, 1, 11):
        prec_at_rec = precision[recall >= thr] if any(recall >= thr) else np.array([0.0])
        ap += np.max(prec_at_rec) / 11.0
    return float(ap)


def _plot_pr_curves(
    pr_data: Dict, aps: Dict, save_path: str
) -> None:
    """Plot PR curves for all classes."""
    colors = plt.cm.Set2(np.linspace(0, 1, len(SHAPE_CLASSES)))
    fig, ax = plt.subplots(figsize=(8, 5))

    for (cls_name, (rec, prec)), color in zip(pr_data.items(), colors):
        ap = aps.get(cls_name, 0.0)
        ax.plot(rec, prec, color=color, lw=1.8,
                label=f"{cls_name} (AP={ap:.3f})")

    ax.set_xlabel("Recall", fontsize=11)
    ax.set_ylabel("Precision", fontsize=11)
    ax.set_title("TinyYOLO Precision-Recall Curves — M2a (Synthetic Data)",
                 fontsize=12, fontweight="bold")
    ax.legend(fontsize=9, loc="upper right")
    ax.set_xlim(0, 1); ax.set_ylim(0, 1.05)
    ax.grid(alpha=0.3)
    fig.tight_layout()
    fig.savefig(save_path, dpi=150, bbox_inches="tight")
    plt.close(fig)


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

def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(
        description="Evaluate M2a Vision models",
        formatter_class=argparse.ArgumentDefaultsHelpFormatter,
    )
    parser.add_argument("--task", choices=["classify", "detect", "both"],
                        default="classify")
    parser.add_argument("--checkpoint", default="checkpoints/best_classifier.pt",
                        help="Model checkpoint path")
    parser.add_argument("--det_checkpoint", default="checkpoints/best_detector.pt")
    parser.add_argument("--clf_checkpoint", default="checkpoints/best_classifier.pt")
    parser.add_argument("--data_dir", default="data/synthetic")
    parser.add_argument("--output_dir", default="assets",
                        help="Directory to save evaluation figures")
    parser.add_argument("--output_json", default=None,
                        help="Save metrics JSON to this path")
    parser.add_argument("--batch_size", type=int, default=32)
    return parser.parse_args()


def main():
    args = parse_args()
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    print(f"Device: {device}\n")

    all_results = {}

    if args.task in ("classify", "both"):
        clf_results = evaluate_classifier(
            checkpoint_path=args.clf_checkpoint if args.task == "both" else args.checkpoint,
            data_dir=args.data_dir,
            output_dir=args.output_dir,
            device=device,
            batch_size=args.batch_size,
        )
        all_results["classifier"] = clf_results

    if args.task in ("detect", "both"):
        det_results = evaluate_detector(
            checkpoint_path=args.det_checkpoint if args.task == "both" else args.checkpoint,
            data_dir=args.data_dir,
            output_dir=args.output_dir,
            device=device,
        )
        all_results["detector"] = det_results

    if args.output_json:
        os.makedirs(os.path.dirname(args.output_json) or ".", exist_ok=True)
        with open(args.output_json, "w") as f:
            json.dump(all_results, f, indent=2)
        print(f"\nResults saved to {args.output_json}")

    return all_results


if __name__ == "__main__":
    main()