from absl.testing import absltest

from test_utils import EnzymeJaxTest


class JAXMD(EnzymeJaxTest):
    def setUp(self):
        from jax import random
        import jax.numpy as jnp
        import jax
        from jax_md import space, energy, simulate, partition

        lattice_constant = 1.37820
        N_rep = 40 if jax.default_backend() == "cpu" else 20
        box_size = N_rep * lattice_constant
        # Using float32 for positions / velocities, but float64 for reductions.
        dtype = jnp.float32

        # Specify the format of the neighbor list.
        # Options are Dense, Sparse, or OrderedSparse.
        format = partition.OrderedSparse

        displacement, shift = space.periodic(box_size)

        R = []
        for i in range(N_rep):
            for j in range(N_rep):
                for k in range(N_rep):
                    R += [[i, j, k]]
        R = jnp.array(R, dtype=dtype) * lattice_constant

        N = R.shape[0]
        phi = N / (lattice_constant * N_rep) ** 3
        print(f"Created a system of {N} LJ particles with number density {phi:.3f}")

        neighbor_fn, energy_fn = energy.lennard_jones_neighbor_list(
            displacement, box_size, r_cutoff=3.0, dr_threshold=1.0, format=format
        )

        init, apply = simulate.nvt_nose_hoover(energy_fn, shift, 5e-3, kT=1.2)

        key = random.PRNGKey(0)

        # We pick an "extra capacity" to ensure ahead of time that the neighbor
        # list will have enough capacity. Since sparse neighbor lists are more
        # robust to changes in the number of particles, in this case we only
        # need to actually add more capacity for dense neighbor lists.
        if format is partition.Dense:
            nbrs = neighbor_fn.allocate(R, extra_capacity=55)
        else:
            nbrs = neighbor_fn.allocate(R)

        state = init(key, R, neighbor=nbrs)

        def step(i, state_and_nbrs):
            state, nbrs = state_and_nbrs
            nbrs = nbrs.update(state.position)
            return apply(state, neighbor=nbrs), nbrs

        iters = 10
        degrees_of_freedom = state.chain.degrees_of_freedom

        def forward(
            position, momentum, force, mass, c_position, c_momentum, c_mass, c_tau, c_KE
        ):
            chain = simulate.NoseHooverChain(
                c_position, c_momentum, c_mass, c_tau, c_KE, degrees_of_freedom
            )
            state = simulate.NVTNoseHooverState(position, momentum, force, mass, chain)
            # new_state, new_nbrs = lax.fori_loop(0, iters, step, (state, nbrs))
            new_state, new_nbrs = step(0, (state, nbrs))
            return (
                new_state.position,
                new_state.momentum,
                new_state.force,
                new_state.mass,
                new_state.chain.position,
                new_state.chain.momentum,
                new_state.chain.mass,
                new_state.chain.tau,
                new_state.chain.kinetic_energy,
            )

        self.fn = forward
        self.name = "jaxmd" + str(N_rep)
        self.repeat = 2

        self.ins = [
            state.position,
            state.momentum,
            state.force,
            state.mass,
            state.chain.position,
            state.chain.momentum,
            state.chain.mass,
            state.chain.tau,
            state.chain.kinetic_energy,
        ]
        # for i, v in enumerate(self.ins):
        #    print("i=", i, v)
        self.dins = [x.copy() for x in self.ins]
        self.douts = [x.copy() for x in self.ins]

        self.atol = 5e-3
        self.rtol = 1e-3


if __name__ == "__main__":
    import platform

    # Deps not available on macos
    # PostRev introduces numerical error -- need to investigate
    if platform.system() != "Darwin" and platform.machine() == "x86_64":
        from test_utils import fix_paths

        fix_paths()
        import jax

        jax.config.update("jax_enable_x64", True)
        absltest.main()
