"""collections of unit tests."""

# Copyright (C) 2024- Kosuke Nakano
# All rights reserved.
#
# Redistribution and use in source and binary forms, with or without
# modification, are permitted provided that the following conditions
# are met:
#
# * Redistributions of source code must retain the above copyright
#   notice, this list of conditions and the following disclaimer.
#
# * Redistributions in binary form must reproduce the above copyright
#   notice, this list of conditions and the following disclaimer in
#   the documentation and/or other materials provided with the
#   distribution.
#
# * Neither the name of the jqmc project nor the names of its
#   contributors may be used to endorse or promote products derived
#   from this software without specific prior written permission.
#
# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS
# "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT
# LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS
# FOR A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE
# COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT,
# INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING,
# BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES;
# LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT
# LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN
# ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE
# POSSIBILITY OF SUCH DAMAGE.

import os
import sys
from pathlib import Path

import jax
import numpy as np
import pytest
from jax import numpy as jnp

project_root = str(Path(__file__).parent.parent)
if project_root not in sys.path:
    sys.path.insert(0, project_root)

from jqmc._setting import (  # noqa: E402
    atol_auto_vs_analytic_deriv,
    atol_auto_vs_numerical_deriv,
    atol_debug_vs_production,
    rtol_auto_vs_analytic_deriv,
    rtol_auto_vs_numerical_deriv,
    rtol_debug_vs_production,
)
from jqmc.determinant import compute_geminal_all_elements  # noqa: E402
from jqmc.jastrow_factor import (  # noqa: E402
    Jastrow_data,
    Jastrow_NN_data,
    Jastrow_three_body_data,
    Jastrow_two_body_data,
)
from jqmc.trexio_wrapper import read_trexio_file  # noqa: E402
from jqmc.wavefunction import (  # noqa: E402
    Wavefunction_data,
    _compute_discretized_kinetic_energy_debug,
    _compute_kinetic_energy_all_elements_auto,
    _compute_kinetic_energy_all_elements_debug,
    _compute_kinetic_energy_all_elements_fast_update_debug,
    _compute_kinetic_energy_auto,
    _compute_kinetic_energy_debug,
    _compute_nodal_distance_debug,
    compute_discretized_kinetic_energy,
    compute_discretized_kinetic_energy_fast_update,
    compute_kinetic_energy,
    compute_kinetic_energy_all_elements,
    compute_kinetic_energy_all_elements_fast_update,
    compute_nodal_distance,
    evaluate_ln_wavefunction,
    evaluate_ln_wavefunction_fast,
    f_epsilon_PW,
)

# JAX float64
jax.config.update("jax_enable_x64", True)
jax.config.update("jax_traceback_filtering", "off")


@pytest.mark.activate_if_skip_heavy
@pytest.mark.parametrize("trexio_file", ["water_ccecp_ccpvqz.h5", "H2_ae_ccpvdz_cart.h5", "N_ae_ccpvdz_cart.h5"])
def test_kinetic_energy_analytic_and_numerical(trexio_file: str):
    """Test the kinetic energy computation."""
    (
        structure_data,
        aos_data,
        _,
        _,
        geminal_mo_data,
        _,
    ) = read_trexio_file(
        trexio_file=os.path.join(os.path.dirname(__file__), "trexio_example_files", trexio_file), store_tuple=True
    )

    jastrow_onebody_data = None
    jastrow_twobody_data = Jastrow_two_body_data.init_jastrow_two_body_data(jastrow_2b_param=0.5, jastrow_2b_type="exp")
    jastrow_threebody_data = Jastrow_three_body_data.init_jastrow_three_body_data(
        orb_data=aos_data, random_init=True, random_scale=1.0e-3
    )
    jastrow_nn_data = Jastrow_NN_data.init_from_structure(structure_data=structure_data, hidden_dim=5, num_layers=2, cutoff=5.0)

    jastrow_data = Jastrow_data(
        jastrow_one_body_data=jastrow_onebody_data,
        jastrow_two_body_data=jastrow_twobody_data,
        jastrow_three_body_data=jastrow_threebody_data,
        jastrow_nn_data=jastrow_nn_data,
    )
    jastrow_data.sanity_check()

    wavefunction_data = Wavefunction_data(geminal_data=geminal_mo_data, jastrow_data=jastrow_data)
    wavefunction_data.sanity_check()

    num_ele_up = geminal_mo_data.num_electron_up
    num_ele_dn = geminal_mo_data.num_electron_dn
    r_cart_min, r_cart_max = -2.0, +2.0
    r_up_carts = (r_cart_max - r_cart_min) * np.random.rand(num_ele_up, 3) + r_cart_min
    r_dn_carts = (r_cart_max - r_cart_min) * np.random.rand(num_ele_dn, 3) + r_cart_min

    K_debug = _compute_kinetic_energy_debug(wavefunction_data=wavefunction_data, r_up_carts=r_up_carts, r_dn_carts=r_dn_carts)
    K_jax = compute_kinetic_energy(wavefunction_data=wavefunction_data, r_up_carts=r_up_carts, r_dn_carts=r_dn_carts)
    assert not np.any(np.isnan(np.asarray(np.asarray(K_debug)))), "NaN detected in first argument"
    assert not np.any(np.isnan(np.asarray(np.asarray(K_jax)))), "NaN detected in second argument"
    np.testing.assert_allclose(
        np.asarray(K_debug),
        np.asarray(K_jax),
        rtol=rtol_auto_vs_numerical_deriv,
        atol=atol_auto_vs_numerical_deriv,
    )


