"""
Hausdorff (fractal) dimension of quantum state trajectories.

The fractal dimension serves as a health metric for the navigation system:
    - d_H ~ 1.0: Smooth geodesic trajectory (error-free unitary evolution)
    - d_H ~ 2.0: Brownian-like trajectory (fully decoherent)
    - 1.0 < d_H < 2.0: Partially coherent (correctable by QEC)

Computed using the box-counting method on the trajectory in density matrix
space, with the Hilbert-Schmidt distance as the metric.

References:
    [1] Feynman & Hibbs, "Quantum Mechanics and Path Integrals" (1965)
        - Quantum paths have Hausdorff dimension d_H = 2
    [2] Grassberger & Procaccia, "Measuring the strangeness of strange attractors"
        Physica D 9:189-208 (1983)
"""

import numpy as np
from typing import List, Optional
from enum import Enum

from .fidelity import bures_distance


class TrajectoryHealth(Enum):
    """Health assessment based on fractal dimension."""
    GEODESIC = "geodesic"              # d_H < 1.1: near-perfect evolution
    COHERENT = "coherent"              # d_H < 1.3: minor decoherence
    PARTIALLY_COHERENT = "partially_coherent"  # d_H < 1.7
    DECOHERENT = "decoherent"          # d_H >= 1.7: heavy decoherence


def hilbert_schmidt_distance(rho1: np.ndarray, rho2: np.ndarray) -> float:
    """
    Hilbert-Schmidt distance between two density matrices.

    d_HS(rho1, rho2) = sqrt(Tr((rho1 - rho2)^dag (rho1 - rho2)))
    """
    diff = rho1 - rho2
    return float(np.sqrt(np.real(np.trace(diff.conj().T @ diff))))


def box_counting_dimension(
    trajectory: List[np.ndarray],
    n_scales: int = 15,
    use_bures: bool = False,
) -> float:
    """
    Estimate the box-counting (Minkowski) dimension of a trajectory.

    d_H = lim_{eps->0} log(N(eps)) / log(1/eps)

    where N(eps) is the number of eps-balls needed to cover the trajectory.

    In practice, we compute N(eps) for several values of eps and fit a line
    to the log-log plot.

    Args:
        trajectory: List of density matrices (the points on the trajectory)
        n_scales: Number of epsilon values to test
        use_bures: If True, use Bures distance; otherwise Hilbert-Schmidt

    Returns:
        Estimated fractal dimension
    """
    if len(trajectory) < 3:
        return 1.0  # Too few points

    # Compute pairwise distances between consecutive points
    n = len(trajectory)
    distances = np.zeros(n - 1)
    for i in range(n - 1):
        if use_bures:
            distances[i] = bures_distance(trajectory[i], trajectory[i + 1])
        else:
            distances[i] = hilbert_schmidt_distance(trajectory[i], trajectory[i + 1])

    if np.max(distances) < 1e-15:
        return 1.0  # Stationary trajectory

    # Generate epsilon values spanning the distance range
    d_min = np.min(distances[distances > 1e-15]) if np.any(distances > 1e-15) else 1e-10
    d_max = np.max(distances)
    # Use total arc length as upper bound
    total_arc = np.sum(distances)

    eps_values = np.logspace(
        np.log10(d_min * 0.5),
        np.log10(total_arc * 0.5),
        n_scales,
    )

    log_eps = []
    log_N = []

    for eps in eps_values:
        # Count boxes: walk along trajectory, start new box when cumulative
        # distance exceeds eps
        n_boxes = 1
        cumulative = 0.0
        for d in distances:
            cumulative += d
            if cumulative >= eps:
                n_boxes += 1
                cumulative = 0.0

        if n_boxes > 1:
            log_eps.append(np.log(1.0 / eps))
            log_N.append(np.log(n_boxes))

    if len(log_eps) < 3:
        return 1.0  # Not enough data points for regression

    # Linear regression: log(N) = d_H * log(1/eps) + const
    log_eps = np.array(log_eps)
    log_N = np.array(log_N)
    coeffs = np.polyfit(log_eps, log_N, 1)
    d_H = coeffs[0]

    # Clamp to physically meaningful range [0.5, 3.0]
    return float(np.clip(d_H, 0.5, 3.0))


def correlation_dimension(
    trajectory: List[np.ndarray],
    max_radius: Optional[float] = None,
    n_radii: int = 15,
) -> float:
    """
    Grassberger-Procaccia correlation dimension.

    D2 = lim_{r->0} log(C(r)) / log(r)

    where C(r) = (2 / (N*(N-1))) * #{(i,j) : d(x_i, x_j) < r}

    This is generally more robust than box-counting for small datasets.

    Args:
        trajectory: List of density matrices
        max_radius: Maximum radius to consider (auto-computed if None)
        n_radii: Number of radius values

    Returns:
        Estimated correlation dimension
    """
    n = len(trajectory)
    if n < 10:
        return 1.0

    # Subsample for efficiency if trajectory is long
    if n > 200:
        indices = np.linspace(0, n - 1, 200, dtype=int)
        trajectory = [trajectory[i] for i in indices]
        n = len(trajectory)

    # Compute all pairwise distances
    dist_matrix = np.zeros((n, n))
    for i in range(n):
        for j in range(i + 1, n):
            d = hilbert_schmidt_distance(trajectory[i], trajectory[j])
            dist_matrix[i, j] = d
            dist_matrix[j, i] = d

    # Remove zeros (self-distances)
    all_dists = dist_matrix[np.triu_indices(n, k=1)]
    all_dists = all_dists[all_dists > 1e-15]

    if len(all_dists) < 10:
        return 1.0

    if max_radius is None:
        max_radius = np.percentile(all_dists, 90)
    min_radius = np.percentile(all_dists, 5)

    if max_radius <= min_radius:
        return 1.0

    radii = np.logspace(np.log10(min_radius), np.log10(max_radius), n_radii)

    log_r = []
    log_C = []

    n_pairs = n * (n - 1) / 2
    for r in radii:
        count = np.sum(all_dists < r)
        C_r = count / n_pairs
        if C_r > 0:
            log_r.append(np.log(r))
            log_C.append(np.log(C_r))

    if len(log_r) < 3:
        return 1.0

    coeffs = np.polyfit(log_r, log_C, 1)
    D2 = coeffs[0]

    return float(np.clip(D2, 0.5, 3.0))


def assess_health(d_H: float) -> TrajectoryHealth:
    """Classify trajectory health based on fractal dimension."""
    if d_H < 1.1:
        return TrajectoryHealth.GEODESIC
    elif d_H < 1.3:
        return TrajectoryHealth.COHERENT
    elif d_H < 1.7:
        return TrajectoryHealth.PARTIALLY_COHERENT
    else:
        return TrajectoryHealth.DECOHERENT