"""Quantum validation utilities for quantum computing research.

This module provides helper functions for validating quantum physical constraints
commonly encountered in quantum computing: unitarity, normalization, density
matrices, CPTP maps, and qubit dimension checks.

Author: Generated by AgentBible
"""

from typing import Optional

import numpy as np
from numpy.typing import NDArray


def check_unitarity(
    matrix: NDArray[np.complexfloating],
    tolerance: float = 1e-10,
    name: Optional[str] = None,
) -> None:
    """Validate that a matrix is unitary (U dagger U = I).

    Args:
        matrix: Square complex matrix to validate.
        tolerance: Numerical tolerance for comparison.
        name: Optional name for error messages.

    Raises:
        ValueError: If matrix is not square or not unitary.

    Example:
        >>> U = np.array([[0, 1], [1, 0]], dtype=complex)  # Pauli X
        >>> check_unitarity(U)  # No error
        >>> check_unitarity(np.array([[1, 1], [0, 1]]))  # Raises ValueError
    """
    label = f" '{name}'" if name else ""

    if matrix.ndim != 2:
        raise ValueError(f"Matrix{label} must be 2D, got shape {matrix.shape}")

    if matrix.shape[0] != matrix.shape[1]:
        raise ValueError(f"Matrix{label} must be square, got shape {matrix.shape}")

    identity = np.eye(matrix.shape[0], dtype=complex)
    product = matrix.conj().T @ matrix

    if not np.allclose(product, identity, atol=tolerance):
        max_error = np.max(np.abs(product - identity))
        raise ValueError(
            f"Matrix{label} is not unitary: max|U dagger U - I| = {max_error:.2e} "
            f"(tolerance: {tolerance:.2e})"
        )


def check_hermitian(
    matrix: NDArray[np.complexfloating],
    tolerance: float = 1e-10,
    name: Optional[str] = None,
) -> None:
    """Validate that a matrix is Hermitian (H = H dagger).

    Args:
        matrix: Square complex matrix to validate.
        tolerance: Numerical tolerance for comparison.
        name: Optional name for error messages.

    Raises:
        ValueError: If matrix is not square or not Hermitian.
    """
    label = f" '{name}'" if name else ""

    if matrix.ndim != 2 or matrix.shape[0] != matrix.shape[1]:
        raise ValueError(f"Matrix{label} must be square, got shape {matrix.shape}")

    if not np.allclose(matrix, matrix.conj().T, atol=tolerance):
        max_error = np.max(np.abs(matrix - matrix.conj().T))
        raise ValueError(
            f"Matrix{label} is not Hermitian: max|H - H dagger| = {max_error:.2e}"
        )


def check_normalized(
    state: NDArray[np.complexfloating],
    tolerance: float = 1e-10,
    name: Optional[str] = None,
) -> None:
    """Validate that a quantum state is normalized (<psi|psi> = 1).

    Args:
        state: Complex state vector.
        tolerance: Numerical tolerance for comparison.
        name: Optional name for error messages.

    Raises:
        ValueError: If state is not normalized.
    """
    label = f" '{name}'" if name else ""
    norm = np.linalg.norm(state)

    if not np.isclose(norm, 1.0, atol=tolerance):
        raise ValueError(
            f"State{label} is not normalized: |<psi|psi>| = {norm:.6f} (expected 1.0)"
        )


def check_density_matrix(
    rho: NDArray[np.complexfloating],
    tolerance: float = 1e-10,
    name: Optional[str] = None,
) -> None:
    """Validate that a matrix is a valid density matrix.

    Checks:
    1. Hermitian: rho = rho dagger
    2. Trace one: Tr(rho) = 1
    3. Positive semidefinite: all eigenvalues >= 0

    Args:
        rho: Square complex matrix to validate.
        tolerance: Numerical tolerance for comparison.
        name: Optional name for error messages.

    Raises:
        ValueError: If any density matrix property is violated.
    """
    label = f" '{name}'" if name else ""

    # Check Hermitian
    check_hermitian(rho, tolerance, name)

    # Check trace = 1
    trace = np.trace(rho)
    if not np.isclose(trace, 1.0, atol=tolerance):
        raise ValueError(f"Density matrix{label} trace = {trace:.6f} (expected 1.0)")

    # Check positive semidefinite
    eigenvalues = np.linalg.eigvalsh(rho)
    min_eigenvalue = np.min(eigenvalues)
    if min_eigenvalue < -tolerance:
        raise ValueError(
            f"Density matrix{label} has negative eigenvalue: {min_eigenvalue:.2e}"
        )