@pytest.mark.parametrize("trexio_file", ["water_ccecp_ccpvqz.h5", "H2_ae_ccpvdz_cart.h5", "N_ae_ccpvdz_cart.h5"])
def test_kinetic_energy_analytic_and_auto(trexio_file: str):
    """Compare analytic and autodiff kinetic energy implementations."""
    (
        _,
        aos_data,
        _,
        _,
        geminal_mo_data,
        _,
    ) = read_trexio_file(
        trexio_file=os.path.join(os.path.dirname(__file__), "trexio_example_files", trexio_file), store_tuple=True
    )

    jastrow_onebody_data = None
    jastrow_twobody_data = Jastrow_two_body_data.init_jastrow_two_body_data(jastrow_2b_param=1.0, jastrow_2b_type="pade")
    jastrow_threebody_data = Jastrow_three_body_data.init_jastrow_three_body_data(
        orb_data=aos_data, random_init=True, random_scale=1.0e-3
    )
    jastrow_data = Jastrow_data(
        jastrow_one_body_data=jastrow_onebody_data,
        jastrow_two_body_data=jastrow_twobody_data,
        jastrow_three_body_data=jastrow_threebody_data,
    )

    wavefunction_data = Wavefunction_data(geminal_data=geminal_mo_data, jastrow_data=jastrow_data)

    num_ele_up = geminal_mo_data.num_electron_up
    num_ele_dn = geminal_mo_data.num_electron_dn
    r_cart_min, r_cart_max = -2.0, +2.0
    r_up_carts = (r_cart_max - r_cart_min) * np.random.rand(num_ele_up, 3) + r_cart_min
    r_dn_carts = (r_cart_max - r_cart_min) * np.random.rand(num_ele_dn, 3) + r_cart_min

    K_analytic = compute_kinetic_energy(wavefunction_data=wavefunction_data, r_up_carts=r_up_carts, r_dn_carts=r_dn_carts)
    K_auto = _compute_kinetic_energy_auto(
        wavefunction_data=wavefunction_data,
        r_up_carts=jnp.asarray(r_up_carts),
        r_dn_carts=jnp.asarray(r_dn_carts),
    )

    assert not np.any(np.isnan(np.asarray(K_analytic))), "NaN detected in first argument"
    assert not np.any(np.isnan(np.asarray(K_auto))), "NaN detected in second argument"
    np.testing.assert_allclose(K_analytic, K_auto, atol=atol_auto_vs_analytic_deriv, rtol=rtol_auto_vs_analytic_deriv)


