import math

import numpy as np
from qibo import Circuit, gates, hamiltonians
from qibo.symbols import Symbol, X, Y, Z

from qibotn.benchmark_cases import exact_pauli_sum
from qibotn.backends.vidal import (
    VidalBackend,
    _can_route_non_adjacent,
    _unsupported_reason,
    _operator_terms_to_mpo,
    _symbolic_hamiltonian_to_operator_terms,
)
from qibotn.backends.vidal_tebd import (
    VidalTEBDExecutor,
    _route_non_adjacent_gates,
    _gate_sites,
)


def build_local_circuit(nqubits=8, nlayers=3, seed=42):
    rng = np.random.default_rng(seed)
    circuit = Circuit(nqubits)
    for layer in range(nlayers):
        for q in range(nqubits):
            circuit.add(gates.RY(q, theta=rng.uniform(-math.pi, math.pi)))
            circuit.add(gates.RZ(q, theta=rng.uniform(-math.pi, math.pi)))
        for q in range(layer % 2, nqubits - 1, 2):
            circuit.add(gates.CNOT(q, q + 1))
    return circuit


def test_vidal_backend_expectation_matches_statevector():
    circuit = build_local_circuit()
    observable = hamiltonians.SymbolicHamiltonian(
        form=0.5 * X(0) * Z(1) + 0.25 * Y(2) * Y(3) - 0.7 * Z(7)
    )
    exact = observable.expectation_from_state(circuit().state(numpy=True))

    backend = VidalBackend()
    backend.configure_tn_simulation(max_bond_dimension=128, tensor_module="torch")
    value = backend.expectation(circuit, observable)

    np.testing.assert_allclose(value, exact, atol=1e-12)


def test_vidal_backend_accepts_unlimited_bond_and_no_cutoff():
    circuit = build_local_circuit(nqubits=6, nlayers=2)
    observable = hamiltonians.SymbolicHamiltonian(
        form=0.5 * X(0) * Z(1) - 0.7 * Z(5)
    )
    exact = observable.expectation_from_state(circuit().state(numpy=True))

    backend = VidalBackend()
    backend.configure_tn_simulation(
        max_bond_dimension=None,
        cut_ratio=None,
        tensor_module="torch",
        fallback=False,
    )
    value = backend.expectation(circuit, observable, preprocess=False)

    np.testing.assert_allclose(value, exact, atol=1e-12)


def test_vidal_backend_fallback_for_non_adjacent_gate():
    """compile_circuit=False (default) → falls back to qmatchatea for non-adjacent."""
    circuit = Circuit(4)
    circuit.add(gates.H(0))
    circuit.add(gates.CNOT(0, 3))
    observable = hamiltonians.SymbolicHamiltonian(form=Z(0) * Z(3))

    backend = VidalBackend()
    backend.configure_tn_simulation(max_bond_dimension=32, tensor_module="torch")
    value = backend.expectation(circuit, observable)

    exact = observable.expectation_from_state(circuit().state(numpy=True))
    np.testing.assert_allclose(value, exact, atol=1e-12)


def test_vidal_backend_routes_non_adjacent_with_compile():
    """Non-adjacent gate with compile_circuit=True goes through Vidal SWAP routing."""
    circuit = Circuit(4)
    circuit.add(gates.H(0))
    circuit.add(gates.CNOT(0, 3))

    observable = hamiltonians.SymbolicHamiltonian(form=Z(0) * Z(3))

    backend = VidalBackend()
    backend.configure_tn_simulation(
        max_bond_dimension=32, tensor_module="torch", compile_circuit=True,
    )
    value = backend.expectation(circuit, observable)

    exact = observable.expectation_from_state(circuit().state(numpy=True))
    np.testing.assert_allclose(value, exact, atol=1e-12)


