"""
Evaluate expert classifier(s) on the Solidity vulnerability evaluation dataset.

For each expert classifier, runs inference on ALL samples in the eval dataset
and reports:
  1. Overall binary metrics (accuracy, F1, precision, recall, AUC)
  2. Per-vulnerability-type breakdown (which vuln types does this expert
     correctly identify vs. false-alarm on?)
  3. Visualizations: grouped bar charts, confusion matrix, score distributions
  4. Trackio dashboard for interactive exploration

Ground truth logic:
  For expert "reentrancy", a sample is "vulnerable" (label=1) if its
  vulnerability_type matches reentrancy. All other types are "safe" (label=0).
  This tests: does the expert fire on its own type and stay quiet on others?

Expected inputs:
  - Eval dataset: jhsu12/solidity-vulnerability-eval-dataset (on Hub)
  - Checkpoints: local folders or Hub repos from train_expert_classifier.py

Usage:
    # Evaluate a single expert:
    python evaluate_classifier.py \
        --checkpoint ./cls-expert-reentrancy/checkpoint-150 \
        --expert reentrancy

    # Evaluate a Hub-hosted expert:
    python evaluate_classifier.py \
        --checkpoint jhsu12/solidity-vuln-cls-reentrancy-v1 \
        --expert reentrancy

    # Evaluate ALL experts at once:
    python evaluate_classifier.py --all --base_dir .

    # Custom eval dataset:
    python evaluate_classifier.py \
        --checkpoint ./cls-expert-reentrancy/checkpoint-150 \
        --expert reentrancy \
        --eval_dataset jhsu12/my-custom-eval-dataset

    # Skip trackio:
    python evaluate_classifier.py --checkpoint ... --expert reentrancy --no_trackio
"""

import argparse
import os
import re
import sys
import json
import io
import numpy as np
import torch
import torch.nn.functional as F

import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt
import matplotlib.ticker as mticker

from PIL import Image as PILImage
from collections import defaultdict
from datasets import load_dataset
from transformers import AutoTokenizer, AutoModelForSequenceClassification, BitsAndBytesConfig
from peft import PeftModel, PeftConfig
from sklearn.metrics import (
    accuracy_score, f1_score, precision_score, recall_score,
    roc_auc_score, confusion_matrix, classification_report,
)

# ══════════════════════════════════════════════════════════════════════════════
# CONFIG
# ══════════════════════════════════════════════════════════════════════════════
BASE_MODEL = "Qwen/Qwen2.5-Coder-3B-Instruct"
EVAL_DATASET = "jhsu12/solidity-vulnerability-eval-dataset"

# Expert slug → list of matching vulnerability_type values in the eval dataset
# (eval dataset uses abbreviated names like "reentrancy (RE)")
EXPERT_VULN_MAPPING = {
    "reentrancy": ["reentrancy (RE)"],
    "access-control": ["dangerous delegatecall (DE)"],
    "integer-overflow-underflow": ["integer overflow (OF)"],
    "timestamp-dependence": [
        "timestamp dependency (TP)",
        "block number dependency (BN)",
    ],
    "unchecked-low-level-calls": ["unchecked external call (UC)"],
}

# All known experts
EXPERTS = {
    "reentrancy": "Reentrancy",
    "access-control": "Access Control",
    "integer-overflow-underflow": "Integer Overflow/Underflow",
    "timestamp-dependence": "Timestamp Dependence",
    "unchecked-low-level-calls": "Unchecked Low-Level Calls",
}


# ══════════════════════════════════════════════════════════════════════════════
# ARGS
# ══════════════════════════════════════════════════════════════════════════════
def parse_args():
    parser = argparse.ArgumentParser(
        description="Evaluate expert classifier(s) on Solidity eval dataset."
    )
    # Single expert mode
    parser.add_argument("--checkpoint", type=str, default=None,
                        help="Path to checkpoint (local folder or Hub ID)")
    parser.add_argument("--expert", type=str, default=None,
                        choices=list(EXPERTS.keys()),
                        help="Expert slug (e.g. 'reentrancy')")

    # Multi-expert mode
    parser.add_argument("--all", action="store_true", default=False,
                        help="Evaluate ALL experts found in --base_dir")
    parser.add_argument("--base_dir", type=str, default=".",
                        help="Base dir containing cls-expert-* folders (for --all mode)")

    # Dataset
    parser.add_argument("--eval_dataset", type=str, default=EVAL_DATASET,
                        help=f"Eval dataset ID (default: {EVAL_DATASET})")
    parser.add_argument("--max_samples", type=int, default=None,
                        help="Limit number of eval samples (for quick testing)")

    # Model
    parser.add_argument("--max_seq_len", type=int, default=1536)
    parser.add_argument("--batch_size", type=int, default=16)
    parser.add_argument("--threshold", type=float, default=0.5)
    parser.add_argument("--load_in_4bit", action="store_true", default=True)
    parser.add_argument("--load_in_8bit", action="store_true", default=False)

    # Output
    parser.add_argument("--output_dir", type=str, default="./eval_results")
    parser.add_argument("--no_trackio", action="store_true", default=False)
    parser.add_argument("--save_predictions", action="store_true", default=False,
                        help="Save per-sample predictions to JSON")

    return parser.parse_args()


