import copy
import dataclasses
from collections.abc import Callable

import jax
import jax.numpy as jnp
import jax.tree_util as jtu
import numpy as np
import numpy.testing as npt
import pytest

from ion import nn, tree


class TestBaseModule:
    def test_call_raises_not_implemented(self):
        """Base Module.__call__ raises NotImplementedError."""

        class Empty(nn.Module):
            pass

        m = Empty()  # type: ignore[reportCallIssue]
        with pytest.raises(NotImplementedError, match="Empty"):
            m(jnp.ones(3))

    def test_iter_yields_field_values(self):
        """Module.__iter__ yields dataclass field values in order."""
        model = nn.Linear(3, 4, key=jax.random.key(0))
        fields = list(model)
        assert len(fields) == 2  # w and b


class TestSubclassTransformation:
    def test_annotations_become_dataclass_fields(self):
        """Annotated fields are recognized by dataclasses.fields()."""

        class Model(nn.Module):
            a: int
            b: float

            def __init__(self, a: int, b: float):
                self.a = a
                self.b = b

        fields = {f.name for f in dataclasses.fields(Model)}  # type: ignore[arg-type]
        assert fields == {"a", "b"}

    def test_custom_init_preserved(self):
        """Subclass with explicit __init__ keeps its constructor logic."""

        class Model(nn.Module):
            w: nn.Param

            def __init__(self, dim: int, *, key: jax.Array):
                self.w = nn.Param(jax.random.normal(key, (dim,)))

        m = Model(4, key=jax.random.key(0))
        assert m.w.shape == (4,)

    def test_generated_init(self):
        """Subclass without explicit __init__ gets one from annotations."""

        class Pair(nn.Module):
            x: int
            y: int

        p = Pair(x=1, y=2)  # type: ignore[call-arg]
        assert p.x == 1
        assert p.y == 2


class TestImmutability:
    def test_setattr_raises_after_init(self):
        """Assigning to a field after construction raises AttributeError."""

        class Model(nn.Module):
            x: int

            def __init__(self, x: int):
                self.x = x

        m = Model(x=1)
        with pytest.raises(AttributeError, match="frozen"):
            m.x = 2

    def test_delattr_raises(self):
        """Deleting a field always raises AttributeError."""

        class Model(nn.Module):
            x: int

            def __init__(self, x: int):
                self.x = x

        m = Model(x=1)
        with pytest.raises(AttributeError, match="Cannot delete"):
            del m.x

    def test_error_message_contents(self):
        """Error messages include the class name and attribute name."""

        class MyLayer(nn.Module):
            w: int

            def __init__(self, w: int):
                self.w = w

        m = MyLayer(w=1)
        with pytest.raises(AttributeError, match="MyLayer"):
            m.w = 2
        with pytest.raises(AttributeError, match="'w'"):
            m.w = 2


class TestPytreeRegistration:
    def test_flatten_unflatten_roundtrip(self):
        """Flatten then unflatten reconstructs the module exactly."""

        class Model(nn.Module):
            w: nn.Param
            b: nn.Param

            def __init__(self, key: jax.Array):
                keys = jax.random.split(key, 2)
                self.w = nn.Param(jax.random.normal(keys[0], (3, 4)))
                self.b = nn.Param(jnp.zeros(4))

        m = Model(key=jax.random.key(0))
        leaves, treedef = jtu.tree_flatten(m)
        reconstructed = treedef.unflatten(leaves)
        npt.assert_array_equal(reconstructed.w._value, m.w._value)
        npt.assert_array_equal(reconstructed.b._value, m.b._value)

    def test_children_follow_field_order(self):
        """Non-array fields become static aux, so leaves are empty, but roundtrip preserves values."""

        class Model(nn.Module):
            first: int
            second: int
            third: int

            def __init__(self, first: int, second: int, third: int):
                self.first = first
                self.second = second
                self.third = third

        m = Model(first=1, second=2, third=3)
        leaves = jax.tree.leaves(m)
        assert leaves == []
        # Roundtrip preserves field values
        reconstructed = jtu.tree_unflatten(*reversed(jtu.tree_flatten(m)))
        assert reconstructed.first == 1
        assert reconstructed.second == 2
        assert reconstructed.third == 3

    def test_unflatten_bypasses_init(self):
        """Unflatten works even when __init__ takes different args than stored fields."""

        class Model(nn.Module):
            w: nn.Param

            def __init__(self, in_dim: int, out_dim: int, *, key: jax.Array):
                self.w = nn.Param(jax.random.normal(key, (in_dim, out_dim)))

        m = Model(3, 4, key=jax.random.key(0))
        leaves, treedef = jtu.tree_flatten(m)
        # This would fail if unflatten tried to call __init__(in_dim, out_dim, key)
        reconstructed = treedef.unflatten(leaves)
        npt.assert_array_equal(reconstructed.w._value, m.w._value)

    def test_tree_map(self):
        """jax.tree.map transforms leaves inside a module."""

        class Model(nn.Module):
            w: nn.Param

            def __init__(self):
                self.w = nn.Param(jnp.array([1.0, 2.0]))

        m = Model()
        doubled = jax.tree.map(lambda x: x * 2, m)
        assert isinstance(doubled, Model)
        npt.assert_array_equal(doubled.w._value, jnp.array([2.0, 4.0]))

    def test_tree_leaves(self):
        """jax.tree.leaves extracts all leaf values."""

        class Model(nn.Module):
            w: nn.Param
            b: nn.Param

            def __init__(self):
                self.w = nn.Param(jnp.array([1.0]))
                self.b = nn.Param(jnp.array([2.0]))

        m = Model()
        leaves = jax.tree.leaves(m)
        assert len(leaves) == 2
        npt.assert_array_equal(leaves[0], jnp.array([1.0]))
        npt.assert_array_equal(leaves[1], jnp.array([2.0]))