@pytest.mark.activate_if_skip_heavy
@pytest.mark.parametrize("trexio_file", ["water_ccecp_ccpvqz.h5", "H2_ae_ccpvdz_cart.h5", "N_ae_ccpvdz_cart.h5"])
def test_debug_and_auto_kinetic_energy_all_elements(trexio_file: str):
    """Debug vs autodiff kinetic energy per-electron arrays."""
    (
        _,
        aos_data,
        _,
        _,
        geminal_mo_data,
        _,
    ) = read_trexio_file(
        trexio_file=os.path.join(os.path.dirname(__file__), "trexio_example_files", trexio_file), store_tuple=True
    )
    jastrow_onebody_data = None
    jastrow_twobody_data = Jastrow_two_body_data.init_jastrow_two_body_data(jastrow_2b_param=1.0, jastrow_2b_type="exp")
    jastrow_threebody_data = Jastrow_three_body_data.init_jastrow_three_body_data(
        orb_data=aos_data, random_init=True, random_scale=1.0e-3
    )
    jastrow_data = Jastrow_data(
        jastrow_one_body_data=jastrow_onebody_data,
        jastrow_two_body_data=jastrow_twobody_data,
        jastrow_three_body_data=jastrow_threebody_data,
    )

    wavefunction_data = Wavefunction_data(geminal_data=geminal_mo_data, jastrow_data=jastrow_data)

    num_ele_up = geminal_mo_data.num_electron_up
    num_ele_dn = geminal_mo_data.num_electron_dn
    rng = np.random.default_rng(42)
    r_up_carts_np = rng.uniform(-2.0, 2.0, size=(num_ele_up, 3))
    r_dn_carts_np = rng.uniform(-2.0, 2.0, size=(num_ele_dn, 3))

    r_up_carts_jnp = jnp.array(r_up_carts_np)
    r_dn_carts_jnp = jnp.array(r_dn_carts_np)

    K_elements_up_debug, K_elements_dn_debug = _compute_kinetic_energy_all_elements_debug(
        wavefunction_data=wavefunction_data, r_up_carts=r_up_carts_np, r_dn_carts=r_dn_carts_np
    )
    K_elements_up_auto, K_elements_dn_auto = _compute_kinetic_energy_all_elements_auto(
        wavefunction_data=wavefunction_data, r_up_carts=r_up_carts_jnp, r_dn_carts=r_dn_carts_jnp
    )

    assert not np.any(np.isnan(np.asarray(K_elements_up_debug))), "NaN detected in first argument"
    assert not np.any(np.isnan(np.asarray(K_elements_up_auto))), "NaN detected in second argument"
    np.testing.assert_allclose(
        K_elements_up_debug, K_elements_up_auto, atol=atol_auto_vs_numerical_deriv, rtol=rtol_auto_vs_numerical_deriv
    )
    assert not np.any(np.isnan(np.asarray(K_elements_dn_debug))), "NaN detected in first argument"
    assert not np.any(np.isnan(np.asarray(K_elements_dn_auto))), "NaN detected in second argument"
    np.testing.assert_allclose(
        K_elements_dn_debug, K_elements_dn_auto, atol=atol_auto_vs_numerical_deriv, rtol=rtol_auto_vs_numerical_deriv
    )

    assert not np.any(np.isnan(np.asarray(np.asarray(K_elements_up_debug)))), "NaN detected in first argument"
    assert not np.any(np.isnan(np.asarray(np.asarray(K_elements_up_auto)))), "NaN detected in second argument"
    np.testing.assert_allclose(
        np.asarray(K_elements_up_debug),
        np.asarray(K_elements_up_auto),
        rtol=rtol_auto_vs_numerical_deriv,
        atol=atol_auto_vs_numerical_deriv,
    )

    assert not np.any(np.isnan(np.asarray(np.asarray(K_elements_dn_debug)))), "NaN detected in first argument"
    assert not np.any(np.isnan(np.asarray(np.asarray(K_elements_dn_auto)))), "NaN detected in second argument"
    np.testing.assert_allclose(
        np.asarray(K_elements_dn_debug),
        np.asarray(K_elements_dn_auto),
        rtol=rtol_auto_vs_numerical_deriv,
        atol=atol_auto_vs_numerical_deriv,
    )