# ══════════════════════════════════════════════════════════════════════════════
# HELPERS
# ══════════════════════════════════════════════════════════════════════════════
def fig_to_trackio_image(fig, caption=""):
    """Convert matplotlib figure to trackio.Image via PIL."""
    import trackio
    buf = io.BytesIO()
    fig.savefig(buf, format="png", bbox_inches="tight", dpi=150)
    buf.seek(0)
    pil_img = PILImage.open(buf).convert("RGB")
    plt.close(fig)
    return trackio.Image(pil_img, caption=caption)


def save_figure(fig, path, caption=""):
    """Save matplotlib figure to file."""
    fig.savefig(path, format="png", bbox_inches="tight", dpi=150)
    plt.close(fig)
    print(f"     📊 Saved: {path}")


def detect_base_model(checkpoint_path):
    """Read base model from adapter config."""
    config_path = os.path.join(checkpoint_path, "adapter_config.json")
    if os.path.isfile(config_path):
        with open(config_path, "r") as f:
            cfg = json.load(f)
        return cfg.get("base_model_name_or_path", BASE_MODEL)
    try:
        peft_config = PeftConfig.from_pretrained(checkpoint_path)
        return peft_config.base_model_name_or_path
    except Exception:
        return BASE_MODEL


def find_best_checkpoint(expert_dir):
    """Find the highest-step checkpoint in an expert directory."""
    if not os.path.isdir(expert_dir):
        return None
    best_step = -1
    best_path = None
    for name in os.listdir(expert_dir):
        match = re.match(r"checkpoint-(\d+)$", name)
        if match:
            step = int(match.group(1))
            if step > best_step:
                best_step = step
                best_path = os.path.join(expert_dir, name)
    # Also check for best_model subfolder
    best_model_path = os.path.join(expert_dir, "best_model")
    if os.path.isdir(best_model_path):
        adapter_file = os.path.join(best_model_path, "adapter_config.json")
        if os.path.isfile(adapter_file):
            return best_model_path
    return best_path


# ══════════════════════════════════════════════════════════════════════════════
# MODEL LOADING
# ══════════════════════════════════════════════════════════════════════════════
def load_classifier(checkpoint_path, load_in_4bit=True, load_in_8bit=False):
    """Load base model + LoRA adapter for classification."""
    base_model_id = detect_base_model(checkpoint_path)
    print(f"     Base model: {base_model_id}")

    has_bf16 = torch.cuda.is_bf16_supported() if torch.cuda.is_available() else False
    compute_dtype = torch.bfloat16 if has_bf16 else torch.float16

    if load_in_8bit:
        bnb_config = BitsAndBytesConfig(load_in_8bit=True)
    elif load_in_4bit:
        bnb_config = BitsAndBytesConfig(
            load_in_4bit=True, bnb_4bit_quant_type="nf4",
            bnb_4bit_compute_dtype=compute_dtype, bnb_4bit_use_double_quant=True,
        )
    else:
        bnb_config = None

    attn_impl = "sdpa"
    try:
        import flash_attn
        attn_impl = "flash_attention_2"
    except ImportError:
        pass

    model = AutoModelForSequenceClassification.from_pretrained(
        base_model_id, num_labels=2,
        id2label={0: "safe", 1: "vulnerable"},
        label2id={"safe": 0, "vulnerable": 1},
        quantization_config=bnb_config, device_map="auto",
        torch_dtype=compute_dtype, trust_remote_code=True,
        attn_implementation=attn_impl, ignore_mismatched_sizes=True,
    )
    model = PeftModel.from_pretrained(model, checkpoint_path)
    model.eval()

    try:
        tokenizer = AutoTokenizer.from_pretrained(checkpoint_path, trust_remote_code=True)
    except Exception:
        tokenizer = AutoTokenizer.from_pretrained(base_model_id, trust_remote_code=True)

    if tokenizer.pad_token is None:
        tokenizer.pad_token = tokenizer.eos_token
    model.config.pad_token_id = tokenizer.pad_token_id

    return model, tokenizer


# ══════════════════════════════════════════════════════════════════════════════
# INFERENCE
# ══════════════════════════════════════════════════════════════════════════════
def run_inference(model, tokenizer, codes, batch_size=16, max_seq_len=1536):
    """Run batched inference on a list of code strings. Returns logits array."""
    all_logits = []

    for i in range(0, len(codes), batch_size):
        batch_codes = codes[i:i + batch_size]
        inputs = tokenizer(
            batch_codes, return_tensors="pt", truncation=True,
            max_length=max_seq_len, padding=True,
        )
        inputs = {k: v.to(model.device) for k, v in inputs.items()}

        with torch.no_grad():
            outputs = model(**inputs)

        all_logits.append(outputs.logits.cpu().float().numpy())

        done = min(i + batch_size, len(codes))
        if (done % (batch_size * 10) == 0) or done == len(codes):
            print(f"       [{done}/{len(codes)}]")

    return np.concatenate(all_logits, axis=0)


