"""Quantum circuit construction utilities.

This module provides helpers for building and simulating quantum circuits
using pure numpy. For more advanced features, consider using Qiskit or Cirq.

Author: Generated by AgentBible
"""

from dataclasses import dataclass
from typing import Optional

import numpy as np
from numpy.typing import NDArray

from .gates import identity, tensor
from .validation import check_normalized, check_unitarity


@dataclass
class QuantumState:
    """Represents a pure quantum state.

    Attributes:
        amplitudes: Complex state vector (normalized).
        n_qubits: Number of qubits.
    """

    amplitudes: NDArray[np.complexfloating]
    n_qubits: int

    def __post_init__(self) -> None:
        """Validate state after initialization."""
        expected_dim = 2**self.n_qubits
        if len(self.amplitudes) != expected_dim:
            raise ValueError(
                f"State dimension {len(self.amplitudes)} does not match "
                f"{self.n_qubits} qubits (expected {expected_dim})"
            )
        check_normalized(self.amplitudes)

    @classmethod
    def zero_state(cls, n_qubits: int) -> "QuantumState":
        """Create |00...0> state for n qubits.

        Args:
            n_qubits: Number of qubits.

        Returns:
            QuantumState initialized to |00...0>.
        """
        amplitudes = np.zeros(2**n_qubits, dtype=complex)
        amplitudes[0] = 1.0
        return cls(amplitudes=amplitudes, n_qubits=n_qubits)

    @classmethod
    def from_label(cls, label: str) -> "QuantumState":
        """Create computational basis state from label.

        Args:
            label: Binary string like "01" or "110".

        Returns:
            QuantumState for the specified basis state.

        Example:
            >>> QuantumState.from_label("01")  # |01> state
            >>> QuantumState.from_label("110")  # |110> state
        """
        n_qubits = len(label)
        if not all(c in "01" for c in label):
            raise ValueError(f"Label must be binary string, got '{label}'")

        index = int(label, 2)
        amplitudes = np.zeros(2**n_qubits, dtype=complex)
        amplitudes[index] = 1.0
        return cls(amplitudes=amplitudes, n_qubits=n_qubits)

    def probabilities(self) -> np.ndarray:
        """Compute measurement probabilities.

        Returns:
            Array of probabilities for each computational basis state.
        """
        return np.abs(self.amplitudes) ** 2

    def measure(self, seed: Optional[int] = None) -> tuple[str, "QuantumState"]:
        """Simulate measurement in computational basis.

        Args:
            seed: Random seed for reproducibility.

        Returns:
            Tuple of (measurement result as binary string, collapsed state).
        """
        if seed is not None:
            np.random.seed(seed)

        probs = self.probabilities()
        outcome = np.random.choice(len(probs), p=probs)

        # Collapse to measured state
        collapsed = np.zeros_like(self.amplitudes)
        collapsed[outcome] = 1.0

        binary_result = format(outcome, f"0{self.n_qubits}b")
        return binary_result, QuantumState(amplitudes=collapsed, n_qubits=self.n_qubits)


def apply_gate(
    state: QuantumState,
    gate: NDArray[np.complexfloating],
    targets: list[int],
    validate: bool = True,
) -> QuantumState:
    """Apply a gate to specific qubits in a quantum state.

    Args:
        state: Input quantum state.
        gate: Unitary gate matrix.
        targets: List of target qubit indices (0-indexed).
        validate: Whether to validate gate unitarity.

    Returns:
        New QuantumState after gate application.

    Raises:
        ValueError: If gate dimensions don't match target qubits.

    Example:
        >>> from src.gates import hadamard, cnot
        >>> state = QuantumState.zero_state(2)
        >>> state = apply_gate(state, hadamard(), [0])  # H on qubit 0
        >>> state = apply_gate(state, cnot(), [0, 1])   # CNOT on qubits 0,1
    """
    if validate:
        check_unitarity(gate)

    n_gate_qubits = int(np.log2(gate.shape[0]))
    if len(targets) != n_gate_qubits:
        raise ValueError(
            f"Gate acts on {n_gate_qubits} qubits but {len(targets)} targets given"
        )

    if any(t >= state.n_qubits or t < 0 for t in targets):
        raise ValueError(
            f"Target qubits {targets} out of range for {state.n_qubits}-qubit state"
        )

    # Build full operator using tensor products
    full_gate = _expand_gate(gate, targets, state.n_qubits)
    new_amplitudes = full_gate @ state.amplitudes

    return QuantumState(amplitudes=new_amplitudes, n_qubits=state.n_qubits)


def _expand_gate(
    gate: NDArray[np.complexfloating],
    targets: list[int],
    n_qubits: int,
) -> NDArray[np.complexfloating]:
    """Expand a gate to act on specific qubits in a larger system.

    This implementation handles arbitrary qubit orderings by constructing
    the appropriate permutation matrices.

    Args:
        gate: Gate matrix.
        targets: Target qubit indices.
        n_qubits: Total number of qubits.

    Returns:
        Full 2^n x 2^n gate matrix.
    """
    n_gate_qubits = len(targets)

    # Simple case: gate acts on consecutive qubits starting from 0
    if targets == list(range(n_gate_qubits)):
        remaining = n_qubits - n_gate_qubits
        if remaining == 0:
            return gate
        return tensor(gate, identity(remaining))

    # Simple case: gate acts on consecutive qubits at the end
    if targets == list(range(n_qubits - n_gate_qubits, n_qubits)):
        prefix = n_qubits - n_gate_qubits
        return tensor(identity(prefix), gate)

    # General case: need to permute qubits
    # For now, use matrix construction approach
    dim = 2**n_qubits
    full_gate = np.zeros((dim, dim), dtype=complex)

    for i in range(dim):
        for j in range(dim):
            # Extract target qubit values
            i_targets = _extract_bits(i, targets, n_qubits)
            j_targets = _extract_bits(j, targets, n_qubits)

            # Check if non-target qubits match
            i_others = _extract_other_bits(i, targets, n_qubits)
            j_others = _extract_other_bits(j, targets, n_qubits)

            if i_others == j_others:
                full_gate[i, j] = gate[i_targets, j_targets]

    return full_gate


def _extract_bits(value: int, positions: list[int], n_bits: int) -> int:
    """Extract bits at specified positions and combine into integer."""
    result = 0
    for k, pos in enumerate(positions):
        bit = (value >> (n_bits - 1 - pos)) & 1
        result |= bit << (len(positions) - 1 - k)
    return result


def _extract_other_bits(value: int, exclude: list[int], n_bits: int) -> int:
    """Extract bits at positions NOT in exclude list."""
    other_positions = [i for i in range(n_bits) if i not in exclude]
    return _extract_bits(value, other_positions, n_bits)
