import os

import numpy as np
import pytest

from pennylane import qchem
from openfermion.hamiltonians import MolecularData

ref_dir = os.path.join(os.path.dirname(os.path.realpath(__file__)), "test_ref_files")

table_1 = np.array(
    [
        [0.0, 0.0, 0.0, 0.0, 0.68238953],
        [0.0, 1.0, 1.0, 0.0, 0.68238953],
        [1.0, 0.0, 0.0, 1.0, 0.68238953],
        [1.0, 1.0, 1.0, 1.0, 0.68238953],
        [0.0, 0.0, 2.0, 2.0, 0.17900058],
        [0.0, 1.0, 3.0, 2.0, 0.17900058],
        [1.0, 0.0, 2.0, 3.0, 0.17900058],
        [1.0, 1.0, 3.0, 3.0, 0.17900058],
        [0.0, 2.0, 0.0, 2.0, 0.17900058],
        [0.0, 3.0, 1.0, 2.0, 0.17900058],
        [1.0, 2.0, 0.0, 3.0, 0.17900058],
        [1.0, 3.0, 1.0, 3.0, 0.17900058],
        [0.0, 2.0, 2.0, 0.0, 0.67073278],
        [0.0, 3.0, 3.0, 0.0, 0.67073278],
        [1.0, 2.0, 2.0, 1.0, 0.67073278],
        [1.0, 3.0, 3.0, 1.0, 0.67073278],
        [2.0, 0.0, 0.0, 2.0, 0.67073278],
        [2.0, 1.0, 1.0, 2.0, 0.67073278],
        [3.0, 0.0, 0.0, 3.0, 0.67073278],
        [3.0, 1.0, 1.0, 3.0, 0.67073278],
        [2.0, 0.0, 2.0, 0.0, 0.17900058],
        [2.0, 1.0, 3.0, 0.0, 0.17900058],
        [3.0, 0.0, 2.0, 1.0, 0.17900058],
        [3.0, 1.0, 3.0, 1.0, 0.17900058],
        [2.0, 2.0, 0.0, 0.0, 0.17900058],
        [2.0, 3.0, 1.0, 0.0, 0.17900058],
        [3.0, 2.0, 0.0, 1.0, 0.17900058],
        [3.0, 3.0, 1.0, 1.0, 0.17900058],
        [2.0, 2.0, 2.0, 2.0, 0.70510563],
        [2.0, 3.0, 3.0, 2.0, 0.70510563],
        [3.0, 2.0, 2.0, 3.0, 0.70510563],
        [3.0, 3.0, 3.0, 3.0, 0.70510563],
    ]
)

table_2 = np.array(
    [
        [0.0, 0.0, 0.0, 0.0, 0.70510563],
        [0.0, 1.0, 1.0, 0.0, 0.70510563],
        [1.0, 0.0, 0.0, 1.0, 0.70510563],
        [1.0, 1.0, 1.0, 1.0, 0.70510563],
    ]
)