class TestReplace:
    def test_single_field(self):
        """Replacing one field returns correct values."""

        class Model(nn.Module):
            a: int
            b: int

            def __init__(self, a: int, b: int):
                self.a = a
                self.b = b

        m = Model(a=1, b=2)
        m2 = m.replace(b=10)
        assert m2.a == 1
        assert m2.b == 10

    def test_multiple_fields(self):
        """Replacing several fields at once."""

        class Model(nn.Module):
            a: int
            b: int
            c: int

            def __init__(self, a: int, b: int, c: int):
                self.a = a
                self.b = b
                self.c = c

        m = Model(a=1, b=2, c=3)
        m2 = m.replace(a=10, c=30)
        assert m2.a == 10
        assert m2.b == 2
        assert m2.c == 30

    def test_returns_new_instance(self):
        """Original module is unchanged after replace."""

        class Model(nn.Module):
            x: int

            def __init__(self, x: int):
                self.x = x

        m = Model(x=1)
        m2 = m.replace(x=2)
        assert m.x == 1
        assert m2.x == 2
        assert m is not m2

    def test_replaced_is_frozen(self):
        """The copy returned by replace is also immutable."""

        class Model(nn.Module):
            x: int

            def __init__(self, x: int):
                self.x = x

        m = Model(x=1)
        m2 = m.replace(x=2)
        with pytest.raises(AttributeError, match="frozen"):
            m2.x = 3

    def test_preserves_type(self):
        """replace returns the same subclass type."""

        class Child(nn.Module):
            x: int

            def __init__(self, x: int):
                self.x = x

        m = Child(x=1)
        m2 = m.replace(x=2)
        assert type(m2) is Child


class TestParams:
    def test_filters_non_param_leaves(self):
        """Non-Param fields are preserved as static metadata."""

        class Model(nn.Module):
            w: nn.Param
            scale: float

            def __init__(self):
                self.w = nn.Param(jnp.array([1.0, 2.0]))
                self.scale = 3.0

        m = Model()
        params = m.params
        assert isinstance(params.w, nn.Param)
        assert params.scale == 3.0

    def test_no_params(self):
        """Module with zero Param fields has no param leaves."""

        class Config(nn.Module):
            size: int
            rate: float

            def __init__(self, size: int, rate: float):
                self.size = size
                self.rate = rate

        m = Config(size=8, rate=0.1)
        leaves = jax.tree.leaves(m.params)
        assert len(leaves) == 0

    def test_nested_module_params(self):
        """Param leaves in nested child modules are included."""
        key = jax.random.key(0)

        class Container(nn.Module):
            layer1: nn.Linear
            layer2: nn.Linear

            def __init__(self, key: jax.Array):
                keys = jax.random.split(key, 2)
                self.layer1 = nn.Linear(4, 8, key=keys[0])
                self.layer2 = nn.Linear(8, 2, key=keys[1])

        m = Container(key=key)
        param_leaves = jax.tree.leaves(m.params)
        # Linear has w (Param) and b (Param), so 2 layers * 2 params = 4
        assert len(param_leaves) == 4
        for leaf in param_leaves:
            assert isinstance(leaf, jnp.ndarray)


class TestNumParams:
    def test_linear(self):
        """Linear(4, 8) has 4*8 + 8 = 40 parameters."""
        model = nn.Linear(4, 8, key=jax.random.key(0))
        assert model.num_params == 40

    def test_no_bias(self):
        """Linear without bias only counts weights."""
        model = nn.Linear(4, 8, bias=False, key=jax.random.key(0))
        assert model.num_params == 32

    def test_nested(self):
        """Nested modules sum all Param leaves."""

        class Net(nn.Module):
            a: nn.Linear
            b: nn.Linear

            def __init__(self, key):
                keys = jax.random.split(key, 2)
                self.a = nn.Linear(4, 8, key=keys[0])
                self.b = nn.Linear(8, 2, key=keys[1])

        model = Net(key=jax.random.key(0))
        # a: 4*8+8=40, b: 8*2+2=18
        assert model.num_params == 58

    def test_no_params(self):
        """Module with no Param fields has 0 parameters."""
        model = nn.Identity()
        assert model.num_params == 0

    def test_frozen_params_counted(self):
        """internals.md: num_params includes both trainable and frozen params."""
        model = nn.Linear(4, 8, key=jax.random.key(0))
        trainable_count = model.num_params
        frozen_count = model.freeze().num_params
        assert trainable_count == frozen_count


class TestRepr:
    def test_param_field(self):
        """repr contains Param wrapper for Param fields."""

        class Model(nn.Module):
            w: nn.Param

            def __init__(self):
                self.w = nn.Param(jnp.zeros((2, 3)))

        r = repr(Model())
        assert "Param(" in r

    def test_array_field(self):
        """repr contains abbreviated dtype and shape for plain array fields."""

        class Model(nn.Module):
            x: jax.Array

            def __init__(self):
                self.x = jnp.zeros((3, 4), dtype=jnp.float32)

        r = repr(Model())
        assert "f32" in r
        assert "[3, 4]" in r

    def test_callable_field(self):
        """repr contains the function __name__ for callable fields."""

        class Model(nn.Module):
            act: Callable

            def __init__(self):
                self.act = jax.nn.relu

        r = repr(Model())
        assert "relu" in r

    def test_nested_modules(self):
        """Nested module tuple indents with open/close brackets."""

        class Container(nn.Module):
            layers: tuple

            def __init__(self):
                self.layers = (nn.Identity(), nn.Identity())

        r = repr(Container())
        assert "(" in r  # tuple open bracket
        assert "Identity()" in r

    def test_empty_module(self):
        """Module with no fields shows ClassName()."""

        class Empty(nn.Module):
            def __init__(self):
                pass

        assert repr(Empty()) == "Empty()"