# ══════════════════════════════════════════════════════════════════════════════
# METRICS
# ══════════════════════════════════════════════════════════════════════════════
def compute_metrics(labels, preds, probs_vuln):
    """Compute binary classification metrics."""
    metrics = {
        "accuracy": accuracy_score(labels, preds),
        "f1": f1_score(labels, preds, average="binary", zero_division=0),
        "precision": precision_score(labels, preds, average="binary", zero_division=0),
        "recall": recall_score(labels, preds, average="binary", zero_division=0),
    }
    if len(set(labels)) > 1:
        try:
            metrics["auc"] = roc_auc_score(labels, probs_vuln)
        except ValueError:
            metrics["auc"] = 0.0
    else:
        metrics["auc"] = 0.0
    return metrics


def compute_per_vuln_metrics(vuln_types, labels, preds, probs_vuln):
    """Compute metrics broken down by vulnerability type."""
    per_vuln = {}
    unique_types = sorted(set(vuln_types))

    for vt in unique_types:
        mask = [i for i, v in enumerate(vuln_types) if v == vt]
        vt_labels = [labels[i] for i in mask]
        vt_preds = [preds[i] for i in mask]
        vt_probs = [probs_vuln[i] for i in mask]

        n_pos = sum(vt_labels)
        n_neg = len(vt_labels) - n_pos

        per_vuln[vt] = {
            "accuracy": accuracy_score(vt_labels, vt_preds),
            "f1": f1_score(vt_labels, vt_preds, average="binary", zero_division=0),
            "precision": precision_score(vt_labels, vt_preds, average="binary", zero_division=0),
            "recall": recall_score(vt_labels, vt_preds, average="binary", zero_division=0),
            "tp": sum(1 for p, l in zip(vt_preds, vt_labels) if p == 1 and l == 1),
            "fp": sum(1 for p, l in zip(vt_preds, vt_labels) if p == 1 and l == 0),
            "tn": sum(1 for p, l in zip(vt_preds, vt_labels) if p == 0 and l == 0),
            "fn": sum(1 for p, l in zip(vt_preds, vt_labels) if p == 0 and l == 1),
            "total": len(vt_labels),
            "n_positive": n_pos,
            "n_negative": n_neg,
            "mean_prob_vuln": float(np.mean(vt_probs)),
        }

    return per_vuln


# ══════════════════════════════════════════════════════════════════════════════
# VISUALIZATIONS
# ══════════════════════════════════════════════════════════════════════════════
COLORS = {
    "precision": "#4C72B0",
    "recall": "#DD8452",
    "f1": "#55A868",
    "accuracy": "#C44E52",
}


def plot_overall_metrics(metrics, expert_name):
    """Bar chart of overall metrics."""
    fig, ax = plt.subplots(figsize=(8, 5))
    metric_names = ["accuracy", "precision", "recall", "f1", "auc"]
    values = [metrics.get(m, 0) for m in metric_names]
    colors = ["#4C72B0", "#DD8452", "#55A868", "#C44E52", "#8172B3"]

    bars = ax.bar(metric_names, values, color=colors, edgecolor="white", linewidth=0.8)
    for bar, val in zip(bars, values):
        ax.text(bar.get_x() + bar.get_width() / 2, bar.get_height() + 0.01,
                f"{val:.3f}", ha="center", va="bottom", fontsize=11, fontweight="bold")

    ax.set_ylim(0, 1.12)
    ax.set_title(f"Overall Metrics — {expert_name}", fontsize=14, fontweight="bold")
    ax.set_ylabel("Score")
    ax.yaxis.set_major_formatter(mticker.FormatStrFormatter("%.2f"))
    ax.grid(axis="y", alpha=0.3)
    fig.tight_layout()
    return fig


def plot_per_vuln_grouped_bar(per_vuln, expert_name):
    """Grouped bar chart: precision/recall/F1 per vulnerability type."""
    vuln_types = list(per_vuln.keys())
    n = len(vuln_types)
    if n == 0:
        return None

    # Shorten labels for display
    short_labels = [vt.split("(")[0].strip()[:20] for vt in vuln_types]

    x = np.arange(n)
    width = 0.25

    prec = [per_vuln[vt]["precision"] for vt in vuln_types]
    rec = [per_vuln[vt]["recall"] for vt in vuln_types]
    f1 = [per_vuln[vt]["f1"] for vt in vuln_types]

    fig, ax = plt.subplots(figsize=(max(10, n * 1.5), 6))
    ax.bar(x - width, prec, width, label="Precision", color=COLORS["precision"], edgecolor="white")
    ax.bar(x, rec, width, label="Recall", color=COLORS["recall"], edgecolor="white")
    ax.bar(x + width, f1, width, label="F1", color=COLORS["f1"], edgecolor="white")

    ax.set_xticks(x)
    ax.set_xticklabels(short_labels, rotation=35, ha="right", fontsize=9)
    ax.set_ylim(0, 1.12)
    ax.set_ylabel("Score")
    ax.set_title(f"Per Vulnerability Type — {expert_name}", fontsize=14, fontweight="bold")
    ax.legend(loc="upper right")
    ax.grid(axis="y", alpha=0.3)
    fig.tight_layout()
    return fig