table_3 = np.array(
    [
        [0.0, 0.0, 0.0, 0.0, 0.48731097],
        [0.0, 1.0, 1.0, 0.0, 0.48731097],
        [1.0, 0.0, 0.0, 1.0, 0.48731097],
        [1.0, 1.0, 1.0, 1.0, 0.48731097],
        [0.0, 0.0, 0.0, 2.0, -0.04857958],
        [0.0, 1.0, 1.0, 2.0, -0.04857958],
        [1.0, 0.0, 0.0, 3.0, -0.04857958],
        [1.0, 1.0, 1.0, 3.0, -0.04857958],
        [0.0, 0.0, 2.0, 0.0, -0.04857958],
        [0.0, 1.0, 3.0, 0.0, -0.04857958],
        [1.0, 0.0, 2.0, 1.0, -0.04857958],
        [1.0, 1.0, 3.0, 1.0, -0.04857958],
        [0.0, 0.0, 2.0, 2.0, 0.01306398],
        [0.0, 1.0, 3.0, 2.0, 0.01306398],
        [1.0, 0.0, 2.0, 3.0, 0.01306398],
        [1.0, 1.0, 3.0, 3.0, 0.01306398],
        [0.0, 2.0, 0.0, 0.0, -0.04857958],
        [0.0, 3.0, 1.0, 0.0, -0.04857958],
        [1.0, 2.0, 0.0, 1.0, -0.04857958],
        [1.0, 3.0, 1.0, 1.0, -0.04857958],
        [0.0, 2.0, 0.0, 2.0, 0.01306398],
        [0.0, 3.0, 1.0, 2.0, 0.01306398],
        [1.0, 2.0, 0.0, 3.0, 0.01306398],
        [1.0, 3.0, 1.0, 3.0, 0.01306398],
        [0.0, 2.0, 2.0, 0.0, 0.22361004],
        [0.0, 3.0, 3.0, 0.0, 0.22361004],
        [1.0, 2.0, 2.0, 1.0, 0.22361004],
        [1.0, 3.0, 3.0, 1.0, 0.22361004],
        [0.0, 2.0, 2.0, 2.0, 0.00748417],
        [0.0, 3.0, 3.0, 2.0, 0.00748417],
        [1.0, 2.0, 2.0, 3.0, 0.00748417],
        [1.0, 3.0, 3.0, 3.0, 0.00748417],
        [2.0, 0.0, 0.0, 0.0, -0.04857958],
        [2.0, 1.0, 1.0, 0.0, -0.04857958],
        [3.0, 0.0, 0.0, 1.0, -0.04857958],
        [3.0, 1.0, 1.0, 1.0, -0.04857958],
        [2.0, 0.0, 0.0, 2.0, 0.22361004],
        [2.0, 1.0, 1.0, 2.0, 0.22361004],
        [3.0, 0.0, 0.0, 3.0, 0.22361004],
        [3.0, 1.0, 1.0, 3.0, 0.22361004],
        [2.0, 0.0, 2.0, 0.0, 0.01306398],
        [2.0, 1.0, 3.0, 0.0, 0.01306398],
        [3.0, 0.0, 2.0, 1.0, 0.01306398],
        [3.0, 1.0, 3.0, 1.0, 0.01306398],
        [2.0, 0.0, 2.0, 2.0, 0.00748417],
        [2.0, 1.0, 3.0, 2.0, 0.00748417],
        [3.0, 0.0, 2.0, 3.0, 0.00748417],
        [3.0, 1.0, 3.0, 3.0, 0.00748417],
        [2.0, 2.0, 0.0, 0.0, 0.01306398],
        [2.0, 3.0, 1.0, 0.0, 0.01306398],
        [3.0, 2.0, 0.0, 1.0, 0.01306398],
        [3.0, 3.0, 1.0, 1.0, 0.01306398],
        [2.0, 2.0, 0.0, 2.0, 0.00748417],
        [2.0, 3.0, 1.0, 2.0, 0.00748417],
        [3.0, 2.0, 0.0, 3.0, 0.00748417],
        [3.0, 3.0, 1.0, 3.0, 0.00748417],
        [2.0, 2.0, 2.0, 0.0, 0.00748417],
        [2.0, 3.0, 3.0, 0.0, 0.00748417],
        [3.0, 2.0, 2.0, 1.0, 0.00748417],
        [3.0, 3.0, 3.0, 1.0, 0.00748417],
        [2.0, 2.0, 2.0, 2.0, 0.33788228],
        [2.0, 3.0, 3.0, 2.0, 0.33788228],
        [3.0, 2.0, 2.0, 3.0, 0.33788228],
        [3.0, 3.0, 3.0, 3.0, 0.33788228],
    ]
)