class TestFreezeUnfreeze:
    def test_freeze_sets_all_params_to_non_trainable(self):
        class Model(nn.Module):
            w: nn.Param
            b: nn.Param

            def __init__(self):
                self.w = nn.Param(jnp.ones(3))
                self.b = nn.Param(jnp.zeros(3))

        m = Model().freeze()
        assert m.w.trainable is False
        assert m.b.trainable is False

    def test_unfreeze_sets_all_params_to_trainable(self):
        class Model(nn.Module):
            w: nn.Param

            def __init__(self):
                self.w = nn.Param(jnp.ones(3), trainable=False)

        m = Model().unfreeze()
        assert m.w.trainable is True

    def test_freeze_unfreeze_roundtrip(self):
        class Model(nn.Module):
            w: nn.Param

            def __init__(self):
                self.w = nn.Param(jnp.ones(3))

        m = Model()
        npt.assert_array_equal(m.freeze().unfreeze().w._value, m.w._value)
        assert m.freeze().unfreeze().w.trainable is True

    def test_freeze_preserves_values(self):
        class Model(nn.Module):
            w: nn.Param

            def __init__(self, key):
                self.w = nn.Param(jax.random.normal(key, (4,)))

        m = Model(key=jax.random.key(0))
        frozen = m.freeze()
        npt.assert_array_equal(frozen.w._value, m.w._value)

    def test_freeze_nested_module(self):
        class Inner(nn.Module):
            w: nn.Param

            def __init__(self):
                self.w = nn.Param(jnp.ones(2))

        class Outer(nn.Module):
            inner: Inner
            b: nn.Param

            def __init__(self):
                self.inner = Inner()
                self.b = nn.Param(jnp.zeros(2))

        m = Outer().freeze()
        assert m.inner.w.trainable is False
        assert m.b.trainable is False

    def test_partial_freeze_via_replace(self):
        """Freeze one sub-module, keep the other trainable."""

        class Model(nn.Module):
            encoder: nn.Linear
            decoder: nn.Linear

            def __init__(self, key):
                keys = jax.random.split(key, 2)
                self.encoder = nn.Linear(4, 8, key=keys[0])
                self.decoder = nn.Linear(8, 4, key=keys[1])

        m = Model(key=jax.random.key(0))
        m = m.replace(encoder=m.encoder.freeze())
        assert m.encoder.w.trainable is False
        assert m.decoder.w.trainable is True

    def test_freeze_idempotent(self):
        class Model(nn.Module):
            w: nn.Param

            def __init__(self):
                self.w = nn.Param(jnp.ones(2), trainable=False)

        m = Model().freeze()
        assert m.w.trainable is False


class TestAstype:
    def test_astype_method(self):
        """Module.astype delegates to tree.astype correctly."""

        class Model(nn.Module):
            w: nn.Param
            buf: jax.Array

            def __init__(self):
                self.w = nn.Param(jnp.ones(2, dtype=jnp.float32))
                self.buf = jnp.zeros(2, dtype=jnp.float32)

        m = Model()
        via_method = m.astype(jnp.bfloat16)
        via_func = tree.astype(m, jnp.bfloat16)
        assert via_method.w._value.dtype == jnp.bfloat16
        assert via_method.buf.dtype == jnp.bfloat16
        npt.assert_allclose(via_method.w._value, via_func.w._value)
        npt.assert_allclose(via_method.buf, via_func.buf)


class TestNoneField:
    def test_module_with_none_field(self):
        """Module can store None as a field value (e.g., optional bias)."""
        linear = nn.Linear(4, 8, bias=False, key=jax.random.key(0))
        assert linear.b is None

    def test_none_field_survives_pytree_roundtrip(self):
        linear = nn.Linear(4, 8, bias=False, key=jax.random.key(0))
        leaves, treedef = jtu.tree_flatten(linear)
        reconstructed = treedef.unflatten(leaves)
        assert reconstructed.b is None
        npt.assert_array_equal(reconstructed.w._value, linear.w._value)

    def test_none_field_works_under_jit(self):
        linear = nn.Linear(4, 8, bias=False, key=jax.random.key(0))
        x = jnp.ones((1, 4))
        eager = linear(x)
        jitted = jax.jit(linear)(x)
        npt.assert_allclose(jitted, eager)


class TestContainerFields:
    def test_tuple_of_modules(self):
        """Module with a tuple of sub-modules."""
        keys = jax.random.split(jax.random.key(0), 2)
        seq = nn.Sequential(nn.Linear(4, 8, key=keys[0]), nn.Linear(8, 2, key=keys[1]))
        x = jnp.ones((1, 4))
        result = jax.jit(seq)(x)
        assert result.shape == (1, 2)

    def test_tuple_of_mixed_callables(self):
        """Sequential with Modules and plain functions."""
        keys = jax.random.split(jax.random.key(0), 2)
        seq = nn.Sequential(
            nn.Linear(4, 8, key=keys[0]),
            jax.nn.relu,
            nn.Linear(8, 2, key=keys[1]),
        )
        x = jnp.ones((1, 4))
        eager = seq(x)
        jitted = jax.jit(seq)(x)
        npt.assert_allclose(jitted, eager, rtol=1e-5, atol=1e-5)

    def test_list_field(self):
        """Module with a list of sub-modules."""

        class Container(nn.Module):
            layers: list

            def __init__(self, key):
                keys = jax.random.split(key, 2)
                self.layers = [nn.Linear(4, 4, key=keys[0]), nn.Linear(4, 4, key=keys[1])]

            def __call__(self, x):
                for layer in self.layers:
                    x = layer(x)
                return x

        m = Container(key=jax.random.key(0))
        x = jnp.ones((1, 4))
        eager = m(x)
        jitted = jax.jit(m)(x)
        npt.assert_allclose(jitted, eager)


