"""v6.1 drift detector retrain — adds `arg_order_inversion_rate`
(23rd feature) to fix the documented arg_swap failure mode.

Feature: for each argument pair where both call values have a detectable
first-appearance position in the prompt, check whether the schema's
declared order of the arg keys disagrees with the prompt-order of their
values. Report fraction of inverted pairs.

Signal probe (scripts/drift_feature_probe_arg_order.py):
  arg_swap mean: 0.4152   gold mean: 0.1662   delta: +0.249
  coverage on arg_swap: 53.3%

Pipeline:
  drift_v0 dataset (n=3700)
  -> 23-dim feature vector per sample
  -> StandardScaler + LogisticRegression (C=1.0, class_weight=balanced)
  -> 5-fold stratified CV, seed=0
  -> pooled AUC + per-drift-type AUC (target: arg_swap >= 0.85)

Output: benchmarks/drift_calibrated_v1.json
"""
from __future__ import annotations

import json
import math
import re
from pathlib import Path
from typing import Dict, List

import numpy as np
from sklearn.linear_model import LogisticRegression
from sklearn.metrics import roc_auc_score
from sklearn.model_selection import StratifiedKFold
from sklearn.preprocessing import StandardScaler

REPO = Path(__file__).resolve().parents[1]
DATA_PATH = REPO / "data" / "drift_v0" / "drift_dataset_v0.jsonl"
OUT_PATH = REPO / "benchmarks" / "drift_calibrated_v1.json"

RANDOM_STATE = 0


def load_dataset():
    rows = []
    with open(DATA_PATH, encoding="utf-8") as f:
        for line in f:
            line = line.strip()
            if line:
                rows.append(json.loads(line))
    return rows


STOPWORDS = {
    "a", "an", "the", "and", "or", "but", "if", "of", "for", "to", "in",
    "on", "at", "by", "with", "from", "is", "are", "was", "were", "be",
    "been", "being", "have", "has", "had", "do", "does", "did", "will",
    "would", "can", "could", "should", "may", "might", "must",
    "this", "that", "these", "those", "i", "you", "he", "she", "it", "we",
    "they", "me", "him", "her", "us", "them", "my", "your", "his", "our",
    "their", "what", "which", "who", "how", "when", "where", "why",
}

WORD_RE = re.compile(r"[A-Za-z0-9_]+")
PLACEHOLDER_RE = re.compile(
    r"^(placeholder|example|test|_[a-z_]+|<[^>]+>)$", re.IGNORECASE
)


def tokens(text):
    return [w.lower() for w in WORD_RE.findall(text or "")]


def content_tokens(text):
    return [t for t in tokens(text) if t not in STOPWORDS and len(t) > 1]


def value_first_position(value, prompt_tokens):
    vtoks = set(tokens(str(value)))
    if not vtoks:
        return None
    for i, tok in enumerate(prompt_tokens):
        if tok in vtoks:
            return i
    return None


def arg_order_inversion_rate(prompt, functions, tool_call):
    args = (tool_call or {}).get("arguments") or {}
    if len(args) < 2:
        return 0.0
    schema = next(
        (fn for fn in (functions or []) if fn.get("name") == tool_call.get("name")),
        None,
    )
    if schema is None:
        return 0.0
    props = (schema.get("parameters") or {}).get("properties") or {}
    if not props:
        return 0.0
    schema_order = {k: i for i, k in enumerate(props.keys())}

    toks = tokens(prompt)
    positions = {}
    for k, v in args.items():
        if k not in schema_order:
            continue
        p = value_first_position(v, toks)
        if p is not None:
            positions[k] = p
    if len(positions) < 2:
        return 0.0

    keys = list(positions.keys())
    inv = 0
    total = 0
    for i in range(len(keys)):
        for j in range(i + 1, len(keys)):
            ki, kj = keys[i], keys[j]
            s_i_before_j = schema_order[ki] < schema_order[kj]
            p_i_before_j = positions[ki] < positions[kj]
            total += 1
            if s_i_before_j != p_i_before_j:
                inv += 1
    return inv / total if total else 0.0