@pytest.mark.parametrize(
    ("name", "core", "active", "table_exp", "v_core_exp"),
    [
        ("h2_pyscf", None, None, table_1, 0),
        ("h2_pyscf", [0], None, table_2, 0.6823895331520422),
        ("h2_pyscf", None, [0, 1], table_1, 0),
        ("h2_pyscf", [0], [1], table_2, 0.6823895331520422),
        ("lih", [0], [1, 2], table_3, 1.6585666870874103),
    ],
)
def test_table_two_particle(name, core, active, table_exp, v_core_exp, tol):
    r"""Test the table of two-particle matrix elements and the contribution of core orbitals
    as implemented in the `two_particle` function of the `obs` module"""

    hf_data = MolecularData(filename=os.path.join(ref_dir, name))

    table, v_core = qchem.two_particle(hf_data.two_body_integrals, core=core, active=active)

    assert np.allclose(table, table_exp, **tol)
    assert np.allclose(v_core, v_core_exp, **tol)


v_me_1D = np.array([1, 2, 3, 4])
v_me_4D = np.full((2, 2, 2, 2), 0.5)


@pytest.mark.parametrize(
    ("v_me", "core", "active", "msg_match"),
    [
        (v_me_1D, [0], None, "'matrix_elements' must be a 4D array"),
        (v_me_4D, [-1, 0, 1, 2], None, "Indices of core orbitals must be between 0 and"),
        (v_me_4D, [0, 1, 2, 3], None, "Indices of core orbitals must be between 0 and"),
        (v_me_4D, None, [-1, 0], "Indices of active orbitals must be between 0 and"),
        (v_me_4D, None, [2, 6], "Indices of active orbitals must be between 0 and"),
    ],
)
def test_exceptions_two_particle(v_me, core, active, msg_match):
    """Test that the function `'two_particle'` throws an exception
    if the dimension of the matrix elements array is not a 4D array or
    if the indices of core and/or active orbitals are out of range."""

    with pytest.raises(ValueError, match=msg_match):
        qchem.two_particle(v_me, core=core, active=active)

import os

import numpy as np
import pytest

from pennylane import qchem
from openfermion.hamiltonians import MolecularData

ref_dir = os.path.join(os.path.dirname(os.path.realpath(__file__)), "test_ref_files")

table_1 = np.array(
    [
        [0.0, 0.0, 0.0, 0.0, 0.68238953],
        [0.0, 1.0, 1.0, 0.0, 0.68238953],
        [1.0, 0.0, 0.0, 1.0, 0.68238953],
        [1.0, 1.0, 1.0, 1.0, 0.68238953],
        [0.0, 0.0, 2.0, 2.0, 0.17900058],
        [0.0, 1.0, 3.0, 2.0, 0.17900058],
        [1.0, 0.0, 2.0, 3.0, 0.17900058],
        [1.0, 1.0, 3.0, 3.0, 0.17900058],
        [0.0, 2.0, 0.0, 2.0, 0.17900058],
        [0.0, 3.0, 1.0, 2.0, 0.17900058],
        [1.0, 2.0, 0.0, 3.0, 0.17900058],
        [1.0, 3.0, 1.0, 3.0, 0.17900058],
        [0.0, 2.0, 2.0, 0.0, 0.67073278],
        [0.0, 3.0, 3.0, 0.0, 0.67073278],
        [1.0, 2.0, 2.0, 1.0, 0.67073278],
        [1.0, 3.0, 3.0, 1.0, 0.67073278],
        [2.0, 0.0, 0.0, 2.0, 0.67073278],
        [2.0, 1.0, 1.0, 2.0, 0.67073278],
        [3.0, 0.0, 0.0, 3.0, 0.67073278],
        [3.0, 1.0, 1.0, 3.0, 0.67073278],
        [2.0, 0.0, 2.0, 0.0, 0.17900058],
        [2.0, 1.0, 3.0, 0.0, 0.17900058],
        [3.0, 0.0, 2.0, 1.0, 0.17900058],
        [3.0, 1.0, 3.0, 1.0, 0.17900058],
        [2.0, 2.0, 0.0, 0.0, 0.17900058],
        [2.0, 3.0, 1.0, 0.0, 0.17900058],
        [3.0, 2.0, 0.0, 1.0, 0.17900058],
        [3.0, 3.0, 1.0, 1.0, 0.17900058],
        [2.0, 2.0, 2.0, 2.0, 0.70510563],
        [2.0, 3.0, 3.0, 2.0, 0.70510563],
        [3.0, 2.0, 2.0, 3.0, 0.70510563],
        [3.0, 3.0, 3.0, 3.0, 0.70510563],
    ]
)

