"""
TDD RED tests for the MLP forward pass in strategies/strategy_mlp_scores.py.

Contract under test:
  - _tanh(x): clamped closed-form tanh, identical formula to the Pine implementation:
        tanh(x) = x >= 20 ?  1.0
                : x <= -20 ? -1.0
                : (exp(2x) - 1) / (exp(2x) + 1)
  - mlp_forward(X, layers): X is a (T, F) float64 matrix; layers is a list of
    (W, b) tuples with W shape (n_out, n_in) and b shape (n_out,).
    Every layer (including the output layer) applies _tanh; the final scalar
    is scaled by 1000 so scores live in [-1000, +1000] like the existing
    activation-score convention.
"""
import numpy as np
import pytest

from strategies.strategy_mlp_scores import _tanh, mlp_forward


class TestTanh:
    def test_matches_numpy_inside_clamp(self):
        x = np.linspace(-19.0, 19.0, 2001)
        assert np.max(np.abs(_tanh(x) - np.tanh(x))) < 1e-12

    def test_clamps_to_exact_one_outside(self):
        # Given the ±20 clamp, When evaluated beyond it, Then output is exactly ±1.
        assert _tanh(np.array([20.0, 25.0, 1e6]))[0] == 1.0
        assert np.all(_tanh(np.array([20.0, 25.0, 1e6])) == 1.0)
        assert np.all(_tanh(np.array([-20.0, -25.0, -1e6])) == -1.0)

    def test_scalar_input(self):
        assert abs(float(_tanh(0.5)) - np.tanh(0.5)) < 1e-12
        assert float(_tanh(0.0)) == 0.0


class TestForward:
    def test_hand_computed_two_layer(self):
        # Given a tiny 3->2->1 net with hand-picked weights,
        # When mlp_forward runs on a single row,
        # Then it reproduces the hand-computed tanh chain to 1e-9.
        W1 = np.array([[0.5, -0.25, 0.1],
                       [0.2,  0.3, -0.4]], dtype=np.float64)
        b1 = np.array([0.1, -0.2], dtype=np.float64)
        W2 = np.array([[0.7, -0.6]], dtype=np.float64)
        b2 = np.array([0.05], dtype=np.float64)
        x = np.array([[1.0, -0.5, 0.25]], dtype=np.float64)

        h = np.tanh(W1 @ x[0] + b1)
        expected = 1000.0 * np.tanh(W2 @ h + b2)[0]

        out = mlp_forward(x, [(W1, b1), (W2, b2)])
        assert out.shape == (1,)
        assert abs(out[0] - expected) < 1e-9

    def test_output_shape_and_range(self):
        rng = np.random.default_rng(42)
        T, arch = 7, [49, 16, 8, 1]
        X = rng.uniform(-1, 1, size=(T, arch[0]))
        layers = [
            (rng.normal(scale=0.5, size=(arch[i + 1], arch[i])),
             rng.normal(scale=0.1, size=arch[i + 1]))
            for i in range(len(arch) - 1)
        ]
        out = mlp_forward(X, layers)
        assert out.shape == (T,)
        assert out.dtype == np.float64
        assert np.all(np.abs(out) <= 1000.0)

    def test_zero_weights_give_zero_score(self):
        X = np.ones((5, 4), dtype=np.float64)
        layers = [(np.zeros((3, 4)), np.zeros(3)), (np.zeros((1, 3)), np.zeros(1))]
        assert np.all(mlp_forward(X, layers) == 0.0)

    def test_matches_torch_reference(self):
        torch = pytest.importorskip("torch")
        rng = np.random.default_rng(7)
        arch = [49, 16, 8, 1]
        X = rng.uniform(-1, 1, size=(64, arch[0]))
        layers = [
            (rng.normal(scale=0.5, size=(arch[i + 1], arch[i])),
             rng.normal(scale=0.1, size=arch[i + 1]))
            for i in range(len(arch) - 1)
        ]
        with torch.no_grad():
            a = torch.tensor(X, dtype=torch.float64)
            for W, b in layers:
                a = torch.tanh(a @ torch.tensor(W.T) + torch.tensor(b))
            ref = (a.squeeze(-1) * 1000.0).numpy()
        assert np.max(np.abs(mlp_forward(X, layers) - ref)) < 1e-9