class TestReplaceEdgeCases:
    def test_replace_unknown_field_raises(self):
        """replace raises ValueError for unknown field names."""

        class Model(nn.Module):
            x: int

            def __init__(self, x: int):
                self.x = x

        m = Model(x=1)
        with pytest.raises(ValueError, match="Unknown field"):
            m.replace(nonexistent=42)

    def test_replace_with_param(self):
        """replace can swap a Param value."""
        linear = nn.Linear(4, 8, key=jax.random.key(0))
        new_w = nn.Param(jnp.zeros_like(linear.w._value))
        replaced = linear.replace(w=new_w)
        npt.assert_array_equal(replaced.w._value, jnp.zeros_like(linear.w._value))
        # Original unchanged
        assert not jnp.array_equal(linear.w._value, replaced.w._value)

    def test_replace_param_with_non_param_breaks_tree_ops(self):
        """Replacing a Param field with a plain value creates a structurally different pytree."""

        class Model(nn.Module):
            w: nn.Param

            def __init__(self):
                self.w = nn.Param(jnp.ones(3))

        m = Model()
        original_def = jtu.tree_flatten(m)[1]

        # Replace Param with plain array
        m2 = m.replace(w=jnp.ones(3))
        new_def = jtu.tree_flatten(m2)[1]

        # The tree structures are different!
        assert original_def != new_def

    def test_replace_changing_structure_works_under_jit(self):
        """internals.md: 'replace() can change pytree structure.' JIT will recompile this."""
        linear = nn.Linear(4, 8, key=jax.random.key(0))
        no_bias = linear.replace(b=None)

        @jax.jit
        def f(m, x):
            return m(x)

        x = jnp.ones((1, 4))
        r1 = f(linear, x)
        r2 = f(no_bias, x)
        assert r1.shape == (1, 8)
        assert r2.shape == (1, 8)

    def test_replace_with_none_changes_structure(self):
        """Replacing a Param with None changes pytree structure (loses a leaf)."""

        class Model(nn.Module):
            w: nn.Param
            b: nn.Param

            def __init__(self):
                self.w = nn.Param(jnp.ones(3))
                self.b = nn.Param(jnp.zeros(3))

        m = Model()
        m2 = m.replace(b=None)
        # Original has 2 leaves (w._value, b._value), replacement has 1
        assert len(jax.tree.leaves(m)) == 2
        assert len(jax.tree.leaves(m2)) == 1


class TestInheritance:
    def test_module_subclass_chain(self):
        """Subclass of a subclass works correctly."""

        class Base(nn.Module):
            x: int

            def __init__(self, x: int):
                self.x = x

        class Child(Base):
            y: int

            def __init__(self, x: int, y: int):
                self.x = x
                self.y = y

        c = Child(x=1, y=2)
        assert c.x == 1
        assert c.y == 2
        # Immutable
        with pytest.raises(AttributeError, match="frozen"):
            c.x = 3
        # Pytree roundtrip
        leaves, treedef = jtu.tree_flatten(c)
        reconstructed = treedef.unflatten(leaves)
        assert reconstructed.x == 1
        assert reconstructed.y == 2

    def test_super_init_does_not_freeze_early(self):
        """Child calling super().__init__() can still set its own fields."""

        class Base(nn.Module):
            w: nn.Param

            def __init__(self, key):
                self.w = nn.Param(jax.random.normal(key, (4,)))

        class Child(Base):
            b: nn.Param

            def __init__(self, key):
                super().__init__(key)
                self.b = nn.Param(jnp.zeros(4))

        c = Child(jax.random.key(0))
        assert c.w.shape == (4,)
        assert c.b.shape == (4,)
        # Frozen after construction
        with pytest.raises(AttributeError, match="frozen"):
            c.b = nn.Param(jnp.ones(4))

    def test_three_level_inheritance_with_super(self):
        """Three levels of inheritance using super().__init__()."""

        class Base(nn.Module):
            w: nn.Param

            def __init__(self, key):
                self.w = nn.Param(jax.random.normal(key, (4,)))

        class Mid(Base):
            b: nn.Param

            def __init__(self, key):
                super().__init__(key)
                self.b = nn.Param(jnp.zeros(4))

        class Top(Mid):
            scale: float

            def __init__(self, key):
                super().__init__(key)
                self.scale = 2.0

            def __call__(self, x):
                return (x @ self.w + self.b) * self.scale

        m = Top(jax.random.key(0))
        assert m.w.shape == (4,)
        assert m.b.shape == (4,)
        assert m.scale == 2.0
        # Works in jit
        result = jax.jit(m)(jnp.ones(4))
        assert result.shape == (4,)
        # Frozen after construction
        with pytest.raises(AttributeError, match="frozen"):
            m.scale = 3.0

    def test_inherited_module_pytree_roundtrip(self):
        """Inherited module survives flatten/unflatten correctly."""

        class Base(nn.Module):
            w: nn.Param

            def __init__(self, key):
                self.w = nn.Param(jax.random.normal(key, (4,)))

        class Child(Base):
            b: nn.Param

            def __init__(self, key):
                super().__init__(key)
                self.b = nn.Param(jnp.zeros(4))

        m = Child(jax.random.key(0))
        leaves, treedef = jtu.tree_flatten(m)
        m2 = treedef.unflatten(leaves)
        npt.assert_array_equal(m2.w._value, m.w._value)
        npt.assert_array_equal(m2.b._value, m.b._value)

    def test_inherited_module_grad(self):
        """jax.grad works on an inherited module."""

        class Base(nn.Module):
            w: nn.Param

            def __init__(self, key):
                self.w = nn.Param(jax.random.normal(key, (4,)))

        class Child(Base):
            b: nn.Param

            def __init__(self, key):
                super().__init__(key)
                self.b = nn.Param(jnp.zeros(4))

        m = Child(jax.random.key(0))
        grads = jax.grad(lambda m: (m.w + m.b).sum())(m)
        npt.assert_allclose(grads.w._value, jnp.ones(4))
        npt.assert_allclose(grads.b._value, jnp.ones(4))