def test_can_route_non_adjacent():
    """_can_route_non_adjacent correctly identifies routable circuits."""
    circuit = Circuit(4)
    circuit.add(gates.H(0))
    circuit.add(gates.CNOT(0, 3))
    assert _can_route_non_adjacent(circuit)

    circuit.add(gates.CNOT(0, 1))
    assert _can_route_non_adjacent(circuit)


def test_cannot_route_multi_qubit():
    """Circuits with 3+ qubit gates cannot be routed."""
    circuit = Circuit(3)
    circuit.add(gates.TOFFOLI(0, 1, 2))
    assert not _can_route_non_adjacent(circuit)


def test_routing_preserves_adjacent_gates():
    """_route_non_adjacent_gates leaves adjacent gates unchanged."""
    circuit = build_local_circuit(nqubits=4, nlayers=2)
    original = list(circuit.queue)
    routed = _route_non_adjacent_gates(original, 4)

    # Count 2Q gates — should be more due to inserted SWAPs, so just
    # check that all 2-site gates ARE adjacent.
    for gate in routed:
        sites = _gate_sites(gate)
        if len(sites) == 2:
            diff = abs(sites[0] - sites[1])
            assert diff == 1, f"Non-adjacent gate after routing: {gate.name} on {sites}"


def test_routing_non_adjacent_cnot():
    """Manually verify SWAP+CNOT+unSWAP for CNOT(0,3)."""
    circuit = Circuit(4)
    circuit.add(gates.H(0))
    circuit.add(gates.H(3))
    circuit.add(gates.CNOT(0, 3))

    routed = _route_non_adjacent_gates(list(circuit.queue), 4)

    # Expected: H(0), H(3), SWAP(2,3), SWAP(1,2), routed CNOT on (0,1), SWAP(1,2), SWAP(2,3)
    names = [getattr(g, "name", g.__class__.__name__) for g in routed]
    assert names == ["h", "h", "swap", "swap", "routed_two_qubit", "swap", "swap"], f"Got {names}"

    # Verify expectation through full pipeline
    observable = hamiltonians.SymbolicHamiltonian(form=Z(0) * Z(3))
    exact = observable.expectation_from_state(circuit().state(numpy=True))

    backend = VidalBackend()
    backend.configure_tn_simulation(
        max_bond_dimension=32, tensor_module="torch", compile_circuit=True,
    )
    value = backend.expectation(circuit, observable)
    np.testing.assert_allclose(value, exact, atol=1e-12)


def test_routing_preserves_reversed_non_adjacent_gate_order():
    circuit = Circuit(6)
    circuit.add(gates.X(5))
    circuit.add(gates.H(0))
    circuit.add(gates.CNOT(5, 0))

    observable = hamiltonians.SymbolicHamiltonian(form=X(0) + Z(5) + Z(0) * Z(5))
    exact = observable.expectation_from_state(circuit().state(numpy=True))

    backend = VidalBackend()
    backend.configure_tn_simulation(
        max_bond_dimension=64,
        tensor_module="torch",
        compile_circuit=True,
        fallback=False,
    )
    value = backend.expectation(circuit, observable, preprocess=False)

    np.testing.assert_allclose(value, exact, atol=1e-12)


def test_vidal_backend_preprocesses_non_adjacent_circuit():
    circuit = Circuit(4)
    circuit.add(gates.H(0))
    circuit.add(gates.CNOT(0, 3))
    observable = hamiltonians.SymbolicHamiltonian(form=Z(0) * Z(3))
    exact = observable.expectation_from_state(circuit().state(numpy=True))

    backend = VidalBackend()
    backend.configure_tn_simulation(
        max_bond_dimension=64,
        tensor_module="torch",
        compile_circuit=True,
        fallback=False,
    )
    value = backend.expectation(circuit, observable, preprocess=True)

    np.testing.assert_allclose(value, exact, atol=1e-12)


