# SPDX-License-Identifier: CC-BY-NC-4.0
# Copyright (c) 2026 Jakub Ćwirlej
"""Unit tests for tethergate.py (CPU, a few seconds):  python test_tethergate.py"""
import math
import sys

import torch

from tethergate import TetherGate, floor_schedule, tethergate_step


def logit(p):
    return math.log(p / (1.0 - p))


def test_floor_schedule():
    kw = dict(ramp_start=100, f_min=0.05, tau=50.0, ramp_end=200)
    ok = floor_schedule(0, **kw) == 1.0 and floor_schedule(99, **kw) == 1.0
    exp_vals = [floor_schedule(t, mode='exp', **kw) for t in range(100, 2000, 7)]
    cos_vals = [floor_schedule(t, mode='cos', **kw) for t in range(100, 400, 3)]
    ok &= all(a >= b for a, b in zip(exp_vals, exp_vals[1:])) and min(exp_vals) > 0.05
    ok &= all(a >= b for a, b in zip(cos_vals, cos_vals[1:])) and abs(floor_schedule(200, mode='cos', **kw) - 0.05) < 1e-12
    return ok


def test_forward_floor():
    g = TetherGate(beta=6.0, theta_init=0.3)
    s = torch.linspace(-2, 2, 11)
    ok = torch.equal(g(s), torch.ones(11))                       # floor 1.0 -> gate fully open
    g.floor.fill_(0.0)
    ref = torch.sigmoid((6.0 * (s - 0.3)).clamp(-8.0, 8.0))
    ok &= torch.allclose(g(s), ref)
    g.floor.fill_(0.2)
    ok &= torch.allclose(g(s), 0.2 + 0.8 * ref)
    return bool(ok)


def test_gradient_alive_below_floor():
    """g_raw below the floor but the argument inside (-8, 8): theta still receives a gradient."""
    g = TetherGate(beta=6.0, theta_init=0.8)                     # arg = 6*(0-0.8) = -4.8 -> g_raw = 0.008 < floor
    g.floor.fill_(0.05)
    g(torch.tensor([0.0])).sum().backward()
    return g.theta.grad is not None and float(g.theta.grad.abs()) > 0


def test_saturation_beyond_clamp_and_spring_recovery():
    """PITFALL (documented): with beta*(s - theta) < arg_min the clamp saturates and the gradient to
    theta is exactly zero. The spring is not gradient-based, so it still pulls theta back."""
    g = TetherGate(beta=6.0, g_target=0.5, theta_init=5.0)      # arg = -30 -> clamped at -8
    g.floor.fill_(0.05)
    g(torch.tensor([0.0])).sum().backward()
    dead = g.theta.grad is not None and float(g.theta.grad.abs()) == 0.0
    s = torch.tensor([0.0])
    for step in range(3000):
        g(s)
        tethergate_step([g], step, lr=1e-2, wd=0.5, ramp_start=0, f_min=0.05, mode='cos', ramp_end=1)
    return dead and abs(g.theta.item() - g.theta_star()) < 1e-3


def test_detach_signal():
    s = torch.tensor([0.5, 1.0], requires_grad=True)
    TetherGate(detach_signal=True)(s).sum().backward()
    a = s.grad is None
    s2 = torch.tensor([0.5, 1.0], requires_grad=True)
    g = TetherGate(detach_signal=False); g.floor.fill_(0.0)
    g(s2).sum().backward()
    return a and s2.grad is not None and float(s2.grad.abs().sum()) > 0


def test_ema_train_eval():
    g = TetherGate(ema_momentum=0.9)
    g.train(); g(torch.full((4,), 2.0))
    a = abs(float(g.ema_s) - 2.0) < 1e-6                         # first call initializes from NaN
    g(torch.full((4,), 0.0))
    b = abs(float(g.ema_s) - 1.8) < 1e-6
    g.eval(); g(torch.full((4,), 100.0))
    return a and b and abs(float(g.ema_s) - 1.8) < 1e-6         # eval does not update


def test_spring_reaches_target():
    """No loss pressure, constant signal: theta -> theta_star and the gate -> g_target."""
    g = TetherGate(beta=6.0, g_target=0.35, theta_init=2.0)
    s = torch.full((8,), 0.7)
    for step in range(3000):
        g(s)
        tethergate_step([g], step, lr=1e-2, wd=0.5, ramp_start=0, f_min=0.0, mode='cos', ramp_end=1)
    g.eval()
    return abs(g.theta.item() - g.theta_star()) < 1e-3 and abs(g(s).mean().item() - 0.35) < 1e-3


def test_equilibrium_independent_of_lr():
    """Constant loss pressure r on theta: theta_eq - theta_star = r / wd, for any LR."""
    r, wd, out = 0.02, 0.5, []
    for lr in (1e-3, 1e-2):
        g = TetherGate(beta=6.0, g_target=0.5, theta_init=0.0)
        s = torch.full((4,), 0.4)
        n = int(30 / (lr * wd))                                  # ~30 time constants
        for step in range(n):
            g(s)
            with torch.no_grad():
                g.theta.add_(lr * r)                             # SGD step with dL/dtheta = -r
            tethergate_step([g], step, lr=lr, wd=wd, ramp_start=0, f_min=0.0, mode='cos', ramp_end=1)
        out.append(g.theta.item() - g.theta_star())
    return all(abs(d - r / wd) < 2e-3 for d in out)


def test_global_vs_local_anchor():
    def run(anchor):
        gs = [TetherGate(beta=6.0, g_target=0.5) for _ in range(2)]
        for step in range(2000):
            gs[0](torch.full((4,), 0.2)); gs[1](torch.full((4,), 1.0))
            tethergate_step(gs, step, lr=1e-2, wd=0.5, anchor=anchor, ramp_start=0, f_min=0.0, mode='cos', ramp_end=1)
        return [g.theta.item() for g in gs]
    glob, loc = run('global'), run('local')
    return abs(glob[0] - glob[1]) < 1e-3 and abs(glob[0] - 0.6) < 1e-3 and abs(loc[0] - 0.2) < 1e-3 and abs(loc[1] - 1.0) < 1e-3


if __name__ == '__main__':
    tests = [test_floor_schedule, test_forward_floor, test_gradient_alive_below_floor,
             test_saturation_beyond_clamp_and_spring_recovery, test_detach_signal,
             test_ema_train_eval, test_spring_reaches_target, test_equilibrium_independent_of_lr,
             test_global_vs_local_anchor]
    res = [(t.__name__, bool(t())) for t in tests]
    for n, r in res:
        print(f'  {n:38s} {"OK" if r else "FAIL"}')
    sys.exit(0 if all(r for _, r in res) else 1)