"""Score paired code with all generated intermediate fields neutralized."""
from __future__ import annotations

import argparse
import hashlib
import json
from pathlib import Path

import numpy as np
import pandas as pd
import torch
from peft import PeftModel
from transformers import AutoModelForCausalLM, AutoTokenizer

from scripts.audit_sft_prefix_sensitivity import score_batch
from scripts.eval_graph_feature_classifier import metrics
from src.contracts.prompt import SYSTEM_PROMPT_V2, build_user_prompt, code_for_sample


NEUTRAL_PREFIX = ("<analysis>\n<explanation></explanation>\n"
                  "<statements></statements>\n<trace>\n</trace>\n"
                  "<cwe>NONE</cwe>\n<label>")


def assert_valid_pair_membership(pair_ids, valid_pair_ids):
    outside_valid = sorted(set(pair_ids) - set(valid_pair_ids))
    if outside_valid:
        raise ValueError(f"source sample pairs are outside chronological VALID: {outside_valid[:5]}")


def run(input_path=Path("results/phase7/prefix_sensitivity_200.json"),
        corpus_path=Path("data/processed/corpus.parquet"),
        checkpoint=Path("checkpoints/sft-contract-v2/checkpoint-500"),
        output_path=Path("results/phase7/code_only_label_signal_200.json"),
        batch_size=4, seed=42, replicates=10000, with_adapter=True):
    prior = json.loads(input_path.read_text())
    sample_ids = [r["sample_id"] for r in prior["per_sample"]]
    frame = pd.read_parquet(corpus_path, columns=["sample_id", "pair_id", "target", "func_vuln", "func_patch"])
    by_sample = frame.set_index("sample_id", drop=False)
    pair_ids = sorted({str(by_sample.loc[sid]["pair_id"]) for sid in sample_ids})
    valid_pair_ids = set(json.loads(Path("data/splits/chronological/valid.json").read_text()))
    assert_valid_pair_membership(pair_ids, valid_pair_ids)
    pair_frame = frame[frame.pair_id.astype(str).isin(pair_ids)]
    pairs = {}
    for pid, group in pair_frame.groupby("pair_id", sort=False):
        if len(group) != 2 or set(group.target.astype(int)) != {0, 1}:
            raise ValueError(f"pair integrity failure: {pid}")
        pairs[str(pid)] = {int(row.target): row for row in group.itertuples(index=False)}

    entries, meta = [], {}
    for pid in pair_ids:
        for target in (0, 1):
            row = pairs[pid][target]
            prompt = f"{SYSTEM_PROMPT_V2}\n\n{build_user_prompt(code_for_sample(row._asdict()), teacher_mode=False)}"
            key = f"{pid}:code{target}"
            entries.append({"key": key, "prompt": prompt, "prefix": NEUTRAL_PREFIX})
            meta[key] = {"pair_id": pid, "target": target}

    base_revision = "1cfa9a7208912126459214e8b04321603b3df60c"
    tokenizer = AutoTokenizer.from_pretrained(
        checkpoint if with_adapter else "Qwen/Qwen3-4B",
        revision=None if with_adapter else base_revision,
        local_files_only=True)
    tokenizer.pad_token = tokenizer.pad_token or tokenizer.eos_token
    tokenizer.padding_side = "left"
    base = AutoModelForCausalLM.from_pretrained(
        "Qwen/Qwen3-4B", torch_dtype=torch.bfloat16,
        revision=base_revision, attn_implementation="sdpa", local_files_only=True).cuda()
    model = (PeftModel.from_pretrained(base, checkpoint, local_files_only=True)
             if with_adapter else base).eval()
    scored = score_batch(model, tokenizer, entries, model.device, batch_size)
    by_pair = {pid: {} for pid in pair_ids}
    for key, info in meta.items():
        values = scored[key]
        by_pair[info["pair_id"]][info["target"]] = {
            "logp_mean_safe": values["SAFE"], "logp_mean_vulnerable": values["VULNERABLE"],
            "vulnerable_margin": values["VULNERABLE"] - values["SAFE"]}

    records = []
    y, pred = [], []
    pair_correct, margins = [], []
    for pid in pair_ids:
        for target in (0, 1):
            margin = by_pair[pid][target]["vulnerable_margin"]
            prediction = int(margin > 0)
            y.append(target); pred.append(prediction)
            records.append({"pair_id": pid, "target": target, "prediction": prediction,
                            "vulnerable_margin": margin})
        difference = (by_pair[pid][1]["vulnerable_margin"] -
                      by_pair[pid][0]["vulnerable_margin"])
        pair_correct.append(difference > 0)
        margins.append(difference)

    rng = np.random.default_rng(seed)
    pair_correct_arr = np.asarray(pair_correct, dtype=float)
    draws = pair_correct_arr[rng.integers(0, len(pair_correct_arr),
                                          size=(replicates, len(pair_correct_arr)))].mean(axis=1)
    result = {"purpose": "frozen-policy inference with neutral explanation/statements/trace/CWE scaffold",
              "policy": "pinned base plus frozen SFT adapter" if with_adapter else "pinned base only",
              "checkpoint": str(checkpoint) if with_adapter else None,
              "base_revision": base_revision, "rows": len(y), "pairs": len(pair_ids),
              "condition_count": len(entries), "score": "mean token log-probability label margin, VULNERABLE minus SAFE",
              "row_metrics": metrics(y, pred, len(pair_ids)),
              "pair_rank_accuracy": float(pair_correct_arr.mean()),
              "pair_rank_bootstrap_95_ci": [float(x) for x in np.quantile(draws, [.025, .975])],
              "mean_pair_code_margin": float(np.mean(margins)),
              "prior_stored_student_accuracy": prior["sampled_student_label_accuracy"],
              "per_row": records,
              "data": {"sample_source_sha256": hashlib.sha256(input_path.read_bytes()).hexdigest(),
                       "corpus_sha256": hashlib.sha256(corpus_path.read_bytes()).hexdigest(),
                       "valid_split_sha256": hashlib.sha256(Path("data/splits/chronological/valid.json").read_bytes()).hexdigest(),
                       "valid_pair_membership_checked": True,
                       "adapter_sha256": hashlib.sha256((checkpoint / "adapter_model.safetensors").read_bytes()).hexdigest() if with_adapter else None,
                       "test_read": False},
              "limitations": ["Neutral empty fields form a counterfactual prompt context and may be out of distribution.",
                              "Same 200-row sample's 187 pairs; one frozen checkpoint; no tuning or training."]}
    output_path.parent.mkdir(parents=True, exist_ok=True)
    output_path.write_text(json.dumps(result, indent=2, sort_keys=True) + "\n")
    return result


def main():
    p = argparse.ArgumentParser(description=__doc__)
    p.add_argument("--output", type=Path, default=Path("results/phase7/code_only_label_signal_200.json"))
    p.add_argument("--batch-size", type=int, default=4)
    p.add_argument("--base-only", action="store_true",
                   help="score with the pinned base model, without the SFT adapter")
    a = p.parse_args()
    if a.base_only and a.output == Path("results/phase7/code_only_label_signal_200.json"):
        a.output = Path("results/phase7/code_only_base_model_200.json")
    print(json.dumps(run(output_path=a.output, batch_size=a.batch_size,
                         with_adapter=not a.base_only), indent=2))


if __name__ == "__main__":
    main()