import ast
import builtins
from contextlib import ExitStack
import io
import inspect
import pickle
import sys
import types
from pathlib import Path
from types import SimpleNamespace
from unittest import mock

import numpy as np
import torch
import torch.nn.functional as F


ROOT = Path(__file__).resolve().parent
TRAINER = ROOT / "casc_eq_train.py"
PROBE = ROOT / "cascade_probe.py"
_REAL_OPEN = builtins.open
_META = pickle.dumps({"vocab_size": 17})


def _fake_open(file, *args, **kwargs):
    normalized = str(file).replace("\\", "/")
    if normalized.endswith("/data/tinystories_bpe/meta.pkl"):
        return io.BytesIO(_META)
    return _REAL_OPEN(file, *args, **kwargs)


def _load_prefix(path, stop_marker, argv, fake_memmap=False):
    source = path.read_text()
    assert stop_marker in source, f"loader marker moved in {path.name}"
    source = source.split(stop_marker, 1)[0]

    module = types.ModuleType(f"_bp_free_audit_{path.stem}")
    module.__file__ = str(path)

    with ExitStack() as stack:
        stack.enter_context(mock.patch.object(sys, "argv", [str(path), *argv]))
        stack.enter_context(mock.patch.object(builtins, "open", _fake_open))
        stack.enter_context(
            mock.patch.object(torch.cuda, "is_available", return_value=False)
        )
        if fake_memmap:
            dummy = np.arange(256, dtype=np.uint16)
            stack.enter_context(
                mock.patch.object(np, "memmap", return_value=dummy)
            )
        exec(compile(source, str(path), "exec"), module.__dict__)

    assert module.dev == "cpu"
    return module


def _load_trainer():
    return _load_prefix(
        TRAINER,
        "\nwb = None",
        [
            "--L", "4",
            "--C", "8",
            "--H", "2",
            "--T", "4",
            "--B", "2",
            "--K", "2",
            "--steps", "-1",
            "--noguard",
            "--gate_every", "0",
        ],
    )


def _load_probe():
    return _load_prefix(
        PROBE,
        "\nflat = lambda",
        [
            "--L", "4",
            "--C", "8",
            "--H", "2",
            "--T", "4",
            "--B", "2",
            "--K", "2",
            "--scheme", "fb",
            "--include_io",
            "--batches", "1",
        ],
        fake_memmap=True,
    )


class _BlockMarker(torch.autograd.Function):
    @staticmethod
    def forward(ctx, x, block_idx):
        ctx.block_idx = int(block_idx)
        return x

    @staticmethod
    def backward(ctx, grad):
        return grad, None


def _tag_blocks(module):
    for idx, block in enumerate(module.blocks):
        original = block.forward

        def tagged(self, z, mask, _original=original, _idx=idx):
            output = _original(z, mask)
            return _BlockMarker.apply(output, _idx)

        block.forward = types.MethodType(tagged, block)


def _tensors(value):
    if isinstance(value, torch.Tensor):
        yield value
    elif isinstance(value, (tuple, list)):
        for item in value:
            yield from _tensors(item)


def _roots(outputs):
    return [
        tensor.grad_fn
        for tensor in _tensors(outputs)
        if tensor.grad_fn is not None
    ]


def _graph_facts(outputs, param_owner):
    """
    Walk every root-to-leaf dependency path.

    A summed theta objective may contain markers from all blocks, but no
    individual path may pass through more than one marker.
    """
    stack = [(root, ()) for root in _roots(outputs)]
    seen = set()
    markers = set()
    max_chain = 0
    foreign_param_paths = []
    unmarked_block_params = []

    while stack:
        node, path = stack.pop()

        block_idx = getattr(node, "block_idx", None)
        if block_idx is not None:
            path = path + ((id(node), int(block_idx)),)
            markers.add(int(block_idx))
            max_chain = max(max_chain, len(path))

        key = (node, tuple(marker_id for marker_id, _ in path))
        if key in seen:
            continue
        seen.add(key)

        variable = getattr(node, "variable", None)
        owner = param_owner.get(id(variable))
        if owner is not None:
            if not path:
                unmarked_block_params.append(owner)
            elif path[-1][1] != owner:
                foreign_param_paths.append(
                    (tuple(idx for _, idx in path), owner)
                )

        for next_node, _ in node.next_functions:
            if next_node is not None:
                stack.append((next_node, path))

    return {
        "markers": markers,
        "max_chain": max_chain,
        "foreign": foreign_param_paths,
        "unmarked": unmarked_block_params,
    }