class TestParamsWithFrozen:
    def test_params_includes_frozen_params(self):
        """params property returns both trainable and frozen Params."""

        class Model(nn.Module):
            w: nn.Param
            b: nn.Param

            def __init__(self):
                self.w = nn.Param(jnp.ones(3))
                self.b = nn.Param(jnp.zeros(3), trainable=False)

        m = Model()
        params = m.params
        assert isinstance(params.w, nn.Param)
        assert isinstance(params.b, nn.Param)
        assert params.w.trainable is True
        assert params.b.trainable is False

    def test_params_on_plain_array_field(self):
        """Plain array (non-Param) fields become None in params."""

        class Model(nn.Module):
            w: nn.Param
            buf: jax.Array

            def __init__(self):
                self.w = nn.Param(jnp.ones(3))
                self.buf = jnp.array([1.0, 2.0, 3.0])

        m = Model()
        params = m.params
        assert isinstance(params.w, nn.Param)
        assert params.buf is None


class TestDeepNesting:
    def test_three_level_nesting(self):
        """Three levels of module nesting work correctly."""

        class Inner(nn.Module):
            w: nn.Param

            def __init__(self, key):
                self.w = nn.Param(jax.random.normal(key, (2,)))

        class Middle(nn.Module):
            inner: Inner

            def __init__(self, key):
                self.inner = Inner(key)

        class Outer(nn.Module):
            middle: Middle

            def __init__(self, key):
                self.middle = Middle(key)

        m = Outer(key=jax.random.key(0))
        leaves = jax.tree.leaves(m)
        assert len(leaves) == 1
        # Roundtrip
        reconstructed = jtu.tree_unflatten(*reversed(jtu.tree_flatten(m)))
        npt.assert_array_equal(reconstructed.middle.inner.w._value, m.middle.inner.w._value)
        # jit
        result = jax.jit(lambda m: jnp.sum(m.middle.inner.w))(m)
        npt.assert_allclose(result, jnp.sum(m.middle.inner.w._value))


class TestStaticWrapping:
    def test_jax_jit_with_callable_field(self):
        """Module with a callable field works under jax.jit."""

        class Model(nn.Module):
            w: nn.Param
            activation: Callable

            def __init__(self, *, key: jax.Array):
                self.w = nn.Param(jax.random.normal(key, (2, 2)))
                self.activation = jax.nn.relu

        def forward(model, x):
            return model.activation(x @ model.w)

        model = Model(key=jax.random.key(0))
        x = jnp.ones((1, 2))
        eager = forward(model, x)
        jitted = jax.jit(forward)(model, x)
        npt.assert_allclose(jitted, eager)

    def test_jax_jit_with_int_field(self):
        """Module with an int field works under jax.jit."""

        class Model(nn.Module):
            w: nn.Param
            scale: int

            def __init__(self, *, key: jax.Array):
                self.w = nn.Param(jax.random.normal(key, (3,)))
                self.scale = 5

        def forward(model):
            return jnp.sum(model.w) * model.scale

        model = Model(key=jax.random.key(0))
        eager = forward(model)
        jitted = jax.jit(forward)(model)
        npt.assert_allclose(jitted, eager)

    def test_jax_grad_with_module(self):
        """jax.grad works directly on a module, producing gradients for all array leaves."""

        class Model(nn.Module):
            w: nn.Param

            def __init__(self):
                self.w = nn.Param(jnp.array([1.0, 2.0, 3.0]))

        def loss(model):
            return jnp.sum(model.w * jnp.array([4.0, 5.0, 6.0]))

        grads = jax.grad(loss)(Model())
        npt.assert_allclose(grads.w._value, jnp.array([4.0, 5.0, 6.0]))


class TestModuleMutableFields:
    def test_list_field_can_be_mutated_in_place(self):
        """Module freezing doesn't prevent in-place mutation of mutable containers."""

        class Model(nn.Module):
            layers: list

            def __init__(self, key):
                self.layers = [nn.Linear(4, 4, key=key)]

        model = Model(key=jax.random.key(0))

        with pytest.raises(AttributeError, match="frozen"):
            model.layers = []

        # But in-place mutation bypasses the freeze
        original_len = len(model.layers)
        model.layers.append(nn.Linear(4, 4, key=jax.random.key(1)))
        assert len(model.layers) == original_len + 1

    def test_numpy_array_field_can_be_mutated(self):
        """numpy arrays in Module fields can be mutated in-place."""

        class Model(nn.Module):
            mask: np.ndarray
            w: nn.Param

            def __init__(self, key):
                self.mask = np.array([1.0, 0.0, 1.0])
                self.w = nn.Param(jax.random.normal(key, (3,)))

        model = Model(key=jax.random.key(0))

        original_val = model.mask[0]
        model.mask[0] = 999.0
        assert model.mask[0] == 999.0
        assert model.mask[0] != original_val


