import math

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

from qibotn.backends.vidal import VidalBackend, _can_route_non_adjacent, _unsupported_reason
from qibotn.backends.vidal_tebd import (
    _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_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_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}"
    )


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)