def _assert_local(outputs, param_owner, label):
    facts = _graph_facts(outputs, param_owner)

    assert facts["max_chain"] <= 1, (
        f"{label}: GLOBAL transformer chain detected: {facts}"
    )
    assert not facts["foreign"], (
        f"{label}: block marker reached another block's parameters: {facts}"
    )
    assert not facts["unmarked"], (
        f"{label}: an unmarked path reached block parameters: {facts}"
    )
    return facts


def _checked_grad(real_grad, param_owner, records):
    def checked(outputs, inputs, *args, **kwargs):
        frame = inspect.currentframe().f_back
        label = (
            f"{Path(frame.f_code.co_filename).name}:"
            f"{frame.f_code.co_name}:{frame.f_lineno}"
        )

        facts = _assert_local(outputs, param_owner, label)

        grad_outputs = kwargs.get(
            "grad_outputs", args[0] if args else None
        )
        assert all(
            tensor.grad_fn is None and not tensor.requires_grad
            for tensor in _tensors(grad_outputs)
        ), f"{label}: VJP seed carries an autograd graph"

        assert all(
            tensor.is_leaf for tensor in _tensors(inputs)
        ), f"{label}: differentiation target is not a leaf"

        result = real_grad(outputs, inputs, *args, **kwargs)

        assert all(
            grad is None
            or (grad.grad_fn is None and not grad.requires_grad)
            for grad in result
        ), f"{label}: autograd.grad returned a higher-order graph"

        records.append((frame.f_code.co_name, frame.f_lineno, facts))
        return result

    return checked


def _dotted_name(node):
    if isinstance(node, ast.Name):
        return node.id
    if isinstance(node, ast.Attribute):
        prefix = _dotted_name(node.value)
        return f"{prefix}.{node.attr}" if prefix else node.attr
    return ""


def _grad_sites(path):
    tree = ast.parse(path.read_text(), filename=str(path))
    sites = []

    class Visitor(ast.NodeVisitor):
        def __init__(self):
            self.functions = []

        def visit_FunctionDef(self, node):
            self.functions.append(node.name)
            self.generic_visit(node)
            self.functions.pop()

        visit_AsyncFunctionDef = visit_FunctionDef

        def visit_Call(self, node):
            if _dotted_name(node.func) == "torch.autograd.grad":
                owner = self.functions[-1] if self.functions else "<module>"
                sites.append((owner, node.lineno))
            self.generic_visit(node)

    Visitor().visit(tree)
    return set(sites)