def check_probabilities(
    probs: NDArray[np.floating],
    tolerance: float = 1e-10,
    name: Optional[str] = None,
) -> None:
    """Validate that an array represents valid probabilities.

    Checks:
    1. All values non-negative
    2. Sum equals 1

    Args:
        probs: Array of probability values.
        tolerance: Numerical tolerance for comparison.
        name: Optional name for error messages.

    Raises:
        ValueError: If probabilities are invalid.
    """
    label = f" '{name}'" if name else ""

    if np.any(probs < -tolerance):
        min_val = np.min(probs)
        raise ValueError(f"Probabilities{label} contain negative value: {min_val:.2e}")

    total = np.sum(probs)
    if not np.isclose(total, 1.0, atol=tolerance):
        raise ValueError(f"Probabilities{label} sum to {total:.6f} (expected 1.0)")


def check_cptp(
    kraus_operators: list[NDArray[np.complexfloating]],
    tolerance: float = 1e-10,
    name: Optional[str] = None,
) -> None:
    """Validate that Kraus operators define a valid CPTP map.

    A quantum channel is Completely Positive Trace Preserving (CPTP) if
    the Kraus operators satisfy: sum_k K_k dagger K_k = I

    Args:
        kraus_operators: List of Kraus operator matrices.
        tolerance: Numerical tolerance for comparison.
        name: Optional name for error messages.

    Raises:
        ValueError: If the Kraus operators don't form a valid CPTP map.
    """
    label = f" '{name}'" if name else ""

    if not kraus_operators:
        raise ValueError(f"CPTP map{label} requires at least one Kraus operator")

    dim = kraus_operators[0].shape[0]
    identity = np.eye(dim, dtype=complex)

    # Compute sum_k K_k dagger K_k
    cptp_sum = np.zeros((dim, dim), dtype=complex)
    for k in kraus_operators:
        if k.shape != (dim, dim):
            raise ValueError(
                f"Kraus operator{label} has inconsistent shape: "
                f"expected ({dim}, {dim}), got {k.shape}"
            )
        cptp_sum += k.conj().T @ k

    if not np.allclose(cptp_sum, identity, atol=tolerance):
        max_error = np.max(np.abs(cptp_sum - identity))
        raise ValueError(
            f"Kraus operators{label} do not form CPTP map: "
            f"max|sum(K dagger K) - I| = {max_error:.2e}"
        )


def check_qubit_dimension(
    matrix: NDArray[np.complexfloating],
    n_qubits: Optional[int] = None,
    name: Optional[str] = None,
) -> int:
    """Validate that matrix dimension is a valid qubit space (2^n).

    Args:
        matrix: Matrix to check dimension.
        n_qubits: If provided, verify matrix is exactly 2^n_qubits dimensional.
        name: Optional name for error messages.

    Returns:
        Number of qubits (n where dimension = 2^n).

    Raises:
        ValueError: If dimension is not a power of 2, or doesn't match n_qubits.
    """
    label = f" '{name}'" if name else ""

    if matrix.ndim != 2 or matrix.shape[0] != matrix.shape[1]:
        raise ValueError(f"Matrix{label} must be square, got shape {matrix.shape}")

    dim = matrix.shape[0]

    # Check if dim is a power of 2
    if dim == 0 or (dim & (dim - 1)) != 0:
        raise ValueError(
            f"Matrix{label} dimension {dim} is not a power of 2 (not a valid qubit space)"
        )

    n = int(np.log2(dim))

    if n_qubits is not None and n != n_qubits:
        raise ValueError(
            f"Matrix{label} has {n} qubits (dim={dim}), expected {n_qubits} qubits "
            f"(dim={2**n_qubits})"
        )

    return n


def check_partial_trace_valid(
    rho: NDArray[np.complexfloating],
    subsystem_dims: list[int],
    name: Optional[str] = None,
) -> None:
    """Validate that density matrix dimensions match subsystem specification.

    Args:
        rho: Density matrix of composite system.
        subsystem_dims: List of dimensions for each subsystem.
        name: Optional name for error messages.

    Raises:
        ValueError: If dimensions don't match.
    """
    label = f" '{name}'" if name else ""

    expected_dim = int(np.prod(subsystem_dims))
    if rho.shape[0] != expected_dim:
        raise ValueError(
            f"Density matrix{label} dimension {rho.shape[0]} does not match "
            f"subsystem dimensions {subsystem_dims} (expected {expected_dim})"
        )
