import itertools

import pytest
import autoray as ar
import numpy as np
from numpy.testing import assert_allclose

import quimb as qu
import quimb.tensor as qtn


class TestPEPSConstruct:
    @pytest.mark.parametrize("Lx", [3, 4, 5])
    @pytest.mark.parametrize("Ly", [3, 4, 5])
    def test_basic_rand(self, Lx, Ly):
        psi = qtn.PEPS.rand(Lx, Ly, bond_dim=4)

        assert psi.max_bond() == 4
        assert psi.Lx == Lx
        assert psi.Ly == Ly
        assert len(psi.tensor_map) == Lx * Ly
        assert psi.site_inds == tuple(
            f"k{i},{j}" for i in range(Lx) for j in range(Ly)
        )
        assert psi.site_tags == tuple(
            f"I{i},{j}" for i in range(Lx) for j in range(Ly)
        )

        assert psi.bond_size((1, 1), (1, 2)) == (4)

        for i in range(Lx):
            assert len(psi.select(f"X{i}").tensor_map) == Ly
        for j in range(Ly):
            assert len(psi.select(f"Y{j}").tensor_map) == Lx

        for i in range(Lx):
            for j in range(Ly):
                assert psi.phys_dim(i, j) == 2
                assert isinstance(psi[i, j], qtn.Tensor)
                assert isinstance(psi[f"I{i},{j}"], qtn.Tensor)

        if Lx == Ly == 3:
            psi_dense = psi.to_qarray(optimize="auto-hq")
            assert psi_dense.shape == (512, 1)

        psi.show()
        assert f"Lx={Lx}" in psi.__str__()
        assert f"Lx={Lx}" in psi.__repr__()

    def test_cyclic_edge_cases(self):
        peps = qtn.PEPS.rand(3, 3, bond_dim=1, cyclic=True)
        assert peps.is_cyclic_x()
        assert peps.is_cyclic_y()
        assert peps.num_indices == peps.num_tensors * 3

    def test_zeros(self):
        peps = qtn.PEPS.zeros(3, 3, cyclic=True, bond_dim=1)
        assert peps.num_tensors == 9
        assert peps.num_indices == 27
        assert_allclose(peps.to_dense(), np.zeros([512, 1]))

    def test_flatten(self):
        psi = qtn.PEPS.rand(3, 5, 3, seed=42)
        norm = psi.H & psi
        assert len(norm.tensors) == 30
        norm.flatten_()
        assert len(norm.tensors) == 15
        assert norm.max_bond() == 9

    def test_add_peps(self):
        pa = qtn.PEPS.rand(3, 4, 2)
        pb = qtn.PEPS.rand(3, 4, 3)
        pc = qtn.PEPS.rand(3, 4, 4)
        pab = pa + pb
        assert pab.max_bond() == 5
        assert pab @ pc == pytest.approx(pa @ pc + pb @ pc)

    @pytest.mark.parametrize("Lx", [3, 4, 5])
    @pytest.mark.parametrize("Ly", [3, 4, 5])
    def test_bond_coordinates(self, Lx, Ly):
        psi = qtn.PEPS.rand(Lx, Ly, bond_dim=1)
        all_bonds = tuple(psi.gen_bond_coos())
        assert len(all_bonds) == 2 * Lx * Ly - Lx - Ly
        he = tuple(psi.gen_horizontal_even_bond_coos())
        ho = tuple(psi.gen_horizontal_odd_bond_coos())
        ve = tuple(psi.gen_vertical_even_bond_coos())
        vo = tuple(psi.gen_vertical_odd_bond_coos())
        for p in (he, ho, ve, vo):
            assert len(set(p)) == len(p)
            # check there is no overlap at all
            sites = tuple(itertools.chain.from_iterable(he))
            assert len(set(sites)) == len(sites)
        # check all coordinates are generated
        assert set(itertools.chain(he, ho, ve, vo)) == set(all_bonds)

    @pytest.mark.parametrize(
        "where",
        [
            [(0, 0)],
            [(0, 1)],
            [(0, 2)],
            [(2, 2)],
            [(3, 2)],
            [(3, 1)],
            [(3, 0)],
            [(2, 0)],
            [(1, 1)],
        ],
    )
    @pytest.mark.parametrize("contract", [False, True])
    def test_gate_2d_single_site(self, where, contract):
        Lx = 4
        Ly = 3
        D = 2

        psi = qtn.PEPS.rand(Lx, Ly, bond_dim=D, seed=42, dtype=complex)
        psi_d = psi.to_qarray()
        G = qu.rand_matrix(2)

        # compute the exact dense reference
        dims = [[2] * Ly] * Lx
        IGI = qu.ikron(G, dims, where, sparse=True)
        xe = (psi_d.H @ IGI @ psi_d).item()

        tn = psi.H & psi.gate(G, where, contract=contract)
        assert len(tn.tensors) == 2 * Lx * Ly + int(not contract)

        assert tn ^ all == pytest.approx(xe)

    @pytest.mark.parametrize(
        "contract", [False, True, "split", "reduce-split"]
    )
    @pytest.mark.parametrize(
        "where",
        [
            [(1, 1), (2, 1)],
            [(3, 2), (2, 2)],
            # XXX: long range, disable for now
            # [(0, 0), (1, 1)],
            # [(3, 1), (1, 2)],
        ],
    )
    def test_gate_2d_two_site(self, where, contract):
        Lx = 4
        Ly = 3
        D = 2

        psi = qtn.PEPS.rand(Lx, Ly, bond_dim=D, seed=42, dtype=complex)
        psi_d = psi.to_qarray()

        # ikron can't tensor operators across non-adjacent subsytems
        # so we explicitly construct the gate as a sum of tensor components
        G_comps = [(qu.rand_matrix(2), qu.rand_matrix(2)) for _ in range(4)]
        G = sum(A & B for A, B in G_comps)

        # compute the exact dense reference
        dims = [[2] * Ly] * Lx
        IGI = sum(
            qu.ikron([A, B], dims, where, sparse=True) for A, B in G_comps
        )

        xe = (psi_d.H @ IGI @ psi_d).item()

        tn = psi.H & psi.gate(G, where, contract=contract)
        change = {False: 1, True: -1, "split": 0, "reduce-split": 0}[contract]
        assert len(tn.tensors) == 2 * Lx * Ly + change

        assert tn ^ all == pytest.approx(xe)

    @pytest.mark.parametrize(
        "propagate_tags", [False, True, "sites", "register"]
    )
    def test_gate_propagate_tags(self, propagate_tags):
        Lx = 4
        Ly = 3
        D = 1
        psi = qtn.PEPS.rand(Lx, Ly, D, tags="PSI0")
        psi.gate_(
            qu.rand_uni(4),
            [(1, 1), (1, 2)],
            tags="G1",
            propagate_tags=propagate_tags,
        )
        psi.gate_(
            qu.rand_uni(4),
            [(1, 2), (3, 2)],
            tags="G2",
            propagate_tags=propagate_tags,
        )
        if propagate_tags is False:
            assert set(psi["G1"].tags) == {"G1"}
            assert set(psi["G2"].tags) == {"G2"}
        if propagate_tags is True:
            tgs1 = {"I1,1", "I1,2", "G1", "PSI0", "Y1", "Y2", "X1"}
            assert set(psi["G1"][0].tags) == tgs1
            assert set(psi["G2"].tags) == tgs1 | {"G2", "I3,2", "X3", "Y2"}
        if propagate_tags == "sites":
            assert set(psi["G1"].tags) == {"G1", "I1,1", "I1,2"}
            assert set(psi["G2"].tags) == {"G2", "I1,1", "I1,2", "I3,2"}
        if propagate_tags == "register":
            assert set(psi["G1"].tags) == {"G1", "I1,1", "I1,2"}
            assert set(psi["G2"].tags) == {"G2", "I1,2", "I3,2"}