table_2 = np.array(
    [
        [0.0, 0.0, 0.0, 0.0, 0.70510563],
        [0.0, 1.0, 1.0, 0.0, 0.70510563],
        [1.0, 0.0, 0.0, 1.0, 0.70510563],
        [1.0, 1.0, 1.0, 1.0, 0.70510563],
    ]
)

table_3 = np.array(
    [
        [0.0, 0.0, 0.0, 0.0, 0.48731097],
        [0.0, 1.0, 1.0, 0.0, 0.48731097],
        [1.0, 0.0, 0.0, 1.0, 0.48731097],
        [1.0, 1.0, 1.0, 1.0, 0.48731097],
        [0.0, 0.0, 0.0, 2.0, -0.04857958],
        [0.0, 1.0, 1.0, 2.0, -0.04857958],
        [1.0, 0.0, 0.0, 3.0, -0.04857958],
        [1.0, 1.0, 1.0, 3.0, -0.04857958],
        [0.0, 0.0, 2.0, 0.0, -0.04857958],
        [0.0, 1.0, 3.0, 0.0, -0.04857958],
        [1.0, 0.0, 2.0, 1.0, -0.04857958],
        [1.0, 1.0, 3.0, 1.0, -0.04857958],
        [0.0, 0.0, 2.0, 2.0, 0.01306398],
        [0.0, 1.0, 3.0, 2.0, 0.01306398],
        [1.0, 0.0, 2.0, 3.0, 0.01306398],
        [1.0, 1.0, 3.0, 3.0, 0.01306398],
        [0.0, 2.0, 0.0, 0.0, -0.04857958],
        [0.0, 3.0, 1.0, 0.0, -0.04857958],
        [1.0, 2.0, 0.0, 1.0, -0.04857958],
        [1.0, 3.0, 1.0, 1.0, -0.04857958],
        [0.0, 2.0, 0.0, 2.0, 0.01306398],
        [0.0, 3.0, 1.0, 2.0, 0.01306398],
        [1.0, 2.0, 0.0, 3.0, 0.01306398],
        [1.0, 3.0, 1.0, 3.0, 0.01306398],
        [0.0, 2.0, 2.0, 0.0, 0.22361004],
        [0.0, 3.0, 3.0, 0.0, 0.22361004],
        [1.0, 2.0, 2.0, 1.0, 0.22361004],
        [1.0, 3.0, 3.0, 1.0, 0.22361004],
        [0.0, 2.0, 2.0, 2.0, 0.00748417],
        [0.0, 3.0, 3.0, 2.0, 0.00748417],
        [1.0, 2.0, 2.0, 3.0, 0.00748417],
        [1.0, 3.0, 3.0, 3.0, 0.00748417],
        [2.0, 0.0, 0.0, 0.0, -0.04857958],
        [2.0, 1.0, 1.0, 0.0, -0.04857958],
        [3.0, 0.0, 0.0, 1.0, -0.04857958],
        [3.0, 1.0, 1.0, 1.0, -0.04857958],
        [2.0, 0.0, 0.0, 2.0, 0.22361004],
        [2.0, 1.0, 1.0, 2.0, 0.22361004],
        [3.0, 0.0, 0.0, 3.0, 0.22361004],
        [3.0, 1.0, 1.0, 3.0, 0.22361004],
        [2.0, 0.0, 2.0, 0.0, 0.01306398],
        [2.0, 1.0, 3.0, 0.0, 0.01306398],
        [3.0, 0.0, 2.0, 1.0, 0.01306398],
        [3.0, 1.0, 3.0, 1.0, 0.01306398],
        [2.0, 0.0, 2.0, 2.0, 0.00748417],
        [2.0, 1.0, 3.0, 2.0, 0.00748417],
        [3.0, 0.0, 2.0, 3.0, 0.00748417],
        [3.0, 1.0, 3.0, 3.0, 0.00748417],
        [2.0, 2.0, 0.0, 0.0, 0.01306398],
        [2.0, 3.0, 1.0, 0.0, 0.01306398],
        [3.0, 2.0, 0.0, 1.0, 0.01306398],
        [3.0, 3.0, 1.0, 1.0, 0.01306398],
        [2.0, 2.0, 0.0, 2.0, 0.00748417],
        [2.0, 3.0, 1.0, 2.0, 0.00748417],
        [3.0, 2.0, 0.0, 3.0, 0.00748417],
        [3.0, 3.0, 1.0, 3.0, 0.00748417],
        [2.0, 2.0, 2.0, 0.0, 0.00748417],
        [2.0, 3.0, 3.0, 0.0, 0.00748417],
        [3.0, 2.0, 2.0, 1.0, 0.00748417],
        [3.0, 3.0, 3.0, 1.0, 0.00748417],
        [2.0, 2.0, 2.0, 2.0, 0.33788228],
        [2.0, 3.0, 3.0, 2.0, 0.33788228],
        [3.0, 2.0, 2.0, 3.0, 0.33788228],
        [3.0, 3.0, 3.0, 3.0, 0.33788228],
    ]
)