class TestModuleInheritanceEdgeCases:
    def test_module_with_no_fields(self):
        """Module subclass with no fields at all."""

        class Empty(nn.Module):
            def __init__(self):
                pass

        m = Empty()
        assert jax.tree.leaves(m) == []
        reconstructed = jtu.tree_unflatten(*reversed(jtu.tree_flatten(m)))
        assert isinstance(reconstructed, Empty)

    def test_abstract_base_then_concrete(self):
        """Abstract base Module with no fields, concrete child with fields."""

        class BaseLayer(nn.Module):
            def forward(self, x):
                raise NotImplementedError

        class ConcreteLayer(BaseLayer):
            w: nn.Param

            def __init__(self, key):
                self.w = nn.Param(jax.random.normal(key, (4,)))

            def forward(self, x):
                return x @ self.w

        m = ConcreteLayer(key=jax.random.key(0))
        assert m.w.shape == (4,)
        leaves, treedef = jtu.tree_flatten(m)
        m2 = treedef.unflatten(leaves)
        npt.assert_array_equal(m2.w._value, m.w._value)

    def test_sibling_classes_dont_interfere(self):
        """Two sibling Module subclasses with same field names don't interfere."""

        class A(nn.Module):
            x: int

            def __init__(self, x):
                self.x = x

        class B(nn.Module):
            x: int

            def __init__(self, x):
                self.x = x

        a = A(x=1)
        b = B(x=2)
        assert a.x == 1
        assert b.x == 2
        assert type(a) is not type(b)


class TestModuleCopy:
    def test_deepcopy_module_preserves_param_wrappers(self):
        """copy.deepcopy on Module preserves Param wrappers."""

        class Model(nn.Module):
            w: nn.Param

            def __init__(self):
                self.w = nn.Param(jnp.ones(3))

        m = Model()
        m2 = copy.deepcopy(m)
        assert isinstance(m2, Model)
        assert isinstance(m2.w, nn.Param)
        npt.assert_array_equal(m2.w._value, m.w._value)
        assert m2.w.trainable == m.w.trainable

    def test_copy_module_works(self):
        """copy.copy on Module works and preserves Param wrappers."""

        class Model(nn.Module):
            w: nn.Param

            def __init__(self):
                self.w = nn.Param(jnp.ones(3))

        m = Model()
        m2 = copy.copy(m)
        assert isinstance(m2.w, nn.Param)
        npt.assert_array_equal(m2.w._value, m.w._value)


class TestModuleWrappingEdgeCases:
    def test_numpy_scalar_treated_as_dynamic_leaf(self):
        """np.ndarray (even 0-d) is a dynamic leaf, not static metadata."""

        class Model(nn.Module):
            w: nn.Param
            buf: np.ndarray

            def __init__(self):
                self.w = nn.Param(jnp.ones(2))
                self.buf = np.array(42.0)

        m = Model()
        leaves = jax.tree.leaves(m)
        assert len(leaves) == 2

    def test_python_int_is_static_metadata(self):
        """Plain Python int becomes static aux (no dynamic leaf)."""

        class Model(nn.Module):
            w: nn.Param
            count: int

            def __init__(self):
                self.w = nn.Param(jnp.ones(2))
                self.count = 42

        m = Model()
        leaves = jax.tree.leaves(m)
        assert len(leaves) == 1

    def test_custom_class_field_becomes_static(self):
        """A custom (non-Module) class instance becomes static aux."""

        class Config:
            def __init__(self, lr):
                self.lr = lr

        class Model(nn.Module):
            w: nn.Param
            config: Config

            def __init__(self, key):
                self.w = nn.Param(jax.random.normal(key, (2,)))
                self.config = Config(lr=0.001)

        m = Model(key=jax.random.key(0))
        leaves = jax.tree.leaves(m)
        assert len(leaves) == 1
        reconstructed = jtu.tree_unflatten(*reversed(jtu.tree_flatten(m)))
        assert reconstructed.config.lr == 0.001

    def test_dict_field_with_arrays_are_dynamic(self):
        """Arrays inside a dict field are dynamic leaves, not static."""

        class Model(nn.Module):
            w: nn.Param
            buffers: dict

            def __init__(self, key):
                self.w = nn.Param(jax.random.normal(key, (2,)))
                self.buffers = {"mask": jnp.array([1.0, 0.0]), "scale": jnp.array(2.0)}

        m = Model(key=jax.random.key(0))
        leaves = jax.tree.leaves(m)
        assert len(leaves) == 3

    def test_dict_field_mixed_static_dynamic(self):
        """Dict with mixed array and non-array values: arrays are dynamic, others are static."""

        class Model(nn.Module):
            w: nn.Param
            config: dict

            def __init__(self, key):
                self.w = nn.Param(jax.random.normal(key, (2,)))
                self.config = {"lr": 0.001, "buf": jnp.array(1.0)}

        m = Model(key=jax.random.key(0))
        leaves = jax.tree.leaves(m)
        assert len(leaves) == 2

        reconstructed = jtu.tree_unflatten(*reversed(jtu.tree_flatten(m)))
        assert reconstructed.config["lr"] == 0.001


class TestReprInsideTransformations:
    """Test that __repr__ doesn't crash inside JAX transformations."""

    def test_repr_inside_jit(self):
        """Calling repr() on a module inside jit doesn't crash."""

        class Model(nn.Module):
            w: nn.Param
            scale: float

            def __init__(self):
                self.w = nn.Param(jnp.ones(3))
                self.scale = 2.0

        m = Model()

        @jax.jit
        def f(m):
            _ = repr(m)
            return jnp.sum(m.w)

        result = f(m)
        npt.assert_allclose(result, 3.0)

    def test_repr_inside_vmap(self):
        """Calling repr() on a module inside vmap doesn't crash."""

        class Model(nn.Module):
            w: nn.Param

            def __init__(self):
                self.w = nn.Param(jnp.ones(3))

        m = Model()

        def f(m):
            _ = repr(m)
            return jnp.sum(m.w)

        # vmap over a batch dim added to the param
        batched_m = jax.tree.map(lambda x: jnp.stack([x, x]), m)
        results = jax.vmap(f)(batched_m)
        npt.assert_allclose(results, jnp.array([3.0, 3.0]))


