#!/usr/bin/env python3
"""Run one preregistered Round 08 train-time uniform-forward cell."""

from __future__ import annotations

import importlib.util
import json
import math
import os
import sys
from pathlib import Path
from typing import Any

import torch
import torch.nn.functional as F


PROTOCOL_ID = "llm-atlas-k3-attnres-forward-training-v1"
PARENT_PROTOCOL_ID = "llm-atlas-k3-attnres-gradient-scale-v1"
VARIANTS = {
    "learned_reference": (),
    "uniform_group_6_forward": tuple(range(40, 48)),
    "uniform_group_7_forward": tuple(range(48, 56)),
    "uniform_groups_6_7_forward": tuple(range(40, 56)),
    "uniform_group_7_mlp_forward": (49, 51, 53, 55),
}
FORMAL_VARIANTS = tuple(name for name in VARIANTS if name != "learned_reference")
EXPECTED_SOURCE_COUNTS = {
    **{40: 6},
    **{index: 7 for index in range(41, 49)},
    **{index: 8 for index in range(49, 56)},
}


def load_parent_module() -> Any:
    path = Path(__file__).resolve().parents[1] / "attnres_gradient" / "train.py"
    spec = importlib.util.spec_from_file_location("k3_attnres_round05_train", path)
    if spec is None or spec.loader is None:
        raise RuntimeError(f"cannot import Round 05 runner from {path}")
    module = importlib.util.module_from_spec(spec)
    sys.modules[spec.name] = module
    spec.loader.exec_module(module)
    return module


parent = load_parent_module()
ACTIVE_VARIANT = "learned_reference"
LAST_MODEL: ForwardInterventionLanguageModel | None = None
LAST_OPTIMIZER: torch.optim.Optimizer | None = None


def extract_wrapper_argument(name: str) -> str:
    try:
        index = sys.argv.index(name)
    except ValueError as error:
        raise ValueError(f"missing required wrapper argument: {name}") from error
    if index + 1 >= len(sys.argv):
        raise ValueError(f"missing value for wrapper argument: {name}")
    value = sys.argv[index + 1]
    del sys.argv[index : index + 2]
    return value


def argument_value(name: str, default: str | None = None) -> str | None:
    try:
        index = sys.argv.index(name)
    except ValueError:
        return default
    if index + 1 >= len(sys.argv):
        raise ValueError(f"missing value for argument: {name}")
    return sys.argv[index + 1]


def parameter_names_for_indices(indices: tuple[int, ...]) -> tuple[str, ...]:
    names = []
    for index in indices:
        names.extend(
            (
                f"mixers.{index}.query",
                f"mixers.{index}.key_norm.weight",
            )
        )
    return tuple(names)


