"""Tests for the intercept decorator."""

import jax
import jax.numpy as jnp
import pytest

from jaxray import intercept


class TestInterceptBasic:
    """Basic functionality tests for the intercept decorator."""

    def test_no_cache_no_patch(self):
        """Decorator should work with no caching or patching."""

        @intercept
        def simple(x, _save):
            x = x * 2
            x = _save("doubled", x)
            return x

        output, cache = simple(jnp.array([1.0, 2.0, 3.0]))
        assert jnp.allclose(output, jnp.array([2.0, 4.0, 6.0]))
        assert cache == {}

    def test_cache_single_point(self):
        """Should cache a single activation point."""

        @intercept
        def mlp(x, _save):
            x = x * 2
            x = _save("layer1", x)
            x = x + 1
            x = _save("layer2", x)
            return x

        output, cache = mlp(jnp.array([1.0, 2.0]), cache={"layer1"})
        assert jnp.allclose(output, jnp.array([3.0, 5.0]))
        assert "layer1" in cache
        assert "layer2" not in cache
        assert jnp.allclose(cache["layer1"], jnp.array([2.0, 4.0]))

    def test_cache_multiple_points(self):
        """Should cache multiple activation points."""

        @intercept
        def mlp(x, _save):
            x = x * 2
            x = _save("layer1", x)
            x = x + 1
            x = _save("layer2", x)
            return x

        output, cache = mlp(jnp.array([1.0]), cache={"layer1", "layer2"})
        assert "layer1" in cache
        assert "layer2" in cache
        assert jnp.allclose(cache["layer1"], jnp.array([2.0]))
        assert jnp.allclose(cache["layer2"], jnp.array([3.0]))

    def test_patch_with_value(self):
        """Should replace activation with a fixed value."""

        @intercept
        def mlp(x, _save):
            x = x * 2
            x = _save("layer1", x)
            x = x + 1
            return x

        # Without patch: [1, 2] -> [2, 4] -> [3, 5]
        output_normal, _ = mlp(jnp.array([1.0, 2.0]))
        assert jnp.allclose(output_normal, jnp.array([3.0, 5.0]))

        # With patch: [1, 2] -> [0, 0] (patched) -> [1, 1]
        output_patched, _ = mlp(
            jnp.array([1.0, 2.0]), patch={"layer1": jnp.array([0.0, 0.0])}
        )
        assert jnp.allclose(output_patched, jnp.array([1.0, 1.0]))

    def test_patch_with_callable(self):
        """Should transform activation with a callable."""

        @intercept
        def mlp(x, _save):
            x = x * 2
            x = _save("layer1", x)
            x = x + 1
            return x

        # Patch: multiply by 0.5 instead of keeping original
        # [1, 2] -> [2, 4] -> [1, 2] (patched: *0.5) -> [2, 3]
        output, _ = mlp(jnp.array([1.0, 2.0]), patch={"layer1": lambda act: act * 0.5})
        assert jnp.allclose(output, jnp.array([2.0, 3.0]))

    def test_cache_and_patch_together(self):
        """Should support caching and patching simultaneously."""

        @intercept
        def mlp(x, _save):
            x = x * 2
            x = _save("layer1", x)
            x = x + 1
            x = _save("layer2", x)
            return x

        output, cache = mlp(
            jnp.array([1.0]),
            cache={"layer1"},
            patch={"layer2": lambda act: act * 10},
        )
        # layer1: 1 * 2 = 2 (cached)
        # layer2: 2 + 1 = 3 -> 30 (patched)
        assert jnp.allclose(cache["layer1"], jnp.array([2.0]))
        assert jnp.allclose(output, jnp.array([30.0]))

    def test_cache_returns_pre_patch_value(self):
        """When caching and patching same point, cache should have pre-patch value."""

        @intercept
        def mlp(x, _save):
            x = x * 2
            x = _save("layer1", x)
            return x

        output, cache = mlp(
            jnp.array([5.0]),
            cache={"layer1"},
            patch={"layer1": jnp.array([0.0])},
        )
        # Original value was 10, but we patched to 0
        assert jnp.allclose(cache["layer1"], jnp.array([10.0]))
        assert jnp.allclose(output, jnp.array([0.0]))


