#!/usr/bin/env python3
"""Deterministic mechanics checks for SHARED_FEEDBACK.md S0."""

import os
import sys

import torch

sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
from sdil.shared_feedback import (
    CONDITIONS, SharedFeedbackConfig, SharedFeedbackNet,
    conditional_selector_data, select_shared_signal, shared_feedback_step,
)


def maximum_difference(left, right):
    return max(float((a - b).abs().max()) for a, b in zip(left, right))


def main():
    torch.set_num_threads(1)
    config = SharedFeedbackConfig(width=8, hidden_layers=2)
    base = SharedFeedbackNet(config, seed=3101, dtype=torch.float64)
    x, z, y = conditional_selector_data(32, 3101)
    x = x.to(torch.float64)

    clones = {condition: base.clone() for condition in CONDITIONS}
    forward_error = max(float((
        clones["oracle"].forward(x, z)["logits"]
        - clones[condition].forward(x, z)["logits"]).abs().max())
                        for condition in CONDITIONS)
    assert forward_error == 0.0

    enabled = base.forward(x, z)
    lesioned = base.forward(x, z, context_enabled=False)
    lesion_context_error = max(float(value.abs().max())
                               for value in lesioned["context"])
    assert lesion_context_error == 0.0
    assert any(not torch.equal(a, b) for a, b in zip(
        enabled["h"][1:-1], lesioned["h"][1:-1]))

    reports = base.fit_neutral_predictor(x, z)
    assert all(row["instruction_observations"] == 0 for row in reports)
    layer = 0
    soma = base.forward(x, z)["h"][1]
    context = base.context_fields(z)[layer]
    instruction = torch.randn_like(context)
    exact, raw, innovation = select_shared_signal(
        base, layer, instruction, context, soma, "exact_subtraction")
    assert torch.equal(raw, instruction + context)
    assert torch.allclose(exact, instruction, atol=1e-15, rtol=1e-15)
    matched, _, _ = select_shared_signal(
        base, layer, instruction, context, soma, "matched_raw")
    assert torch.allclose(matched.norm(dim=1), innovation.norm(dim=1),
                          atol=1e-12, rtol=1e-12)

    # With context and predictor exactly zero, all four signal rules coincide.
    zero = base.clone()
    for value in zero.C + zero.P + zero.P_bias:
        value.zero_()
    zero_clones = {condition: zero.clone() for condition in CONDITIONS}
    for condition, net in zero_clones.items():
        shared_feedback_step(net, x, z, y, condition)
    reference = zero_clones["oracle"]
    zero_context_update_error = max(
        maximum_difference(reference.W, net.W)
        + maximum_difference(reference.Q[1:], net.Q[1:])
        for condition, net in zero_clones.items() if condition != "oracle")
    assert zero_context_update_error < 1e-14

    # The reciprocal change is independently applied from the same local
    # direction; it is not copied from the updated forward tensor.
    kp = base.clone()
    before_w = [value.clone() for value in kp.W]
    before_q = [None] + [value.clone() for value in kp.Q[1:]]
    _, auxiliary = shared_feedback_step(kp, x, z, y, "oracle")
    reciprocal_direction_error = 0.0
    for layer in range(1, len(kp.W)):
        expected_w = ((kp.W[layer] - before_w[layer])
                      / kp.config.learning_rate)
        expected_q = ((kp.Q[layer] - before_q[layer])
                      / kp.config.reciprocal_learning_rate)
        reciprocal_direction_error = max(
            reciprocal_direction_error,
            float((expected_w - expected_q
                   - kp.config.weight_decay
                   * (before_q[layer] - before_w[layer])).abs().max()))
    assert reciprocal_direction_error < 1e-12
    assert all(not value.requires_grad for value in kp.W + kp.Q[1:])
    assert all(direction.grad_fn is None for direction in auxiliary["directions"])

    print({
        "forward_identity_error": forward_error,
        "lesion_context_error": lesion_context_error,
        "zero_context_update_error": zero_context_update_error,
        "reciprocal_direction_error": reciprocal_direction_error,
        "neutral_instruction_observations": max(
            row["instruction_observations"] for row in reports),
    })


if __name__ == "__main__":
    main()