"""
Evaluate classifier expert(s) on the COMBINED training data from ALL expert datasets.

Loads train+test from all 5 expert datasets, merges them into one pool, then
runs each classifier against the full set. For each classifier, its own
expert type = label 1 (vulnerable), everything else = label 0 (safe).

This gives you per-vulnerability-type breakdown using the same data distribution
the models were trained on — no length mismatch issues.

Usage:
    # Evaluate one classifier against all combined data:
    python evaluate_on_all_train_data.py \
        --checkpoint ./cls-expert-reentrancy/checkpoint-480 \
        --expert reentrancy

    # Evaluate ALL classifiers at once:
    python evaluate_on_all_train_data.py --all --base_dir .

    # Quick test:
    python evaluate_on_all_train_data.py --all --base_dir . --max_samples_per_dataset 100
"""

import argparse
import os
import re
import json
import numpy as np
import torch
from collections import Counter
from datasets import load_dataset, concatenate_datasets
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

# ── Config ────────────────────────────────────────────────────────────────────
BASE_MODEL = "Qwen/Qwen2.5-Coder-3B-Instruct"

EXPERT_DATASETS = {
    "reentrancy": "jhsu12/solidity-vuln-expert-reentrancy",
    "access-control": "jhsu12/solidity-vuln-expert-access-control",
    "integer-overflow-underflow": "jhsu12/solidity-vuln-expert-integer-overflow-underflow",
    "timestamp-dependence": "jhsu12/solidity-vuln-expert-timestamp-dependence",
    "unchecked-low-level-calls": "jhsu12/solidity-vuln-expert-unchecked-low-level-calls",
}

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",
}


def parse_args():
    parser = argparse.ArgumentParser(
        description="Evaluate classifier(s) on combined data from all expert datasets."
    )
    # Single expert mode
    parser.add_argument("--checkpoint", type=str, default=None)
    parser.add_argument("--expert", type=str, default=None, choices=list(EXPERTS.keys()))

    # Multi-expert mode
    parser.add_argument("--all", action="store_true", default=False)
    parser.add_argument("--base_dir", type=str, default=".")

    # Options
    parser.add_argument("--max_samples_per_dataset", type=int, default=None,
                        help="Limit samples per expert dataset (for quick testing)")
    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("--preview", type=int, default=5)
    return parser.parse_args()


def extract_user_code(messages):
    for msg in messages:
        if msg["role"] == "user":
            return msg["content"]
    return ""


def detect_base_model(checkpoint_path):
    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):
    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)
    best_model_path = os.path.join(expert_dir, "best_model")
    if os.path.isdir(best_model_path):
        if os.path.isfile(os.path.join(best_model_path, "adapter_config.json")):
            return best_model_path
    return best_path


def load_classifier(checkpoint_path):
    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

    bnb_config = BitsAndBytesConfig(
        load_in_4bit=True, bnb_4bit_quant_type="nf4",
        bnb_4bit_compute_dtype=compute_dtype, bnb_4bit_use_double_quant=True,
    )

    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


def load_all_data(max_samples_per_dataset=None):
    """Load train+test from all 5 expert datasets, tag each sample with its source."""
    all_codes = []
    all_source_experts = []  # which expert dataset this sample came from
    all_is_expert = []       # is_expert_type from the original dataset

    for slug, dataset_id in EXPERT_DATASETS.items():
        print(f"  Loading {slug}...", end=" ")
        ds = load_dataset(dataset_id)

        for split_name in ["train", "test"]:
            split = ds[split_name]
            if max_samples_per_dataset:
                n = min(max_samples_per_dataset, len(split))
                split = split.select(range(n))

            for row in split:
                code = extract_user_code(row["messages"])
                all_codes.append(code)
                all_source_experts.append(slug)
                all_is_expert.append(int(row["is_expert_type"]))

        total = len(ds["train"]) + len(ds["test"])
        if max_samples_per_dataset:
            total = min(max_samples_per_dataset, len(ds["train"])) + min(max_samples_per_dataset, len(ds["test"]))
        print(f"{total} samples")

    return all_codes, all_source_experts, all_is_expert


def run_inference(model, tokenizer, codes, batch_size=16, max_seq_len=1536):
    all_logits = []
    for i in range(0, len(codes), batch_size):
        batch = codes[i:i + batch_size]
        inputs = tokenizer(batch, 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 * 20) == 0 or done == len(codes):
            print(f"     [{done}/{len(codes)}]")

    return np.concatenate(all_logits, axis=0)