@pytest.mark.parametrize("trexio_file", ["water_ccecp_ccpvqz.h5", "H2_ae_ccpvdz_cart.h5", "N_ae_ccpvdz_cart.h5"])
def test_auto_and_analytic_kinetic_energy_all_elements(trexio_file: str):
    """Autodiff vs analytic kinetic energy per-electron arrays."""
    (
        _,
        aos_data,
        _,
        _,
        geminal_mo_data,
        _,
    ) = read_trexio_file(
        trexio_file=os.path.join(os.path.dirname(__file__), "trexio_example_files", trexio_file), store_tuple=True
    )

    jastrow_onebody_data = None
    jastrow_twobody_data = Jastrow_two_body_data.init_jastrow_two_body_data(jastrow_2b_param=1.0, jastrow_2b_type="pade")
    jastrow_threebody_data = Jastrow_three_body_data.init_jastrow_three_body_data(
        orb_data=aos_data, random_init=True, random_scale=1.0e-3
    )
    jastrow_data = Jastrow_data(
        jastrow_one_body_data=jastrow_onebody_data,
        jastrow_two_body_data=jastrow_twobody_data,
        jastrow_three_body_data=jastrow_threebody_data,
    )

    wavefunction_data = Wavefunction_data(geminal_data=geminal_mo_data, jastrow_data=jastrow_data)

    num_ele_up = geminal_mo_data.num_electron_up
    num_ele_dn = geminal_mo_data.num_electron_dn
    rng = np.random.default_rng(43)
    r_up_carts_np = rng.uniform(-2.0, 2.0, size=(num_ele_up, 3))
    r_dn_carts_np = rng.uniform(-2.0, 2.0, size=(num_ele_dn, 3))

    r_up_carts_jnp = jnp.array(r_up_carts_np)
    r_dn_carts_jnp = jnp.array(r_dn_carts_np)

    K_elements_up_auto, K_elements_dn_auto = _compute_kinetic_energy_all_elements_auto(
        wavefunction_data=wavefunction_data, r_up_carts=r_up_carts_jnp, r_dn_carts=r_dn_carts_jnp
    )
    K_elements_up_analytic, K_elements_dn_analytic = compute_kinetic_energy_all_elements(
        wavefunction_data=wavefunction_data, r_up_carts=r_up_carts_jnp, r_dn_carts=r_dn_carts_jnp
    )

    assert not np.any(np.isnan(np.asarray(K_elements_up_auto))), "NaN detected in first argument"
    assert not np.any(np.isnan(np.asarray(K_elements_up_analytic))), "NaN detected in second argument"
    np.testing.assert_allclose(
        K_elements_up_auto, K_elements_up_analytic, atol=atol_auto_vs_analytic_deriv, rtol=rtol_auto_vs_analytic_deriv
    )
    assert not np.any(np.isnan(np.asarray(K_elements_dn_auto))), "NaN detected in first argument"
    assert not np.any(np.isnan(np.asarray(K_elements_dn_analytic))), "NaN detected in second argument"
    np.testing.assert_allclose(
        K_elements_dn_auto, K_elements_dn_analytic, atol=atol_auto_vs_analytic_deriv, rtol=rtol_auto_vs_analytic_deriv
    )


@pytest.mark.parametrize("trexio_file", ["water_ccecp_ccpvqz.h5", "H2_ae_ccpvdz_cart.h5", "N_ae_ccpvdz_cart.h5"])
def test_fast_update_kinetic_energy_all_elements(trexio_file: str):
    """Fast-update per-electron kinetic energy should match the standard analytic path."""
    (
        _,
        aos_data,
        _,
        _,
        geminal_mo_data,
        _,
    ) = read_trexio_file(
        trexio_file=os.path.join(os.path.dirname(__file__), "trexio_example_files", trexio_file),
        store_tuple=True,
    )

    jastrow_onebody_data = None
    jastrow_twobody_data = Jastrow_two_body_data.init_jastrow_two_body_data(jastrow_2b_param=1.0, jastrow_2b_type="exp")
    jastrow_threebody_data = Jastrow_three_body_data.init_jastrow_three_body_data(
        orb_data=aos_data, random_init=True, random_scale=1.0e-3
    )
    jastrow_data = Jastrow_data(
        jastrow_one_body_data=jastrow_onebody_data,
        jastrow_two_body_data=jastrow_twobody_data,
        jastrow_three_body_data=jastrow_threebody_data,
    )

    wavefunction_data = Wavefunction_data(geminal_data=geminal_mo_data, jastrow_data=jastrow_data)

    num_ele_up = geminal_mo_data.num_electron_up
    num_ele_dn = geminal_mo_data.num_electron_dn
    r_cart_min, r_cart_max = -2.0, +2.0
    r_up_carts = (r_cart_max - r_cart_min) * np.random.rand(num_ele_up, 3) + r_cart_min
    r_dn_carts = (r_cart_max - r_cart_min) * np.random.rand(num_ele_dn, 3) + r_cart_min

    r_up_carts_jnp = jnp.asarray(r_up_carts)
    r_dn_carts_jnp = jnp.asarray(r_dn_carts)

    # Standard analytic per-electron kinetic energy
    ke_up_debug, ke_dn_debug = _compute_kinetic_energy_all_elements_fast_update_debug(
        wavefunction_data=wavefunction_data,
        r_up_carts=r_up_carts_jnp,
        r_dn_carts=r_dn_carts_jnp,
    )

    # Build geminal inverse explicitly for the fast path
    A = compute_geminal_all_elements(
        geminal_data=geminal_mo_data,
        r_up_carts=r_up_carts_jnp,
        r_dn_carts=r_dn_carts_jnp,
    )
    A_inv = jnp.asarray(np.linalg.inv(np.array(A)))

    ke_up_fast, ke_dn_fast = compute_kinetic_energy_all_elements_fast_update(
        wavefunction_data=wavefunction_data,
        r_up_carts=r_up_carts_jnp,
        r_dn_carts=r_dn_carts_jnp,
        geminal_inverse=A_inv,
    )

    assert not np.any(np.isnan(np.asarray(ke_up_fast))), "NaN detected in first argument"
    assert not np.any(np.isnan(np.asarray(ke_up_debug))), "NaN detected in second argument"
    np.testing.assert_allclose(ke_up_fast, ke_up_debug, atol=atol_debug_vs_production, rtol=rtol_debug_vs_production)
    assert not np.any(np.isnan(np.asarray(ke_dn_fast))), "NaN detected in first argument"
    assert not np.any(np.isnan(np.asarray(ke_dn_debug))), "NaN detected in second argument"
    np.testing.assert_allclose(ke_dn_fast, ke_dn_debug, atol=atol_debug_vs_production, rtol=rtol_debug_vs_production)