def plot_per_vuln_detection_rate(per_vuln, expert_name, expert_vuln_types):
    """
    Stacked bar: for each vuln type, show the fraction predicted as vulnerable.
    Highlight the expert's own type vs others.
    """
    vuln_types = list(per_vuln.keys())
    n = len(vuln_types)
    if n == 0:
        return None

    short_labels = [vt.split("(")[0].strip()[:20] for vt in vuln_types]
    detection_rates = [per_vuln[vt]["mean_prob_vuln"] for vt in vuln_types]

    # Color: expert's own type in green, others in blue
    colors = []
    for vt in vuln_types:
        if vt in expert_vuln_types:
            colors.append("#55A868")  # green — should detect
        else:
            colors.append("#4C72B0")  # blue — should NOT detect

    fig, ax = plt.subplots(figsize=(max(10, n * 1.5), 6))
    bars = ax.bar(range(n), detection_rates, color=colors, edgecolor="white", linewidth=0.8)

    for bar, val in zip(bars, detection_rates):
        ax.text(bar.get_x() + bar.get_width() / 2, bar.get_height() + 0.01,
                f"{val:.2f}", ha="center", va="bottom", fontsize=9)

    ax.set_xticks(range(n))
    ax.set_xticklabels(short_labels, rotation=35, ha="right", fontsize=9)
    ax.set_ylim(0, 1.12)
    ax.set_ylabel("Mean P(vulnerable)")
    ax.set_title(f"Detection Confidence by Vuln Type — {expert_name}\n"
                 f"(green = expert's target type, blue = other types)",
                 fontsize=12, fontweight="bold")
    ax.axhline(y=0.5, color="red", linestyle="--", alpha=0.5, label="threshold=0.5")
    ax.legend()
    ax.grid(axis="y", alpha=0.3)
    fig.tight_layout()
    return fig


def plot_confusion_matrix(labels, preds, expert_name):
    """Binary confusion matrix heatmap."""
    cm = confusion_matrix(labels, preds, labels=[0, 1])
    fig, ax = plt.subplots(figsize=(6, 5))

    im = ax.imshow(cm, cmap="Blues", aspect="auto")
    ax.set_xticks([0, 1])
    ax.set_yticks([0, 1])
    ax.set_xticklabels(["Safe", "Vulnerable"])
    ax.set_yticklabels(["Safe", "Vulnerable"])
    ax.set_xlabel("Predicted")
    ax.set_ylabel("Actual")
    ax.set_title(f"Confusion Matrix — {expert_name}", fontsize=13, fontweight="bold")

    # Annotate cells
    for i in range(2):
        for j in range(2):
            val = cm[i, j]
            color = "white" if val > cm.max() / 2 else "black"
            ax.text(j, i, str(val), ha="center", va="center",
                    fontsize=16, fontweight="bold", color=color)

    fig.colorbar(im, ax=ax, shrink=0.8)
    fig.tight_layout()
    return fig


def plot_score_distribution(probs_vuln, labels, expert_name):
    """Histogram of predicted P(vulnerable) split by actual label."""
    fig, ax = plt.subplots(figsize=(9, 5))

    probs_safe = [p for p, l in zip(probs_vuln, labels) if l == 0]
    probs_vuln_actual = [p for p, l in zip(probs_vuln, labels) if l == 1]

    ax.hist(probs_safe, bins=50, alpha=0.6, color="#4C72B0", label=f"Actually Safe (n={len(probs_safe)})",
            density=True, edgecolor="white")
    if probs_vuln_actual:
        ax.hist(probs_vuln_actual, bins=50, alpha=0.6, color="#C44E52",
                label=f"Actually Vulnerable (n={len(probs_vuln_actual)})",
                density=True, edgecolor="white")

    ax.axvline(x=0.5, color="black", linestyle="--", alpha=0.7, label="threshold=0.5")
    ax.set_xlabel("P(vulnerable)")
    ax.set_ylabel("Density")
    ax.set_title(f"Prediction Score Distribution — {expert_name}", fontsize=13, fontweight="bold")
    ax.legend()
    ax.grid(axis="y", alpha=0.3)
    fig.tight_layout()
    return fig


def plot_false_alarm_by_type(per_vuln, expert_name, expert_vuln_types):
    """Bar chart showing false positive RATE for each non-target vuln type."""
    non_target = {vt: m for vt, m in per_vuln.items() if vt not in expert_vuln_types}
    if not non_target:
        return None

    vuln_types = list(non_target.keys())
    short_labels = [vt.split("(")[0].strip()[:20] for vt in vuln_types]
    fp_rates = []
    for vt in vuln_types:
        m = non_target[vt]
        total_neg = m["tn"] + m["fp"]
        fp_rates.append(m["fp"] / total_neg if total_neg > 0 else 0)

    fig, ax = plt.subplots(figsize=(max(10, len(vuln_types) * 1.5), 5))
    bars = ax.bar(range(len(vuln_types)), fp_rates, color="#DD8452", edgecolor="white")

    for bar, val in zip(bars, fp_rates):
        ax.text(bar.get_x() + bar.get_width() / 2, bar.get_height() + 0.01,
                f"{val:.1%}", ha="center", va="bottom", fontsize=9)

    ax.set_xticks(range(len(vuln_types)))
    ax.set_xticklabels(short_labels, rotation=35, ha="right", fontsize=9)
    ax.set_ylim(0, max(fp_rates) * 1.3 + 0.05 if fp_rates else 1)
    ax.set_ylabel("False Positive Rate")
    ax.set_title(f"False Alarm Rate on Non-Target Types — {expert_name}",
                 fontsize=12, fontweight="bold")
    ax.grid(axis="y", alpha=0.3)
    fig.tight_layout()
    return fig


