"""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
from experiments.import_native_baseline import (summarize_burst_history,
                                                summarize_dualprop_hist)


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))

    # Exact author update: contrast weakly-clamped and free correlations,
    # divide by beta and minibatch size, then apply the layer-specific rate.
    rates = [0.01, 0.005]
    old = [w.clone() for w in net.W]
    old_biases = [b.clone() for b in net.b]
    expected_free = net._settle(x, beta=0.0, T=net.T_free)
    expected_nudged = net._settle(
        x, y, beta=net.beta, s=[v.clone() for v in expected_free], T=net.T_nudge)
    expected_weights, expected_biases = [], []
    for k in range(net.L):
        below_free = x if k == 0 else net.rho(expected_free[k - 1])
        below_nudged = x if k == 0 else net.rho(expected_nudged[k - 1])
        correlation_delta = (
            net.rho(expected_nudged[k]).t() @ below_nudged
            - net.rho(expected_free[k]).t() @ below_free) / (net.beta * x.shape[0])
        bias_delta = (net.rho(expected_nudged[k])
                      - net.rho(expected_free[k])).mean(0) / net.beta
        expected_weights.append(old[k] + rates[k] * correlation_delta)
        expected_biases.append(old_biases[k] + rates[k] * bias_delta)

    _, free_state = net.train_step(x, y.argmax(1), y, eta=rates,
                                   return_free_state=True)
    assert all(torch.allclose(actual, expected, atol=1e-7, rtol=1e-5)
               for actual, expected in zip(net.W, expected_weights))
    assert all(torch.allclose(actual, expected, atol=1e-7, rtol=1e-5)
               for actual, expected in zip(net.b, expected_biases))
    assert all(torch.equal(actual, expected)
               for actual, expected in zip(free_state, expected_free))
    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 contrastive parameter update: exact")
    print("EP free particles persist across presentations")


def check_native_result_selection():
    burst_rows = [
        {"epoch": 0, "epoch/top1_error/test": 30.0,
         "epoch/top1_error/val": 28.0, "_runtime": 10.0},
        {"epoch": 1, "epoch/top1_error/test": 20.0,
         "epoch/top1_error/val": 22.0, "_runtime": 20.0},
        {"epoch": 2, "epoch/top1_error/test": 21.0,
         "epoch/top1_error/val": 18.0, "_runtime": 30.0},
    ]
    burst = summarize_burst_history(burst_rows, 3)
    assert burst["primary_final_test_accuracy_percent"] == 79.0
    assert burst["validation_selected_epoch"] == 2
    assert burst["validation_selected_test_accuracy_percent"] == 79.0
    assert burst["test_selected_epoch"] == 1
    assert burst["test_selected_test_accuracy_percent"] == 80.0

    dual = summarize_dualprop_hist({
        "train_loss": [2.0, 1.0], "train_accuracy": [30.0, 70.0],
        "train_time": [4.0, 5.0], "val_loss": [1.5, 1.2],
        "val_accuracy": [60.0, 65.0], "val_top5accuracy": [90.0, 92.0],
        "val_time": [1.0, 1.5], "test_loss": 1.1,
        "test_accuracy": 64.0, "test_top5accuracy": 93.0, "test_time": 2.0,
    }, 2)
    assert dual["validation_selected_epoch"] == 2
    assert dual["primary_test_accuracy_percent"] == 64.0
    assert dual["test_evaluations"] == 1
    assert dual["wall_s"] == 13.5
    print("native final/validation-selected/test-selected semantics: exact")


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