class TestStaticFieldTypes:
    """Test that various field types (list, dict, tuple, callable) work through jit."""

    def test_list_of_ints_through_jit(self):
        class Model(nn.Module):
            w: nn.Param
            dims: list

            def __init__(self):
                self.w = nn.Param(jnp.ones(3))
                self.dims = [1, 2, 3]

            def __call__(self, x):
                return x * self.dims[0]

        m = Model()
        x = jnp.array(5.0)
        npt.assert_allclose(jax.jit(m)(x), 5.0)

    def test_dict_of_scalars_through_jit(self):
        class Model(nn.Module):
            w: nn.Param
            config: dict

            def __init__(self):
                self.w = nn.Param(jnp.ones(3))
                self.config = {"scale": 2.0, "offset": 1.0}

            def __call__(self, x):
                return x * self.config["scale"] + self.config["offset"]

        m = Model()
        x = jnp.array(3.0)
        npt.assert_allclose(jax.jit(m)(x), 7.0)

    def test_tuple_of_callables_through_jit(self):
        class Model(nn.Module):
            w: nn.Param
            activations: tuple

            def __init__(self):
                self.w = nn.Param(jnp.ones(3))
                self.activations = (jax.nn.relu, jax.nn.sigmoid)

            def __call__(self, x):
                for act in self.activations:
                    x = act(x)
                return x

        m = Model()
        x = jnp.array([-1.0, 0.0, 1.0])
        eager = m(x)
        jitted = jax.jit(m)(x)
        npt.assert_allclose(jitted, eager)

    def test_nested_dict_with_arrays_through_jit(self):
        """Dict mixing arrays and scalars: arrays are dynamic, scalars are static."""

        class Model(nn.Module):
            w: nn.Param
            meta: dict

            def __init__(self):
                self.w = nn.Param(jnp.ones(3))
                self.meta = {"scale": 2.0, "mask": jnp.array([1.0, 0.0, 1.0])}

            def __call__(self, x):
                return x * self.meta["mask"] * self.meta["scale"]

        m = Model()
        x = jnp.array([1.0, 2.0, 3.0])
        eager = m(x)
        jitted = jax.jit(m)(x)
        npt.assert_allclose(jitted, eager)

    def test_mutating_static_list_uses_stale_cache(self):
        """Sharp edge: mutating a mutable field in-place does NOT trigger retrace.
        JAX caches by object identity, so the same list object (now mutated)
        still hits the cached trace with the old value baked in."""

        class Model(nn.Module):
            w: nn.Param
            dims: list

            def __init__(self):
                self.w = nn.Param(jnp.ones(3))
                self.dims = [1]

            def __call__(self, x):
                return x * len(self.dims)

        m = Model()
        x = jnp.array(1.0)

        result1 = jax.jit(m)(x)
        npt.assert_allclose(result1, 1.0)

        # Mutate the list in-place, JAX does NOT retrace
        m.dims.append(2)
        result2 = jax.jit(m)(x)
        # Stale! Still returns 1.0, not 2.0, because the old trace is cached
        npt.assert_allclose(result2, 1.0)


class TestPytreeSharedReferences:
    """internals.md: 'JAX pytrees are trees, not graphs'.'"""

    def test_shared_module_is_duplicated_through_flatten_unflatten(self):
        """Shared sub-module reference becomes two independent copies after roundtrip."""

        class Inner(nn.Module):
            w: nn.Param

            def __init__(self, key):
                self.w = nn.Param(jax.random.normal(key, (2,)))

        class Shared(nn.Module):
            a: Inner
            b: Inner

            def __init__(self, layer):
                self.a = layer
                self.b = layer

        layer = Inner(jax.random.key(0))
        model = Shared(layer)
        assert model.a is model.b

        rebuilt = jtu.tree_unflatten(*reversed(jtu.tree_flatten(model)))
        assert rebuilt.a is not rebuilt.b
        npt.assert_array_equal(rebuilt.a.w._value, rebuilt.b.w._value)

    def test_shared_param_is_duplicated(self):
        """Shared Param reference becomes two independent copies after roundtrip."""

        class SharedParam(nn.Module):
            a: nn.Param
            b: nn.Param

            def __init__(self, p):
                self.a = p
                self.b = p

        p = nn.Param(jnp.ones(3))
        model = SharedParam(p)
        assert model.a is model.b

        rebuilt = jtu.tree_unflatten(*reversed(jtu.tree_flatten(model)))
        assert rebuilt.a is not rebuilt.b

    def test_weight_tying_via_array_reference(self):
        """Weight tying via direct array reference works and grads flow through."""

        class TiedModel(nn.Module):
            embed: nn.Param
            hidden: nn.Param

            def __init__(self, key: jax.Array):
                self.embed = nn.Param(jax.random.normal(key, (4, 8)))
                self.hidden = nn.Param(jnp.zeros(8))

            def decode(self, h):
                return h @ self.embed.T

        model = TiedModel(key=jax.random.key(0))
        h = jnp.ones(8)
        out = model.decode(h)
        assert out.shape == (4,)

        grads = jax.grad(lambda m, h: jnp.sum(m.decode(h)))(model, h)
        assert jnp.any(grads.embed._value != 0)


