import numpy as np
import pytest
from pytest_lazyfixture import lazy_fixture as lf
import tensorflow as tf
import optax

try:
    import openfermion as _
except ModuleNotFoundError:
    pytestmark = pytest.mark.skip("skip fgs test due to missing openfermion module")

import tensorcircuit as tc

F = tc.fgs.FGSSimulator
FT = tc.fgs.FGSTestSimulator


@pytest.mark.parametrize("backend", [lf("npb"), lf("tfb"), lf("jaxb")])
def test_cmatrix(backend, highp):
    c = tc.fgs.FGSSimulator(4, [0])
    c1 = tc.fgs.FGSTestSimulator(4, [0])
    c.evol_hp(0, 1, 1.2)
    c1.evol_hp(0, 1, 1.2)
    c.evol_icp(2, 0.5)
    c1.evol_icp(2, 0.5)
    c.evol_cp(0, -0.8)
    c1.evol_cp(0, -0.8)
    c.evol_sp(1, 0, 0.5)
    c1.evol_sp(1, 0, 0.5)
    c.evol_sp(1, 2, 1.5)
    c1.evol_sp(1, 2, 1.5)
    c.evol_ihamiltonian(c.chemical_potential(1.0, 1, 4))
    c1.evol_ihamiltonian(c1.chemical_potential_jw(1.0, 1, 4))
    np.testing.assert_allclose(
        c1.get_cmatrix_majorana(), c.get_cmatrix_majorana(), atol=1e-5
    )
    np.testing.assert_allclose(c1.get_cmatrix(), c.get_cmatrix(), atol=1e-5)
    np.testing.assert_allclose(c1.entropy([0, 2]), c.entropy([0, 2]), atol=1e-5)
    np.testing.assert_allclose(
        c1.renyi_entropy(2, [1]), c.renyi_entropy(2, [1]), atol=1e-5
    )


@pytest.mark.parametrize("backend", [lf("tfb"), lf("jaxb")])
def test_fgs_ad(backend, highp):
    N = 18

    def f(chi):
        c = tc.fgs.FGSSimulator(N, [0])
        for i in range(N - 1):
            c.evol_hp(i, i + 1, chi[i])
        cm = c.get_cmatrix()
        return -tc.backend.real(1 - cm[N, N])

    chi_vg = tc.backend.jit(tc.backend.value_and_grad(f))

    if tc.backend.name == "tensorflow":
        opt = tc.backend.optimizer(tf.keras.optimizers.Adam(1e-2))
    else:
        opt = tc.backend.optimizer(optax.adam(1e-2))

    param = tc.backend.ones([N - 1], dtype="float64")
    for _ in range(300):
        vs, gs = chi_vg(param)
        param = opt.update(gs, param)

    np.testing.assert_allclose(vs, -0.9986, atol=1e-2)


@pytest.mark.parametrize("backend", [lf("npb"), lf("tfb"), lf("jaxb")])
def test_hamiltonian_generation(backend):
    hc = F.chemical_potential(0.8, 0, 3)
    h1 = FT.get_hmatrix(hc, 3) + 0.4 * tc.backend.eye(2**3)
    # constant shift
    h2 = FT.chemical_potential_jw(0.8, 0, 3)
    np.testing.assert_allclose(h1, h2, atol=1e-5)

    hc = F.hopping(0.8j, 0, 1, 3)
    h1 = FT.get_hmatrix(hc, 3)
    h2 = FT.hopping_jw(0.8j, 0, 1, 3)
    np.testing.assert_allclose(h1, h2, atol=1e-5)

    hc = F.sc_pairing(0.8, 1, 0, 3)
    h1 = FT.get_hmatrix(hc, 3)
    h2 = FT.sc_pairing_jw(0.8, 1, 0, 3)
    np.testing.assert_allclose(h1, h2, atol=1e-5)


@pytest.mark.parametrize("backend", [lf("npb"), lf("tfb"), lf("jaxb")])
def test_ground_state(backend, highp):
    N = 3
    hc = (
        F.hopping(1.0, 0, 1, N)
        + F.hopping(1.0, 1, 2, N)
        + F.chemical_potential(0.4, 0, N)
        - F.chemical_potential(0.4, 1, N)
        + F.sc_pairing(0.8, 0, 1, N)
    )

    c = F(N, hc=hc)
    c1 = FT(N, hc=hc)

    np.testing.assert_allclose(c.get_cmatrix(), c1.get_cmatrix(), atol=1e-6)
    np.testing.assert_allclose(c1.entropy([0]), c.entropy([0]), atol=1e-5)