class ForwardInterventionLanguageModel(parent.GradientLanguageModel):
    """Round 05 model with one frozen selector and parameter-free uniform mixers."""

    def __init__(self, architecture: str):
        super().__init__(architecture)
        global LAST_MODEL
        if architecture != "block":
            raise ValueError("Round 08 only permits the block architecture")
        if ACTIVE_VARIANT not in VARIANTS:
            raise ValueError(f"unknown Round 08 variant: {ACTIVE_VARIANT}")
        self.forward_variant = ACTIVE_VARIANT
        self.selected_indices = tuple(VARIANTS[ACTIVE_VARIANT])
        self.selected_set = frozenset(self.selected_indices)
        self.forward_calls = 0
        self.depth_visits = [0] * len(self.mixers)
        self.output_visits = 0
        self.source_counts: dict[int, set[int]] = {
            index: set() for index in range(len(self.mixers))
        }
        self.uniform_weight_max_abs_error = 0.0
        selected_names = parameter_names_for_indices(self.selected_indices)
        named_parameters = dict(self.named_parameters())
        self.selected_initial_tensors = {
            name: named_parameters[name].detach().cpu().clone()
            for name in selected_names
        }
        self.gradient_hook_calls = {
            name: 0
            for name in named_parameters
            if name.startswith("mixers.") or name.startswith("output_mixer.")
        }
        self._gradient_hooks = []
        for name, parameter in named_parameters.items():
            if name not in self.gradient_hook_calls:
                continue

            def count_hook(
                gradient: torch.Tensor, *, parameter_name: str = name
            ) -> torch.Tensor:
                self.gradient_hook_calls[parameter_name] += 1
                return gradient

            self._gradient_hooks.append(parameter.register_hook(count_hook))
        LAST_MODEL = self

    def mix(
        self,
        mixer_index: int,
        sources: list[torch.Tensor],
        capture: bool,
    ) -> tuple[torch.Tensor, dict[str, Any] | None]:
        self.depth_visits[mixer_index] += 1
        self.source_counts[mixer_index].add(len(sources))
        if mixer_index not in self.selected_set:
            return self.mixers[mixer_index](sources, capture)

        values = torch.stack(sources, dim=0)
        logits = torch.zeros(
            values.shape[0],
            values.shape[1],
            values.shape[2],
            dtype=torch.float32,
            device=values.device,
        )
        weights = torch.softmax(logits, dim=0)
        expected = torch.tensor(
            1.0 / len(sources), dtype=weights.dtype, device=weights.device
        )
        error = (weights - expected).abs().max().detach().cpu().item()
        self.uniform_weight_max_abs_error = max(
            self.uniform_weight_max_abs_error, error
        )
        output = torch.einsum(
            "nbt,nbtd->btd", weights, values.float()
        ).to(values.dtype)
        if not capture:
            return output, None
        entropy = -(weights * torch.log(weights.clamp_min(1e-30))).sum(dim=0)
        return output, {
            "mean_weights": weights.mean(dim=(1, 2)).detach().cpu().tolist(),
            "entropy_mean": entropy.mean().detach().cpu().item(),
            "sources": len(sources),
        }

    def forward(
        self, input_ids: torch.Tensor, capture: bool = False
    ) -> tuple[torch.Tensor, parent.ActivationTrace | None]:
        if not self.selected_indices:
            return super().forward(input_ids, capture)

        self.forward_calls += 1
        embedded = self.embed(input_ids)
        trace = parent.ActivationTrace([], [], [], [], []) if capture else None
        completed = [embedded]
        partial: torch.Tensor | None = None
        mixer_index = 0
        for block in self.blocks:
            for branch_index in range(2):
                sources = completed + ([] if partial is None else [partial])
                branch_input, weights = self.mix(
                    mixer_index, sources, capture
                )
                mixer_index += 1
                if branch_index == 0:
                    branch_output = block.attention(
                        block.attention_norm(branch_input)
                    )
                else:
                    branch_output = block.mlp(block.mlp_norm(branch_input))
                branch_for_residual = branch_output.float()
                partial = (
                    branch_for_residual
                    if partial is None
                    else partial + branch_for_residual
                )
                if trace is not None:
                    trace.layer_input_rms.append(parent.rms(branch_input))
                    trace.branch_output_rms.append(parent.rms(branch_output))
                    trace.stream_state_rms.append(parent.rms(partial))
                    trace.depth_weights.append(weights or {})
                    if branch_index == 1:
                        partial.retain_grad()
                        trace.block_outputs.append(partial)
                if mixer_index % parent.round04.SUBLAYERS_PER_BLOCK == 0:
                    completed.append(partial)
                    partial = None
        if partial is not None or len(completed) != parent.BLOCK_GROUPS + 1:
            raise RuntimeError("Round 08 Block AttnRes aggregation failed")
        if self.output_mixer is None:
            raise RuntimeError("Round 08 output mixer missing")
        self.output_visits += 1
        hidden, output_weights = self.output_mixer(completed, capture)
        if trace is not None:
            trace.output_weights = output_weights
        normalized = self.final_norm(hidden)
        logits = F.linear(normalized, self.token_embedding.weight)
        return logits, trace


