"""Reusable benchmark circuits and observables for expectation runs."""

from __future__ import annotations

import math

import numpy as np
from qibo import Circuit, gates


CIRCUITS = (
    "brickwall_cnot",
    "reversed_cnot",
    "shifted_cz",
    "rxx_rzz",
    "swap_scramble",
    "ghz_ladder",
)

OBSERVABLES = (
    "ring_xz",
    "open_zz",
    "mixed_local",
    "range2_xx",
    "long_z_string",
)


def parse_names(raw, valid, label):
    if raw == ["all"]:
        return list(valid)
    unknown = sorted(set(raw) - set(valid))
    if unknown:
        raise ValueError(f"Unknown {label}: {', '.join(unknown)}")
    return raw


def build_circuit(kind, nqubits, nlayers, seed):
    rng = np.random.default_rng(seed)
    circuit = Circuit(nqubits)

    if kind == "ghz_ladder":
        circuit.add(gates.H(0))
        for qubit in range(nqubits - 1):
            circuit.add(gates.CNOT(qubit, qubit + 1))
        return circuit

    for layer in range(nlayers):
        for qubit in range(nqubits):
            circuit.add(gates.RY(qubit, theta=rng.uniform(-math.pi, math.pi)))
            circuit.add(gates.RZ(qubit, theta=rng.uniform(-math.pi, math.pi)))
            if kind in ("rxx_rzz", "swap_scramble"):
                circuit.add(gates.RX(qubit, theta=rng.uniform(-math.pi, math.pi)))

        if kind == "brickwall_cnot":
            add_brickwall(circuit, nqubits, gates.CNOT, layer, reverse=False)
        elif kind == "reversed_cnot":
            add_brickwall(circuit, nqubits, gates.CNOT, layer, reverse=True)
        elif kind == "shifted_cz":
            for qubit in range(layer % 2, nqubits - 1, 2):
                circuit.add(gates.CZ(qubit, qubit + 1))
        elif kind == "rxx_rzz":
            for qubit in range(layer % 2, nqubits - 1, 2):
                circuit.add(gates.RXX(qubit, qubit + 1, theta=rng.uniform(-0.7, 0.7)))
                circuit.add(gates.RZZ(qubit, qubit + 1, theta=rng.uniform(-0.7, 0.7)))
        elif kind == "swap_scramble":
            for qubit in range(layer % 2, nqubits - 1, 2):
                circuit.add(gates.CZ(qubit, qubit + 1))
                if layer % 4 == 3:
                    circuit.add(gates.SWAP(qubit, qubit + 1))
        else:
            raise ValueError(f"Unknown circuit kind {kind!r}.")

    return circuit


def add_brickwall(circuit, nqubits, gate, layer, reverse):
    for qubit in range(0, nqubits - 1, 2):
        if reverse and layer % 2:
            circuit.add(gate(qubit + 1, qubit))
        else:
            circuit.add(gate(qubit, qubit + 1))
    for qubit in range(1, nqubits - 1, 2):
        if reverse and not layer % 2:
            circuit.add(gate(qubit + 1, qubit))
        else:
            circuit.add(gate(qubit, qubit + 1))


def observable_terms(kind, nqubits):
    if kind == "ring_xz":
        return [
            (0.5, (("X", site), ("Z", (site + 1) % nqubits)))
            for site in range(nqubits)
        ]
    if kind == "open_zz":
        return [
            (1.0 / (nqubits - 1), (("Z", site), ("Z", site + 1)))
            for site in range(nqubits - 1)
        ]
    if kind == "mixed_local":
        terms = [(0.25, (("X", 0),)), (-0.5, (("Z", nqubits - 1),))]
        terms += [
            (0.125, (("Y", site), ("Y", site + 1)))
            for site in range(0, nqubits - 1, 3)
        ]
        return terms
    if kind == "range2_xx":
        return [
            (1.0 / max(1, nqubits - 2), (("X", site), ("X", site + 2)))
            for site in range(nqubits - 2)
        ]
    if kind == "long_z_string":
        stride = max(1, nqubits // 16)
        return [(1.0, tuple(("Z", site) for site in range(0, nqubits, stride)))]
    raise ValueError(f"Unknown observable kind {kind!r}.")


def terms_to_dict(terms):
    return {
        "terms": [
            {
                "coefficient": float(np.real(coeff)),
                "operators": [(name, int(site)) for name, site in ops],
            }
            for coeff, ops in terms
        ]
    }


def exact_pauli_sum(circuit, terms, nqubits):
    state = circuit().state(numpy=True).reshape(-1)
    indices = np.arange(state.size, dtype=np.int64)
    value = 0.0 + 0.0j
    for coeff, ops in terms:
        flipped = indices.copy()
        phase = np.ones(state.size, dtype=np.complex128)
        for name, site in ops:
            shift = nqubits - 1 - site
            bit = (indices >> shift) & 1
            if name == "X":
                flipped ^= 1 << shift
            elif name == "Y":
                flipped ^= 1 << shift
                phase *= 1j * (1 - 2 * bit)
            elif name == "Z":
                phase *= 1 - 2 * bit
            elif name != "I":
                raise ValueError(f"Unsupported Pauli {name!r}.")
        value += coeff * np.vdot(state[flipped], phase * state)
    return float(value.real)