class Test2DContract:
    @pytest.mark.parametrize("mode", ["mps", "projector", "full-bond"])
    def test_contract_boundary(self, mode):
        # make a large but cheap and easy (mostly positive) TN
        rng = np.random.default_rng(42)
        tn = qtn.TN2D_from_fill_fn(
            lambda shape: rng.uniform(low=-0.1, size=shape),
            Lx=8,
            Ly=8,
            D=2,
        )
        Zex = tn.contract(...)
        Z = tn.contract_boundary(max_bond=4, mode=mode)
        assert Z == pytest.approx(Zex, rel=1e-3)

    def test_contract_2d_one_layer_boundary(self):
        psi = qtn.PEPS.rand(4, 4, 3, seed=42)
        norm = psi.make_norm()
        xe = norm.contract(all, optimize="auto-hq")
        xt = norm.contract_boundary(max_bond=9)
        assert xt == pytest.approx(xe, rel=1e-2)

    def test_contract_2d_two_layer_boundary(self):
        psi = qtn.PEPS.rand(4, 4, 3, seed=42, tags="KET")
        norm = psi.make_norm()
        xe = norm.contract(all, optimize="auto-hq")
        xt = norm.contract_boundary(max_bond=27, layer_tags=["KET", "BRA"])
        assert xt == pytest.approx(xe, rel=1e-2)

    def test_contract_2d_full_bond(self):
        psi = qtn.PEPS.rand(4, 4, 3, seed=42, tags="KET")
        norm = psi.make_norm()
        xe = norm.contract(all, optimize="auto-hq")
        xt = norm.contract_boundary(max_bond=27, mode="full-bond")
        assert xt == pytest.approx(xe, rel=1e-2)

    @pytest.mark.parametrize("dims", [(10, 4), (4, 10)])
    def test_contract_boundary_stopping_criterion(self, dims):
        tn = qtn.TN2D_from_fill_fn(
            lambda shape: ar.lazy.Variable(shape=shape, backend="numpy"),
            *dims,
            D=2,
        )
        tn.contract_ctmrg_(4, cutoff=0.0, final_contract=False, progbar=True)
        assert tn.max_bond() == 4
        assert 16 <= tn.num_tensors <= 20

    @pytest.mark.parametrize("lazy", [False, True])
    def test_coarse_grain_basics(self, lazy):
        tn = qtn.TN2D_from_fill_fn(
            lambda shape: ar.lazy.Variable(shape, backend="numpy"),
            Lx=6,
            Ly=7,
            D=2,
        )
        tncg = tn.coarse_grain_hotrg("x", max_bond=3, cutoff=0.0, lazy=lazy)
        assert (tncg.Lx, tncg.Ly) == (3, 7)
        assert not tncg.outer_inds()
        assert tncg.max_bond() == 3
        assert "I4,0" not in tncg.tag_map
        assert "X5" not in tncg.tag_map

        tncg = tn.coarse_grain_hotrg("y", max_bond=3, cutoff=0.0, lazy=lazy)
        assert (tncg.Lx, tncg.Ly) == (6, 4)
        assert not tncg.outer_inds()
        assert tncg.max_bond() == 3
        assert "I0,5" not in tncg.tag_map
        assert "Y6" not in tncg.tag_map

    def test_contract_hotrg(self):
        tn = qtn.TN2D_classical_ising_partition_function(16, 16, 0.44)
        tn.contract_hotrg_(max_bond=5, progbar=True, equalize_norms=1.0)
        Zap = tn.item() * 10**tn.exponent
        assert Zap == pytest.approx(8.459419593253275e100, rel=2e-3)

    def test_contract_hotrg_two_layer_rand_peps(self):
        rng = np.random.default_rng(42)
        psi = qtn.PEPS.from_fill_fn(
            lambda shape: rng.uniform(low=-0.1, size=shape),
            Lx=7,
            Ly=5,
            bond_dim=2,
        )
        norm = psi.make_norm()
        xe = norm.contract(all, optimize="auto-hq")
        xt = norm.contract_hotrg(max_bond=5)
        assert xt == pytest.approx(xe, rel=1e-4)

    def test_ising_accuracy_regression(self):
        tn = qtn.TN2D_classical_ising_partition_function(16, 16, 0.44)
        for s in [("xmin",), ("xmax",), ("ymin",), ("ymax",)]:
            Zap = tn.contract_boundary(max_bond=8, sequence=s)
            assert Zap == pytest.approx(8.459419593253275e100, rel=2.2e-7)
        for s in [("xmin", "xmax"), ("ymin", "ymax")]:
            Zap = tn.contract_boundary(max_bond=8, sequence=s)
            assert Zap == pytest.approx(8.459419593253275e100, rel=3.9e-9)

    @pytest.mark.parametrize("mode", ["mps", "ctmrg", "hotrg"])
    def test_cdl_rand_large(self, mode):
        tn = qtn.TN2D_rand_hidden_loop(10, 10, seed=42, contract_sites=False)
        Zex = tn.contract(...)
        tn = qtn.TN2D_rand_hidden_loop(10, 10, seed=42, contract_sites=True)

        if mode == "mps":
            Z = tn.contract_boundary(max_bond=16)
        elif mode == "ctmrg":
            Z = tn.contract_ctmrg(max_bond=16)
        elif mode == "hotrg":
            Z = tn.contract_hotrg(max_bond=16)

        assert Z == pytest.approx(Zex, rel=1e-1)

    @pytest.mark.parametrize(
        "mode,two_layer",
        [
            ("mps", False),
            ("mps", True),
            ("full-bond", False),
        ],
    )
    def test_compute_x_envs(self, mode, two_layer):
        psi = qtn.PEPS.rand(5, 4, 2, seed=42, tags="KET")
        norm = psi.make_norm()
        ex = norm.contract(all)

        if two_layer:
            compress_opts = {
                "cutoff": 1e-6,
                "max_bond": 12,
                "mode": mode,
                "layer_tags": ["KET", "BRA"],
            }
        else:
            compress_opts = {"cutoff": 1e-6, "max_bond": 8, "mode": mode}
        row_envs = norm.compute_x_environments(**compress_opts)

        for i in range(norm.Lx):
            norm_i = (
                row_envs["xmin", i]
                & norm.select(norm.x_tag(i))
                & row_envs["xmax", i]
            )
            x = norm_i.contract(all)
            assert x == pytest.approx(ex, rel=1e-2)

    @pytest.mark.parametrize(
        "mode,two_layer",
        [
            ("mps", False),
            ("mps", True),
            ("full-bond", False),
        ],
    )
    def test_compute_y_envs(self, mode, two_layer):
        psi = qtn.PEPS.rand(4, 5, 2, seed=42, tags="KET")
        norm = psi.retag({"KET": "BRA"}).H | psi
        ex = norm.contract(all)

        if two_layer:
            compress_opts = {
                "cutoff": 1e-6,
                "max_bond": 12,
                "mode": mode,
                "layer_tags": ["KET", "BRA"],
            }
        else:
            compress_opts = {"cutoff": 1e-6, "max_bond": 8, "mode": mode}
        col_envs = norm.compute_y_environments(**compress_opts)

        for j in range(norm.Lx):
            norm_j = (
                col_envs["ymin", j]
                & norm.select(norm.y_tag(j))
                & col_envs["ymax", j]
            )
            x = norm_j.contract(all)
            assert x == pytest.approx(ex, rel=1e-2)

    def test_normalize(self):
        psi = qtn.PEPS.rand(4, 5, 2, seed=42)
        norm = (psi.H | psi).contract(all)
        assert norm != pytest.approx(1.0)
        psi.normalize_(balance_bonds=True, equalize_norms=True, cutoff=2e-3)
        norm = (psi.H | psi).contract(all)
        assert norm == pytest.approx(1.0, rel=0.01)

    @pytest.mark.parametrize("normalized", [False, True])
    @pytest.mark.parametrize("mode", ["mps", "full-bond"])
    def test_compute_local_expectation_one_sites(self, mode, normalized):
        peps = qtn.PEPS.rand(4, 3, 2, seed=42, dtype="complex")

        # reference
        k = peps.to_qarray()
        if normalized:
            qu.normalize(k)
        coos = list(itertools.product([0, 2, 3], [0, 1, 2]))
        terms = {coo: qu.rand_matrix(2) for coo in coos}
        dims = [[2] * 3] * 4
        A = sum(
            qu.ikron(A, dims, [coo], sparse=True) for coo, A in terms.items()
        )
        ex = qu.expec(A, k)

        opts = dict(cutoff=2e-3, max_bond=9, contract_optimize="auto-hq")
        e = peps.compute_local_expectation(
            terms, mode=mode, normalized=normalized, **opts
        )

        assert e == pytest.approx(ex, rel=1e-2)

    @pytest.mark.parametrize("normalized", [False, True])
    @pytest.mark.parametrize("mode", ["mps", "full-bond"])
    def test_compute_local_expectation_two_sites(self, mode, normalized):
        H = qu.ham_heis_2D(4, 3, sparse=True)
        Hij = qu.ham_heis(2, cyclic=False)

        peps = qtn.PEPS.rand(4, 3, 2, seed=42)
        k = peps.to_qarray()

        if normalized:
            qu.normalize(k)
        ex = qu.expec(H, k)

        opts = dict(
            mode=mode,
            normalized=normalized,
            cutoff=2e-3,
            max_bond=16,
            contract_optimize="auto-hq",
        )

        # compute 2x1 and 1x2 plaquettes separately
        hterms = {coos: Hij for coos in peps.gen_horizontal_bond_coos()}
        vterms = {coos: Hij for coos in peps.gen_vertical_bond_coos()}

        he = peps.compute_local_expectation(hterms, **opts)
        ve = peps.compute_local_expectation(vterms, **opts)

        assert he + ve == pytest.approx(ex, rel=1e-2)

        # compute all terms in 2x2 plaquettes
        terms_all = {**hterms, **vterms}
        e = peps.compute_local_expectation(terms_all, autogroup=False, **opts)

        assert e == pytest.approx(ex, rel=1e-2)

    def test_cyclic_basic(self):
        tn = qtn.TN2D_rand(Lx=3, Ly=4, D=2, cyclic=True)
        assert tn.is_cyclic_x()
        assert tn.is_cyclic_y()
        assert tn.num_indices == 2 * 3 * 4
        tn = qtn.TN2D_rand(Lx=3, Ly=4, D=2, cyclic=(False, True))
        assert not tn.is_cyclic_x()
        assert tn.is_cyclic_y()
        assert tn.num_indices == 2 * 3 * 4 - 4
        tn = qtn.TN2D_rand(Lx=3, Ly=4, D=2, cyclic=(True, False))
        assert tn.is_cyclic_x()
        assert not tn.is_cyclic_y()
        assert tn.num_indices == 2 * 3 * 4 - 3
        tn = qtn.TN2D_rand(Lx=3, Ly=4, D=2, cyclic=(False, False))
        assert not tn.is_cyclic_x()
        assert not tn.is_cyclic_y()
        assert tn.num_indices == 2 * 3 * 4 - 7

    @pytest.mark.parametrize("cyclicx", [False, True])
    @pytest.mark.parametrize("cyclicy", [False, True])
    @pytest.mark.parametrize("mode", ["mps", "hotrg", "ctmrg"])
    def test_cyclic_contract(self, cyclicx, cyclicy, mode):
        Lx = 5
        Ly = 6
        D = 2
        chi = 3
        tn = qtn.TN2D_rand(
            Lx,
            Ly,
            D,
            cyclic=(cyclicx, cyclicy),
            seed=42,
            dist="uniform",
        )
        Zex = tn.contract(...)
        if mode == "hotrg":
            Z = tn.contract_hotrg(chi)
        elif mode == "ctmrg":
            Z = tn.contract_ctmrg(chi)
        else:
            Z = tn.contract_boundary(chi, mode=mode)
        assert abs(1 - Z / Zex) < 1e-3


