"""Independent verification of v2 ACO specialist models.

Recreates test splits with same seed, loads v1 and v2 models,
computes metrics, compares delta.

Key fix: v1 (DistilBERT) has max_position_embeddings=512, v2 (ModernBERT) has 8192.
We detect the limit from model config to avoid overflow.

Usage via hf_jobs:
    hf_jobs run --script verify_v2.py --deps transformers,torch,datasets,scikit-learn,huggingface_hub --hardware a10g-large --timeout 2h
"""
import torch, numpy as np, json, os, sys
from datasets import Dataset, load_dataset
from transformers import AutoTokenizer, AutoModelForSequenceClassification, AutoConfig
from sklearn.metrics import accuracy_score, f1_score, classification_report, precision_recall_fscore_support
from torch.utils.data import DataLoader

# ═══════════════════════════════════════════
# Constants
# ═══════════════════════════════════════════

V1_MODELS = {
    "tier_router": "narcolepticchicken/aco-specialists-tier-router",
    "tool_gater": "narcolepticchicken/aco-specialists-tool-gater",
    "verifier_gater": "narcolepticchicken/aco-specialists-verifier-gater",
}
V2_MODELS = {
    "tier_router": "narcolepticchicken/aco-specialists-tier-router-v2",
    "tool_gater": "narcolepticchicken/aco-specialists-tool-gater-v2",
    "verifier_gater": "narcolepticchicken/aco-specialists-verifier-gater-v2",
}
NUM_LABELS_MAP = {"tier_router": 3, "tool_gater": 2, "verifier_gater": 2}
TASK_NAMES = ["tier_router", "tool_gater", "verifier_gater"]

# ═══════════════════════════════════════════
# Dataset loaders (SAME as training script)
# ═══════════════════════════════════════════

def load_tool_gater():
    import re
    ds = load_dataset("lockon/ToolACE", split="train")
    t, l = [], []
    for row in ds:
        conv = row.get("conversations", [])
        q = ""
        for turn in conv:
            if turn.get("from") == "user":
                q = turn.get("value", "")[:1500]
                break
        if not q:
            continue
        called = any(re.search(r'\[[A-Z][a-zA-Z]+\s*\(', turn["value"]) 
                    for turn in conv if turn.get("from") == "assistant")
        text = f"Query: {q}"
        if row.get("system"):
            text = f"System: {row['system'][:500]}\n\n{text}"
        t.append(text[:2000])
        l.append(1 if called else 0)
    ds = Dataset.from_dict({"text": t, "labels": l})
    return ds.train_test_split(test_size=0.15, seed=42)

def load_tier_router():
    ds = load_dataset("RouteWorks/RouterArena", "default", split="full")
    tmap = {"easy": 0, "medium": 1, "hard": 2}
    t, l = [], []
    for row in ds:
        d = row.get("Difficulty", "").strip().lower()
        if d not in tmap:
            continue
        parts = []
        if row.get("Domain"):
            parts.append(f"[{row['Domain']}]")
        if row.get("Context"):
            parts.append(f"Context: {row['Context']}")
        parts.append(row.get("Question", ""))
        o = row.get("Options", "")
        if o:
            parts.append(f"Options: {'; '.join(o) if isinstance(o, list) else o}")
        t.append(" ".join(parts)[:2000])
        l.append(tmap[d])
    ds = Dataset.from_dict({"text": t, "labels": l})
    return ds.train_test_split(test_size=0.15, seed=42)

def load_verifier_gater():
    import re
    ds = load_dataset("R2E-Gym/R2EGym-Verifier-Trajectories", split="train")
    t, l = [], []
    for row in ds:
        messages = row["messages"]
        fl = messages[1]["content"] if len(messages) > 1 else ""
        task_text = ""
        for msg in messages:
            if msg["role"] == "user" and "INTERACTION LOG" in msg["content"]:
                m = re.search(r'<github_issue>(.*?)</github_issue>', msg["content"], re.DOTALL)
                if m:
                    task_text = m.group(1).strip()[:1000]
                    break
        if not task_text:
            for msg in messages:
                if msg["role"] == "system":
                    task_text = msg["content"][:500]
                    break
        ab = re.findall(r'\[ASSISTANT\](.*?)(?:\[USER\]|\[STEP\]|$)', fl, re.DOTALL)
        agent_sum = " ".join(b.strip()[:200] for b in ab[-3:])
        pm = re.search(r'=== FINAL PATCH ===\s*\n(.*?)\n=== END FINAL PATCH ===', fl, re.DOTALL)
        patch = pm.group(1)[:500] if pm else ""
        text = f"TASK: {task_text[:600]}\nAGENT_ACTIONS: {agent_sum[:600]}\nPATCH: {patch[:400]}"
        t.append(text[:2000])
        l.append(1 if row["rewards"] >= 1.0 else 0)
    ds = Dataset.from_dict({"text": t, "labels": l})
    return ds.train_test_split(test_size=0.15, seed=42)

