"""Observable helpers shared by tensor-network backends and benchmarks."""

from qibo import hamiltonians
from qibo.symbols import I, X, Y, Z


def check_observable(observable, circuit_nqubit):
    """Checks the type of observable and returns the appropriate Hamiltonian."""
    if observable is None:
        return build_observable(circuit_nqubit)
    if isinstance(observable, dict):
        return create_hamiltonian_from_dict(observable, circuit_nqubit)
    if isinstance(observable, hamiltonians.SymbolicHamiltonian):
        return observable
    raise TypeError("Invalid observable type.")


def build_observable(circuit_nqubit):
    """Construct the default benchmark observable used by qibotn."""
    hamiltonian_form = 0
    for i in range(circuit_nqubit):
        hamiltonian_form += 0.5 * X(i % circuit_nqubit) * Z((i + 1) % circuit_nqubit)

    return hamiltonians.SymbolicHamiltonian(form=hamiltonian_form)


def create_hamiltonian_from_dict(data, circuit_nqubit):
    """Create a Qibo SymbolicHamiltonian from the qibotn dict representation."""
    pauli_gates = {"X": X, "Y": Y, "Z": Z}
    terms = []

    for term in data["terms"]:
        coeff = term["coefficient"]
        operators = term["operators"]
        operator_dict = {q: pauli_gates[g] for g, q in operators}

        full_term_expr = [
            operator_dict[q](q) if q in operator_dict else I(q)
            for q in range(circuit_nqubit)
        ]

        term_expr = full_term_expr[0]
        for op in full_term_expr[1:]:
            term_expr *= op

        terms.append(coeff * term_expr)

    if not terms:
        raise ValueError("No valid Hamiltonian terms were added.")

    return hamiltonians.SymbolicHamiltonian(sum(terms))


def build_random_circuit(nqubits, nlayers, seed=42):
    """Build a random circuit with RY+RZ+CNOT layers for benchmarks."""
    import numpy as np
    from qibo import Circuit, gates
    np.random.seed(seed)
    c = Circuit(nqubits)
    for _ in range(nlayers):
        for q in range(nqubits):
            c.add(gates.RY(q, theta=np.random.uniform(0, 2*np.pi)))
            c.add(gates.RZ(q, theta=np.random.uniform(0, 2*np.pi)))
        for q in range(nqubits):
            c.add(gates.CNOT(q % nqubits, (q + 1) % nqubits))
    return c


def extract_gates_and_qubits(hamiltonian):
    """Extract per-term Pauli factors from a Qibo SymbolicHamiltonian.

    Returns list of terms, where each term is (coefficient, [(qubit, gate_name), ...]).
    """
    extracted_terms = []

    if not isinstance(hamiltonian, hamiltonians.SymbolicHamiltonian):
        raise ValueError(
            "Unsupported Hamiltonian type. Must be SymbolicHamiltonian or Hamiltonian."
        )

    for term in hamiltonian.terms:
        coeff = term.coefficient
        factors = [(int(str(f)[1:]), str(f)[0]) for f in term.factors]
        extracted_terms.append((coeff, factors))

    return extracted_terms