def test_vidal_backend_preprocesses_toffoli_locally():
    circuit = Circuit(4)
    circuit.add(gates.H(0))
    circuit.add(gates.H(1))
    circuit.add(gates.TOFFOLI(0, 1, 3))
    observable = hamiltonians.SymbolicHamiltonian(form=Z(0) * Z(3))
    exact = observable.expectation_from_state(circuit().state(numpy=True))

    backend = VidalBackend()
    backend.configure_tn_simulation(
        max_bond_dimension=128,
        tensor_module="torch",
        compile_circuit=True,
        fallback=False,
    )
    value = backend.expectation(circuit, observable, preprocess=True)

    np.testing.assert_allclose(value, exact, atol=1e-12)


def test_vidal_expectation_preserves_complex_coefficients():
    circuit = Circuit(1)
    observable = hamiltonians.SymbolicHamiltonian(form=(1.0 + 2.0j) * Z(0))

    backend = VidalBackend()
    backend.configure_tn_simulation(
        max_bond_dimension=8,
        tensor_module="torch",
        fallback=False,
    )
    value = backend.expectation(circuit, observable, preprocess=False)

    np.testing.assert_allclose(value, 1.0 + 2.0j, atol=1e-12)


def test_vidal_expectation_supports_custom_local_symbols():
    circuit = build_local_circuit(nqubits=4, nlayers=2)
    a0 = Symbol(0, np.array([[0.2, 1.0], [1.0, -0.3]], dtype=complex), name="A")
    b2 = Symbol(2, np.array([[0.7, -0.4j], [0.4j, 0.1]], dtype=complex), name="B")
    a3 = Symbol(3, np.array([[0.5, 0.2], [0.2, -0.8]], dtype=complex), name="A")
    observable = hamiltonians.SymbolicHamiltonian(form=0.7 * a0 * b2 - 0.4 * a3)
    exact = observable.expectation_from_state(circuit().state(numpy=True))

    backend = VidalBackend()
    backend.configure_tn_simulation(
        max_bond_dimension=64,
        tensor_module="torch",
        fallback=False,
    )
    value = backend.expectation(circuit, observable, preprocess=False)

    np.testing.assert_allclose(value, exact, atol=1e-12)


def test_vidal_executor_mpo_expectation_matches_pauli_sum():
    circuit = build_local_circuit(nqubits=4, nlayers=2)
    executor = VidalTEBDExecutor(
        nqubits=circuit.nqubits,
        max_bond=64,
        tensor_module="torch",
    )
    executor.run_circuit(circuit)

    x = np.array([[0, 1], [1, 0]], dtype=complex)
    z = np.array([[1, 0], [0, -1]], dtype=complex)
    i2 = np.eye(2, dtype=complex)
    mpo = [
        x.reshape(1, 2, 2, 1),
        z.reshape(1, 2, 2, 1),
        i2.reshape(1, 2, 2, 1),
        i2.reshape(1, 2, 2, 1),
    ]
    mpo_value = executor.expectation_mpo(mpo)
    pauli_value = executor.expectation_pauli_sum([(1.0, (("X", 0), ("Z", 1)))])

    np.testing.assert_allclose(mpo_value, pauli_value, atol=1e-12)


def test_vidal_backend_accepts_mpo_observable_dict():
    circuit = build_local_circuit(nqubits=4, nlayers=2)
    x = np.array([[0, 1], [1, 0]], dtype=complex)
    z = np.array([[1, 0], [0, -1]], dtype=complex)
    i2 = np.eye(2, dtype=complex)
    mpo = [
        x.reshape(1, 2, 2, 1),
        z.reshape(1, 2, 2, 1),
        i2.reshape(1, 2, 2, 1),
        i2.reshape(1, 2, 2, 1),
    ]
    exact = exact_pauli_sum(circuit, [(1.0, (("X", 0), ("Z", 1)))], 4)

    backend = VidalBackend()
    backend.configure_tn_simulation(
        max_bond_dimension=64,
        tensor_module="torch",
        fallback=False,
    )
    value = backend.expectation(circuit, {"mpo_tensors": mpo}, preprocess=False)

    np.testing.assert_allclose(value, exact, atol=1e-12)