@pytest.mark.parametrize("trexio_file", ["water_ccecp_ccpvqz.h5", "H2_ae_ccpvdz_cart.h5", "N_ae_ccpvdz_cart.h5"])
def test_debug_and_jax_discretized_kinetic_energy(trexio_file: str):
    """Test the discretized kinetic energy computation."""
    (
        _,
        aos_data,
        _,
        _,
        geminal_mo_data,
        _,
    ) = read_trexio_file(
        trexio_file=os.path.join(os.path.dirname(__file__), "trexio_example_files", trexio_file), store_tuple=True
    )

    jastrow_onebody_data = None
    jastrow_twobody_data = Jastrow_two_body_data.init_jastrow_two_body_data(jastrow_2b_param=1.0, jastrow_2b_type="pade")
    jastrow_threebody_data = Jastrow_three_body_data.init_jastrow_three_body_data(
        orb_data=aos_data, random_init=True, random_scale=1.0e-3
    )

    jastrow_data = Jastrow_data(
        jastrow_one_body_data=jastrow_onebody_data,
        jastrow_two_body_data=jastrow_twobody_data,
        jastrow_three_body_data=jastrow_threebody_data,
    )
    jastrow_data.sanity_check()

    wavefunction_data = Wavefunction_data(geminal_data=geminal_mo_data, jastrow_data=jastrow_data)
    wavefunction_data.sanity_check()

    num_ele_up = geminal_mo_data.num_electron_up
    num_ele_dn = geminal_mo_data.num_electron_dn
    rng = np.random.default_rng(44)
    r_up_carts_np = rng.uniform(-2.0, 2.0, size=(num_ele_up, 3))
    r_dn_carts_np = rng.uniform(-2.0, 2.0, size=(num_ele_dn, 3))

    r_up_carts_jnp = jnp.array(r_up_carts_np)
    r_dn_carts_jnp = jnp.array(r_dn_carts_np)

    alat = 0.05
    RT = np.eye(3)
    mesh_kinetic_part_r_up_carts_debug, mesh_kinetic_part_r_dn_carts_debug, elements_kinetic_part_debug = (
        _compute_discretized_kinetic_energy_debug(
            alat=alat, wavefunction_data=wavefunction_data, r_up_carts=r_up_carts_np, r_dn_carts=r_dn_carts_np
        )
    )

    # elements_kinetic_part_debug_all = np.array(elements_kinetic_part_debug).reshape(-1, 6)
    # print(np.array(elements_kinetic_part_debug))
    # print(elements_kinetic_part_debug_all.shape)
    # print(elements_kinetic_part_debug_all)

    mesh_kinetic_part_r_up_carts_jax, mesh_kinetic_part_r_dn_carts_jax, elements_kinetic_part_jax = (
        compute_discretized_kinetic_energy(
            alat=alat, wavefunction_data=wavefunction_data, r_up_carts=r_up_carts_jnp, r_dn_carts=r_dn_carts_jnp, RT=RT
        )
    )

    A = compute_geminal_all_elements(geminal_data=geminal_mo_data, r_up_carts=r_up_carts_jnp, r_dn_carts=r_dn_carts_jnp)
    A_old_inv = np.linalg.inv(A)
    (
        mesh_kinetic_part_r_up_carts_jax_fast_update,
        mesh_kinetic_part_r_dn_carts_jax_fast_update,
        elements_kinetic_part_jax_fast_update,
    ) = compute_discretized_kinetic_energy_fast_update(
        alat=alat,
        A_old_inv=A_old_inv,
        wavefunction_data=wavefunction_data,
        r_up_carts=r_up_carts_jnp,
        r_dn_carts=r_dn_carts_jnp,
        RT=RT,
    )

    assert not np.any(np.isnan(np.asarray(mesh_kinetic_part_r_up_carts_jax))), "NaN detected in first argument"
    assert not np.any(np.isnan(np.asarray(mesh_kinetic_part_r_up_carts_debug))), "NaN detected in second argument"
    np.testing.assert_allclose(
        mesh_kinetic_part_r_up_carts_jax,
        mesh_kinetic_part_r_up_carts_debug,
        atol=atol_debug_vs_production,
        rtol=rtol_debug_vs_production,
    )
    assert not np.any(np.isnan(np.asarray(mesh_kinetic_part_r_dn_carts_jax))), "NaN detected in first argument"
    assert not np.any(np.isnan(np.asarray(mesh_kinetic_part_r_dn_carts_debug))), "NaN detected in second argument"
    np.testing.assert_allclose(
        mesh_kinetic_part_r_dn_carts_jax,
        mesh_kinetic_part_r_dn_carts_debug,
        atol=atol_debug_vs_production,
        rtol=rtol_debug_vs_production,
    )
    assert not np.any(np.isnan(np.asarray(mesh_kinetic_part_r_up_carts_jax_fast_update))), "NaN detected in first argument"
    assert not np.any(np.isnan(np.asarray(mesh_kinetic_part_r_up_carts_debug))), "NaN detected in second argument"
    np.testing.assert_allclose(
        mesh_kinetic_part_r_up_carts_jax_fast_update,
        mesh_kinetic_part_r_up_carts_debug,
        atol=atol_debug_vs_production,
        rtol=rtol_debug_vs_production,
    )
    assert not np.any(np.isnan(np.asarray(mesh_kinetic_part_r_dn_carts_jax_fast_update))), "NaN detected in first argument"
    assert not np.any(np.isnan(np.asarray(mesh_kinetic_part_r_dn_carts_debug))), "NaN detected in second argument"
    np.testing.assert_allclose(
        mesh_kinetic_part_r_dn_carts_jax_fast_update,
        mesh_kinetic_part_r_dn_carts_debug,
        atol=atol_debug_vs_production,
        rtol=rtol_debug_vs_production,
    )
    assert not np.any(np.isnan(np.asarray(elements_kinetic_part_jax))), "NaN detected in first argument"
    assert not np.any(np.isnan(np.asarray(elements_kinetic_part_debug))), "NaN detected in second argument"
    np.testing.assert_allclose(
        elements_kinetic_part_jax, elements_kinetic_part_debug, atol=atol_debug_vs_production, rtol=rtol_debug_vs_production
    )
    assert not np.any(np.isnan(np.asarray(elements_kinetic_part_jax_fast_update))), "NaN detected in first argument"
    assert not np.any(np.isnan(np.asarray(elements_kinetic_part_debug))), "NaN detected in second argument"
    np.testing.assert_allclose(
        elements_kinetic_part_jax_fast_update,
        elements_kinetic_part_debug,
        atol=atol_debug_vs_production,
        rtol=rtol_debug_vs_production,
    )


