import runpy
from pathlib import Path

import equinox as eqx
import jax
import jax.numpy as jnp
import jax.random as jrandom

SCRIPT_PATH = Path(__file__).parents[1] / "scripts" / "train.py"


def test_initialize_network_restores_serialized_weights(tmp_path):
    script = runpy.run_path(str(SCRIPT_PATH), run_name="train_script")
    initialize_network = script["initialize_network"]

    original = initialize_network(jrandom.PRNGKey(0))
    checkpoint_path = tmp_path / "model.eqx"
    eqx.tree_serialise_leaves(checkpoint_path, original)

    restored = initialize_network(jrandom.PRNGKey(1), str(checkpoint_path))
    original_leaves = jax.tree.leaves(eqx.filter(original, eqx.is_array))
    restored_leaves = jax.tree.leaves(eqx.filter(restored, eqx.is_array))

    assert len(original_leaves) == len(restored_leaves)
    assert all(
        jnp.array_equal(original_leaf, restored_leaf)
        for original_leaf, restored_leaf in zip(original_leaves, restored_leaves, strict=True)
    )