def extract_features(r):
    prompt = r.get("prompt", "") or ""
    call = r.get("tool_call", {}) or {}
    call_name = (call.get("name") or "").lower()
    call_args = call.get("arguments") or {}
    funcs = r.get("functions", []) or []

    prompt_tokens = set(tokens(prompt))
    prompt_content = set(content_tokens(prompt))

    f_tool_in_prompt = 1.0 if call_name and call_name in prompt_tokens else 0.0
    tool_parts = set(call_name.split("_"))
    f_tool_parts_in_prompt = (
        len(tool_parts & prompt_content) / max(1, len(tool_parts))
        if tool_parts else 0.0
    )
    call_text_tokens = set()
    for v in call_args.values():
        call_text_tokens.update(tokens(str(v)))
    call_text_tokens.update(tokens(call_name))
    union = prompt_content | call_text_tokens
    f_overlap_jaccard = (
        len(prompt_content & call_text_tokens) / max(1, len(union))
        if union else 0.0
    )
    f_prompt_coverage = (
        len(prompt_content & call_text_tokens) / max(1, len(prompt_content))
    )
    if call_args:
        hits = total = 0
        for v in call_args.values():
            vt = set(tokens(str(v)))
            if not vt:
                continue
            total += 1
            if vt & prompt_tokens:
                hits += 1
        f_arg_verbatim_rate = hits / total if total else 0.0
    else:
        f_arg_verbatim_rate = 0.0

    schema = next((f for f in funcs if f.get("name") == call.get("name")), None)
    if schema is None:
        f_tool_in_schema = 0.0
        f_missing_required_frac = 1.0
        f_spurious_arg_frac = 1.0
        f_type_mismatch_frac = 1.0
        f_arg_count_zscore = 0.0
        f_required_count = 0.0
    else:
        f_tool_in_schema = 1.0
        props = (schema.get("parameters") or {}).get("properties") or {}
        required = (schema.get("parameters") or {}).get("required") or []
        spec_args = set(props.keys())
        req_args = set(required)
        called = set(call_args.keys())
        f_missing_required_frac = len(req_args - called) / max(1, len(req_args))
        f_spurious_arg_frac = len(called - spec_args) / max(1, len(spec_args))
        mismatches = 0
        checks = 0
        for k, v in call_args.items():
            if k not in props:
                continue
            checks += 1
            declared = (props[k] or {}).get("type", "")
            if declared in ("integer", "int"):
                if isinstance(v, bool) or not isinstance(v, int):
                    try:
                        int(str(v).replace(",", ""))
                    except (ValueError, TypeError):
                        mismatches += 1
            elif declared in ("number", "float"):
                try:
                    float(str(v).replace(",", ""))
                except (ValueError, TypeError):
                    mismatches += 1
            elif declared == "boolean":
                if not isinstance(v, bool):
                    mismatches += 1
            elif declared in ("array", "list"):
                if not isinstance(v, list):
                    mismatches += 1
            elif declared in ("object", "dict"):
                if not isinstance(v, dict):
                    mismatches += 1
        f_type_mismatch_frac = mismatches / max(1, checks)
        f_arg_count_zscore = (
            (len(called) - len(spec_args))
            / max(1.0, math.sqrt(max(1, len(spec_args))))
        )
        f_required_count = float(len(req_args))

    f_arg_order_inversion = arg_order_inversion_rate(prompt, funcs, call)

    if call_args:
        placeholders = sum(
            1 for v in call_args.values()
            if isinstance(v, str) and PLACEHOLDER_RE.match(v.strip())
        )
        f_placeholder_frac = placeholders / len(call_args)
    else:
        f_placeholder_frac = 0.0
    f_tool_name_len = math.log(max(1, len(call_name)))
    f_tool_in_any_schema = (
        1.0 if call_name in {f.get("name", "").lower() for f in funcs} else 0.0
    )
    f_n_available_tools = math.log(max(1, len(funcs)))

    f_n_args_called = float(len(call_args))
    f_prompt_len = math.log(max(1, len(prompt)))
    if call_args:
        avg_arg_len = sum(len(str(v)) for v in call_args.values()) / len(call_args)
    else:
        avg_arg_len = 0.0
    f_avg_arg_len = math.log(max(1.0, avg_arg_len + 1))
    f_has_nested = 1.0 if any(isinstance(v, dict) for v in call_args.values()) else 0.0
    f_has_list = 1.0 if any(isinstance(v, list) for v in call_args.values()) else 0.0
    f_prompt_is_question = 1.0 if "?" in prompt else 0.0
    f_prompt_imperative = (
        1.0 if (prompt[:1].isupper()
                and not prompt.startswith(
                    ("What", "Why", "How", "When", "Where", "Who", "Which",
                     "Can ", "Could ", "Would ", "Should ")
                ))
        else 0.0
    )

    return [
        # A — semantic alignment (5)
        f_tool_in_prompt, f_tool_parts_in_prompt, f_overlap_jaccard,
        f_prompt_coverage, f_arg_verbatim_rate,
        # B — schema conformance (7, +1 new)
        f_tool_in_schema, f_missing_required_frac, f_spurious_arg_frac,
        f_type_mismatch_frac, f_arg_count_zscore, f_required_count,
        f_arg_order_inversion,  # NEW
        # C — lexical drift (4)
        f_placeholder_frac, f_tool_name_len, f_tool_in_any_schema,
        f_n_available_tools,
        # D — structural (7)
        f_n_args_called, f_prompt_len, f_avg_arg_len, f_has_nested,
        f_has_list, f_prompt_is_question, f_prompt_imperative,
    ]


FEATURE_NAMES = [
    "tool_in_prompt", "tool_parts_in_prompt", "overlap_jaccard",
    "prompt_coverage", "arg_verbatim_rate",
    "tool_in_schema", "missing_required_frac", "spurious_arg_frac",
    "type_mismatch_frac", "arg_count_zscore", "required_count",
    "arg_order_inversion",
    "placeholder_frac", "tool_name_len", "tool_in_any_schema",
    "n_available_tools",
    "n_args_called", "prompt_len", "avg_arg_len", "has_nested",
    "has_list", "prompt_is_question", "prompt_imperative",
]


