"""
Quantum state fidelity measures.

Implements Uhlmann-Jozsa fidelity and related metrics.
All functions require an explicit target state -- no implicit defaults.

References:
    [1] Uhlmann, "The transition probability in the state space" (1976)
    [2] Jozsa, "Fidelity for mixed quantum states" (1994)
    [3] Nielsen & Chuang, "Quantum Computation and Quantum Information" Ch. 9
"""

import numpy as np
from scipy import linalg


def uhlmann_fidelity(rho: np.ndarray, sigma: np.ndarray) -> float:
    """
    Uhlmann-Jozsa fidelity between two density matrices.

    F(rho, sigma) = (Tr sqrt(sqrt(rho) * sigma * sqrt(rho)))^2

    Properties:
        - F(rho, sigma) in [0, 1]
        - F(rho, sigma) = 1 iff rho = sigma
        - F(rho, sigma) = F(sigma, rho) (symmetric)
        - For pure states: F(|a><a|, |b><b|) = |<a|b>|^2

    Args:
        rho: First density matrix (d x d, Hermitian, PSD, trace 1)
        sigma: Second density matrix (d x d, Hermitian, PSD, trace 1)

    Returns:
        Fidelity in [0, 1]
    """
    # Compute sqrt(rho) via eigendecomposition
    eigvals_rho, eigvecs_rho = linalg.eigh(rho)
    eigvals_rho = np.maximum(eigvals_rho, 0.0)
    sqrt_rho = eigvecs_rho @ np.diag(np.sqrt(eigvals_rho)) @ eigvecs_rho.conj().T

    # Compute sqrt(rho) * sigma * sqrt(rho)
    product = sqrt_rho @ sigma @ sqrt_rho

    # Eigenvalues of the product (should be real and non-negative)
    eigvals_product = linalg.eigvalsh(product)
    eigvals_product = np.maximum(eigvals_product, 0.0)

    # F = (sum of sqrt of eigenvalues)^2
    fidelity = np.sum(np.sqrt(eigvals_product)) ** 2

    # Clamp to [0, 1] for numerical safety
    return float(np.clip(np.real(fidelity), 0.0, 1.0))


def infidelity(rho: np.ndarray, sigma: np.ndarray) -> float:
    """
    Infidelity (error measure): 1 - F(rho, sigma).

    This is the natural error metric: 0 means perfect, 1 means orthogonal.
    """
    return 1.0 - uhlmann_fidelity(rho, sigma)


def bures_distance(rho: np.ndarray, sigma: np.ndarray) -> float:
    """
    Bures distance between two density matrices.

    d_B(rho, sigma) = sqrt(2 * (1 - sqrt(F(rho, sigma))))

    This is a proper metric on the space of density matrices.
    """
    f = uhlmann_fidelity(rho, sigma)
    return float(np.sqrt(2.0 * (1.0 - np.sqrt(max(f, 0.0)))))


def pure_state_fidelity(psi: np.ndarray, rho: np.ndarray) -> float:
    """
    Fidelity between a pure state |psi> and a density matrix rho.

    F(|psi><psi|, rho) = <psi|rho|psi>

    This is much cheaper to compute than the general formula.

    Args:
        psi: State vector (d,), normalized
        rho: Density matrix (d x d)

    Returns:
        Fidelity in [0, 1]
    """
    return float(np.real(psi.conj() @ rho @ psi))


def von_neumann_entropy(rho: np.ndarray) -> float:
    """
    Von Neumann entropy: S(rho) = -Tr(rho log rho).

    For a d-dimensional system:
        - S = 0 for pure states
        - S = log(d) for maximally mixed state
    """
    eigvals = linalg.eigvalsh(rho)
    eigvals = eigvals[eigvals > 1e-15]  # Remove numerical zeros
    return float(-np.sum(eigvals * np.log(eigvals)))


def purity(rho: np.ndarray) -> float:
    """
    Purity: Tr(rho^2).

    - Purity = 1 for pure states
    - Purity = 1/d for maximally mixed state of dimension d
    """
    return float(np.real(np.trace(rho @ rho)))


def linear_entropy(rho: np.ndarray) -> float:
    """Linear entropy: S_L = 1 - Tr(rho^2)."""
    return 1.0 - purity(rho)


def validate_density_matrix(rho: np.ndarray, tol: float = 1e-8) -> dict:
    """
    Check that a matrix is a valid density matrix.

    Returns dict with 'valid' bool and individual check results.
    """
    d = rho.shape[0]
    checks = {}

    # Square matrix
    checks['square'] = rho.shape[0] == rho.shape[1]

    # Hermitian: rho = rho^dagger
    checks['hermitian'] = np.allclose(rho, rho.conj().T, atol=tol)

    # Trace = 1
    checks['trace_one'] = abs(np.trace(rho) - 1.0) < tol

    # Positive semidefinite: all eigenvalues >= 0
    eigvals = linalg.eigvalsh(rho)
    checks['positive_semidefinite'] = np.all(eigvals >= -tol)

    checks['valid'] = all(checks.values())
    return checks