def test_vidal_symbolic_hamiltonian_auto_mpo_matches_operator_sum():
    circuit = build_local_circuit(nqubits=5, nlayers=2)
    observable = hamiltonians.SymbolicHamiltonian(
        form=0.3 * X(0) * Z(1) - 0.2j * Y(2) + 0.7 * Z(3) * X(4)
    )

    executor = VidalTEBDExecutor(
        nqubits=circuit.nqubits,
        max_bond=64,
        tensor_module="torch",
    )
    executor.run_circuit(circuit)
    terms = _symbolic_hamiltonian_to_operator_terms(observable)

    term_value = executor.expectation_operator_sum(terms)
    mpo_value = executor.expectation_mpo(_operator_terms_to_mpo(terms, circuit.nqubits))

    np.testing.assert_allclose(mpo_value, term_value, atol=1e-12)


def test_vidal_backend_accepts_dense_two_qubit_observable():
    circuit = Circuit(2)
    circuit.add(gates.H(0))
    circuit.add(gates.CNOT(0, 1))

    bell = np.zeros((4, 4), dtype=complex)
    bell[0, 0] = bell[0, 3] = bell[3, 0] = bell[3, 3] = 0.5
    observable = {"matrix": bell, "qubits": [0, 1]}

    backend = VidalBackend()
    backend.configure_tn_simulation(
        max_bond_dimension=16,
        tensor_module="torch",
        fallback=False,
    )
    value = backend.expectation(circuit, observable, preprocess=False)

    np.testing.assert_allclose(value, 1.0, atol=1e-12)


def test_vidal_backend_dense_observable_preserves_complex_value():
    circuit = Circuit(2)
    circuit.add(gates.H(0))
    circuit.add(gates.H(1))

    op = np.zeros((4, 4), dtype=complex)
    op[0, 3] = 1.0
    observable = {"coefficient": 1.0j, "matrix": op, "qubits": [0, 1]}

    backend = VidalBackend()
    backend.configure_tn_simulation(
        max_bond_dimension=16,
        tensor_module="torch",
        fallback=False,
    )
    value = backend.expectation(circuit, observable, preprocess=False)

    np.testing.assert_allclose(value, 0.25j, atol=1e-12)


def test_truncation_error_no_truncation():
    """With large bond, truncation error should be essentially zero."""
    circuit = build_local_circuit(nqubits=6, nlayers=2)
    observable = hamiltonians.SymbolicHamiltonian(form=0.5 * X(0) * Z(1))

    backend = VidalBackend()
    backend.configure_tn_simulation(max_bond_dimension=256, tensor_module="torch")
    value = backend.expectation(circuit, observable)
    _ = value  # ensure computation runs

    assert backend.last_truncation_error < 1e-14, (
        f"Expected near-zero truncation error, got {backend.last_truncation_error}"
    )
    assert backend.last_max_truncation_error < 1e-14, (
        "Expected near-zero max truncation error, got "
        f"{backend.last_max_truncation_error}"
    )


def test_vidal_backend_matches_statevector_multiterm():
    """Multi-term observable with non-adjacent gates, compile_circuit=True."""
    circuit = Circuit(5)
    for q in range(5):
        circuit.add(gates.RY(q, theta=0.7))
        circuit.add(gates.RZ(q, theta=0.3))
    circuit.add(gates.CNOT(0, 2))
    circuit.add(gates.CNOT(1, 4))

    observable = hamiltonians.SymbolicHamiltonian(
        form=(0.3 * X(0) * Z(2) + 0.7 * Y(1) * Y(4) - 0.5 * Z(0) * X(4))
    )

    exact_state = circuit().state(numpy=True)
    exact = observable.expectation_from_state(exact_state)

    backend = VidalBackend()
    backend.configure_tn_simulation(
        max_bond_dimension=64, tensor_module="torch", compile_circuit=True,
    )
    value = backend.expectation(circuit, observable)
    np.testing.assert_allclose(value, exact, atol=1e-10)