import math

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

from qibotn.backends.cpu import CpuTensorNet
from qibotn.benchmark_cases import (
    build_circuit as build_benchmark_circuit,
    exact_pauli_sum,
)


def build_circuit(nqubits=6):
    circuit = Circuit(nqubits)
    for qubit in range(nqubits):
        circuit.add(gates.RY(qubit, theta=0.1 * (qubit + 1)))
        circuit.add(gates.RZ(qubit, theta=-0.05 * (qubit + 1)))
    for qubit in range(nqubits - 1):
        circuit.add(gates.CNOT(qubit, qubit + 1))
    return circuit


def build_observable(nqubits):
    form = 0
    for qubit in range(nqubits):
        form += 0.5 * X(qubit) * Z((qubit + 1) % nqubits)
    return hamiltonians.SymbolicHamiltonian(form=form)


def test_cpu_generic_tn_expectation_matches_statevector():
    circuit = build_circuit()
    observable = build_observable(circuit.nqubits)
    exact = observable.expectation_from_state(circuit().state(numpy=True))

    backend = CpuTensorNet(
        {
            "MPI_enabled": False,
            "MPS_enabled": False,
            "NCCL_enabled": False,
            "expectation_enabled": observable,
        }
    )
    value = backend.execute_circuit(circuit)[0]

    assert math.isclose(value, exact, abs_tol=1e-12)


def test_cpu_mps_expectation_matches_statevector():
    circuit = build_circuit()
    observable = build_observable(circuit.nqubits)
    exact = observable.expectation_from_state(circuit().state(numpy=True))

    backend = CpuTensorNet(
        {
            "MPI_enabled": False,
            "MPS_enabled": True,
            "NCCL_enabled": False,
            "expectation_enabled": observable,
            "max_bond_dimension": 64,
            "tensor_module": "torch",
            "torch_threads": 1,
        }
    )
    value = backend.execute_circuit(circuit)[0]

    assert math.isclose(value, exact, abs_tol=1e-12)


def test_cpu_runcard_pauli_pattern_matches_statevector():
    circuit = build_circuit()
    observable = {"pauli_string_pattern": "IXZ"}
    exact_hamiltonian = hamiltonians.SymbolicHamiltonian(
        form=X(1) * Z(2) * X(4) * Z(5)
    )
    exact = exact_hamiltonian.expectation_from_state(circuit().state(numpy=True))

    for mps_enabled in (False, True):
        backend = CpuTensorNet(
            {
                "MPI_enabled": False,
                "MPS_enabled": mps_enabled,
                "NCCL_enabled": False,
                "expectation_enabled": observable,
                "max_bond_dimension": 64,
                "tensor_module": "torch",
                "torch_threads": 1,
            }
        )
        value = backend.execute_circuit(circuit)[0]

        assert math.isclose(value, exact, abs_tol=1e-12)


def test_cpu_mps_sampling_uses_nshots():
    circuit = Circuit(4)
    circuit.add(gates.H(0))
    for qubit in range(3):
        circuit.add(gates.CNOT(qubit, qubit + 1))

    backend = CpuTensorNet(
        {
            "MPI_enabled": False,
            "MPS_enabled": True,
            "NCCL_enabled": False,
            "expectation_enabled": False,
        }
    )
    result = backend.execute_circuit(circuit, nshots=100)

    assert sum(result.frequencies().values()) == 100
    assert set(result.frequencies()) <= {"0000", "1111"}


def test_cpu_mps_mpo_expectation_matches_statevector():
    circuit = build_circuit(nqubits=4)
    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 = CpuTensorNet(
        {
            "MPI_enabled": False,
            "MPS_enabled": True,
            "NCCL_enabled": False,
            "expectation_enabled": {"mpo_tensors": mpo},
            "max_bond_dimension": 64,
            "tensor_module": "torch",
            "torch_threads": 1,
        }
    )
    value = backend.execute_circuit(circuit)[0]

    assert math.isclose(value, exact, abs_tol=1e-12)


def test_cpu_mps_dense_observable_dict_matches_known_value():
    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

    backend = CpuTensorNet(
        {
            "MPI_enabled": False,
            "MPS_enabled": True,
            "NCCL_enabled": False,
            "expectation_enabled": {"matrix": bell, "qubits": [0, 1]},
            "max_bond_dimension": 16,
            "tensor_module": "torch",
            "torch_threads": 1,
        }
    )
    value = backend.execute_circuit(circuit)[0]

    assert math.isclose(value, 1.0, abs_tol=1e-12)


def test_cpu_generic_tn_long_pauli_string_matches_statevector():
    circuit = build_benchmark_circuit("rxx_rzz", 10, 2, 42)
    observable = {"pauli_string_pattern": "XZ"}
    exact_hamiltonian = hamiltonians.SymbolicHamiltonian(
        form=X(0) * Z(1) * X(2) * Z(3) * X(4) * Z(5) * X(6) * Z(7) * X(8) * Z(9)
    )
    exact = exact_hamiltonian.expectation_from_state(circuit().state(numpy=True))

    backend = CpuTensorNet(
        {
            "MPI_enabled": False,
            "MPS_enabled": False,
            "NCCL_enabled": False,
            "expectation_enabled": observable,
        }
    )
    value = backend.execute_circuit(circuit)[0]

    assert math.isclose(value, exact, abs_tol=1e-12)