def evaluate_one_expert(checkpoint_path, expert_slug, all_codes, all_source_experts,
                        all_is_expert, args):
    """Evaluate one classifier on the full combined dataset."""
    expert_name = EXPERTS[expert_slug]

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

    # ── Ground truth: for THIS expert, label=1 if sample came from this expert's
    #    dataset AND is_expert_type=True. Everything else is label=0.
    labels = []
    for i in range(len(all_codes)):
        if all_source_experts[i] == expert_slug and all_is_expert[i] == 1:
            labels.append(1)
        else:
            labels.append(0)

    n_pos = sum(labels)
    n_neg = len(labels) - n_pos
    print(f"\n     Total: {len(labels)} samples")
    print(f"     Positive (this expert's vuln): {n_pos}")
    print(f"     Negative (everything else):    {n_neg}")

    # ── Load model ────────────────────────────────────────────────────────────
    print(f"\n     Loading model...")
    model, tokenizer = load_classifier(checkpoint_path)

    # ── Inference ─────────────────────────────────────────────────────────────
    print(f"\n     Running inference on {len(all_codes)} samples...")
    logits = run_inference(model, tokenizer, all_codes,
                           batch_size=args.batch_size, max_seq_len=args.max_seq_len)

    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]

    # ── Preview first N samples ───────────────────────────────────────────────
    n_preview = min(args.preview, len(all_codes))
    print(f"\n     {'─' * 50}")
    print(f"     FIRST {n_preview} SAMPLES")
    print(f"     {'─' * 50}")
    for idx in range(n_preview):
        code_preview = all_codes[idx].replace('\n', ' ')[:80]
        gt = "VULNERABLE" if labels[idx] == 1 else "SAFE"
        pred = "VULNERABLE" if preds[idx] == 1 else "SAFE"
        match = "✅" if labels[idx] == preds[idx] else "❌"
        print(f"\n     [{idx+1}] {match}")
        print(f"       Source:       {all_source_experts[idx]} (is_expert={all_is_expert[idx]})")
        print(f"       Code:         {code_preview}...")
        print(f"       Ground truth: {gt}  |  Prediction: {pred}")
        print(f"       P(vuln):      {probs_vuln[idx]:.6f}  |  Logits: safe={logits[idx][0]:.4f} vuln={logits[idx][1]:.4f}")

    # ── Overall metrics ───────────────────────────────────────────────────────
    acc = accuracy_score(labels, preds)
    f1 = f1_score(labels, preds, average="binary", zero_division=0)
    prec = precision_score(labels, preds, average="binary", zero_division=0)
    rec = recall_score(labels, preds, average="binary", zero_division=0)
    try:
        auc = roc_auc_score(labels, probs_vuln) if len(set(labels)) > 1 else 0.0
    except ValueError:
        auc = 0.0

    cm = confusion_matrix(labels, preds, labels=[0, 1])

    print(f"\n     {'─' * 50}")
    print(f"     OVERALL METRICS")
    print(f"     {'─' * 50}")
    print(f"     Accuracy:  {acc:.4f}")
    print(f"     F1:        {f1:.4f}")
    print(f"     Precision: {prec:.4f}")
    print(f"     Recall:    {rec:.4f}")
    print(f"     AUC:       {auc:.4f}")
    print(f"\n     Confusion Matrix:")
    print(f"                      Pred SAFE  Pred VULN")
    print(f"     Actual SAFE      {cm[0][0]:>8}  {cm[0][1]:>8}")
    print(f"     Actual VULN      {cm[1][0]:>8}  {cm[1][1]:>8}")

    # ── Per source-expert breakdown ───────────────────────────────────────────
    print(f"\n     {'─' * 50}")
    print(f"     PER SOURCE DATASET BREAKDOWN")
    print(f"     {'─' * 50}")
    print(f"     {'Source':<35} {'N':>5} {'GT=1':>5} {'Pred=1':>6} {'TP':>4} {'FP':>4} {'TN':>4} {'FN':>4} {'Acc':>6} {'F1':>6} {'MeanP':>6}")
    print(f"     {'─' * 95}")

    per_source = {}
    for source_slug in EXPERT_DATASETS.keys():
        mask = [i for i, s in enumerate(all_source_experts) if s == source_slug]
        if not mask:
            continue

        s_labels = [labels[i] for i in mask]
        s_preds = [preds[i] for i in mask]
        s_probs = [probs_vuln[i] for i in mask]

        s_n = len(mask)
        s_gt1 = sum(s_labels)
        s_pred1 = sum(s_preds)
        s_tp = sum(1 for p, l in zip(s_preds, s_labels) if p == 1 and l == 1)
        s_fp = sum(1 for p, l in zip(s_preds, s_labels) if p == 1 and l == 0)
        s_tn = sum(1 for p, l in zip(s_preds, s_labels) if p == 0 and l == 0)
        s_fn = sum(1 for p, l in zip(s_preds, s_labels) if p == 0 and l == 1)
        s_acc = accuracy_score(s_labels, s_preds)
        s_f1 = f1_score(s_labels, s_preds, average="binary", zero_division=0)
        s_mean_p = np.mean(s_probs)

        is_self = " ◀" if source_slug == expert_slug else ""
        source_name = EXPERTS.get(source_slug, source_slug)[:33]
        print(f"     {source_name:<35} {s_n:>5} {s_gt1:>5} {s_pred1:>6} "
              f"{s_tp:>4} {s_fp:>4} {s_tn:>4} {s_fn:>4} "
              f"{s_acc:>6.3f} {s_f1:>6.3f} {s_mean_p:>6.3f}{is_self}")

        per_source[source_slug] = {
            "n": s_n, "gt_positive": s_gt1, "pred_positive": s_pred1,
            "tp": s_tp, "fp": s_fp, "tn": s_tn, "fn": s_fn,
            "accuracy": s_acc, "f1": s_f1, "mean_prob_vuln": float(s_mean_p),
        }

    # ── Diagnosis ─────────────────────────────────────────────────────────────
    self_data = per_source.get(expert_slug, {})
    print(f"\n     {'─' * 50}")
    print(f"     DIAGNOSIS for {expert_name}")
    print(f"     {'─' * 50}")

    if self_data:
        self_recall = self_data["tp"] / (self_data["tp"] + self_data["fn"]) if (self_data["tp"] + self_data["fn"]) > 0 else 0
        print(f"     Own type recall:   {self_recall:.3f} ({self_data['tp']}/{self_data['tp']+self_data['fn']} detected)")

    # False alarms on other types
    other_fp_total = sum(per_source[s]["fp"] for s in per_source if s != expert_slug)
    other_total = sum(per_source[s]["n"] for s in per_source if s != expert_slug)
    if other_total > 0:
        fp_rate = other_fp_total / other_total
        print(f"     False alarm rate on other types: {fp_rate:.3f} ({other_fp_total}/{other_total})")

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

    return {
        "expert_slug": expert_slug,
        "expert_name": expert_name,
        "checkpoint": checkpoint_path,
        "overall": {"accuracy": acc, "f1": f1, "precision": prec, "recall": rec, "auc": auc},
        "per_source": per_source,
    }


