"""Standard quantum gates for quantum computing.

This module provides commonly used quantum gates as numpy arrays, with
validation and composition utilities.

All gates are returned as complex numpy arrays and are verified for unitarity.

Author: Generated by AgentBible
"""

from typing import Optional

import numpy as np
from numpy.typing import NDArray

# =============================================================================
# Pauli Matrices
# =============================================================================


def pauli_x() -> NDArray[np.complexfloating]:
    """Pauli X gate (bit flip).

    X = [[0, 1], [1, 0]]

    Returns:
        2x2 unitary matrix.
    """
    return np.array([[0, 1], [1, 0]], dtype=complex)


def pauli_y() -> NDArray[np.complexfloating]:
    """Pauli Y gate.

    Y = [[0, -i], [i, 0]]

    Returns:
        2x2 unitary matrix.
    """
    return np.array([[0, -1j], [1j, 0]], dtype=complex)


def pauli_z() -> NDArray[np.complexfloating]:
    """Pauli Z gate (phase flip).

    Z = [[1, 0], [0, -1]]

    Returns:
        2x2 unitary matrix.
    """
    return np.array([[1, 0], [0, -1]], dtype=complex)


def identity(n: int = 1) -> NDArray[np.complexfloating]:
    """Identity gate for n qubits.

    Args:
        n: Number of qubits.

    Returns:
        2^n x 2^n identity matrix.
    """
    return np.eye(2**n, dtype=complex)


# =============================================================================
# Common Single-Qubit Gates
# =============================================================================


def hadamard() -> NDArray[np.complexfloating]:
    """Hadamard gate.

    H = (1/sqrt(2)) * [[1, 1], [1, -1]]

    Creates equal superposition: H|0> = |+>, H|1> = |->

    Returns:
        2x2 unitary matrix.
    """
    return np.array([[1, 1], [1, -1]], dtype=complex) / np.sqrt(2)


def phase(phi: float = np.pi / 2) -> NDArray[np.complexfloating]:
    """Phase gate S (default) or general phase gate.

    P(phi) = [[1, 0], [0, exp(i*phi)]]
    S = P(pi/2) (default)

    Args:
        phi: Phase angle in radians. Default is pi/2 (S gate).

    Returns:
        2x2 unitary matrix.
    """
    return np.array([[1, 0], [0, np.exp(1j * phi)]], dtype=complex)


def t_gate() -> NDArray[np.complexfloating]:
    """T gate (pi/8 gate).

    T = [[1, 0], [0, exp(i*pi/4)]]

    Returns:
        2x2 unitary matrix.
    """
    return phase(np.pi / 4)


def s_gate() -> NDArray[np.complexfloating]:
    """S gate (phase gate).

    S = [[1, 0], [0, i]]

    Returns:
        2x2 unitary matrix.
    """
    return phase(np.pi / 2)


# =============================================================================
# Rotation Gates
# =============================================================================


def rx(theta: float) -> NDArray[np.complexfloating]:
    """Rotation around X-axis.

    Rx(theta) = exp(-i * theta/2 * X)
              = [[cos(theta/2), -i*sin(theta/2)],
                 [-i*sin(theta/2), cos(theta/2)]]

    Args:
        theta: Rotation angle in radians.

    Returns:
        2x2 unitary matrix.
    """
    c = np.cos(theta / 2)
    s = np.sin(theta / 2)
    return np.array([[c, -1j * s], [-1j * s, c]], dtype=complex)


def ry(theta: float) -> NDArray[np.complexfloating]:
    """Rotation around Y-axis.

    Ry(theta) = exp(-i * theta/2 * Y)
              = [[cos(theta/2), -sin(theta/2)],
                 [sin(theta/2), cos(theta/2)]]

    Args:
        theta: Rotation angle in radians.

    Returns:
        2x2 unitary matrix.
    """
    c = np.cos(theta / 2)
    s = np.sin(theta / 2)
    return np.array([[c, -s], [s, c]], dtype=complex)


