Source code for uniqc.algorithmics.circuits.entangled_states

"""Entangled state preparation circuits: GHZ, W, and Cluster states."""

__all__ = [
    "ghz_state",
    "w_state",
    "cluster_state",
]

from typing import List, Optional, Tuple

from uniqc.circuit_builder import Circuit
from uniqc.algorithmics.circuits.dicke_state import dicke_state_circuit




def ghz_state(
    circuit: Circuit,
    qubits: Optional[List[int]] = None,
) -> None:
    r"""Prepare a GHZ (Greenberger–Horne–Zeilinger) state.

    Produces the state:

    .. math::

        \frac{1}{\sqrt{2}}(|00\ldots0\rangle + |11\ldots1\rangle)

    Implementation:

    1. Hadamard on the first qubit: :math:`\frac{1}{\sqrt{2}}(|0\rangle + |1\rangle)`
    2. Chain of CNOT gates: ``CNOT(q[0], q[1])``, ``CNOT(q[1], q[2])``, ...

    Args:
        circuit: Quantum circuit to operate on (mutated in-place).
        qubits: Qubit indices. ``None`` means all qubits of *circuit*.

    Raises:
        ValueError: Fewer than 2 qubits.

    Example:
        >>> from uniqc.circuit_builder import Circuit
        >>> from uniqc.algorithmics.circuits import ghz_state
        >>> c = Circuit(3)
        >>> ghz_state(c)
    """
    if qubits is None:
        qubits = list(range(circuit.qubit_num))

    n = len(qubits)
    if n < 2:
        raise ValueError("ghz_state requires at least 2 qubits")

    circuit.h(qubits[0])
    for i in range(n - 1):
        circuit.cnot(qubits[i], qubits[i + 1])





def w_state(
    circuit: Circuit,
    qubits: Optional[List[int]] = None,
) -> None:
    r"""Prepare a W state.

    Produces the state:

    .. math::

        \frac{1}{\sqrt{n}}(|10\ldots0\rangle + |01\ldots0\rangle +
        \ldots + |00\ldots1\rangle)

    Implemented as a Dicke state :math:`|D(n,1)\rangle` (equal superposition
    of all single-excitation computational basis states) via
    :func:`dicke_state_circuit` with ``k=1``.

    Args:
        circuit: Quantum circuit to operate on (mutated in-place).
        qubits: Qubit indices. ``None`` means all qubits of *circuit*.

    Raises:
        ValueError: Fewer than 2 qubits.

    Example:
        >>> from uniqc.circuit_builder import Circuit
        >>> from uniqc.algorithmics.circuits import w_state
        >>> c = Circuit(4)
        >>> w_state(c)
    """
    if qubits is None:
        qubits = list(range(circuit.qubit_num))

    if len(qubits) < 2:
        raise ValueError("w_state requires at least 2 qubits")

    dicke_state_circuit(circuit, k=1, qubits=qubits)





def cluster_state(
    circuit: Circuit,
    qubits: Optional[List[int]] = None,
    edges: Optional[List[Tuple[int, int]]] = None,
) -> None:
    r"""Prepare a cluster state (graph state).

    A cluster state is prepared by:

    1. Applying Hadamard to all qubits: :math:`H^{\otimes n}`
    2. Applying CZ (controlled-Z) on each edge of the graph

    The resulting state is:

    .. math::

        \frac{1}{\sqrt{2^n}} \sum_{x \in \{0,1\}^n} (-1)^{\sum_{(i,j) \in E}
        x_i x_j} |x\rangle

    Args:
        circuit: Quantum circuit to operate on (mutated in-place).
        qubits: Qubit indices. ``None`` means all qubits of *circuit*.
        edges: List of (source, target) pairs defining the entanglement graph.
            Indices refer to positions in *qubits*. ``None`` uses a linear
            nearest-neighbor chain: ``(0,1), (1,2), (2,3), ...``.

    Raises:
        ValueError: Fewer than 1 qubit.

    Example:
        >>> from uniqc.circuit_builder import Circuit
        >>> from uniqc.algorithmics.circuits import cluster_state
        >>> c = Circuit(4)
        >>> cluster_state(c)  # linear chain
        >>> # Custom graph (square)
        >>> c2 = Circuit(4)
        >>> cluster_state(c2, edges=[(0,1), (1,2), (2,3), (3,0)])
    """
    if qubits is None:
        qubits = list(range(circuit.qubit_num))

    n = len(qubits)
    if n < 1:
        raise ValueError("cluster_state requires at least 1 qubit")

    # Step 1: Hadamard on all qubits
    for q in qubits:
        circuit.h(q)

    # Step 2: CZ on each edge
    if edges is None:
        # Linear chain
        edges = [(i, i + 1) for i in range(n - 1)]

    for src_idx, tgt_idx in edges:
        if src_idx >= n or tgt_idx >= n:
            raise ValueError(
                f"Edge ({src_idx}, {tgt_idx}) out of range for {n} qubits"
            )
        circuit.cz(qubits[src_idx], qubits[tgt_idx])