def main():
    args = parse_args()

    print("=" * 60)
    print("  Evaluate Classifiers on Combined Training Data")
    print("=" * 60)

    # ── Determine which experts to evaluate ───────────────────────────────────
    eval_tasks = []

    if args.all:
        base_dir = os.path.abspath(args.base_dir)
        print(f"\n🔍 Scanning: {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("❌ Provide --checkpoint + --expert, or use --all")
        return

    if not eval_tasks:
        print("❌ No checkpoints found!")
        return

    # ── Load ALL data ─────────────────────────────────────────────────────────
    print(f"\n📦 Loading all expert datasets (train + test)...")
    all_codes, all_source_experts, all_is_expert = load_all_data(args.max_samples_per_dataset)

    print(f"\n  Combined dataset: {len(all_codes)} total samples")
    source_dist = Counter(all_source_experts)
    for slug, count in source_dist.most_common():
        n_expert = sum(1 for s, e in zip(all_source_experts, all_is_expert) if s == slug and e == 1)
        print(f"    {EXPERTS[slug]:<35} {count:>6} samples ({n_expert} positive)")

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

    for checkpoint_path, expert_slug in eval_tasks:
        result = evaluate_one_expert(
            checkpoint_path, expert_slug,
            all_codes, all_source_experts, all_is_expert, args,
        )
        all_results.append(result)

    # ── Final summary ─────────────────────────────────────────────────────────
    if len(all_results) > 1:
        print(f"\n{'━' * 60}")
        print(f"  SUMMARY — ALL EXPERTS")
        print(f"{'━' * 60}")
        print(f"  {'Expert':<30} {'Acc':>6} {'F1':>6} {'Prec':>6} {'Rec':>6} {'AUC':>6}")
        print(f"  {'─' * 66}")
        for r in all_results:
            o = r["overall"]
            print(f"  {r['expert_name']:<30} {o['accuracy']:>6.3f} {o['f1']:>6.3f} "
                  f"{o['precision']:>6.3f} {o['recall']:>6.3f} {o['auc']:>6.3f}")

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


if __name__ == "__main__":
    main()