@pytest.mark.parametrize(
    ("name", "core", "active", "table_exp", "v_core_exp"),
    [
        ("h2_pyscf", None, None, table_1, 0),
        ("h2_pyscf", [0], None, table_2, 0.6823895331520422),
        ("h2_pyscf", None, [0, 1], table_1, 0),
        ("h2_pyscf", [0], [1], table_2, 0.6823895331520422),
        ("lih", [0], [1, 2], table_3, 1.6585666870874103),
    ],
)
def test_table_two_particle(name, core, active, table_exp, v_core_exp, tol):
    r"""Test the table of two-particle matrix elements and the contribution of core orbitals
    as implemented in the `two_particle` function of the `obs` module"""

    hf_data = MolecularData(filename=os.path.join(ref_dir, name))

    table, v_core = qchem.two_particle(hf_data.two_body_integrals, core=core, active=active)

    assert np.allclose(table, table_exp, **tol)
    assert np.allclose(v_core, v_core_exp, **tol)


v_me_1D = np.array([1, 2, 3, 4])
v_me_4D = np.full((2, 2, 2, 2), 0.5)


@pytest.mark.parametrize(
    ("v_me", "core", "active", "msg_match"),
    [
        (v_me_1D, [0], None, "'matrix_elements' must be a 4D array"),
        (v_me_4D, [-1, 0, 1, 2], None, "Indices of core orbitals must be between 0 and"),
        (v_me_4D, [0, 1, 2, 3], None, "Indices of core orbitals must be between 0 and"),
        (v_me_4D, None, [-1, 0], "Indices of active orbitals must be between 0 and"),
        (v_me_4D, None, [2, 6], "Indices of active orbitals must be between 0 and"),
    ],
)
def test_exceptions_two_particle(v_me, core, active, msg_match):
    """Test that the function `'two_particle'` throws an exception
    if the dimension of the matrix elements array is not a 4D array or
    if the indices of core and/or active orbitals are out of range."""

    with pytest.raises(ValueError, match=msg_match):
        qchem.two_particle(v_me, core=core, active=active)

Test the table of two-particle matrix elements and the contribution of core orbitals     as implemented in the `two_particle` function of the `obs` module Test that the function `'two_particle'` throws an exception     if the dimension of the matrix elements array is not a 4D array or     if the indices of core and/or active orbitals are out of range.