@pytest.mark.parametrize("trexio_file", ["water_ccecp_ccpvqz.h5", "H2_ae_ccpvdz_cart.h5", "N_ae_ccpvdz_cart.h5"])
def test_nodal_distance_analytic_vs_debug(trexio_file: str):
    """Analytic compute_nodal_distance should match _compute_nodal_distance_debug."""
    (
        structure_data,
        aos_data,
        _,
        _,
        geminal_mo_data,
        _,
    ) = read_trexio_file(
        trexio_file=os.path.join(os.path.dirname(__file__), "trexio_example_files", trexio_file), store_tuple=True
    )

    jastrow_onebody_data = None
    jastrow_twobody_data = Jastrow_two_body_data.init_jastrow_two_body_data(jastrow_2b_param=0.5, jastrow_2b_type="exp")
    jastrow_threebody_data = Jastrow_three_body_data.init_jastrow_three_body_data(
        orb_data=aos_data, random_init=True, random_scale=1.0e-3
    )
    jastrow_nn_data = Jastrow_NN_data.init_from_structure(structure_data=structure_data, hidden_dim=5, num_layers=2, cutoff=5.0)

    jastrow_data = Jastrow_data(
        jastrow_one_body_data=jastrow_onebody_data,
        jastrow_two_body_data=jastrow_twobody_data,
        jastrow_three_body_data=jastrow_threebody_data,
        jastrow_nn_data=jastrow_nn_data,
    )

    wavefunction_data = Wavefunction_data(geminal_data=geminal_mo_data, jastrow_data=jastrow_data)

    num_ele_up = geminal_mo_data.num_electron_up
    num_ele_dn = geminal_mo_data.num_electron_dn
    r_cart_min, r_cart_max = -3.0, +3.0
    np.random.seed(42)
    r_up_carts = (r_cart_max - r_cart_min) * np.random.rand(num_ele_up, 3) + r_cart_min
    r_dn_carts = (r_cart_max - r_cart_min) * np.random.rand(num_ele_dn, 3) + r_cart_min

    r_up_carts_jnp = jnp.asarray(r_up_carts)
    r_dn_carts_jnp = jnp.asarray(r_dn_carts)

    # Analytic path
    nd_analytic = compute_nodal_distance(
        wavefunction_data=wavefunction_data,
        r_up_carts=r_up_carts_jnp,
        r_dn_carts=r_dn_carts_jnp,
    )

    # Debug path (paper formula)
    nd_debug = _compute_nodal_distance_debug(
        wavefunction_data=wavefunction_data,
        r_up_carts=r_up_carts_jnp,
        r_dn_carts=r_dn_carts_jnp,
    )

    # They should be identical up to numerical noise
    np.testing.assert_allclose(
        np.asarray(nd_analytic),
        np.asarray(nd_debug),
        rtol=rtol_auto_vs_numerical_deriv,
        atol=atol_auto_vs_numerical_deriv,
    )

    # Sanity: nodal distance should be positive
    assert float(nd_analytic) > 0.0
    assert float(nd_debug) > 0.0