class TestInterceptWithParams:
    """Tests with model parameters (more realistic usage)."""

    def test_mlp_with_params(self):
        """Should work with parameterized models."""
        params = {
            "w1": jnp.array([[1.0, 2.0], [3.0, 4.0]]),
            "b1": jnp.array([0.1, 0.2]),
            "w2": jnp.array([[1.0], [1.0]]),
            "b2": jnp.array([0.0]),
        }

        @intercept
        def mlp(params, x, _save):
            x = jnp.dot(x, params["w1"]) + params["b1"]
            x = _save("pre_relu", x)
            x = jax.nn.relu(x)
            x = _save("post_relu", x)
            x = jnp.dot(x, params["w2"]) + params["b2"]
            return x

        x = jnp.array([[1.0, 0.0]])
        output, cache = mlp(params, x, cache={"pre_relu", "post_relu"})

        assert "pre_relu" in cache
        assert "post_relu" in cache
        assert cache["pre_relu"].shape == (1, 2)
        assert cache["post_relu"].shape == (1, 2)


class TestJAXTransforms:
    """Tests for compatibility with JAX transforms."""

    def test_jit_compatibility(self):
        """Should work with jax.jit when cache/patch are fixed via partial."""
        from functools import partial

        @intercept
        def mlp(x, _save):
            x = x * 2
            x = _save("layer1", x)
            x = x + 1
            return x

        # Fix cache/patch via partial, then jit the result
        mlp_with_cache = partial(mlp, cache={"layer1"})
        jitted = jax.jit(mlp_with_cache)
        output, cache = jitted(jnp.array([1.0, 2.0]))

        assert jnp.allclose(output, jnp.array([3.0, 5.0]))
        assert jnp.allclose(cache["layer1"], jnp.array([2.0, 4.0]))

    def test_grad_compatibility(self):
        """Should work with jax.grad."""

        @intercept
        def mlp(x, _save):
            x = x * 2
            x = _save("layer1", x)
            x = x**2
            return x.sum()

        # grad with respect to first argument (x)
        def loss_fn(x):
            output, _ = mlp(x)
            return output

        grad_fn = jax.grad(loss_fn)
        grads = grad_fn(jnp.array([1.0, 2.0, 3.0]))

        # d/dx of sum((2x)^2) = d/dx of sum(4x^2) = 8x
        expected = jnp.array([8.0, 16.0, 24.0])
        assert jnp.allclose(grads, expected)

    def test_vmap_compatibility(self):
        """Should work with jax.vmap when cache/patch are excluded from mapping."""

        @intercept
        def mlp(x, _save):
            x = x * 2
            x = _save("layer1", x)
            x = x + 1
            return x

        # Use functools.partial to fix cache/patch, then vmap over x only
        from functools import partial

        mlp_with_cache = partial(mlp, cache={"layer1"})
        batched = jax.vmap(mlp_with_cache)
        batch_input = jnp.array([[1.0, 2.0], [3.0, 4.0], [5.0, 6.0]])

        output, cache = batched(batch_input)

        assert output.shape == (3, 2)
        assert cache["layer1"].shape == (3, 2)
        assert jnp.allclose(output, batch_input * 2 + 1)

    def test_jit_and_vmap_together(self):
        """Should work with combined jit and vmap."""
        from functools import partial

        @intercept
        def mlp(x, _save):
            x = x * 2
            x = _save("doubled", x)
            x = jax.nn.relu(x - 5)
            x = _save("relu_out", x)
            return x

        mlp_with_cache = partial(mlp, cache={"doubled", "relu_out"})
        jit_vmap = jax.jit(jax.vmap(mlp_with_cache))
        batch = jnp.array([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]])

        output, cache = jit_vmap(batch)

        assert output.shape == (2, 3)
        assert cache["doubled"].shape == (2, 3)
        assert cache["relu_out"].shape == (2, 3)

    def test_grad_through_patched_activation(self):
        """Gradients should flow through patched activations."""

        @intercept
        def mlp(x, _save):
            x = x * 2
            x = _save("layer1", x)
            x = x**2
            return x.sum()

        def loss_fn(x, scale):
            # Patch layer1 by scaling it
            output, _ = mlp(x, patch={"layer1": lambda act: act * scale})
            return output

        # Gradient w.r.t. scale
        grad_scale = jax.grad(loss_fn, argnums=1)
        x = jnp.array([1.0, 2.0])
        scale = 2.0

        # layer1 = 2x, patched = 2x * scale, output = sum((2x*scale)^2)
        # d/d(scale) = sum(2 * (2x*scale) * 2x) = sum(8 * x^2 * scale)
        # At scale=2, x=[1,2]: 8 * (1 + 4) * 2 = 80
        # Wait, let me recalculate:
        # output = sum((2x * scale)^2) = sum(4 * x^2 * scale^2)
        # d/d(scale) = sum(8 * x^2 * scale) = 8 * scale * sum(x^2)
        # At scale=2, x=[1,2]: 8 * 2 * (1 + 4) = 80
        g = grad_scale(x, scale)
        assert jnp.allclose(g, 80.0)