def rz(theta: float) -> NDArray[np.complexfloating]:
    """Rotation around Z-axis.

    Rz(theta) = exp(-i * theta/2 * Z)
              = [[exp(-i*theta/2), 0],
                 [0, exp(i*theta/2)]]

    Args:
        theta: Rotation angle in radians.

    Returns:
        2x2 unitary matrix.
    """
    return np.array(
        [[np.exp(-1j * theta / 2), 0], [0, np.exp(1j * theta / 2)]], dtype=complex
    )


# =============================================================================
# Two-Qubit Gates
# =============================================================================


def cnot() -> NDArray[np.complexfloating]:
    """Controlled-NOT (CNOT) gate.

    CNOT = [[1, 0, 0, 0],
            [0, 1, 0, 0],
            [0, 0, 0, 1],
            [0, 0, 1, 0]]

    Control on first qubit, target on second qubit.

    Returns:
        4x4 unitary matrix.
    """
    return np.array(
        [[1, 0, 0, 0], [0, 1, 0, 0], [0, 0, 0, 1], [0, 0, 1, 0]], dtype=complex
    )


def cz() -> NDArray[np.complexfloating]:
    """Controlled-Z (CZ) gate.

    CZ = [[1, 0, 0, 0],
          [0, 1, 0, 0],
          [0, 0, 1, 0],
          [0, 0, 0, -1]]

    Returns:
        4x4 unitary matrix.
    """
    return np.array(
        [[1, 0, 0, 0], [0, 1, 0, 0], [0, 0, 1, 0], [0, 0, 0, -1]], dtype=complex
    )


def swap() -> NDArray[np.complexfloating]:
    """SWAP gate.

    SWAP = [[1, 0, 0, 0],
            [0, 0, 1, 0],
            [0, 1, 0, 0],
            [0, 0, 0, 1]]

    Returns:
        4x4 unitary matrix.
    """
    return np.array(
        [[1, 0, 0, 0], [0, 0, 1, 0], [0, 1, 0, 0], [0, 0, 0, 1]], dtype=complex
    )


def iswap() -> NDArray[np.complexfloating]:
    """iSWAP gate.

    iSWAP = [[1, 0, 0, 0],
             [0, 0, i, 0],
             [0, i, 0, 0],
             [0, 0, 0, 1]]

    Returns:
        4x4 unitary matrix.
    """
    return np.array(
        [[1, 0, 0, 0], [0, 0, 1j, 0], [0, 1j, 0, 0], [0, 0, 0, 1]], dtype=complex
    )


# =============================================================================
# Gate Composition Utilities
# =============================================================================


def controlled(
    gate: NDArray[np.complexfloating],
    name: Optional[str] = None,
) -> NDArray[np.complexfloating]:
    """Create a controlled version of a single-qubit gate.

    The control qubit is the first qubit, target is the second.

    Args:
        gate: 2x2 unitary gate to make controlled.
        name: Optional name for error messages.

    Returns:
        4x4 controlled gate matrix.

    Raises:
        ValueError: If input is not a 2x2 matrix.

    Example:
        >>> controlled(pauli_x())  # Returns CNOT
        >>> controlled(pauli_z())  # Returns CZ
    """
    label = f" '{name}'" if name else ""

    if gate.shape != (2, 2):
        raise ValueError(f"Gate{label} must be 2x2, got shape {gate.shape}")

    result = np.eye(4, dtype=complex)
    result[2:4, 2:4] = gate
    return result


def tensor(*gates: NDArray[np.complexfloating]) -> NDArray[np.complexfloating]:
    """Compute tensor product of multiple gates.

    Args:
        *gates: Variable number of gate matrices.

    Returns:
        Tensor product of all gates.

    Raises:
        ValueError: If no gates provided.

    Example:
        >>> tensor(hadamard(), pauli_x())  # H (x) X
        >>> tensor(identity(), identity(), pauli_z())  # I (x) I (x) Z
    """
    if not gates:
        raise ValueError("tensor() requires at least one gate")

    result = gates[0]
    for gate in gates[1:]:
        result = np.kron(result, gate)
    return result