def test_f_epsilon_PW_boundary_values():
    """Test f_epsilon_PW boundary conditions and properties."""
    epsilon_PW = 0.5

    # f(0) = 0
    np.testing.assert_allclose(float(f_epsilon_PW(jnp.array(0.0), epsilon_PW)), 0.0, atol=1e-14)

    # f(epsilon) = 1 (i.e., |t|=1)
    np.testing.assert_allclose(float(f_epsilon_PW(jnp.array(epsilon_PW), epsilon_PW)), 1.0, atol=1e-14)

    # f(x) = 1 for |x| >= epsilon
    np.testing.assert_allclose(float(f_epsilon_PW(jnp.array(2.0 * epsilon_PW), epsilon_PW)), 1.0, atol=1e-14)
    np.testing.assert_allclose(float(f_epsilon_PW(jnp.array(10.0 * epsilon_PW), epsilon_PW)), 1.0, atol=1e-14)

    # f'(0) = 0 and f'(1) = 0 (smooth at boundaries)
    # Check numerically: f(small) ~ 0 (derivative at origin is zero)
    delta = 1.0e-8
    f_at_delta = float(f_epsilon_PW(jnp.array(delta), epsilon_PW))
    # Should scale as t^2, so f(delta/eps) ~ 9*(delta/eps)^2
    assert f_at_delta < 1.0e-10  # very close to 0

    # f(x) >= 0 for interior
    t_vals = jnp.linspace(0.0, epsilon_PW, 100)
    f_vals = f_epsilon_PW(t_vals, epsilon_PW)
    assert jnp.all(f_vals >= -1e-14)