LOADERS = {"tool_gater": load_tool_gater, "tier_router": load_tier_router, "verifier_gater": load_verifier_gater}

# ═══════════════════════════════════════════
# Evaluation
# ═══════════════════════════════════════════

def get_max_length(model_config):
    """Detect model's maximum position embeddings from config."""
    if hasattr(model_config, "max_position_embeddings"):
        return model_config.max_position_embeddings
    if hasattr(model_config, "n_positions"):
        return model_config.n_positions
    return 512  # safe default

def evaluate_model(model_name, task_name, num_labels):
    print(f"\n{'='*60}")
    print(f"EVALUATING: {model_name} [{task_name}]")
    print(f"{'='*60}")

    # Load data
    ds = LOADERS[task_name]()
    test_ds = ds["test"]
    print(f"  Test samples: {len(test_ds)}")

    # Class distribution
    lc = {}
    for lb in test_ds["labels"]:
        lc[lb] = lc.get(lb, 0) + 1
    print(f"  Class dist: {lc}")

    # Load model and detect max token length
    try:
        config = AutoConfig.from_pretrained(model_name)
        max_len = get_max_length(config)
        tokenizer = AutoTokenizer.from_pretrained(model_name)
        if tokenizer.pad_token is None:
            tokenizer.pad_token = tokenizer.eos_token
        model = AutoModelForSequenceClassification.from_pretrained(
            model_name, num_labels=num_labels, ignore_mismatched_sizes=True)
        model.eval()
        device = "cuda" if torch.cuda.is_available() else "cpu"
        model.to(device)
        print(f"  Model loaded on {device}, max_len={max_len}")
    except Exception as e:
        print(f"  FAILED to load model: {e}")
        import traceback; traceback.print_exc()
        return None

    # Extract threshold from config
    threshold = getattr(model.config, "threshold", 0.5)
    print(f"  Threshold from config: {threshold}")

    # Tokenize with model-specific max length
    texts = list(test_ds["text"])
    labels_list = list(test_ds["labels"])
    encodings = tokenizer(texts, truncation=True, max_length=max_len, padding=True)

    # Batch inference with DataLoader
    input_ids = torch.tensor(encodings["input_ids"])
    attention_mask = torch.tensor(encodings["attention_mask"])
    labels_t = torch.tensor(labels_list)
    ds_tensor = torch.utils.data.TensorDataset(input_ids, attention_mask, labels_t)
    loader = DataLoader(ds_tensor, batch_size=32, shuffle=False)

    all_probs = []
    all_labels = []
    with torch.no_grad():
        for batch in loader:
            b_input_ids, b_attention_mask, b_labels = [x.to(device) for x in batch]
            logits = model(input_ids=b_input_ids, attention_mask=b_attention_mask).logits
            probs = torch.softmax(logits, dim=-1).cpu().numpy()
            all_probs.append(probs)
            all_labels.extend(b_labels.cpu().numpy().tolist())

    probs = np.vstack(all_probs)
    labels = np.array(all_labels)

    # Default predictions
    preds_default = np.argmax(probs, axis=-1)
    acc_default = accuracy_score(labels, preds_default)
    f1_default = f1_score(labels, preds_default, average="macro", zero_division=0)

    print(f"  Default:  acc={acc_default:.4f}, f1_macro={f1_default:.4f}")

    if num_labels == 2:
        # Calibrated predictions
        preds_cal = (probs[:, 1] >= threshold).astype(int)
        acc_cal = accuracy_score(labels, preds_cal)
        f1_cal = f1_score(labels, preds_cal, average="macro", zero_division=0)

        # Per-class precision/recall
        p, r, f1_per, support_per = precision_recall_fscore_support(labels, preds_cal, zero_division=0)

        print(f"  Calibrated (t={threshold:.3f}): acc={acc_cal:.4f}, f1_macro={f1_cal:.4f}")

        # DETECT COLLAPSE
        unique_preds = np.unique(preds_cal)
        if len(unique_preds) == 1:
            majority_pct = (labels == unique_preds[0]).mean()
            print(f"  ⚠️  MAJORITY-CLASS COLLAPSE: predicts only class {unique_preds[0]} "
                  f"(base rate={majority_pct:.1%})")

        print(f"\n  Classification Report (calibrated):")
        print(f"  {classification_report(labels, preds_cal, target_names=['neg','pos'], zero_division=0, digits=4)}")
        print(f"  Per-class: neg P={p[0]:.4f} R={r[0]:.4f} F1={f1_per[0]:.4f} | pos P={p[1]:.4f} R={r[1]:.4f} F1={f1_per[1]:.4f}")

        return {
            "accuracy": acc_cal, "f1_macro": f1_cal,
            "accuracy_default": acc_default, "f1_default": f1_default,
            "threshold": threshold,
            "per_class": {
                "neg": {"precision": float(p[0]), "recall": float(r[0]), "f1": float(f1_per[0]), "support": int(support_per[0])},
                "pos": {"precision": float(p[1]), "recall": float(r[1]), "f1": float(f1_per[1]), "support": int(support_per[1])},
            },
            "collapsed": len(unique_preds) == 1,
            "class_dist": lc,
        }
    else:
        # Multi-class
        p, r, f1_per, support_per = precision_recall_fscore_support(labels, preds_default, zero_division=0)
        print(f"\n  Classification Report:")
        print(f"  {classification_report(labels, preds_default, zero_division=0, digits=4)}")

        per_class = {}
        for i in range(num_labels):
            per_class[str(i)] = {"precision": float(p[i]), "recall": float(r[i]), "f1": float(f1_per[i]), "support": int(support_per[i])}

        return {
            "accuracy": acc_default, "f1_macro": f1_default,
            "threshold": None,
            "per_class": per_class,
            "collapsed": np.unique(preds_default).size == 1,
            "class_dist": lc,
        }

