"""Paired answer-token extraction at exactly the decoder outputs we later steer."""

from __future__ import annotations

import hashlib
import json
from collections.abc import Iterable, Sequence
import torch
from .layers import decoder_layers, hidden_tensor
from .vectors import AffectVector


def answer_tokens(tokenizer, situation: str, answer: str, maximum: int = 1024):
    if not answer.strip():
        raise ValueError("An empty answer has no extraction region")
    prompt = tokenizer.apply_chat_template(
        [{"role": "user", "content": situation}],
        tokenize=False,
        add_generation_prompt=True,
        enable_thinking=False,
    )
    prefix = tokenizer(prompt, add_special_tokens=False)["input_ids"]
    full = tokenizer(prompt + answer, add_special_tokens=False)["input_ids"]
    # A BPE merge at the prompt/answer boundary can change the final prefix
    # token. Start at the first changed token, never an assumed character index.
    start = 0
    while start < min(len(prefix), len(full)) and prefix[start] == full[start]:
        start += 1
    if start >= len(full) or len(full) > maximum:
        raise ValueError(
            "No answer tokens or extraction example exceeds maximum length"
        )
    return full, start


@torch.no_grad()
def answer_statistics(model, tokenizer, situation, answer, layers):
    ids, start = answer_tokens(tokenizer, situation, answer)
    blocks = decoder_layers(model)
    captured = {}
    handles = []

    def hook(layer):
        def record(_module, _inputs, output):
            h = hidden_tensor(output)[0, start:].detach().float()
            captured[layer] = (h.mean(dim=0).cpu(), h.norm(dim=-1).mean().item())

        return record

    try:
        for layer in layers:
            handles.append(blocks[layer].register_forward_hook(hook(layer)))
        device = model.get_input_embeddings().weight.device
        model(input_ids=torch.tensor([ids], device=device), use_cache=False)
    finally:
        for handle in handles:
            handle.remove()
    if set(captured) != set(layers):
        raise RuntimeError("Decoder adapter did not observe every requested layer")
    return captured


def extract_vector(
    model,
    tokenizer,
    pairs: Iterable[dict],
    layers: Sequence[int],
    name="pain",
    log=None,
):
    layers = list(layers)
    blocks = decoder_layers(model)
    if (
        not layers
        or len(set(layers)) != len(layers)
        or any(i < 0 or i >= len(blocks) for i in layers)
    ):
        raise ValueError("Invalid extraction layer selection")
    sums = {}
    squares = {}
    norms = {i: 0.0 for i in layers}
    count = 0
    digest = hashlib.sha256()
    was_training = model.training
    model.eval()
    try:
        for pair in pairs:
            if pair.get("split", "train") != "train":
                raise ValueError("Extraction only accepts the training split")
            target = pair.get("target", pair.get(name))
            if target is None:
                raise ValueError(f"Missing contrastive answer for {name}")
            if pair.get("channel", name) != name:
                raise ValueError("Mixed channels in one extraction run")
            record = {
                "situation": pair["situation"],
                "calm": pair["calm"],
                "target": target,
            }
            digest.update(
                (json.dumps(record, sort_keys=True, ensure_ascii=False) + "\n").encode()
            )
            baseline = answer_statistics(
                model, tokenizer, record["situation"], record["calm"], layers
            )
            affected = answer_statistics(
                model, tokenizer, record["situation"], target, layers
            )
            count += 1
            for layer in layers:
                delta = affected[layer][0] - baseline[layer][0]
                sums[layer] = sums.get(layer, torch.zeros_like(delta)) + delta
                squares[layer] = (
                    squares.get(layer, torch.zeros_like(delta)) + delta.square()
                )
                norms[layer] += (baseline[layer][1] + affected[layer][1]) * 0.5
            if log and count % 10 == 0:
                log(f"{name}: {count} contrastive pairs")
    finally:
        model.train(was_training)
    if not count:
        raise ValueError("No training pairs")
    directions = {}
    stats = {}
    for layer in layers:
        mean = sums[layer] / count
        magnitude = mean.norm().item()
        if magnitude < 1e-8:
            raise ValueError(f"No contrastive signal at layer {layer}")
        directions[layer] = mean / magnitude
        norms[layer] /= count
        variance = (squares[layer] / count - mean.square()).clamp_min(0)
        stats[str(layer)] = {
            "mean_difference_norm": magnitude,
            "coordinate_variance_mean": variance.mean().item(),
            "standard_error_l2": (variance.sum() / count).sqrt().item(),
        }
    config = model.config
    metadata = {
        "release": "CODELOCK-8Bbeta-v0.2",
        "signal": name,
        "pairs": count,
        "split": "train",
        "base_model": getattr(config, "_name_or_path", "unknown"),
        "model_revision": getattr(config, "_commit_hash", None),
        "hidden_size": next(iter(directions.values())).numel(),
        "num_hidden_layers": len(blocks),
        "layers": layers,
        "data_sha256": digest.hexdigest(),
        "capture_site": "decoder block output, before any final model norm",
        "method": "paired mean difference over answer tokens, unit direction per layer",
        "reference_norm": "mean per-answer-token L2 activation norm",
        "statistics": stats,
    }
    return AffectVector(directions, norms, metadata).validate()