# ══════════════════════════════════════════════════════════════════════════════
# EVALUATE ONE EXPERT
# ══════════════════════════════════════════════════════════════════════════════
def evaluate_expert(checkpoint_path, expert_slug, dataset, args, use_trackio=True):
    """Run full evaluation for one expert. Returns results dict."""
    expert_name = EXPERTS.get(expert_slug, expert_slug)
    expert_vuln_types = EXPERT_VULN_MAPPING.get(expert_slug, [])

    print(f"\n{'━' * 60}")
    print(f"  🔬 Evaluating: {expert_name}")
    print(f"     Checkpoint:  {checkpoint_path}")
    print(f"     Target types: {expert_vuln_types}")
    print(f"{'━' * 60}")

    # ── Create ground truth labels ────────────────────────────────────────────
    vuln_types = dataset["vulnerability_type"]
    labels = [1 if vt in expert_vuln_types else 0 for vt in vuln_types]
    n_pos = sum(labels)
    n_neg = len(labels) - n_pos
    print(f"\n     Ground truth: {n_pos} positive, {n_neg} negative "
          f"({n_pos / len(labels):.1%} positive rate)")

    # ── Load model ────────────────────────────────────────────────────────────
    print(f"\n     Loading model...")
    model, tokenizer = load_classifier(
        checkpoint_path,
        load_in_4bit=args.load_in_4bit and not args.load_in_8bit,
        load_in_8bit=args.load_in_8bit,
    )

    # ── Run inference ─────────────────────────────────────────────────────────
    codes = dataset["code"]
    print(f"\n     Running inference on {len(codes)} samples (batch_size={args.batch_size})...")
    logits = run_inference(model, tokenizer, codes,
                           batch_size=args.batch_size, max_seq_len=args.max_seq_len)

    # Compute probabilities and predictions
    probs = torch.softmax(torch.tensor(logits), dim=-1).numpy()
    probs_vuln = probs[:, 1].tolist()
    preds = [1 if p >= args.threshold else 0 for p in probs_vuln]

    # ── Print first N examples for sanity check ─────────────────────────────
    n_preview = min(5, len(codes))
    print(f"\n     {'─' * 40}")
    print(f"     FIRST {n_preview} SAMPLES (sanity check)")
    print(f"     {'─' * 40}")
    for idx in range(n_preview):
        code_preview = codes[idx].replace('\n', ' ')[:80]
        gt_label = "VULNERABLE" if labels[idx] == 1 else "SAFE"
        pred_label = "VULNERABLE" if preds[idx] == 1 else "SAFE"
        match = "✅" if labels[idx] == preds[idx] else "❌"
        print(f"\n     [{idx+1}/{n_preview}] {match}")
        print(f"       File:         {dataset['filepath'][idx]}")
        print(f"       Vuln type:    {vuln_types[idx]}")
        print(f"       Code:         {code_preview}...")
        print(f"       Ground truth: {gt_label} (label={labels[idx]})")
        print(f"       Prediction:   {pred_label} (label={preds[idx]})")
        print(f"       P(safe):      {probs[idx][0]:.6f}")
        print(f"       P(vulnerable):{probs_vuln[idx]:.6f}")
        print(f"       Logits:       safe={logits[idx][0]:.4f}  vuln={logits[idx][1]:.4f}")

    # ── Compute metrics ───────────────────────────────────────────────────────
    print(f"\n     Computing metrics...")
    overall = compute_metrics(labels, preds, probs_vuln)
    per_vuln = compute_per_vuln_metrics(vuln_types, labels, preds, probs_vuln)

    # Print overall
    print(f"\n     {'─' * 40}")
    print(f"     OVERALL METRICS")
    print(f"     {'─' * 40}")
    print(f"     Accuracy:  {overall['accuracy']:.4f}")
    print(f"     F1:        {overall['f1']:.4f}")
    print(f"     Precision: {overall['precision']:.4f}")
    print(f"     Recall:    {overall['recall']:.4f}")
    print(f"     AUC:       {overall['auc']:.4f}")

    # Print per-vuln breakdown
    print(f"\n     {'─' * 40}")
    print(f"     PER VULNERABILITY TYPE")
    print(f"     {'─' * 40}")
    print(f"     {'Type':<30} {'Acc':>5} {'F1':>5} {'Prec':>5} {'Rec':>5} {'TP':>4} {'FP':>4} {'TN':>4} {'FN':>4} {'N':>5}")
    print(f"     {'─' * 87}")
    for vt in sorted(per_vuln.keys()):
        m = per_vuln[vt]
        marker = " ◀" if vt in expert_vuln_types else ""
        short_name = vt[:28]
        print(f"     {short_name:<30} {m['accuracy']:>5.3f} {m['f1']:>5.3f} "
              f"{m['precision']:>5.3f} {m['recall']:>5.3f} "
              f"{m['tp']:>4} {m['fp']:>4} {m['tn']:>4} {m['fn']:>4} {m['total']:>5}{marker}")

    # ── Create output dir ─────────────────────────────────────────────────────
    expert_output_dir = os.path.join(args.output_dir, expert_slug)
    os.makedirs(expert_output_dir, exist_ok=True)

    # ── Generate visualizations ───────────────────────────────────────────────
    print(f"\n     Generating visualizations...")

    fig_overall = plot_overall_metrics(overall, expert_name)
    fig_grouped = plot_per_vuln_grouped_bar(per_vuln, expert_name)
    fig_detection = plot_per_vuln_detection_rate(per_vuln, expert_name, expert_vuln_types)
    fig_cm = plot_confusion_matrix(labels, preds, expert_name)
    fig_dist = plot_score_distribution(probs_vuln, labels, expert_name)
    fig_fp = plot_false_alarm_by_type(per_vuln, expert_name, expert_vuln_types)

    # Save all figures locally
    figures = {
        "overall_metrics": fig_overall,
        "per_vuln_grouped_bar": fig_grouped,
        "detection_confidence": fig_detection,
        "confusion_matrix": fig_cm,
        "score_distribution": fig_dist,
        "false_alarm_rate": fig_fp,
    }
    for name, fig in figures.items():
        if fig is not None:
            save_figure(fig, os.path.join(expert_output_dir, f"{name}.png"), name)

    # ── Log to Trackio ────────────────────────────────────────────────────────
    if use_trackio:
        import trackio
        import pandas as pd

        # Overall scalars
        trackio.log({
            f"{expert_slug}/overall/accuracy": overall["accuracy"],
            f"{expert_slug}/overall/f1": overall["f1"],
            f"{expert_slug}/overall/precision": overall["precision"],
            f"{expert_slug}/overall/recall": overall["recall"],
            f"{expert_slug}/overall/auc": overall["auc"],
        })

        # Per-vuln scalars
        for vt, m in per_vuln.items():
            safe_vt = vt.replace("/", "-").replace(" ", "_")
            trackio.log({
                f"{expert_slug}/per_vuln/{safe_vt}/accuracy": m["accuracy"],
                f"{expert_slug}/per_vuln/{safe_vt}/f1": m["f1"],
                f"{expert_slug}/per_vuln/{safe_vt}/precision": m["precision"],
                f"{expert_slug}/per_vuln/{safe_vt}/recall": m["recall"],
                f"{expert_slug}/per_vuln/{safe_vt}/mean_prob_vuln": m["mean_prob_vuln"],
            })

        # Per-vuln breakdown table
        table_data = []
        for vt in sorted(per_vuln.keys()):
            m = per_vuln[vt]
            is_target = "YES" if vt in expert_vuln_types else ""
            table_data.append([
                vt, is_target, m["total"], m["n_positive"], m["n_negative"],
                round(m["accuracy"], 4), round(m["precision"], 4),
                round(m["recall"], 4), round(m["f1"], 4),
                m["tp"], m["fp"], m["tn"], m["fn"],
                round(m["mean_prob_vuln"], 4),
            ])

        df = pd.DataFrame(table_data, columns=[
            "vulnerability_type", "is_target", "total", "n_pos", "n_neg",
            "accuracy", "precision", "recall", "f1",
            "tp", "fp", "tn", "fn", "mean_prob_vuln",
        ])
        trackio.log({
            f"{expert_slug}/tables/per_vuln_breakdown": trackio.Table(dataframe=df),
        })

        # Figures
        for name, fig_path in [
            ("overall_metrics", os.path.join(expert_output_dir, "overall_metrics.png")),
            ("per_vuln_grouped_bar", os.path.join(expert_output_dir, "per_vuln_grouped_bar.png")),
            ("detection_confidence", os.path.join(expert_output_dir, "detection_confidence.png")),
            ("confusion_matrix", os.path.join(expert_output_dir, "confusion_matrix.png")),
            ("score_distribution", os.path.join(expert_output_dir, "score_distribution.png")),
            ("false_alarm_rate", os.path.join(expert_output_dir, "false_alarm_rate.png")),
        ]:
            if os.path.isfile(fig_path):
                trackio.log({
                    f"{expert_slug}/charts/{name}": trackio.Image(fig_path, caption=name),
                })

        # Score distribution histogram
        trackio.log({
            f"{expert_slug}/distributions/prob_vulnerable": trackio.Histogram(
                np.array(probs_vuln), num_bins=50
            ),
        })

        # Markdown summary
        best_vt = max(per_vuln.items(), key=lambda x: x[1]["f1"])
        worst_vt = min(per_vuln.items(), key=lambda x: x[1]["f1"])
        trackio.log({
            f"{expert_slug}/reports/summary": trackio.Markdown(f"""
# {expert_name} — Evaluation Summary

| Metric | Score |
|--------|-------|
| Accuracy | {overall['accuracy']:.4f} |
| F1 | {overall['f1']:.4f} |
| Precision | {overall['precision']:.4f} |
| Recall | {overall['recall']:.4f} |
| AUC | {overall['auc']:.4f} |

**Samples**: {len(labels)} ({n_pos} positive, {n_neg} negative)
**Threshold**: {args.threshold}

**Best per-type F1**: {best_vt[0]} → {best_vt[1]['f1']:.4f}
**Worst per-type F1**: {worst_vt[0]} → {worst_vt[1]['f1']:.4f}
"""),
        })

    # ── Save results JSON ─────────────────────────────────────────────────────
    results = {
        "expert_slug": expert_slug,
        "expert_name": expert_name,
        "checkpoint": checkpoint_path,
        "threshold": args.threshold,
        "n_samples": len(labels),
        "n_positive": n_pos,
        "n_negative": n_neg,
        "overall": overall,
        "per_vulnerability_type": per_vuln,
    }

    results_path = os.path.join(expert_output_dir, "results.json")
    with open(results_path, "w") as f:
        json.dump(results, f, indent=2)
    print(f"\n     💾 Results saved: {results_path}")

    # Save per-sample predictions if requested
    if args.save_predictions:
        predictions = []
        for i in range(len(codes)):
            predictions.append({
                "filepath": dataset["filepath"][i],
                "vulnerability_type": vuln_types[i],
                "ground_truth": labels[i],
                "prediction": preds[i],
                "prob_vulnerable": round(probs_vuln[i], 6),
                "prob_safe": round(probs[i][0].item(), 6),
            })
        pred_path = os.path.join(expert_output_dir, "predictions.json")
        with open(pred_path, "w") as f:
            json.dump(predictions, f, indent=2)
        print(f"     💾 Predictions saved: {pred_path}")

    # Free GPU memory
    del model
    if torch.cuda.is_available():
        torch.cuda.empty_cache()

    return results