class TestPEPO:
    @pytest.mark.parametrize("Lx", [3, 4, 5])
    @pytest.mark.parametrize("Ly", [3, 4, 5])
    def test_basic_rand(self, Lx, Ly):
        X = qtn.PEPO.rand_herm(Lx, Ly, bond_dim=4)

        assert X.max_bond() == 4
        assert X.Lx == Lx
        assert X.Ly == Ly
        assert len(X.tensor_map) == Lx * Ly
        assert X.upper_inds == tuple(
            f"k{i},{j}" for i in range(Lx) for j in range(Ly)
        )
        assert X.lower_inds == tuple(
            f"b{i},{j}" for i in range(Lx) for j in range(Ly)
        )
        assert X.site_tags == tuple(
            f"I{i},{j}" for i in range(Lx) for j in range(Ly)
        )

        assert X.bond_size((1, 1), (1, 2)) == (4)

        for i in range(Lx):
            assert len(X.select(f"X{i}").tensor_map) == Ly
        for j in range(Ly):
            assert len(X.select(f"Y{j}").tensor_map) == Lx

        for i in range(Lx):
            for j in range(Ly):
                assert X.phys_dim(i, j) == 2
                assert isinstance(X[i, j], qtn.Tensor)
                assert isinstance(X[f"I{i},{j}"], qtn.Tensor)

        if Lx == Ly == 3:
            X_dense = X.to_qarray(optimize="auto-hq")
            assert X_dense.shape == (512, 512)
            assert qu.isherm(X_dense)

        X.show()
        assert f"Lx={Lx}" in X.__str__()
        assert f"Lx={Lx}" in X.__repr__()

    def test_add_pepo(self):
        pa = qtn.PEPO.rand(3, 4, 2)
        pb = qtn.PEPO.rand(3, 4, 3)
        pc = qtn.PEPO.rand(3, 4, 4)
        pab = pa + pb
        assert pab.max_bond() == 5
        assert pab @ pc == pytest.approx(pa @ pc + pb @ pc)

    def test_apply_pepo(self):
        A = qtn.PEPO.rand(Lx=3, Ly=2, bond_dim=2, seed=1)
        x = qtn.PEPS.rand(Lx=3, Ly=2, bond_dim=2, seed=0)
        y = A.apply(x)
        assert y.num_indices == x.num_indices
        Ad = A.to_qarray()
        xd = x.to_qarray()
        yd = y.to_qarray()
        assert_allclose(Ad @ xd, yd)
        yc = A.apply(x, compress=True, max_bond=3)
        assert yc.max_bond() == 3


