"""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_debug_vs_production,
    rtol_debug_vs_production,
)
from jqmc.coulomb_potential import (  # noqa: E402
    _compute_bare_coulomb_potential_debug,
    _compute_bare_coulomb_potential_el_ion_element_wise_debug,
    _compute_discretized_bare_coulomb_potential_el_ion_element_wise_debug,
    _compute_ecp_local_parts_all_pairs_debug,
    _compute_ecp_non_local_parts_all_pairs_debug,
    _compute_ecp_non_local_parts_nearest_neighbors_debug,
    compute_bare_coulomb_potential,
    compute_bare_coulomb_potential_el_ion_element_wise,
    compute_coulomb_potential,
    compute_coulomb_potential_fast,
    compute_discretized_bare_coulomb_potential_el_ion_element_wise,
    compute_ecp_coulomb_potential,
    compute_ecp_coulomb_potential_fast,
    compute_ecp_local_parts_all_pairs,
    compute_ecp_non_local_parts_all_pairs,
    compute_ecp_non_local_parts_nearest_neighbors,
    compute_ecp_non_local_parts_nearest_neighbors_fast_update,
)
from jqmc.determinant import compute_geminal_all_elements  # noqa: E402
from jqmc.jastrow_factor import (  # noqa: E402
    Jastrow_data,
    Jastrow_one_body_data,
    Jastrow_three_body_data,
    Jastrow_two_body_data,
)
from jqmc.trexio_wrapper import read_trexio_file  # noqa: E402
from jqmc.wavefunction import Wavefunction_data  # noqa: E402

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

# params
angle_values = [
    (0.0, 0.0, 0.0),
    (+1.0 / 18.0 * np.pi, -1.0 / 12.0 * np.pi, +1.0 / 95.0 * np.pi),
    (+1.0 / 12.0 * np.pi, -1.0 / 17.0 * np.pi, +1.0 / 15.0 * np.pi),
    (+1.0 / 22.0 * np.pi, +1.0 / 54.0 * np.pi, +1.0 / 75.0 * np.pi),
    (-1.0 / 56.0 * np.pi, -1.0 / 13.0 * np.pi, +1.0 / 25.0 * np.pi),
    (+np.pi, +np.pi, +np.pi),
    (-np.pi, -np.pi, -np.pi),
    (+2.0 * np.pi, +2.0 * np.pi, +2.0 * np.pi),
    (-2.0 * np.pi, -2.0 * np.pi, -2.0 * np.pi),
]

# angle parameters
angle_params = [
    pytest.param(alpha, beta, gamma, id=f"(alpha,beta,gamma)=({alpha:.2f},{beta:.2f},{gamma:.2f})")
    for alpha, beta, gamma in angle_values
]

# Nv parameters
Nv_params = [pytest.param(Nv, id=f"Nv={Nv}") for Nv in (4, 6, 12, 18)]


@pytest.mark.parametrize("trexio_file", ["water_ccecp_ccpvqz.h5"])
def test_debug_and_jax_bare_coulomb(trexio_file: str):
    """Test the bare coulomb potential computation."""
    (
        _,
        _,
        _,
        _,
        _,
        coulomb_potential_data,
    ) = read_trexio_file(
        trexio_file=os.path.join(os.path.dirname(__file__), "trexio_example_files", trexio_file), store_tuple=True
    )

    # check its input
    coulomb_potential_data.sanity_check()

    r_up_carts_np = np.array(
        [
            [0.64878536, -0.83275288, 0.33532629],
            [0.55271273, 0.72310605, 0.93443775],
            [0.66767275, 0.1206456, -0.36521208],
            [-0.93165236, -0.0120386, 0.33003036],
        ]
    )
    r_dn_carts_np = np.array(
        [
            [1.0347816, 1.26162081, 0.42301735],
            [-0.57843435, 1.03651987, -0.55091542],
            [-1.56091964, -0.58952149, -0.99268141],
            [0.61863233, -0.14903326, 0.51962683],
        ]
    )

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

    # bare coulomb
    vpot_bare_jax = compute_bare_coulomb_potential(
        coulomb_potential_data=coulomb_potential_data,
        r_up_carts=r_up_carts_jnp,
        r_dn_carts=r_dn_carts_jnp,
    )

    vpot_bare_debug = _compute_bare_coulomb_potential_debug(
        coulomb_potential_data=coulomb_potential_data,
        r_up_carts=r_up_carts_np,
        r_dn_carts=r_dn_carts_np,
    )

    # print(f"vpot_bare_jax = {vpot_bare_jax}")
    # print(f"vpot_bare_debug = {vpot_bare_debug}")
    assert not np.any(np.isnan(np.asarray(vpot_bare_jax))), "NaN detected in first argument"
    assert not np.any(np.isnan(np.asarray(vpot_bare_debug))), "NaN detected in second argument"
    np.testing.assert_allclose(vpot_bare_jax, vpot_bare_debug, atol=atol_debug_vs_production, rtol=rtol_debug_vs_production)


@pytest.mark.parametrize("trexio_file", ["water_ccecp_ccpvqz.h5"])
def test_debug_and_jax_ecp_local(trexio_file: str):
    """Test the local ECP potential computation."""
    (
        _,
        _,
        _,
        _,
        _,
        coulomb_potential_data,
    ) = read_trexio_file(
        trexio_file=os.path.join(os.path.dirname(__file__), "trexio_example_files", trexio_file), store_tuple=True
    )

    # check its input
    coulomb_potential_data.sanity_check()

    r_up_carts_np = np.array(
        [
            [0.64878536, -0.83275288, 0.33532629],
            [0.55271273, 0.72310605, 0.93443775],
            [0.66767275, 0.1206456, -0.36521208],
            [-0.93165236, -0.0120386, 0.33003036],
        ]
    )
    r_dn_carts_np = np.array(
        [
            [1.0347816, 1.26162081, 0.42301735],
            [-0.57843435, 1.03651987, -0.55091542],
            [-1.56091964, -0.58952149, -0.99268141],
            [0.61863233, -0.14903326, 0.51962683],
        ]
    )

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

    # ecp local
    vpot_ecp_local_full_NN_jax = compute_ecp_local_parts_all_pairs(
        coulomb_potential_data=coulomb_potential_data,
        r_up_carts=r_up_carts_jnp,
        r_dn_carts=r_dn_carts_jnp,
    )

    vpot_ecp_local_full_NN_debug = _compute_ecp_local_parts_all_pairs_debug(
        coulomb_potential_data=coulomb_potential_data,
        r_up_carts=r_up_carts_np,
        r_dn_carts=r_dn_carts_np,
    )

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


@pytest.mark.activate_if_skip_heavy
@pytest.mark.parametrize("Nv", Nv_params)
@pytest.mark.parametrize("alpha, beta, gamma", angle_params)
def test_debug_and_jax_ecp_non_local_full_NN(Nv, alpha, beta, gamma):
    """Test the non-local ECP potential computation with the full neibohrs."""
    (
        structure_data,
        _,
        _,
        _,
        geminal_mo_data,
        coulomb_potential_data,
    ) = read_trexio_file(
        trexio_file=os.path.join(os.path.dirname(__file__), "trexio_example_files", "water_ccecp_ccpvqz.h5"), store_tuple=True
    )

    # check its input
    coulomb_potential_data.sanity_check()

    # n_atom
    n_atom = structure_data.natom

    # define data
    jastrow_data = Jastrow_data(
        jastrow_one_body_data=None,
        jastrow_two_body_data=None,
        jastrow_three_body_data=None,
    )  # no jastrow for the time-being.

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

    r_up_carts_np = np.array(
        [
            [0.64878536, -0.83275288, 0.33532629],
            [0.55271273, 0.72310605, 0.93443775],
            [0.66767275, 0.1206456, -0.36521208],
            [-0.93165236, -0.0120386, 0.33003036],
        ]
    )
    r_dn_carts_np = np.array(
        [
            [1.0347816, 1.26162081, 0.42301735],
            [-0.57843435, 1.03651987, -0.55091542],
            [-1.56091964, -0.58952149, -0.99268141],
            [0.61863233, -0.14903326, 0.51962683],
        ]
    )

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

    cos_a, sin_a = np.cos(alpha), np.sin(alpha)
    cos_b, sin_b = np.cos(beta), np.sin(beta)
    cos_g, sin_g = np.cos(gamma), np.sin(gamma)

    R_np = np.array(
        [
            [cos_b * cos_g, cos_g * sin_a * sin_b - cos_a * sin_g, sin_a * sin_g + cos_a * cos_g * sin_b],
            [cos_b * sin_g, cos_a * cos_g + sin_a * sin_b * sin_g, cos_a * sin_b * sin_g - cos_g * sin_a],
            [-sin_b, cos_b * sin_a, cos_a * cos_b],
        ]
    )
    R_jnp = jnp.array(R_np)

    # ecp non-local (full_NN)
    (
        mesh_non_local_ecp_part_r_up_carts_full_NN_jax,
        mesh_non_local_ecp_part_r_dn_carts_full_NN_jax,
        V_nonlocal_full_NN_jax,
        sum_V_nonlocal_full_NN_jax,
    ) = compute_ecp_non_local_parts_all_pairs(
        coulomb_potential_data=coulomb_potential_data,
        wavefunction_data=wavefunction_data,
        r_up_carts=r_up_carts_jnp,
        r_dn_carts=r_dn_carts_jnp,
        Nv=Nv,
        RT=R_jnp.T,
    )

    (
        mesh_non_local_ecp_part_r_up_carts_full_NN_debug,
        mesh_non_local_ecp_part_r_dn_carts_full_NN_debug,
        V_nonlocal_full_NN_debug,
        sum_V_nonlocal_full_NN_debug,
    ) = _compute_ecp_non_local_parts_all_pairs_debug(
        coulomb_potential_data=coulomb_potential_data,
        wavefunction_data=wavefunction_data,
        r_up_carts=r_up_carts_np,
        r_dn_carts=r_dn_carts_np,
        Nv=Nv,
        RT=R_np.T,
    )

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

    mesh_non_local_r_up_carts_max_full_NN_jax = mesh_non_local_ecp_part_r_up_carts_full_NN_jax[
        np.argmax(V_nonlocal_full_NN_jax)
    ]
    mesh_non_local_r_up_carts_max_full_NN_debug = mesh_non_local_ecp_part_r_up_carts_full_NN_debug[
        np.argmax(V_nonlocal_full_NN_debug)
    ]
    mesh_non_local_r_dn_carts_max_full_NN_jax = mesh_non_local_ecp_part_r_dn_carts_full_NN_jax[
        np.argmax(V_nonlocal_full_NN_jax)
    ]
    mesh_non_local_r_dn_carts_max_full_NN_debug = mesh_non_local_ecp_part_r_dn_carts_full_NN_debug[
        np.argmax(V_nonlocal_full_NN_debug)
    ]
    V_ecp_non_local_max_full_NN_jax = V_nonlocal_full_NN_jax[np.argmax(V_nonlocal_full_NN_jax)]
    V_ecp_non_local_max_full_NN_debug = V_nonlocal_full_NN_debug[np.argmax(V_nonlocal_full_NN_debug)]

    assert not np.any(np.isnan(np.asarray(V_ecp_non_local_max_full_NN_jax))), "NaN detected in first argument"
    assert not np.any(np.isnan(np.asarray(V_ecp_non_local_max_full_NN_debug))), "NaN detected in second argument"
    np.testing.assert_allclose(
        V_ecp_non_local_max_full_NN_jax,
        V_ecp_non_local_max_full_NN_debug,
        atol=atol_debug_vs_production,
        rtol=rtol_debug_vs_production,
    )
    assert not np.any(np.isnan(np.asarray(mesh_non_local_r_up_carts_max_full_NN_jax))), "NaN detected in first argument"
    assert not np.any(np.isnan(np.asarray(mesh_non_local_r_up_carts_max_full_NN_debug))), "NaN detected in second argument"
    np.testing.assert_allclose(
        mesh_non_local_r_up_carts_max_full_NN_jax,
        mesh_non_local_r_up_carts_max_full_NN_debug,
        atol=atol_debug_vs_production,
        rtol=rtol_debug_vs_production,
    )
    assert not np.any(np.isnan(np.asarray(mesh_non_local_r_dn_carts_max_full_NN_jax))), "NaN detected in first argument"
    assert not np.any(np.isnan(np.asarray(mesh_non_local_r_dn_carts_max_full_NN_debug))), "NaN detected in second argument"
    np.testing.assert_allclose(
        mesh_non_local_r_dn_carts_max_full_NN_jax,
        mesh_non_local_r_dn_carts_max_full_NN_debug,
        atol=atol_debug_vs_production,
        rtol=rtol_debug_vs_production,
    )

    # ecp non-local (NN, N=max)
    (
        mesh_non_local_ecp_part_r_up_carts_NN_check_jax,
        mesh_non_local_ecp_part_r_dn_carts_NN_check_jax,
        V_nonlocal_NN_check_jax,
        sum_V_nonlocal_NN_check_jax,
    ) = compute_ecp_non_local_parts_nearest_neighbors(
        coulomb_potential_data=coulomb_potential_data,
        wavefunction_data=wavefunction_data,
        NN=n_atom,
        Nv=Nv,
        RT=R_jnp.T,
        r_up_carts=r_up_carts_jnp,
        r_dn_carts=r_dn_carts_jnp,
    )

    (
        mesh_non_local_ecp_part_r_up_carts_NN_check_debug,
        mesh_non_local_ecp_part_r_dn_carts_NN_check_debug,
        V_nonlocal_NN_check_debug,
        sum_V_nonlocal_NN_check_debug,
    ) = _compute_ecp_non_local_parts_nearest_neighbors_debug(
        coulomb_potential_data=coulomb_potential_data,
        wavefunction_data=wavefunction_data,
        NN=n_atom,
        Nv=Nv,
        RT=R_np.T,
        r_up_carts=r_up_carts_np,
        r_dn_carts=r_dn_carts_np,
    )

    # debug, full-NN vs check-NN
    assert not np.any(np.isnan(np.asarray(sum_V_nonlocal_full_NN_debug))), "NaN detected in first argument"
    assert not np.any(np.isnan(np.asarray(sum_V_nonlocal_NN_check_debug))), "NaN detected in second argument"
    np.testing.assert_allclose(
        sum_V_nonlocal_full_NN_debug,
        sum_V_nonlocal_NN_check_debug,
        atol=atol_debug_vs_production,
        rtol=rtol_debug_vs_production,
    )

    # jax, full-NN vs check-NN
    assert not np.any(np.isnan(np.asarray(sum_V_nonlocal_full_NN_jax))), "NaN detected in first argument"
    assert not np.any(np.isnan(np.asarray(sum_V_nonlocal_NN_check_jax))), "NaN detected in second argument"
    np.testing.assert_allclose(
        sum_V_nonlocal_full_NN_jax, sum_V_nonlocal_NN_check_jax, atol=atol_debug_vs_production, rtol=rtol_debug_vs_production
    )

    mesh_non_local_r_up_carts_max_NN_check_jax = mesh_non_local_ecp_part_r_up_carts_NN_check_jax[
        np.argmax(V_nonlocal_NN_check_jax)
    ]
    mesh_non_local_r_up_carts_max_NN_check_debug = mesh_non_local_ecp_part_r_up_carts_NN_check_debug[
        np.argmax(V_nonlocal_NN_check_debug)
    ]
    mesh_non_local_r_dn_carts_max_NN_check_jax = mesh_non_local_ecp_part_r_dn_carts_NN_check_jax[
        np.argmax(V_nonlocal_NN_check_jax)
    ]
    mesh_non_local_r_dn_carts_max_NN_check_debug = mesh_non_local_ecp_part_r_dn_carts_NN_check_debug[
        np.argmax(V_nonlocal_NN_check_debug)
    ]
    V_ecp_non_local_max_NN_check_jax = V_nonlocal_NN_check_jax[np.argmax(V_nonlocal_NN_check_jax)]
    V_ecp_non_local_max_NN_check_debug = V_nonlocal_NN_check_debug[np.argmax(V_nonlocal_NN_check_debug)]

    # debug, full-NN vs check-NN
    assert not np.any(np.isnan(np.asarray(V_ecp_non_local_max_full_NN_debug))), "NaN detected in first argument"
    assert not np.any(np.isnan(np.asarray(V_ecp_non_local_max_NN_check_debug))), "NaN detected in second argument"
    np.testing.assert_allclose(
        V_ecp_non_local_max_full_NN_debug,
        V_ecp_non_local_max_NN_check_debug,
        atol=atol_debug_vs_production,
        rtol=rtol_debug_vs_production,
    )
    assert not np.any(np.isnan(np.asarray(mesh_non_local_r_up_carts_max_full_NN_debug))), "NaN detected in first argument"
    assert not np.any(np.isnan(np.asarray(mesh_non_local_r_up_carts_max_NN_check_debug))), "NaN detected in second argument"
    np.testing.assert_allclose(
        mesh_non_local_r_up_carts_max_full_NN_debug,
        mesh_non_local_r_up_carts_max_NN_check_debug,
        atol=atol_debug_vs_production,
        rtol=rtol_debug_vs_production,
    )
    assert not np.any(np.isnan(np.asarray(mesh_non_local_r_dn_carts_max_full_NN_debug))), "NaN detected in first argument"
    assert not np.any(np.isnan(np.asarray(mesh_non_local_r_dn_carts_max_NN_check_debug))), "NaN detected in second argument"
    np.testing.assert_allclose(
        mesh_non_local_r_dn_carts_max_full_NN_debug,
        mesh_non_local_r_dn_carts_max_NN_check_debug,
        atol=atol_debug_vs_production,
        rtol=rtol_debug_vs_production,
    )

    # jax, full-NN vs check-NN
    assert not np.any(np.isnan(np.asarray(V_ecp_non_local_max_full_NN_jax))), "NaN detected in first argument"
    assert not np.any(np.isnan(np.asarray(V_ecp_non_local_max_NN_check_jax))), "NaN detected in second argument"
    np.testing.assert_allclose(
        V_ecp_non_local_max_full_NN_jax,
        V_ecp_non_local_max_NN_check_jax,
        atol=atol_debug_vs_production,
        rtol=rtol_debug_vs_production,
    )
    assert not np.any(np.isnan(np.asarray(mesh_non_local_r_up_carts_max_full_NN_jax))), "NaN detected in first argument"
    assert not np.any(np.isnan(np.asarray(mesh_non_local_r_up_carts_max_NN_check_jax))), "NaN detected in second argument"
    np.testing.assert_allclose(
        mesh_non_local_r_up_carts_max_full_NN_jax,
        mesh_non_local_r_up_carts_max_NN_check_jax,
        atol=atol_debug_vs_production,
        rtol=rtol_debug_vs_production,
    )
    assert not np.any(np.isnan(np.asarray(mesh_non_local_r_dn_carts_max_full_NN_jax))), "NaN detected in first argument"
    assert not np.any(np.isnan(np.asarray(mesh_non_local_r_dn_carts_max_NN_check_jax))), "NaN detected in second argument"
    np.testing.assert_allclose(
        mesh_non_local_r_dn_carts_max_full_NN_jax,
        mesh_non_local_r_dn_carts_max_NN_check_jax,
        atol=atol_debug_vs_production,
        rtol=rtol_debug_vs_production,
    )


@pytest.mark.activate_if_skip_heavy
@pytest.mark.parametrize("Nv", Nv_params)
@pytest.mark.parametrize("alpha, beta, gamma", angle_params)
def test_debug_and_jax_ecp_non_local_partial_NN(Nv, alpha, beta, gamma):
    """Test the non-local ECP potential computation with partial neibohrs."""
    (
        structure_data,
        _,
        _,
        _,
        geminal_mo_data,
        coulomb_potential_data,
    ) = read_trexio_file(
        trexio_file=os.path.join(os.path.dirname(__file__), "trexio_example_files", "water_ccecp_ccpvqz.h5"), store_tuple=True
    )

    # check its input
    coulomb_potential_data.sanity_check()

    # n_atom
    n_atom = structure_data.natom

    # define data
    jastrow_data = Jastrow_data(
        jastrow_one_body_data=None,
        jastrow_two_body_data=None,
        jastrow_three_body_data=None,
    )  # no jastrow for the time-being.

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

    r_up_carts_np = np.array(
        [
            [0.64878536, -0.83275288, 0.33532629],
            [0.55271273, 0.72310605, 0.93443775],
            [0.66767275, 0.1206456, -0.36521208],
            [-0.93165236, -0.0120386, 0.33003036],
        ]
    )
    r_dn_carts_np = np.array(
        [
            [1.0347816, 1.26162081, 0.42301735],
            [-0.57843435, 1.03651987, -0.55091542],
            [-1.56091964, -0.58952149, -0.99268141],
            [0.61863233, -0.14903326, 0.51962683],
        ]
    )

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

    cos_a, sin_a = np.cos(alpha), np.sin(alpha)
    cos_b, sin_b = np.cos(beta), np.sin(beta)
    cos_g, sin_g = np.cos(gamma), np.sin(gamma)

    R_np = np.array(
        [
            [cos_b * cos_g, cos_g * sin_a * sin_b - cos_a * sin_g, sin_a * sin_g + cos_a * cos_g * sin_b],
            [cos_b * sin_g, cos_a * cos_g + sin_a * sin_b * sin_g, cos_a * sin_b * sin_g - cos_g * sin_a],
            [-sin_b, cos_b * sin_a, cos_a * cos_b],
        ]
    )
    R_jnp = jnp.array(R_np)

    for NN in range(1, n_atom):
        # ecp non-local (NN, NN=NN)
        (
            mesh_non_local_ecp_part_r_up_carts_NN_jax,
            mesh_non_local_ecp_part_r_dn_carts_NN_jax,
            V_nonlocal_NN_jax,
            sum_V_nonlocal_NN_jax,
        ) = compute_ecp_non_local_parts_nearest_neighbors(
            coulomb_potential_data=coulomb_potential_data,
            wavefunction_data=wavefunction_data,
            r_up_carts=r_up_carts_jnp,
            r_dn_carts=r_dn_carts_jnp,
            Nv=Nv,
            NN=NN,
            RT=R_jnp.T,
        )

        (
            mesh_non_local_ecp_part_r_up_carts_NN_debug,
            mesh_non_local_ecp_part_r_dn_carts_NN_debug,
            V_nonlocal_NN_debug,
            sum_V_nonlocal_NN_debug,
        ) = _compute_ecp_non_local_parts_nearest_neighbors_debug(
            coulomb_potential_data=coulomb_potential_data,
            wavefunction_data=wavefunction_data,
            r_up_carts=r_up_carts_np,
            r_dn_carts=r_dn_carts_np,
            Nv=Nv,
            NN=NN,
            RT=R_np.T,
        )

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

        mesh_non_local_r_up_carts_max_NN_jax = mesh_non_local_ecp_part_r_up_carts_NN_jax[np.argmax(V_nonlocal_NN_jax)]
        mesh_non_local_r_up_carts_max_NN_debug = mesh_non_local_ecp_part_r_up_carts_NN_debug[np.argmax(V_nonlocal_NN_debug)]
        mesh_non_local_r_dn_carts_max_NN_jax = mesh_non_local_ecp_part_r_dn_carts_NN_jax[np.argmax(V_nonlocal_NN_jax)]
        mesh_non_local_r_dn_carts_max_NN_debug = mesh_non_local_ecp_part_r_dn_carts_NN_debug[np.argmax(V_nonlocal_NN_debug)]
        V_ecp_non_local_max_NN_jax = V_nonlocal_NN_jax[np.argmax(V_nonlocal_NN_jax)]
        V_ecp_non_local_max_NN_debug = V_nonlocal_NN_debug[np.argmax(V_nonlocal_NN_debug)]

        assert not np.any(np.isnan(np.asarray(V_ecp_non_local_max_NN_jax))), "NaN detected in first argument"
        assert not np.any(np.isnan(np.asarray(V_ecp_non_local_max_NN_debug))), "NaN detected in second argument"
        np.testing.assert_allclose(
            V_ecp_non_local_max_NN_jax,
            V_ecp_non_local_max_NN_debug,
            atol=atol_debug_vs_production,
            rtol=rtol_debug_vs_production,
        )
        assert not np.any(np.isnan(np.asarray(mesh_non_local_r_up_carts_max_NN_jax))), "NaN detected in first argument"
        assert not np.any(np.isnan(np.asarray(mesh_non_local_r_up_carts_max_NN_debug))), "NaN detected in second argument"
        np.testing.assert_allclose(
            mesh_non_local_r_up_carts_max_NN_jax,
            mesh_non_local_r_up_carts_max_NN_debug,
            atol=atol_debug_vs_production,
            rtol=rtol_debug_vs_production,
        )
        assert not np.any(np.isnan(np.asarray(mesh_non_local_r_dn_carts_max_NN_jax))), "NaN detected in first argument"
        assert not np.any(np.isnan(np.asarray(mesh_non_local_r_dn_carts_max_NN_debug))), "NaN detected in second argument"
        np.testing.assert_allclose(
            mesh_non_local_r_dn_carts_max_NN_jax,
            mesh_non_local_r_dn_carts_max_NN_debug,
            atol=atol_debug_vs_production,
            rtol=rtol_debug_vs_production,
        )


def _build_full_jastrow_data(structure_data, geminal_mo_data, coulomb_potential_data, rng: np.random.Generator) -> Jastrow_data:
    core_electrons = tuple(coulomb_potential_data.z_cores)

    jastrow_one_body_data = Jastrow_one_body_data(
        jastrow_1b_param=1.0,
        jastrow_1b_type="exp",
        structure_data=structure_data,
        core_electrons=core_electrons,
    )

    jastrow_two_body_data = Jastrow_two_body_data(jastrow_2b_param=1.0, jastrow_2b_type="exp")

    num_orb = geminal_mo_data.orb_data_up_spin._num_orb
    j_matrix = rng.random((num_orb, num_orb + 1))
    jastrow_three_body_data = Jastrow_three_body_data(
        orb_data=geminal_mo_data.orb_data_up_spin,
        j_matrix=j_matrix,
    )

    return Jastrow_data(
        jastrow_one_body_data=jastrow_one_body_data,
        jastrow_two_body_data=jastrow_two_body_data,
        jastrow_three_body_data=jastrow_three_body_data,
    )


@pytest.mark.parametrize("trexio_file", ["water_ccecp_ccpvqz.h5"])
def test_fast_update_ecp_non_local_partial_NN(trexio_file: str):
    """Fast-update nearest-neighbor ECP matches reference with all Jastrow terms enabled."""
    (
        structure_data,
        _aos_data,
        _mos_data_up,
        _mos_data_dn,
        geminal_mo_data,
        coulomb_potential_data,
    ) = read_trexio_file(
        trexio_file=os.path.join(os.path.dirname(__file__), "trexio_example_files", trexio_file),
        store_tuple=True,
    )

    rng = np.random.default_rng(0)
    jastrow_data = _build_full_jastrow_data(structure_data, geminal_mo_data, coulomb_potential_data, rng)
    wavefunction_data = Wavefunction_data(geminal_data=geminal_mo_data, jastrow_data=jastrow_data)

    r_up_carts = rng.uniform(-1.0, 1.0, size=(geminal_mo_data.num_electron_up, 3))
    r_dn_carts = rng.uniform(-1.0, 1.0, size=(geminal_mo_data.num_electron_dn, 3))

    geminal_old = compute_geminal_all_elements(
        geminal_data=geminal_mo_data,
        r_up_carts=r_up_carts,
        r_dn_carts=r_dn_carts,
    )
    A_old_inv = np.linalg.inv(np.array(geminal_old))

    for NN in range(1, structure_data.natom):
        # ecp non-local (NN, NN=NN)
        Nv = 6
        RT = np.eye(3)

        (
            mesh_up_ref,
            mesh_dn_ref,
            V_ref,
            sum_V_ref,
        ) = compute_ecp_non_local_parts_nearest_neighbors(
            coulomb_potential_data=coulomb_potential_data,
            wavefunction_data=wavefunction_data,
            r_up_carts=r_up_carts,
            r_dn_carts=r_dn_carts,
            NN=NN,
            Nv=Nv,
            RT=RT,
        )

        (
            mesh_up_fast,
            mesh_dn_fast,
            V_fast,
            sum_V_fast,
        ) = compute_ecp_non_local_parts_nearest_neighbors_fast_update(
            coulomb_potential_data=coulomb_potential_data,
            wavefunction_data=wavefunction_data,
            r_up_carts=r_up_carts,
            r_dn_carts=r_dn_carts,
            RT=RT,
            A_old_inv=A_old_inv,
            NN=NN,
            Nv=Nv,
            flag_determinant_only=False,
        )

        assert not np.any(np.isnan(np.asarray(np.asarray(sum_V_fast)))), "NaN detected in first argument"
        assert not np.any(np.isnan(np.asarray(np.asarray(sum_V_ref)))), "NaN detected in second argument"
        np.testing.assert_allclose(
            np.asarray(sum_V_fast), np.asarray(sum_V_ref), atol=atol_debug_vs_production, rtol=rtol_debug_vs_production
        )
        assert not np.any(np.isnan(np.asarray(np.asarray(V_fast)))), "NaN detected in first argument"
        assert not np.any(np.isnan(np.asarray(np.asarray(V_ref)))), "NaN detected in second argument"
        np.testing.assert_allclose(
            np.asarray(V_fast), np.asarray(V_ref), atol=atol_debug_vs_production, rtol=rtol_debug_vs_production
        )
        assert not np.any(np.isnan(np.asarray(np.asarray(mesh_up_fast)))), "NaN detected in first argument"
        assert not np.any(np.isnan(np.asarray(np.asarray(mesh_up_ref)))), "NaN detected in second argument"
        np.testing.assert_allclose(
            np.asarray(mesh_up_fast), np.asarray(mesh_up_ref), atol=atol_debug_vs_production, rtol=rtol_debug_vs_production
        )
        assert not np.any(np.isnan(np.asarray(np.asarray(mesh_dn_fast)))), "NaN detected in first argument"
        assert not np.any(np.isnan(np.asarray(np.asarray(mesh_dn_ref)))), "NaN detected in second argument"
        np.testing.assert_allclose(
            np.asarray(mesh_dn_fast), np.asarray(mesh_dn_ref), atol=atol_debug_vs_production, rtol=rtol_debug_vs_production
        )


@pytest.mark.parametrize("trexio_file", ["water_ccecp_ccpvqz.h5"])
def test_debug_and_jax_bare_el_ion_elements(trexio_file: str):
    """Test the bare couloumb potential computation."""
    (
        structure_data,
        _,
        _,
        _,
        geminal_mo_data,
        coulomb_potential_data,
    ) = read_trexio_file(
        trexio_file=os.path.join(os.path.dirname(__file__), "trexio_example_files", trexio_file), store_tuple=True
    )

    # check its input
    coulomb_potential_data.sanity_check()

    r_up_carts_np = np.array(
        [
            [0.64878536, -0.83275288, 0.33532629],
            [0.55271273, 0.72310605, 0.93443775],
            [0.66767275, 0.1206456, -0.36521208],
            [-0.93165236, -0.0120386, 0.33003036],
        ]
    )
    r_dn_carts_np = np.array(
        [
            [1.0347816, 1.26162081, 0.42301735],
            [-0.57843435, 1.03651987, -0.55091542],
            [-1.56091964, -0.58952149, -0.99268141],
            [0.61863233, -0.14903326, 0.51962683],
        ]
    )

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

    interactions_R_r_up_debug, interactions_R_r_dn_debug = _compute_bare_coulomb_potential_el_ion_element_wise_debug(
        coulomb_potential_data=coulomb_potential_data,
        r_up_carts=r_up_carts_np,
        r_dn_carts=r_dn_carts_np,
    )

    interactions_R_r_up_jax, interactions_R_r_dn_jax = compute_bare_coulomb_potential_el_ion_element_wise(
        coulomb_potential_data=coulomb_potential_data,
        r_up_carts=r_up_carts_jnp,
        r_dn_carts=r_dn_carts_jnp,
    )

    assert not np.any(np.isnan(np.asarray(interactions_R_r_up_debug))), "NaN detected in first argument"
    assert not np.any(np.isnan(np.asarray(interactions_R_r_up_jax))), "NaN detected in second argument"
    np.testing.assert_allclose(
        interactions_R_r_up_debug, interactions_R_r_up_jax, atol=atol_debug_vs_production, rtol=rtol_debug_vs_production
    )
    assert not np.any(np.isnan(np.asarray(interactions_R_r_dn_debug))), "NaN detected in first argument"
    assert not np.any(np.isnan(np.asarray(interactions_R_r_dn_jax))), "NaN detected in second argument"
    np.testing.assert_allclose(
        interactions_R_r_dn_debug, interactions_R_r_dn_jax, atol=atol_debug_vs_production, rtol=rtol_debug_vs_production
    )


@pytest.mark.parametrize("trexio_file", ["water_ccecp_ccpvqz.h5"])
def test_debug_and_jax_discretized_bare_el_ion_elements(trexio_file: str):
    """Test the bare couloumb potential computation."""
    (
        structure_data,
        _,
        _,
        _,
        geminal_mo_data,
        coulomb_potential_data,
    ) = read_trexio_file(
        trexio_file=os.path.join(os.path.dirname(__file__), "trexio_example_files", trexio_file), store_tuple=True
    )

    # check its input
    coulomb_potential_data.sanity_check()

    r_up_carts_np = np.array(
        [
            [0.64878536, -0.83275288, 0.33532629],
            [0.55271273, 0.72310605, 0.93443775],
            [0.66767275, 0.1206456, -0.36521208],
            [-0.93165236, -0.0120386, 0.33003036],
        ]
    )
    r_dn_carts_np = np.array(
        [
            [1.0347816, 1.26162081, 0.42301735],
            [-0.57843435, 1.03651987, -0.55091542],
            [-1.56091964, -0.58952149, -0.99268141],
            [0.61863233, -0.14903326, 0.51962683],
        ]
    )

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

    alat = 0.5

    interactions_R_r_up_debug, interactions_R_r_dn_debug = (
        _compute_discretized_bare_coulomb_potential_el_ion_element_wise_debug(
            coulomb_potential_data=coulomb_potential_data,
            r_up_carts=r_up_carts_np,
            r_dn_carts=r_dn_carts_np,
            alat=alat,
        )
    )

    interactions_R_r_up_jax, interactions_R_r_dn_jax = compute_discretized_bare_coulomb_potential_el_ion_element_wise(
        coulomb_potential_data=coulomb_potential_data,
        r_up_carts=r_up_carts_jnp,
        r_dn_carts=r_dn_carts_jnp,
        alat=alat,
    )

    assert not np.any(np.isnan(np.asarray(interactions_R_r_up_debug))), "NaN detected in first argument"
    assert not np.any(np.isnan(np.asarray(interactions_R_r_up_jax))), "NaN detected in second argument"
    np.testing.assert_allclose(
        interactions_R_r_up_debug, interactions_R_r_up_jax, atol=atol_debug_vs_production, rtol=rtol_debug_vs_production
    )
    assert not np.any(np.isnan(np.asarray(interactions_R_r_dn_debug))), "NaN detected in first argument"
    assert not np.any(np.isnan(np.asarray(interactions_R_r_dn_jax))), "NaN detected in second argument"
    np.testing.assert_allclose(
        interactions_R_r_dn_debug, interactions_R_r_dn_jax, atol=atol_debug_vs_production, rtol=rtol_debug_vs_production
    )


@pytest.mark.parametrize("Nv", Nv_params)
@pytest.mark.parametrize("alpha, beta, gamma", angle_params)
def test_compute_ecp_coulomb_potential_fast(Nv, alpha, beta, gamma):
    """compute_ecp_coulomb_potential_fast matches compute_ecp_coulomb_potential when A_old_inv is exact.

    With a freshly computed LU inverse both functions must agree to full double
    precision because the only difference is which code path owns the LU solve.
    """
    (
        _structure_data,
        _aos_data,
        _mos_data_up,
        _mos_data_dn,
        geminal_mo_data,
        coulomb_potential_data,
    ) = read_trexio_file(
        trexio_file=os.path.join(os.path.dirname(__file__), "trexio_example_files", "water_ccecp_ccpvqz.h5"),
        store_tuple=True,
    )
    rng = np.random.default_rng(0)
    wavefunction_data = Wavefunction_data(
        geminal_data=geminal_mo_data,
        jastrow_data=Jastrow_data(jastrow_one_body_data=None, jastrow_two_body_data=None, jastrow_three_body_data=None),
    )
    r_up_carts = jnp.array(rng.uniform(-1.0, 1.0, size=(geminal_mo_data.num_electron_up, 3)))
    r_dn_carts = jnp.array(rng.uniform(-1.0, 1.0, size=(geminal_mo_data.num_electron_dn, 3)))
    G = compute_geminal_all_elements(geminal_data=geminal_mo_data, r_up_carts=r_up_carts, r_dn_carts=r_dn_carts)
    A_old_inv = jnp.linalg.inv(G)

    cos_a, sin_a = jnp.cos(alpha), jnp.sin(alpha)
    cos_b, sin_b = jnp.cos(beta), jnp.sin(beta)
    cos_g, sin_g = jnp.cos(gamma), jnp.sin(gamma)
    R = jnp.array(
        [
            [cos_b * cos_g, cos_g * sin_a * sin_b - cos_a * sin_g, sin_a * sin_g + cos_a * cos_g * sin_b],
            [cos_b * sin_g, cos_a * cos_g + sin_a * sin_b * sin_g, cos_a * sin_b * sin_g - cos_g * sin_a],
            [-sin_b, cos_b * sin_a, cos_a * cos_b],
        ]
    )
    RT = R.T

    # Reference: compute_ecp_coulomb_potential performs its own LU internally
    V_ref = compute_ecp_coulomb_potential(
        coulomb_potential_data=coulomb_potential_data,
        wavefunction_data=wavefunction_data,
        r_up_carts=r_up_carts,
        r_dn_carts=r_dn_carts,
        RT=RT,
        Nv=Nv,
    )

    # Fast path: accepts pre-computed A_old_inv, skips LU
    V_fast = compute_ecp_coulomb_potential_fast(
        coulomb_potential_data=coulomb_potential_data,
        wavefunction_data=wavefunction_data,
        r_up_carts=r_up_carts,
        r_dn_carts=r_dn_carts,
        RT=RT,
        A_old_inv=A_old_inv,
        Nv=Nv,
    )

    assert not np.any(np.isnan(np.asarray(V_ref))), "NaN detected in V_ref"
    assert not np.any(np.isnan(np.asarray(V_fast))), "NaN detected in V_fast"
    np.testing.assert_allclose(
        np.asarray(V_fast), np.asarray(V_ref), atol=atol_debug_vs_production, rtol=rtol_debug_vs_production
    )


@pytest.mark.parametrize("Nv", Nv_params)
@pytest.mark.parametrize("alpha, beta, gamma", angle_params)
def test_compute_coulomb_potential_fast(Nv, alpha, beta, gamma):
    """compute_coulomb_potential_fast matches compute_coulomb_potential when A_old_inv is exact.

    Covers both bare-Coulomb and ECP branches.  For ECP systems the non-local
    Psi(r')/Psi(r) ratios must be identical between the two paths because the
    same A_old_inv (exact LU inverse) is used.
    """
    (
        _structure_data,
        _aos_data,
        _mos_data_up,
        _mos_data_dn,
        geminal_mo_data,
        coulomb_potential_data,
    ) = read_trexio_file(
        trexio_file=os.path.join(os.path.dirname(__file__), "trexio_example_files", "water_ccecp_ccpvqz.h5"),
        store_tuple=True,
    )
    rng = np.random.default_rng(0)
    wavefunction_data = Wavefunction_data(
        geminal_data=geminal_mo_data,
        jastrow_data=Jastrow_data(jastrow_one_body_data=None, jastrow_two_body_data=None, jastrow_three_body_data=None),
    )
    r_up_carts = jnp.array(rng.uniform(-1.0, 1.0, size=(geminal_mo_data.num_electron_up, 3)))
    r_dn_carts = jnp.array(rng.uniform(-1.0, 1.0, size=(geminal_mo_data.num_electron_dn, 3)))
    G = compute_geminal_all_elements(geminal_data=geminal_mo_data, r_up_carts=r_up_carts, r_dn_carts=r_dn_carts)
    A_old_inv = jnp.linalg.inv(G)

    cos_a, sin_a = jnp.cos(alpha), jnp.sin(alpha)
    cos_b, sin_b = jnp.cos(beta), jnp.sin(beta)
    cos_g, sin_g = jnp.cos(gamma), jnp.sin(gamma)
    R = jnp.array(
        [
            [cos_b * cos_g, cos_g * sin_a * sin_b - cos_a * sin_g, sin_a * sin_g + cos_a * cos_g * sin_b],
            [cos_b * sin_g, cos_a * cos_g + sin_a * sin_b * sin_g, cos_a * sin_b * sin_g - cos_g * sin_a],
            [-sin_b, cos_b * sin_a, cos_a * cos_b],
        ]
    )
    RT = R.T

    # Reference: standard path, computes LU internally for ECP non-local part
    V_ref = compute_coulomb_potential(
        coulomb_potential_data=coulomb_potential_data,
        r_up_carts=r_up_carts,
        r_dn_carts=r_dn_carts,
        RT=RT,
        wavefunction_data=wavefunction_data,
        Nv=Nv,
    )

    # Fast path: accepts caller-supplied inverse
    V_fast = compute_coulomb_potential_fast(
        coulomb_potential_data=coulomb_potential_data,
        r_up_carts=r_up_carts,
        r_dn_carts=r_dn_carts,
        RT=RT,
        A_old_inv=A_old_inv,
        wavefunction_data=wavefunction_data,
        Nv=Nv,
    )

    assert not np.any(np.isnan(np.asarray(V_ref))), "NaN detected in V_ref"
    assert not np.any(np.isnan(np.asarray(V_fast))), "NaN detected in V_fast"
    np.testing.assert_allclose(
        np.asarray(V_fast), np.asarray(V_ref), atol=atol_debug_vs_production, rtol=rtol_debug_vs_production
    )


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)
