"""Execute the authors' construction notebooks without modifying their bytes.

The arXiv source links to https://github.com/SprocketLab/hybrid-expressivity.
``source/official-code`` is a detached checkout of the last repository commit
before arXiv:2603.08859v1.  This runner reads the code cells directly from the
two authored notebooks, executes their setup/weight-assignment cells unchanged
on CPU, and evaluates the constructed models on deterministic batches generated
by the authors' own Dataset implementation.  Notebook display/plot cells are
intentionally not executed because they are not part of the forward pass.

This is construction validation only.  It is not an attempted rerun of the
GPU-only learned-model training entry point, which hard-codes ``cuda``.
"""

from __future__ import annotations

import copy
import hashlib
import json
import os
import sys
from pathlib import Path
from typing import Any

import numpy as np
import torch


ROOT = Path(__file__).resolve().parent
CODE = ROOT / "source" / "official-code"
NOTEBOOKS = {
    "claim3_native": (CODE / "constructions" / "construction_var_copy.ipynb", 18),
    "claim4_native": (CODE / "constructions" / "construction_decode_recall.ipynb", 27),
}
OUT = ROOT / "outputs"


def sha256(path: Path) -> str:
    return hashlib.sha256(path.read_bytes()).hexdigest()


def execute_notebook_prefix(path: Path, last_cell: int) -> dict[str, Any]:
    """Execute exact code-cell text through the construction's model cell."""
    notebook = json.loads(path.read_text())
    cells = notebook["cells"]
    assert cells[0]["cell_type"] == "code"
    # The two notebook magics only enable autoreload; skipping that display
    # cell preserves every model/data/weight-assignment cell verbatim.
    selected = [i for i in range(1, last_cell + 1) if cells[i]["cell_type"] == "code"]
    namespace: dict[str, Any] = {"__name__": "__main__", "__file__": str(path)}
    old_cwd = Path.cwd()
    old_path = list(sys.path)
    try:
        os.chdir(path.parent)
        sys.path.insert(0, str(path.parent))
        np.random.seed(1729)
        torch.manual_seed(1729)
        for i in selected:
            source = "".join(cells[i]["source"])
            exec(compile(source, f"{path}::cell-{i}", "exec"), namespace)
    finally:
        os.chdir(old_cwd)
        sys.path[:] = old_path
    namespace["_executed_cells"] = selected
    return namespace


def evaluate_last_position(ns: dict[str, Any], n_batches: int = 256) -> dict[str, Any]:
    """Use the already-authored Dataset/model objects and score final tokens."""
    model = ns["model"]
    dataset = ns["train_dataset"]
    model.eval()
    eligible = correct = 0
    with torch.no_grad():
        for i in range(n_batches):
            batch = dataset[i]
            logits = model(batch["input_ids"])
            pred = torch.argmax(logits[:, -1], dim=-1)
            valid = batch["mask"][:, -1].bool()
            eligible += int(valid.sum())
            correct += int((pred[valid] == batch["output_ids"][:, -1][valid]).sum())
    return {
        "n_batches": n_batches,
        "examples_per_batch": int(ns["args"].train_batch_size),
        "eligible_last_position_examples": eligible,
        "correct_last_position_examples": correct,
        "last_position_accuracy": correct / eligible if eligible else None,
    }


def destructive_no_first_ssm(ns: dict[str, Any], n_batches: int = 256) -> dict[str, Any]:
    """Ablate the first authored SSM gate in-memory; source bytes stay intact."""
    model = copy.deepcopy(ns["model"])
    model.layers[0].layer.Delta.data.zero_()
    dataset = ns["train_dataset"]
    model.eval()
    eligible = correct = 0
    # Re-seed data generation so the control sees exactly the baseline inputs.
    np.random.seed(1729 + 1)
    with torch.no_grad():
        for i in range(n_batches):
            batch = dataset[i]
            logits = model(batch["input_ids"])
            pred = torch.argmax(logits[:, -1], dim=-1)
            valid = batch["mask"][:, -1].bool()
            eligible += int(valid.sum())
            correct += int((pred[valid] == batch["output_ids"][:, -1][valid]).sum())
    return {
        "ablation": "first SimpleSSM Delta vector set to zero in memory",
        "eligible_last_position_examples": eligible,
        "correct_last_position_examples": correct,
        "last_position_accuracy": correct / eligible if eligible else None,
    }


def main() -> None:
    OUT.mkdir(exist_ok=True)
    git_head = (CODE / ".git").exists()
    for output_name, (notebook, last_cell) in NOTEBOOKS.items():
        # Baseline and destructive control receive identical generated batches.
        ns = execute_notebook_prefix(notebook, last_cell)
        np.random.seed(1729 + 1)
        baseline = evaluate_last_position(ns)
        control = destructive_no_first_ssm(ns)
        payload = {
            "kind": "direct_native_notebook_execution",
            "official_repository": "https://github.com/SprocketLab/hybrid-expressivity",
            "repository_checkout": "6be8f8fbc2169290af6f4ba5e4bd53a5c6485f7b",
            "repository_is_git_checkout": git_head,
            "notebook": str(notebook.relative_to(ROOT)),
            "notebook_sha256": sha256(notebook),
            "executed_code_cells": ns["_executed_cells"],
            "device": "cpu",
            "torch_version": torch.__version__,
            "baseline": baseline,
            "destructive_control": control,
            "control_is_lower": (
                control["last_position_accuracy"] < baseline["last_position_accuracy"]
                if baseline["last_position_accuracy"] is not None
                else None
            ),
            "scope": (
                "Construction cells plus deterministic last-position batches; "
                "not a learned-model training rerun."
            ),
        }
        (OUT / f"{output_name}.json").write_text(json.dumps(payload, indent=2) + "\n")
        print(
            f"{output_name}: baseline={baseline['last_position_accuracy']:.6f} "
            f"control={control['last_position_accuracy']:.6f} n={baseline['eligible_last_position_examples']}",
            flush=True,
        )


if __name__ == "__main__":
    main()