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


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_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)