def tensor_exact(left: torch.Tensor, right: torch.Tensor) -> bool:
    return (
        left.dtype == right.dtype
        and tuple(left.shape) == tuple(right.shape)
        and torch.equal(left.detach().cpu(), right.detach().cpu())
    )


def build_intervention_audit(
    model: ForwardInterventionLanguageModel,
    optimizer: torch.optim.Optimizer,
    study_manifest: dict[str, Any],
) -> dict[str, Any]:
    selected = tuple(model.selected_indices)
    selected_names = set(parameter_names_for_indices(selected))
    mixer_parameters = {
        name: parameter
        for name, parameter in model.named_parameters()
        if name.startswith("mixers.") or name.startswith("output_mixer.")
    }
    optimizer_parameters = {
        parameter
        for group in optimizer.param_groups
        for parameter in group["params"]
    }
    selected_parameter_checks = {}
    for name in sorted(selected_names):
        parameter = mixer_parameters[name]
        selected_parameter_checks[name] = {
            "gradient_hook_calls": model.gradient_hook_calls[name],
            "in_optimizer_param_group": parameter in optimizer_parameters,
            "optimizer_state_present": parameter in optimizer.state,
            "final_equals_initial": tensor_exact(
                parameter, model.selected_initial_tensors[name]
            ),
        }
    unselected_parameter_checks = {}
    for name, parameter in sorted(mixer_parameters.items()):
        if name in selected_names:
            continue
        unselected_parameter_checks[name] = {
            "gradient_hook_calls": model.gradient_hook_calls[name],
            "in_optimizer_param_group": parameter in optimizer_parameters,
            "optimizer_state_present": parameter in optimizer.state,
        }

    source_counts = {
        str(index): sorted(values)
        for index, values in model.source_counts.items()
    }
    selected_source_gate = {
        str(index): (
            source_counts[str(index)]
            == [study_manifest["selected_source_counts"][str(index)]]
        )
        for index in selected
    }
    visit_gate = (
        all(value == model.forward_calls for value in model.depth_visits)
        and model.output_visits == model.forward_calls
    )
    selected_parameter_gate = all(
        check["gradient_hook_calls"] == 0
        and check["in_optimizer_param_group"]
        and not check["optimizer_state_present"]
        and check["final_equals_initial"]
        for check in selected_parameter_checks.values()
    )
    unselected_parameter_gate = all(
        check["gradient_hook_calls"] > 0
        and check["in_optimizer_param_group"]
        and check["optimizer_state_present"]
        for check in unselected_parameter_checks.values()
    )
    expected_selected = tuple(
        study_manifest["variants"]
        .get(model.forward_variant, {"selected_depth_indices": []})[
            "selected_depth_indices"
        ]
    )
    selector_gate = (
        selected == expected_selected
        and 64 not in selected
        and selected_source_gate == {
            str(index): True for index in selected
        }
    )
    threshold = study_manifest["thresholds"][
        "uniform_weight_max_abs_error"
    ]
    uniform_gate = model.uniform_weight_max_abs_error <= threshold
    passed = (
        visit_gate
        and selector_gate
        and selected_parameter_gate
        and unselected_parameter_gate
        and uniform_gate
    )
    return {
        "passed": passed,
        "variant": model.forward_variant,
        "selected_depth_indices": list(selected),
        "output_mixer_selected": False,
        "forward_calls": model.forward_calls,
        "depth_visit_counts": model.depth_visits,
        "output_visit_count": model.output_visits,
        "visit_gate": visit_gate,
        "source_counts_by_depth_index": source_counts,
        "selected_source_count_checks": selected_source_gate,
        "selector_gate": selector_gate,
        "uniform_weight_max_abs_error": model.uniform_weight_max_abs_error,
        "uniform_weight_threshold": threshold,
        "uniform_weight_gate": uniform_gate,
        "selected_parameters": selected_parameter_checks,
        "selected_parameter_reachability_gate": selected_parameter_gate,
        "unselected_parameters": unselected_parameter_checks,
        "unselected_parameter_reachability_gate": unselected_parameter_gate,
        "semantics": (
            "selected depth mixers use parameter-free constant-zero logits "
            "with the parent softmax+einsum arithmetic kernel"
        ),
    }