def main():
    rows = load_dataset()
    print(f"loaded {len(rows)} samples")

    X = np.array([extract_features(r) for r in rows])
    y = np.array([r["drift"] for r in rows])
    sources = np.array([r["source"] for r in rows])
    drift_types = np.array([r["drift_type"] for r in rows])

    print(f"  feature matrix: {X.shape}")
    print(f"  n_features: {len(FEATURE_NAMES)}")
    assert X.shape[1] == len(FEATURE_NAMES), "feature-count mismatch"

    skf = StratifiedKFold(n_splits=5, shuffle=True, random_state=RANDOM_STATE)
    fold_aucs = []
    all_scores = []
    all_labels = []
    all_idx = []

    for fold, (tr, te) in enumerate(skf.split(X, y)):
        sc = StandardScaler()
        X_tr = sc.fit_transform(X[tr])
        X_te = sc.transform(X[te])
        clf = LogisticRegression(
            C=1.0, max_iter=2000, random_state=RANDOM_STATE,
            class_weight="balanced",
        )
        clf.fit(X_tr, y[tr])
        p = clf.predict_proba(X_te)[:, 1]
        auc = roc_auc_score(y[te], p)
        fold_aucs.append(auc)
        all_scores.extend(p)
        all_labels.extend(y[te])
        all_idx.extend(te)
        print(f"  fold {fold+1}: AUC {auc:.4f}")

    mean_auc = float(np.mean(fold_aucs))
    std_auc = float(np.std(fold_aucs))
    pooled = roc_auc_score(all_labels, all_scores)
    print()
    print(f"  MEAN fold AUC: {mean_auc:.4f} ± {std_auc:.4f}")
    print(f"  POOLED AUC:    {pooled:.4f}")

    scores_arr = np.array(all_scores)
    labels_arr = np.array(all_labels)
    idx_arr = np.array(all_idx)

    source_aucs = {}
    for src in sorted(set(sources)):
        mask = sources[idx_arr] == src
        if mask.sum() == 0:
            continue
        sub_l = labels_arr[mask]
        sub_s = scores_arr[mask]
        if len(set(sub_l)) < 2:
            source_aucs[src] = None
            continue
        source_aucs[src] = float(roc_auc_score(sub_l, sub_s))
    print()
    print("  per-source AUC:")
    for src, a in source_aucs.items():
        print(f"    {src:<20s}  {a:.4f}" if a is not None else f"    {src:<20s}  n/a")

    gold_mask = drift_types == "gold"
    type_aucs = {}
    for dt in sorted(set(drift_types)):
        if dt == "gold":
            continue
        mask = (drift_types[idx_arr] == dt) | gold_mask[idx_arr]
        sub_l = labels_arr[mask]
        sub_s = scores_arr[mask]
        if len(set(sub_l)) < 2:
            type_aucs[dt] = None
            continue
        type_aucs[dt] = float(roc_auc_score(sub_l, sub_s))
    print()
    print("  per-drift-type AUC (vs gold):")
    for dt, a in type_aucs.items():
        val = f"{a:.4f}" if a is not None else "n/a"
        print(f"    {dt:<22s}  {val}")

    # Fit final on all for export
    sc = StandardScaler()
    X_all = sc.fit_transform(X)
    clf = LogisticRegression(
        C=1.0, max_iter=2000, random_state=RANDOM_STATE,
        class_weight="balanced",
    )
    clf.fit(X_all, y)

    print()
    print("  feature importance (scaled, |coef| desc):")
    pairs = sorted(zip(FEATURE_NAMES, clf.coef_[0]), key=lambda kv: -abs(kv[1]))
    for name, c in pairs:
        print(f"    {name:<25s}  {c:+.3f}")
    print(f"    intercept               {clf.intercept_[0]:+.3f}")

    result = {
        "methodology": "drift v1 — 23-feature calibrated LR, 5-fold stratified CV, seed=0",
        "new_feature": "arg_order_inversion — positional inversion rate of call values in prompt vs schema arg order",
        "n_samples": len(rows),
        "n_features": len(FEATURE_NAMES),
        "feature_names": FEATURE_NAMES,
        "fold_aucs": fold_aucs,
        "mean_auc": mean_auc,
        "std_auc": std_auc,
        "pooled_auc": float(pooled),
        "per_source_auc": source_aucs,
        "per_drift_type_auc": type_aucs,
        "coefficients": dict(zip(FEATURE_NAMES, clf.coef_[0].tolist())),
        "intercept": float(clf.intercept_[0]),
        "scaler_mean": sc.mean_.tolist(),
        "scaler_scale": sc.scale_.tolist(),
    }
    OUT_PATH.parent.mkdir(parents=True, exist_ok=True)
    OUT_PATH.write_text(json.dumps(result, indent=2), encoding="utf-8")
    print()
    print(f"wrote -> {OUT_PATH.relative_to(REPO)}")


if __name__ == "__main__":
    main()