def test_casc_eq_training_graphs_are_local_and_gate_is_read_only():
    module = _load_trainer()
    _tag_blocks(module)

    param_owner = {
        id(param): idx
        for idx, block in enumerate(module.blocks)
        for param in block.parameters()
    }

    generator = torch.Generator().manual_seed(123)
    x = torch.randint(
        module.vocab,
        (module.args.B, module.args.T),
        generator=generator,
    )
    y = torch.randint(
        module.vocab,
        (module.args.B, module.args.T),
        generator=generator,
    )

    # Free-state firewalls.
    z0, states, ins, outs = module.free_states_graphed(x)
    assert z0.grad_fn is None and not z0.requires_grad
    assert all(
        state.grad_fn is None and not state.requires_grad
        for state in states
    )
    assert all(item.is_leaf and item.requires_grad for item in ins)

    for idx, output in enumerate(outs):
        facts = _assert_local(output, param_owner, f"free output {idx}")
        assert facts["markers"] == {idx}

    real_grad = torch.autograd.grad
    records = []
    checked = _checked_grad(real_grad, param_owner, records)

    with mock.patch.object(torch.autograd, "grad", checked):
        relaxed, last_outs = module.relax(
            z0,
            states,
            ins,
            outs,
            y,
            module.args.beta,
            module.args.K,
            x,
        )

        assert all(
            state.grad_fn is None and not state.requires_grad
            for state in relaxed
        )

        for idx, output in enumerate(last_outs):
            facts = _assert_local(
                output, param_owner, f"last-round output {idx}"
            )
            assert facts["markers"] == {idx}

        # Exercise the otherwise-unused helper call site too.
        module.dFdtheta(relaxed, x, y, module.args.beta)

        _, _, _, ok = module.ep_step(x, y)
        assert ok

    # Every non-gate autograd.grad call currently in the source must have
    # been exercised by the instrumented tiny run.
    sites = _grad_sites(TRAINER)
    local_sites = {
        site for site in sites if site[0] != "bp_gate"
    }
    observed = {(name, line) for name, line, _ in records}
    assert local_sites <= observed

    # The theta objective should contain all block markers, but as a
    # disconnected forest with one marker maximum per path.
    assert any(
        facts["markers"] == set(range(module.args.L))
        and facts["max_chain"] == 1
        for _, _, facts in records
    )

    # Known global exception: prove both that it is truly global and that
    # it cannot overwrite the EP gradients stored in p.grad.
    for idx, param in enumerate(module.all_params):
        param.grad = torch.full_like(param, (idx + 1) / 1000)

    before = [param.grad.clone() for param in module.all_params]
    gate_facts = []

    def gate_grad(outputs, inputs, *args, **kwargs):
        gate_facts.append(_graph_facts(outputs, param_owner))
        return real_grad(outputs, inputs, *args, **kwargs)

    with mock.patch.object(torch.autograd, "grad", gate_grad):
        gbp = module.bp_gate(x, y)

    assert len(gate_facts) == 1
    assert gate_facts[0]["max_chain"] == module.args.L
    assert all(
        torch.equal(old, param.grad)
        for old, param in zip(before, module.all_params)
    ), "bp_gate populated or changed p.grad"
    assert all(
        grad is None
        or (grad.grad_fn is None and not grad.requires_grad)
        for grad in gbp
    )

    # The trainer should not acquire any Tensor.backward/autograd.backward
    # calls later.
    tree = ast.parse(TRAINER.read_text())
    backward_lines = [
        node.lineno
        for node in ast.walk(tree)
        if isinstance(node, ast.Call)
        and _dotted_name(node.func).endswith(".backward")
    ]
    assert not backward_lines, (
        f"new .backward() calls at lines {backward_lines}"
    )


def test_theta_readout_has_zero_distant_block_influence():
    module = _load_trainer()
    generator = torch.Generator().manual_seed(456)

    x = torch.randint(
        module.vocab,
        (module.args.B, module.args.T),
        generator=generator,
    )
    y = torch.randint(
        module.vocab,
        (module.args.B, module.args.T),
        generator=generator,
    )
    states = [
        torch.randn(
            module.args.B,
            module.args.T,
            module.args.C,
            generator=generator,
        )
        for _ in range(module.args.L)
    ]

    baseline = module.dFdtheta(states, x, y, module.args.beta)
    param_index = {
        id(param): idx
        for idx, param in enumerate(module.all_params)
    }

    for changed_state in range(module.args.L):
        altered = [state.clone() for state in states]
        altered[changed_state].add_(
            0.125
            * torch.randn(
                altered[changed_state].shape,
                generator=generator,
            )
        )
        got = module.dFdtheta(altered, x, y, module.args.beta)

        for block_idx, block in enumerate(module.blocks):
            # E_i may depend numerically only on z_i and z_{i-1}.
            if changed_state in {block_idx - 1, block_idx}:
                continue

            for param in block.parameters():
                idx = param_index[id(param)]
                assert torch.equal(baseline[idx], got[idx]), (
                    f"block {block_idx} gradient changed after distant "
                    f"state {changed_state} changed"
                )


