"""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
    try:
        return hamiltonians.SymbolicHamiltonian(form=observable)
    except Exception as exc:
        raise TypeError("Invalid observable type.") from exc


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."""
    if "pauli_string_pattern" in data:
        return create_hamiltonian_from_pauli_pattern(
            data["pauli_string_pattern"], circuit_nqubit
        )

    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 create_hamiltonian_from_pauli_pattern(pattern, circuit_nqubit):
    """Create a single Pauli-string Hamiltonian by repeating ``pattern``.

    Example: pattern ``"IXZ"`` on 5 qubits becomes ``I0 * X1 * Z2 * I3 * X4``.
    Identity factors are omitted except for the all-identity case.
    """
    if not isinstance(pattern, str) or not pattern:
        raise ValueError("pauli_string_pattern must be a non-empty string.")

    pauli_gates = {"X": X, "Y": Y, "Z": Z}
    pattern = pattern.upper()
    invalid = sorted(set(pattern) - {"I", "X", "Y", "Z"})
    if invalid:
        raise ValueError(
            "pauli_string_pattern characters must be one of I/X/Y/Z; "
            f"got {''.join(invalid)!r}."
        )

    expr = None
    for qubit in range(circuit_nqubit):
        name = pattern[qubit % len(pattern)]
        if name == "I":
            continue
        factor = pauli_gates[name](qubit)
        expr = factor if expr is None else expr * factor

    if expr is None:
        expr = I(0)

    return hamiltonians.SymbolicHamiltonian(form=expr)


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