@pytest.mark.activate_if_skip_heavy
@pytest.mark.parametrize("trexio_file", ["water_ccecp_ccpvqz.h5", "H2_ae_ccpvdz_cart.h5", "N_ae_ccpvdz_cart.h5"])
def test_evaluate_ln_wavefunction_fast_forward(trexio_file):
    """Forward value of evaluate_ln_wavefunction_fast must match evaluate_ln_wavefunction."""
    _, _, _, _, geminal_data, _ = read_trexio_file(
        trexio_file=os.path.join(os.path.dirname(__file__), "trexio_example_files", trexio_file),
        store_tuple=True,
    )
    wavefunction_data = Wavefunction_data(geminal_data=geminal_data)
    rng = np.random.default_rng(2)
    n_up = geminal_data.num_electron_up
    n_dn = geminal_data.num_electron_dn

    for _ in range(10):
        r_up = jnp.array(rng.standard_normal((n_up, 3)) * 1.2, dtype=jnp.float64)
        r_dn = jnp.array(rng.standard_normal((n_dn, 3)) * 1.2, dtype=jnp.float64)
        G = compute_geminal_all_elements(geminal_data, r_up, r_dn)
        G_inv = jnp.linalg.inv(G)

        val_ref = float(evaluate_ln_wavefunction(wavefunction_data, r_up, r_dn))
        val_fast = float(evaluate_ln_wavefunction_fast(wavefunction_data, r_up, r_dn, G_inv))

        assert np.isfinite(val_ref), f"Reference value is not finite: {val_ref}"
        assert np.isfinite(val_fast), f"Fast value is not finite: {val_fast}"
        np.testing.assert_allclose(
            val_fast,
            val_ref,
            atol=atol_debug_vs_production,
            rtol=rtol_debug_vs_production,
            err_msg=f"Forward mismatch: fast={val_fast:.15f}, ref={val_ref:.15f}",
        )


@pytest.mark.activate_if_skip_heavy
@pytest.mark.parametrize("trexio_file", ["water_ccecp_ccpvqz.h5", "H2_ae_ccpvdz_cart.h5", "N_ae_ccpvdz_cart.h5"])
def test_evaluate_ln_wavefunction_fast_backward(trexio_file):
    """Gradient of evaluate_ln_wavefunction_fast w.r.t. wavefunction_data must match evaluate_ln_wavefunction."""
    _, _, _, _, geminal_data, _ = read_trexio_file(
        trexio_file=os.path.join(os.path.dirname(__file__), "trexio_example_files", trexio_file),
        store_tuple=True,
    )
    wavefunction_data = Wavefunction_data(geminal_data=geminal_data)
    rng = np.random.default_rng(3)
    n_up = geminal_data.num_electron_up
    n_dn = geminal_data.num_electron_dn

    grad_ref_fn = jax.grad(evaluate_ln_wavefunction, argnums=0)
    grad_fast_fn = jax.grad(evaluate_ln_wavefunction_fast, argnums=0)

    for _ in range(10):
        r_up = jnp.array(rng.standard_normal((n_up, 3)) * 1.2, dtype=jnp.float64)
        r_dn = jnp.array(rng.standard_normal((n_dn, 3)) * 1.2, dtype=jnp.float64)
        G = compute_geminal_all_elements(geminal_data, r_up, r_dn)
        G_inv = jnp.linalg.inv(G)

        grad_ref = grad_ref_fn(wavefunction_data, r_up, r_dn)
        grad_fast = grad_fast_fn(wavefunction_data, r_up, r_dn, G_inv)

        jax.tree_util.tree_map(
            lambda a, b: np.testing.assert_allclose(
                np.asarray(a),
                np.asarray(b),
                atol=atol_debug_vs_production,
                rtol=rtol_debug_vs_production,
                err_msg="Backward mismatch in evaluate_ln_wavefunction_fast",
            ),
            grad_ref,
            grad_fast,
        )


if __name__ == "__main__":
    from logging import Formatter, StreamHandler, getLogger

    logger = getLogger("jqmc")
    logger.setLevel("INFO")
    stream_handler = StreamHandler()
    stream_handler.setLevel("INFO")
    handler_format = Formatter("%(name)s - %(levelname)s - %(lineno)d - %(message)s")
    stream_handler.setFormatter(handler_format)
    logger.addHandler(stream_handler)