# ═══════════════════════════════════════════
# Main
# ═══════════════════════════════════════════

def main():
    results = {}

    for task_name in TASK_NAMES:
        num_labels = NUM_LABELS_MAP[task_name]

        # Evaluate v2
        print(f"\n{'#'*60}")
        print(f"### V2 MODEL: {task_name}")
        print(f"{'#'*60}")
        v2_res = evaluate_model(V2_MODELS[task_name], task_name, num_labels)

        # Evaluate v1
        print(f"\n{'#'*60}")
        print(f"### V1 MODEL: {task_name} (baseline)")
        print(f"{'#'*60}")
        v1_res = evaluate_model(V1_MODELS[task_name], task_name, num_labels)

        if v2_res is not None and v1_res is not None:
            delta_f1 = v2_res["f1_macro"] - v1_res["f1_macro"]
            delta_acc = v2_res["accuracy"] - v1_res["accuracy"]
            print(f"\n  >>> v1 → v2 delta: F1 {v1_res['f1_macro']:.4f} → {v2_res['f1_macro']:.4f} = {delta_f1:+.4f}")
            print(f"  >>> v1 → v2 delta: Acc {v1_res['accuracy']:.4f} → {v2_res['accuracy']:.4f} = {delta_acc:+.4f}")
            results[task_name] = {"v1": v1_res, "v2": v2_res, "delta_f1": delta_f1, "delta_acc": delta_acc}
        else:
            results[task_name] = {"v1": v1_res, "v2": v2_res, "error": "One or both models failed"}

    # Final summary
    print(f"\n{'='*60}")
    print("FINAL COMPARISON")
    print(f"{'='*60}")

    for tn in TASK_NAMES:
        r = results.get(tn, {})
        v1_c = r.get("v1", {}).get("collapsed", True) if r.get("v1") else True
        v2_c = r.get("v2", {}).get("collapsed", True) if r.get("v2") else True
        delta = r.get("delta_f1", float("nan"))

        status = "OK"
        if v2_c:
            status = "⚠️  V2 COLLAPSED"
        elif v1_c:
            status = "⚠️  V1 COLLAPSED"

        print(f"  {tn:<20} v1_f1={r.get('v1',{}).get('f1_macro',float('nan')):.4f}  "
              f"v2_f1={r.get('v2',{}).get('f1_macro',float('nan')):.4f}  "
              f"delta={delta:+.4f}  {status}")

    print(json.dumps(results, indent=2, default=str))

    # Save results
    with open("/tmp/v2_verification_results.json", "w") as f:
        json.dump(results, f, indent=2, default=str)

    # Push results
    from huggingface_hub import HfApi
    api = HfApi()
    api.upload_file(
        path_or_fileobj="/tmp/v2_verification_results.json",
        path_in_repo="v2_verification_results.json",
        repo_id="narcolepticchicken/agent-cost-optimizer",
        repo_type="model",
    )
    print("\nResults pushed to narcolepticchicken/agent-cost-optimizer")

if __name__ == "__main__":
    main()