def test_cascade_probe_fb_ep_side_is_local():
    module = _load_probe()
    _tag_blocks(module)

    param_owner = {
        id(param): idx
        for idx, block in enumerate(module.blocks)
        for param in block.parameters()
    }

    generator = torch.Generator().manual_seed(789)
    x = torch.randint(
        module.vocab,
        (module.args.B, module.args.T),
        generator=generator,
    )
    y = torch.randint(
        module.vocab,
        (module.args.B, module.args.T),
        generator=generator,
    )
    positions = torch.arange(module.args.T)
    z0 = (module.tok(x) + module.pos(positions)[None]).detach()

    real_grad = torch.autograd.grad
    records = []
    checked = _checked_grad(real_grad, param_owner, records)

    with (
        mock.patch.object(torch.autograd, "grad", checked),
        mock.patch.object(
            torch.Tensor,
            "backward",
            side_effect=AssertionError("fb called backward()"),
        ),
    ):
        relaxed = module.relax(z0, y, module.args.beta)
        module.dFdtheta(
            relaxed,
            z0,
            y,
            module.args.beta,
            module.gate_params,
            x,
        )

    assert all(
        state.grad_fn is None and not state.requires_grad
        for state in relaxed
    )

    # Per round: one top-force call plus L-1 feedback calls.
    # Then one disconnected theta-readout forest.
    assert len(records) == module.args.K * module.args.L + 1
    assert sum(
        not facts["markers"] for _, _, facts in records
    ) == module.args.K
    assert any(
        facts["markers"] == set(range(module.args.L))
        and facts["max_chain"] == 1
        for _, _, facts in records
    )

    # Positive control corresponding to the probe's intended BP reference.
    for idx, param in enumerate(module.gate_params):
        param.grad = torch.full_like(param, (idx + 1) / 1000)
    before = [param.grad.clone() for param in module.gate_params]

    z = module.tok(x) + module.pos(positions)[None]
    for block in module.blocks:
        z = block(z, module.mask)
    ce = F.cross_entropy(
        module.readout(z).reshape(-1, module.vocab),
        y.reshape(-1),
    )

    facts = _graph_facts(ce, param_owner)
    assert facts["max_chain"] == module.args.L

    gbp = real_grad(ce, module.gate_params, allow_unused=True)
    assert all(
        torch.equal(old, param.grad)
        for old, param in zip(before, module.gate_params)
    )
    assert all(grad is None or not grad.requires_grad for grad in gbp)


def test_gate_every_zero_is_a_true_off_switch():
    """
    This intentionally fails on the audited source. It passes after changing
    the loop condition to guard gate_every > 0 before taking the modulo.
    """
    tree = ast.parse(TRAINER.read_text(), filename=str(TRAINER))
    gate_ifs = []

    for node in ast.walk(tree):
        if isinstance(node, ast.If) and any(
            isinstance(child, ast.Call)
            and _dotted_name(child.func) == "bp_gate"
            for child in ast.walk(node)
        ):
            gate_ifs.append(node)

    assert len(gate_ifs) == 1
    expression = ast.Expression(gate_ifs[0].test)
    ast.fix_missing_locations(expression)

    try:
        enabled = eval(
            compile(expression, str(TRAINER), "eval"),
            {},
            {
                "args": SimpleNamespace(gate_every=0),
                "step": 0,
                "ok": True,
            },
        )
    except ZeroDivisionError:
        raise AssertionError(
            "--gate_every 0 crashes; guard modulo with "
            "args.gate_every > 0"
        ) from None

    assert enabled is False, (
        "--gate_every 0 must prevent the global bp_gate call"
    )