"""Mechanism checks for the publication baselines (CPU, no dataset needed)."""
import os
import sys

import torch

sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
from sdil.data import onehot
from sdil.local_baselines import FANet, PEPITANet, FFNet, EPNet
from sdil import probes


def check_fa_residual_transport():
    torch.manual_seed(1)
    x = torch.randn(32, 8)
    y = torch.randint(0, 3, (32,))
    yoh = onehot(y, 3)
    residual = FANet([8, 6, 6, 6, 3], seed=2, residual=True)
    plain = FANet([8, 6, 6, 6, 3], seed=2, residual=False)
    for net in (residual, plain):
        for l in (1, 2):
            net.B[l].zero_()
    before_res = [w.clone() for w in residual.W]
    before_plain = [w.clone() for w in plain.W]
    residual.fa_step(x, y, yoh, eta=0.01)
    plain.fa_step(x, y, yoh, eta=0.01)
    res_changes = [(a - b).norm().item() for a, b in zip(residual.W, before_res)]
    plain_changes = [(a - b).norm().item() for a, b in zip(plain.W, before_plain)]
    assert all(v > 0 for v in res_changes)
    assert plain_changes[0] == 0 and plain_changes[1] == 0
    assert plain_changes[2] > 0 and plain_changes[3] > 0
    alignment = probes.fa_alignment_report(residual, x, y, yoh)["cos_fa_negg"]
    assert len(alignment) == 3 and all(torch.isfinite(torch.tensor(alignment)))
    print("FA residual identity transport:", res_changes)
    print("FA measurable hidden alignment:", alignment)


def check_pepita_output_rule():
    torch.manual_seed(2)
    net = PEPITANet([5, 4, 3], act="relu", seed=3, keep_prob=1.0)
    x = torch.rand(16, 5)
    y = torch.randint(0, 3, (16,))
    yoh = onehot(y, 3)
    with torch.no_grad():
        clean = net._pepita_forward(x, None)
        error = torch.softmax(clean[-1], 1) - yoh
        mod = net._pepita_forward(x + error @ net.Fproj.t(), None)
        mod_error = torch.softmax(mod[-1], 1) - yoh
        expected = -(mod_error.t() @ mod[-2]) / x.shape[0]
        old = net.W[-1].clone()
    net.pepita_step(x, y, yoh, eta=0.1, momentum=0.0)
    assert torch.allclose(net.W[-1] - old, 0.1 * expected, atol=1e-7, rtol=1e-5)
    assert net.Fproj.abs().max() <= (6.0 / 5) ** 0.5 * 0.05 + 1e-7
    print("PEPITA modulated-output delta: exact")


def ff_loss(net, layer, x, y, neg):
    hp = net._inputs_to_layer(layer, net._overlay(x, y))
    hn = net._inputs_to_layer(layer, net._overlay(x, neg))
    gp = net._layer_forward(layer, hp).pow(2).mean(1)
    gn = net._layer_forward(layer, hn).pow(2).mean(1)
    return (torch.nn.functional.softplus(-gp + net.thr)
            + torch.nn.functional.softplus(gn - net.thr)).mean().item()


def check_ff_local_optimization():
    torch.manual_seed(3)
    net = FFNet([20, 32, 32], seed=4)
    x = torch.randn(128, 20)
    y = torch.randint(0, 10, (128,))
    neg = (y + 1) % 10
    initial = ff_loss(net, 0, x, y, neg)
    for _ in range(80):
        net.train_layer(0, x, y, eta=0.03, negative_labels=neg)
    final = ff_loss(net, 0, x, y, neg)
    assert final < initial - 0.1
    assert net._layer_forward(0, x).shape == (128, 32)
    print(f"FF layer-local loss: {initial:.4f} -> {final:.4f}")


def check_ep_energy_dynamics():
    torch.manual_seed(4)
    net = EPNet([4, 3, 2], seed=5, dt=0.2, random_beta_sign=False)
    x = torch.rand(7, 4)
    y = onehot(torch.randint(0, 2, (7,)), 2)
    s = [torch.rand(7, 3) * 0.8 + 0.1, torch.rand(7, 2) * 0.8 + 0.1]
    beta = 0.3
    manual = []
    for k in range(net.L):
        below = x if k == 0 else s[k - 1]
        drive = -s[k] + below @ net.W[k].t() + net.b[k]
        if k < net.L - 1:
            drive = drive + s[k + 1] @ net.W[k + 1]
        else:
            drive = drive + 2 * beta * (y - s[k])
        manual.append((s[k] + net.dt * drive).clamp(0, 1))
    actual = net._settle(x, y, beta=beta, s=[v.clone() for v in s], T=1)
    assert all(torch.allclose(a, b, atol=1e-7) for a, b in zip(actual, manual))
    old = [w.clone() for w in net.W]
    _, free_state = net.train_step(x, y.argmax(1), y, eta=[0.01, 0.005],
                                   return_free_state=True)
    assert all(torch.isfinite(w).all() for w in net.W)
    assert any(not torch.equal(a, b) for a, b in zip(net.W, old))
    assert len(free_state) == net.L
    continued = net._settle(x, s=free_state, T=1)
    restarted = net._settle(x, T=1)
    assert any(not torch.equal(a, b) for a, b in zip(continued, restarted))
    print("EP one-step -d(E+beta*C)/ds dynamics: exact")
    print("EP free particles persist across presentations")


if __name__ == "__main__":
    check_fa_residual_transport()
    check_pepita_output_rule()
    check_ff_local_optimization()
    check_ep_energy_dynamics()
    print("ALL BASELINE SMOKE CHECKS PASSED")