class TestParamsPropertyEdgeCases:
    def test_params_replaces_plain_array_with_none_but_keeps_static(self):
        """Plain arrays become None but static fields are preserved (by design)."""

        class Model(nn.Module):
            w: nn.Param
            buf: jax.Array
            scale: float

            def __init__(self):
                self.w = nn.Param(jnp.ones(3))
                self.buf = jnp.ones(3)
                self.scale = 2.0

        m = Model()
        p = m.params
        assert isinstance(p.w, nn.Param)
        assert p.buf is None
        assert p.scale == 2.0

    def test_params_with_callable_field(self):
        """internals.md: 'Module.params preserves static fields alongside Param leaves.
        Non-array fields remain unchanged.'"""

        class Model(nn.Module):
            w: nn.Param
            act: Callable

            def __init__(self):
                self.w = nn.Param(jnp.ones(3))
                self.act = jax.nn.relu

        m = Model()
        p = m.params
        assert isinstance(p.w, nn.Param)
        assert p.act is jax.nn.relu

    def test_params_with_nested_module_plain_array(self):
        """Nested module's plain array fields also become None."""

        class Inner(nn.Module):
            w: nn.Param
            buf: jax.Array

            def __init__(self, key):
                self.w = nn.Param(jax.random.normal(key, (2,)))
                self.buf = jnp.ones(2)

        class Outer(nn.Module):
            inner: Inner

            def __init__(self, key):
                self.inner = Inner(key)

        m = Outer(key=jax.random.key(0))
        p = m.params
        assert isinstance(p.inner.w, nn.Param)
        assert p.inner.buf is None


class TestFieldPartitioning:
    """Tests for the dynamic child vs static aux partitioning in _register_module."""

    def test_replace_param_with_none_changes_treedef(self):
        """Replacing a Param field with None moves it from dynamic child to static aux."""
        model = nn.Linear(3, 4, key=jax.random.key(0))
        _, treedef_with_bias = jtu.tree_flatten(model)

        model_no_bias = model.replace(b=None)
        leaves, treedef_no_bias = jtu.tree_flatten(model_no_bias)

        # Different treedef (b moved from child to static aux)
        assert treedef_with_bias != treedef_no_bias
        # Only w remains as a leaf
        assert len(leaves) == 1
        # Roundtrip preserves None
        rebuilt = treedef_no_bias.unflatten(leaves)
        assert rebuilt.b is None
        npt.assert_array_equal(rebuilt.w._value, model_no_bias.w._value)

    def test_replace_none_with_param_changes_treedef(self):
        """Replacing None with a Param moves it from static aux to dynamic child."""
        model = nn.Linear(3, 4, bias=False, key=jax.random.key(0))
        assert model.b is None
        leaves_before = jax.tree.leaves(model)
        assert len(leaves_before) == 1

        new_bias = nn.Param(jnp.zeros(4))
        model_with_bias = model.replace(b=new_bias)
        leaves_after = jax.tree.leaves(model_with_bias)
        assert len(leaves_after) == 2

    def test_surgery_both_versions_work_through_jit(self):
        """Both pre- and post-surgery models work through jit."""
        model = nn.Linear(3, 4, key=jax.random.key(0))
        x = jnp.ones((2, 3))

        result_with_bias = jax.jit(model)(x)
        assert result_with_bias.shape == (2, 4)

        model_no_bias = model.replace(b=None)
        result_no_bias = jax.jit(model_no_bias)(x)
        assert result_no_bias.shape == (2, 4)

    def test_mixed_tuple_modules_and_callable_through_jit_grad(self):
        """Tuple mixing Modules and callables: Modules are dynamic, callables wrapped in _Static."""

        class Model(nn.Module):
            layers: tuple

            def __init__(self, key):
                self.layers = (nn.Linear(3, 3, key=key), jax.nn.relu)

            def __call__(self, x):
                for layer in self.layers:
                    x = layer(x)
                return x

        m = Model(key=jax.random.key(0))
        x = jnp.ones((2, 3))

        # Works through jit
        eager = m(x)
        jitted = jax.jit(m)(x)
        npt.assert_allclose(jitted, eager)

        # Works through grad
        grads = jax.grad(lambda m, x: jnp.sum(m(x)))(m, x)
        assert isinstance(grads.layers[0], nn.Linear)
        # Callable preserved through roundtrip
        assert grads.layers[1] is jax.nn.relu

    def test_empty_tuple_field_is_static(self):
        """Empty tuple has no array-like elements, so it goes to static aux."""

        class Model(nn.Module):
            w: nn.Param
            items: tuple

            def __init__(self):
                self.w = nn.Param(jnp.ones(3))
                self.items = ()

        m = Model()
        assert len(jax.tree.leaves(m)) == 1
        rebuilt = jtu.tree_unflatten(*reversed(jtu.tree_flatten(m)))
        assert rebuilt.items == ()

    def test_empty_list_field_is_static(self):
        """Empty list has no array-like elements, so it goes to static aux."""

        class Model(nn.Module):
            w: nn.Param
            items: list

            def __init__(self):
                self.w = nn.Param(jnp.ones(3))
                self.items = []

        m = Model()
        assert len(jax.tree.leaves(m)) == 1
        rebuilt = jtu.tree_unflatten(*reversed(jtu.tree_flatten(m)))
        assert rebuilt.items == []

    def test_empty_dict_field_is_static(self):
        """Empty dict has no array-like values, so it goes to static aux."""

        class Model(nn.Module):
            w: nn.Param
            meta: dict

            def __init__(self):
                self.w = nn.Param(jnp.ones(3))
                self.meta = {}

        m = Model()
        assert len(jax.tree.leaves(m)) == 1
        rebuilt = jtu.tree_unflatten(*reversed(jtu.tree_flatten(m)))
        assert rebuilt.meta == {}

    def test_none_field_is_static(self):
        """None field goes to static aux, not dynamic children."""

        class Model(nn.Module):
            w: nn.Param
            optional: None

            def __init__(self):
                self.w = nn.Param(jnp.ones(3))
                self.optional = None

        m = Model()
        assert len(jax.tree.leaves(m)) == 1
        rebuilt = jtu.tree_unflatten(*reversed(jtu.tree_flatten(m)))
        assert rebuilt.optional is None
        npt.assert_allclose(jax.jit(lambda m: jnp.sum(m.w))(m), 3.0)
