from functools import partial

import jax
import jax.numpy as jnp
from flax import nnx, struct


# Case 1: Standard JAX Pytree (flax.struct.dataclass)
@struct.dataclass
class StateAsDataclass:
	data: jax.Array


# Case 2: NNX Pytree
class StateAsPytree(nnx.Pytree):
	def __init__(self, data: jax.Array):
		self.data = data


class Model(nnx.Module):
	def _step(self, state, sow: bool):
		new_data = state.data + 1

		if sow:
			self.sow(nnx.Intermediate, "data", new_data)

		# Return the next carry, which must have the same type
		if isinstance(state, StateAsDataclass):
			return state.replace(data=new_data)
		else:
			state.data = new_data
			return state

	@partial(nnx.jit, static_argnames=("num_steps", "sow"))
	def __call__(self, state, *, num_steps: int, sow: bool = False):
		state_axes = nnx.StateAxes({nnx.Intermediate: 0, ...: nnx.Carry})
		state_final = nnx.scan(
			partial(Model._step, sow=sow),
			in_axes=(state_axes, nnx.Carry), 
			out_axes=nnx.Carry,
			length=num_steps,
		)(self, state)

		return state_final


NUM_STEPS = 5

# --- Test Case 1: StateAsDataclass (This works) ---
print("--- Testing with StateAsDataclass ---")
model = Model()
state = StateAsDataclass(data=jnp.array(0.0))

state_final = model(state, num_steps=NUM_STEPS, sow=True)
intermediates = nnx.pop(model, nnx.Intermediate)

print(f"Final data: {state_final.data}")
print(f"Sowed data shape: {intermediates['data'].value[0].shape}")
assert intermediates["data"].value[0].shape == (NUM_STEPS,)
print("SUCCESS: `sow` worked as expected.\n")


# --- Test Case 2: StateAsPytree (This fails) ---
print("--- Testing with StateAsPytree ---")
model = Model()
state = StateAsPytree(data=jnp.array(0.0))

state_final = model(state, num_steps=NUM_STEPS, sow=True)
intermdiates = nnx.pop(model, nnx.Intermediate)

print(f"Final data: {state_final.data}")
print(f"Sowed data shape: {intermediates['data'].value[0].shape}")
assert intermediates["data"].value[0].shape == (NUM_STEPS,)
print("SUCCESS\n")