# ══════════════════════════════════════════════════════════════════════════════
# CROSS-EXPERT COMPARISON
# ══════════════════════════════════════════════════════════════════════════════
def plot_cross_expert_comparison(all_results, output_dir, use_trackio=True):
    """Compare all experts side-by-side."""
    if len(all_results) < 2:
        return

    print(f"\n{'━' * 60}")
    print(f"  📊 Cross-Expert Comparison")
    print(f"{'━' * 60}")

    experts = [r["expert_name"] for r in all_results]
    short_experts = [e[:15] for e in experts]

    # Overall metrics comparison
    fig, ax = plt.subplots(figsize=(max(10, len(experts) * 2), 6))
    x = np.arange(len(experts))
    width = 0.2

    metrics_to_plot = ["accuracy", "precision", "recall", "f1"]
    colors = [COLORS[m] for m in metrics_to_plot]

    for i, (metric, color) in enumerate(zip(metrics_to_plot, colors)):
        values = [r["overall"][metric] for r in all_results]
        offset = (i - len(metrics_to_plot) / 2 + 0.5) * width
        bars = ax.bar(x + offset, values, width, label=metric.capitalize(), color=color, edgecolor="white")

    ax.set_xticks(x)
    ax.set_xticklabels(short_experts, rotation=20, ha="right")
    ax.set_ylim(0, 1.12)
    ax.set_ylabel("Score")
    ax.set_title("Cross-Expert Comparison", fontsize=14, fontweight="bold")
    ax.legend()
    ax.grid(axis="y", alpha=0.3)
    fig.tight_layout()
    save_figure(fig, os.path.join(output_dir, "cross_expert_comparison.png"))

    # Heatmap: each expert × each vuln type → F1
    all_vuln_types = sorted(set(
        vt for r in all_results for vt in r["per_vulnerability_type"].keys()
    ))
    heatmap_data = np.zeros((len(all_results), len(all_vuln_types)))
    for i, r in enumerate(all_results):
        for j, vt in enumerate(all_vuln_types):
            if vt in r["per_vulnerability_type"]:
                heatmap_data[i, j] = r["per_vulnerability_type"][vt]["f1"]

    fig, ax = plt.subplots(figsize=(max(12, len(all_vuln_types) * 1.5), max(5, len(experts) * 0.8)))
    short_vuln = [vt.split("(")[0].strip()[:18] for vt in all_vuln_types]

    im = ax.imshow(heatmap_data, cmap="RdYlGn", aspect="auto", vmin=0, vmax=1)
    ax.set_xticks(range(len(all_vuln_types)))
    ax.set_xticklabels(short_vuln, rotation=40, ha="right", fontsize=9)
    ax.set_yticks(range(len(experts)))
    ax.set_yticklabels(short_experts, fontsize=10)
    ax.set_title("F1 Score: Expert × Vulnerability Type", fontsize=14, fontweight="bold")

    for i in range(len(all_results)):
        for j in range(len(all_vuln_types)):
            val = heatmap_data[i, j]
            color = "white" if val < 0.5 else "black"
            ax.text(j, i, f"{val:.2f}", ha="center", va="center",
                    fontsize=9, color=color, fontweight="bold")

    fig.colorbar(im, ax=ax, shrink=0.8, label="F1 Score")
    fig.tight_layout()
    save_figure(fig, os.path.join(output_dir, "expert_vuln_heatmap.png"))

    if use_trackio:
        import trackio
        import pandas as pd

        trackio.log({
            "comparison/charts/cross_expert": trackio.Image(
                os.path.join(output_dir, "cross_expert_comparison.png"),
                caption="Cross-Expert Comparison"),
            "comparison/charts/heatmap": trackio.Image(
                os.path.join(output_dir, "expert_vuln_heatmap.png"),
                caption="Expert × Vuln Type F1 Heatmap"),
        })

        # Summary table
        summary_data = []
        for r in all_results:
            summary_data.append([
                r["expert_name"],
                round(r["overall"]["accuracy"], 4),
                round(r["overall"]["precision"], 4),
                round(r["overall"]["recall"], 4),
                round(r["overall"]["f1"], 4),
                round(r["overall"]["auc"], 4),
                r["n_positive"],
                r["n_negative"],
            ])
        df = pd.DataFrame(summary_data, columns=[
            "expert", "accuracy", "precision", "recall", "f1", "auc", "n_pos", "n_neg"
        ])
        trackio.log({"comparison/tables/summary": trackio.Table(dataframe=df)})

        trackio.log({
            "comparison/reports/summary": trackio.Markdown(
                "# Cross-Expert Evaluation Summary\n\n"
                + df.to_markdown(index=False)
            ),
        })