class TestMisc:
    def test_calc_plaquette_sizes(self):
        from quimb.tensor.tensor_2d import calc_plaquette_sizes

        H2 = {None: qu.ham_heis(2)}
        ham = qtn.LocalHam2D(10, 10, H2)
        assert calc_plaquette_sizes(ham.terms.keys()) == ((1, 2), (2, 1))
        assert calc_plaquette_sizes(ham.terms.keys(), autogroup=False) == (
            (2, 2),
        )
        H2[(1, 1), (2, 2)] = 0.5 * qu.ham_heis(2)
        ham = qtn.LocalHam2D(10, 10, H2)
        assert calc_plaquette_sizes(ham.terms.keys()) == ((2, 2),)
        H2[(2, 2), (2, 4)] = 0.25 * qu.ham_heis(2)
        H2[(2, 4), (4, 4)] = 0.25 * qu.ham_heis(2)
        ham = qtn.LocalHam2D(10, 10, H2)
        assert calc_plaquette_sizes(ham.terms.keys()) == (
            (1, 3),
            (2, 2),
            (3, 1),
        )
        assert calc_plaquette_sizes(ham.terms.keys(), autogroup=False) == (
            (3, 3),
        )

    def test_calc_plaquette_map(self):
        from quimb.tensor.tensor_2d import calc_plaquette_map

        plaquettes = [
            # 2x2 plaquette covering all sites
            ((0, 0), (2, 2)),
            # horizontal plaquettes
            ((0, 0), (1, 2)),
            ((1, 0), (1, 2)),
            # vertical plaquettes
            ((0, 0), (2, 1)),
            ((0, 1), (2, 1)),
        ]
        assert calc_plaquette_map(plaquettes) == {
            (0, 0): ((0, 0), (2, 1)),
            (0, 1): ((0, 1), (2, 1)),
            (1, 0): ((1, 0), (1, 2)),
            (1, 1): ((1, 0), (1, 2)),
            ((0, 0), (0, 1)): ((0, 0), (1, 2)),
            ((0, 0), (1, 0)): ((0, 0), (2, 1)),
            ((0, 0), (1, 1)): ((0, 0), (2, 2)),
            ((0, 1), (1, 0)): ((0, 0), (2, 2)),
            ((0, 1), (1, 1)): ((0, 1), (2, 1)),
            ((1, 0), (1, 1)): ((1, 0), (1, 2)),
        }