def rewrite_result(
    output_path: Path,
    study_manifest_path: Path,
    study_manifest: dict[str, Any],
) -> None:
    if LAST_MODEL is None or LAST_OPTIMIZER is None:
        raise RuntimeError("runner capture state missing")
    result = json.loads(output_path.read_text())
    parent_self_hash = result.pop("canonical_sha256_without_self")
    if result["protocol_id"] != PARENT_PROTOCOL_ID:
        raise RuntimeError("parent runner protocol drift")
    result["schema_version"] = 2
    result["protocol_id"] = PROTOCOL_ID
    result["parent_protocol_id"] = PARENT_PROTOCOL_ID
    result["variant"] = ACTIVE_VARIANT
    result["parent_runner_canonical_sha256"] = parent_self_hash
    result["study_manifest"] = {
        "path": str(study_manifest_path),
        "file_sha256": parent.file_sha256(study_manifest_path),
        "status": study_manifest["status"],
    }
    result["forward_intervention"] = build_intervention_audit(
        LAST_MODEL, LAST_OPTIMIZER, study_manifest
    )
    result["canonical_sha256_without_self"] = parent.canonical_sha256(result)
    temporary = output_path.with_suffix(output_path.suffix + ".round08.tmp")
    temporary.write_text(
        json.dumps(result, ensure_ascii=False, indent=2, sort_keys=True) + "\n"
    )
    os.replace(temporary, output_path)
    if not result["forward_intervention"]["passed"]:
        raise RuntimeError(
            f"forward intervention audit failed: "
            f"{result['forward_intervention']}"
        )


def main() -> None:
    global ACTIVE_VARIANT, LAST_OPTIMIZER
    variant = extract_wrapper_argument("--variant")
    study_manifest_path = Path(
        extract_wrapper_argument("--study-manifest")
    ).resolve()
    if variant not in VARIANTS:
        raise ValueError(f"unknown variant: {variant}")
    run_kind = argument_value("--run-kind", "formal")
    if run_kind in ("formal", "replay") and variant not in FORMAL_VARIANTS:
        raise ValueError("learned_reference is smoke-only")
    if argument_value("--architecture") != "block":
        raise ValueError("Round 08 requires --architecture block")
    if argument_value("--depth") != "32":
        raise ValueError("Round 08 requires --depth 32")
    if run_kind == "replay" and variant != "uniform_groups_6_7_forward":
        raise ValueError("the frozen replay uses the primary joint variant")

    study_manifest = json.loads(study_manifest_path.read_text())
    if (
        study_manifest["protocol_id"] != PROTOCOL_ID
        or study_manifest["status"] != "frozen-before-model-output"
    ):
        raise ValueError("study manifest is not the frozen Round 08 contract")
    expected = tuple(
        study_manifest["variants"]
        .get(variant, {"selected_depth_indices": []})[
            "selected_depth_indices"
        ]
    )
    if expected != VARIANTS[variant]:
        raise ValueError("study manifest selector drift")

    output_value = argument_value("--output")
    if output_value is None:
        raise ValueError("--output is required")
    output_path = Path(output_value).resolve()
    ACTIVE_VARIANT = variant
    parent.GradientLanguageModel = ForwardInterventionLanguageModel

    original_adamw = torch.optim.AdamW

    def capture_adamw(*args: Any, **kwargs: Any) -> torch.optim.Optimizer:
        global LAST_OPTIMIZER
        LAST_OPTIMIZER = original_adamw(*args, **kwargs)
        return LAST_OPTIMIZER

    torch.optim.AdamW = capture_adamw  # type: ignore[assignment]
    try:
        parent.main()
    finally:
        torch.optim.AdamW = original_adamw  # type: ignore[assignment]
    rewrite_result(output_path, study_manifest_path, study_manifest)


if __name__ == "__main__":
    main()