# ══════════════════════════════════════════════════════════════════════════════
# MAIN
# ══════════════════════════════════════════════════════════════════════════════
def main():
    args = parse_args()

    print("=" * 60)
    print("  Expert Classifier Evaluation")
    print("=" * 60)

    # ── Determine which experts to evaluate ───────────────────────────────────
    eval_tasks = []  # list of (checkpoint_path, expert_slug)

    if args.all:
        base_dir = os.path.abspath(args.base_dir)
        print(f"\n🔍 Scanning for experts in: {base_dir}")
        for slug in EXPERTS:
            expert_dir = os.path.join(base_dir, f"cls-expert-{slug}")
            ckpt = find_best_checkpoint(expert_dir)
            if ckpt:
                eval_tasks.append((ckpt, slug))
                print(f"     ✅ {slug}: {ckpt}")
            else:
                print(f"     ⬜ {slug}: not found")
    elif args.checkpoint and args.expert:
        eval_tasks.append((args.checkpoint, args.expert))
    else:
        print("\n❌ Provide --checkpoint + --expert, or use --all")
        sys.exit(1)

    if not eval_tasks:
        print("\n❌ No expert checkpoints found!")
        sys.exit(1)

    print(f"\n   Will evaluate {len(eval_tasks)} expert(s)")

    # ── Load eval dataset ─────────────────────────────────────────────────────
    print(f"\n📦 Loading eval dataset: {args.eval_dataset}")
    dataset = load_dataset(args.eval_dataset, split="train")
    if args.max_samples:
        dataset = dataset.select(range(min(args.max_samples, len(dataset))))
    print(f"   {len(dataset)} samples")

    from collections import Counter
    vuln_dist = Counter(dataset["vulnerability_type"])
    print(f"   {len(vuln_dist)} vulnerability types:")
    for vt, count in vuln_dist.most_common():
        print(f"     {vt}: {count}")

    # ── Init Trackio ──────────────────────────────────────────────────────────
    use_trackio = not args.no_trackio
    if use_trackio:
        import trackio
        trackio.init(
            project="solidity-cls-expert-eval",
            name=f"eval-{'all' if args.all else args.expert}",
            config={
                "eval_dataset": args.eval_dataset,
                "n_samples": len(dataset),
                "threshold": args.threshold,
                "max_seq_len": args.max_seq_len,
                "experts": [slug for _, slug in eval_tasks],
            },
        )
        print(f"\n📈 Trackio initialized: project='solidity-cls-expert-eval'")

    # ── Create output dir ─────────────────────────────────────────────────────
    os.makedirs(args.output_dir, exist_ok=True)

    # ── Evaluate each expert ──────────────────────────────────────────────────
    all_results = []

    for checkpoint_path, expert_slug in eval_tasks:
        results = evaluate_expert(
            checkpoint_path, expert_slug, dataset, args,
            use_trackio=use_trackio,
        )
        all_results.append(results)

    # ── Cross-expert comparison ───────────────────────────────────────────────
    if len(all_results) > 1:
        plot_cross_expert_comparison(all_results, args.output_dir, use_trackio=use_trackio)

    # ── Final summary ─────────────────────────────────────────────────────────
    if use_trackio:
        import trackio
        trackio.finish()

    print(f"\n{'=' * 60}")
    print(f"  EVALUATION COMPLETE")
    print(f"{'=' * 60}")
    print(f"\n  Results saved in: {args.output_dir}/")
    for _, slug in eval_tasks:
        r = next(r for r in all_results if r["expert_slug"] == slug)
        print(f"\n  {EXPERTS[slug]}:")
        print(f"    Accuracy:  {r['overall']['accuracy']:.4f}")
        print(f"    F1:        {r['overall']['f1']:.4f}")
        print(f"    Precision: {r['overall']['precision']:.4f}")
        print(f"    Recall:    {r['overall']['recall']:.4f}")
        print(f"    AUC:       {r['overall']['auc']:.4f}")

    if use_trackio:
        print(f"\n  📈 Trackio dashboard: check your trackio space")

    print(f"\n{'=' * 60}")


if __name__ == "__main__":
    main()