"""Tests for E13's Git Re-Basin weight-matching + barrier (pure NumPy, always runnable)."""

from __future__ import annotations

import numpy as np
import pytest

from neural.rebasin import apply_perms, barrier, canonicalise_scale, interpolate, weight_matching


def _mlp(sizes, rng):
    return [(rng.standard_normal((sizes[i + 1], sizes[i])), rng.standard_normal(sizes[i + 1]))
            for i in range(len(sizes) - 1)]


def _forward(params, X):
    h = X
    for i, (W, b) in enumerate(params):
        h = h @ W.T + b
        if i < len(params) - 1:
            h = np.maximum(h, 0)
    return h


def test_weight_matching_recovers_a_known_permutation():
    # A random-but-functionally-identical permuted copy is in a different "basis"; weight matching must
    # recover the permutation, so realigning makes the copy functionally identical to the original again.
    rng = np.random.default_rng(0)
    A = _mlp([8, 16, 16, 3], rng)
    p1, p2 = rng.permutation(16), rng.permutation(16)
    B = apply_perms(A, [p1, p2])                                # functionally identical to A, permuted basis
    X = rng.standard_normal((32, 8))
    assert np.allclose(_forward(A, X), _forward(B, X))         # permutation preserves the function
    perms = weight_matching(A, B, np.random.default_rng(1))
    B_realigned = apply_perms(B, perms)
    assert np.allclose(_forward(A, X), _forward(B_realigned, X), atol=1e-6)  # recovered -> function matches


def test_weight_matching_is_deterministic():
    rng = np.random.default_rng(2)
    A, B = _mlp([6, 10, 4], rng), _mlp([6, 10, 4], rng)
    p1 = weight_matching(A, B, np.random.default_rng(3))
    p2 = weight_matching(A, B, np.random.default_rng(3))
    assert all(np.array_equal(a, b) for a, b in zip(p1, p2))   # deterministic given inputs + seed


def test_interpolate_endpoints_and_barrier_zero_for_identical():
    rng = np.random.default_rng(4)
    A = _mlp([5, 8, 2], rng)
    B = _mlp([5, 8, 2], rng)
    assert np.allclose(interpolate(A, B, 0.0)[0][0], A[0][0])
    assert np.allclose(interpolate(A, B, 1.0)[0][0], B[0][0])
    X = rng.standard_normal((20, 5)); tgt = rng.standard_normal((20, 2))

    def loss_fn(P):
        L = float(((_forward(P, X) - tgt) ** 2).mean()); return L, L
    b = barrier(A, A, loss_fn)                                  # a model with itself: no barrier
    assert b["loss_barrier"] == pytest.approx(0.0, abs=1e-9)


def test_apply_perms_preserves_function():
    rng = np.random.default_rng(5)
    A = _mlp([4, 7, 7, 3], rng)
    perms = [rng.permutation(7), rng.permutation(7)]
    X = rng.standard_normal((16, 4))
    assert np.allclose(_forward(A, X), _forward(apply_perms(A, perms), X))


def _rescale(params, scales_per_layer):
    # Apply the positive per-unit rescaling symmetry: unit i of hidden layer k scaled by c>0.
    out = [(W.copy(), b.copy()) for W, b in params]
    for k, scales in enumerate(scales_per_layer):
        W, b = out[k]
        out[k] = (W * scales[:, None], b * scales)
        Wn, bn = out[k + 1]
        out[k + 1] = (Wn / scales[None, :], bn)
    return out


def test_canonicalise_scale_preserves_function_and_normalises():
    rng = np.random.default_rng(6)
    A = _mlp([5, 9, 9, 2], rng)
    C = canonicalise_scale(A)
    X = rng.standard_normal((24, 5))
    assert np.allclose(_forward(A, X), _forward(C, X), atol=1e-8)   # function-preserving (ReLU homogeneity)
    for k in range(len(C) - 1):                                     # every hidden unit's (W, b) is unit-norm
        W, b = C[k]
        assert np.allclose(np.sqrt((W ** 2).sum(axis=1) + b ** 2), 1.0)


def test_weight_matching_recovers_permutation_and_rescaling():
    # A permuted AND positively-rescaled copy is functionally identical; permutation-only matching can
    # miss it, but canonicalise-then-match must realign it to functional identity — the full ReLU
    # symmetry group (the E13c referee-proofing gate).
    rng = np.random.default_rng(7)
    A = _mlp([6, 12, 12, 3], rng)
    B = apply_perms(_rescale(A, [np.exp(rng.uniform(-2, 2, 12)), np.exp(rng.uniform(-2, 2, 12))]),
                    [rng.permutation(12), rng.permutation(12)])
    X = rng.standard_normal((32, 6))
    assert np.allclose(_forward(A, X), _forward(B, X), atol=1e-6)   # symmetry-equivalent copy
    cA, cB = canonicalise_scale(A), canonicalise_scale(B)
    perms = weight_matching(cA, cB, np.random.default_rng(8))
    B_aligned = apply_perms(cB, perms)
    assert np.allclose(_forward(cA, X), _forward(B_aligned, X), atol=1e-5)  # realigned exactly
    # and the aligned weights themselves coincide (canonical form is unique up to permutation)
    for (Wa, ba), (Wb, bb) in zip(cA, B_aligned):
        assert np.allclose(Wa, Wb, atol=1e-6) and np.allclose(ba, bb, atol=1e-6)