"""
Multi-clock trajectory divergence analyzer.

Evolves N offset copies of the quantum state ("clocks") started at
different times to empirically measure trajectory properties that
complement the analytical Frenet-Serret approach.

The original Fractal Time Navigation Compass used 5 offset clocks at
t = -2, -1, 0, +1, +2 to sample "adjacent timelines" and track their
divergence. This module implements that concept rigorously.

Capabilities:
  1. Empirical curvature from finite-difference divergence (model-free)
  2. Divergence tensor for directional drift sensitivity (per 11D dimension)
  3. Non-Markovianity detection via BLP measure (reconvergence = memory)
  4. Coherence width (maximum offset before fidelity drops)
  5. Phase dispersion across the clock ensemble

References:
    Breuer, Laine, Piilo, "Measure for the degree of non-Markovian
    behavior of quantum processes in open systems", PRL 103, 210401 (2009)
"""

import numpy as np
from dataclasses import dataclass, field
from typing import List, Optional, Dict, Tuple, Any
from scipy.linalg import expm

from .fidelity import uhlmann_fidelity, validate_density_matrix
from .lindblad import (
    LindbladEvolver, build_lindblad_operators,
    compute_decoherence_rates, unitary_evolve, normalize_density_matrix,
)
from .frenet_serret import QuantumFrenetSerret
from .hausdorff import hilbert_schmidt_distance


# ---------------------------------------------------------------------------
# Utility functions
# ---------------------------------------------------------------------------

def fubini_study_distance(rho1: np.ndarray, rho2: np.ndarray) -> float:
    """
    Fubini-Study distance on the space of density matrices.

    d_FS = arccos(sqrt(F(rho1, rho2)))

    This is the natural geodesic distance on projective Hilbert space.
    Returns a value in [0, pi/2].
    """
    F = uhlmann_fidelity(rho1, rho2)
    F = np.clip(F, 0.0, 1.0)
    return np.arccos(np.sqrt(F))


def compute_pairwise_distances(
    states: List[np.ndarray],
    metric: str = 'fubini_study',
) -> np.ndarray:
    """
    Compute distance matrix between all states.

    Args:
        states: List of N density matrices.
        metric: 'fubini_study', 'hilbert_schmidt', or 'bures'.

    Returns:
        (N, N) symmetric distance matrix with zeros on diagonal.
    """
    N = len(states)
    D = np.zeros((N, N))

    if metric == 'fubini_study':
        dist_fn = fubini_study_distance
    elif metric == 'hilbert_schmidt':
        dist_fn = hilbert_schmidt_distance
    elif metric == 'bures':
        def dist_fn(a, b):
            return np.sqrt(2 * (1 - np.sqrt(uhlmann_fidelity(a, b))))
    else:
        raise ValueError(f"Unknown metric: {metric}")

    for i in range(N):
        for j in range(i + 1, N):
            d = dist_fn(states[i], states[j])
            D[i, j] = d
            D[j, i] = d

    return D


def fit_divergence_parabola(
    offsets: np.ndarray,
    distances: np.ndarray,
) -> Tuple[float, float, float]:
    """
    Fit d(offset) = a + b*offset + c*offset^2.

    The quadratic coefficient c is proportional to the curvature:
    large c means the trajectory is highly sensitive to small perturbations.

    Args:
        offsets: Array of time offsets (e.g., [-2, -1, 0, 1, 2]).
        distances: Array of distances from the central clock.

    Returns:
        (a, b, c) polynomial coefficients. c ~ empirical curvature.
    """
    if len(offsets) < 3:
        return (0.0, 0.0, 0.0)

    coeffs = np.polyfit(offsets, distances, 2)
    # polyfit returns [c, b, a] for c*x^2 + b*x + a
    return (float(coeffs[2]), float(coeffs[1]), float(coeffs[0]))


def map_divergence_to_dimensions(
    divergence_rates: np.ndarray,
    scaling_factors: np.ndarray,
) -> np.ndarray:
    """
    Project scalar divergence rates onto 11D navigation dimensions.

    Each dimension's sensitivity = divergence_rate * scaling_factor.
    Spacetime dimensions (large scaling) get high sensitivity.
    Compact dimensions (suppressed scaling) get low sensitivity.

    Args:
        divergence_rates: Array of divergence rates per clock pair.
        scaling_factors: Array of 11 dimensional scaling factors.

    Returns:
        (11,) array of per-dimension divergence sensitivity.
    """
    mean_divergence = np.mean(divergence_rates) if len(divergence_rates) > 0 else 0.0
    sensitivity = mean_divergence * scaling_factors
    max_val = np.max(np.abs(sensitivity))
    if max_val > 0:
        sensitivity = sensitivity / max_val
    return sensitivity


# ---------------------------------------------------------------------------
# Dataclasses
# ---------------------------------------------------------------------------

@dataclass
class MultiClockConfig:
    """Configuration for multi-clock analysis."""
    n_clocks: int = 5
    offset_multiplier: float = 1.0  # Scale offsets by this * dt
    metric: str = 'fubini_study'    # Distance metric
    coherence_threshold: float = 0.9  # Fidelity threshold for coherence width
    blp_tolerance: float = 1e-6     # Threshold for positive dF/dt
    auto_scale: bool = False         # Auto-compute offset_multiplier from Hamiltonian


@dataclass
class ClockState:
    """State of a single clock in the ensemble."""
    clock_id: int
    time_offset: float
    density_matrices: List[np.ndarray]
    times: np.ndarray
    fidelity_to_central: np.ndarray


@dataclass
class DivergenceTensor:
    """Pairwise divergence analysis between all clocks."""
    # (n_clocks, n_clocks, n_steps+1) distance at each timestep
    distance_history: np.ndarray
    # (n_clocks, n_clocks) mean rate of distance growth
    divergence_rates: np.ndarray
    # (11,) per-dimension sensitivity
    dimension_sensitivity: np.ndarray
    # Maximum offset where coherence is maintained
    max_stable_offset: float


@dataclass
class NonMarkovianityMeasure:
    """BLP non-Markovianity quantification."""
    blp_measure: float               # Integral of positive dF/dt
    reconvergence_times: List[float]  # Times where clocks reconverge
    is_markovian: bool                # True if monotonic fidelity decay
    fidelity_flow: np.ndarray         # dF/dt at each timestep


@dataclass
class MultiClockResult:
    """Complete result of a multi-clock analysis."""
    # Clock data
    clocks: List[ClockState]
    central_index: int               # Index of the t=0 clock
    offsets: np.ndarray              # Time offsets for all clocks
    times: np.ndarray                # Common time grid

    # Empirical curvature from finite differences
    empirical_curvatures: np.ndarray   # (n_steps+1,)
    analytical_curvatures: np.ndarray  # (n_steps+1,) from Frenet-Serret
    curvature_correlation: float       # Pearson r between them

    # Divergence
    divergence_tensor: DivergenceTensor
    coherence_width: float             # Max offset keeping F > threshold

    # Phase dispersion
    phase_dispersions: np.ndarray      # (n_steps+1,) variance of phases

    # Non-Markovianity
    non_markovianity: NonMarkovianityMeasure


@dataclass
class ClockDisplacementResult:
    """Per-dimension displacement estimated from 5-clock divergence.

    The 5 offset clocks sample the quantum state at slightly different
    evolution times. By comparing KK coordinates (x_d = Tr(P_d @ rho) * R_d)
    across clocks and between initial and final states, we get:
    - displacement: how far the state has drifted in each KK dimension
    - velocity: dx/dt estimated via finite differences across clocks
    """
    clock_coordinates: np.ndarray       # (n_clocks, n_dims) final KK coordinates
    central_coordinates: np.ndarray     # (n_dims,) central clock final coords
    initial_coordinates: np.ndarray     # (n_dims,) coords before evolution
    displacements: np.ndarray           # (n_dims,) = central_final - initial
    clock_velocity_estimates: np.ndarray  # (n_dims,) dx/dt via finite differences
    clock_fidelities: np.ndarray        # (n_clocks,) fidelity of each clock to central
    clock_result: MultiClockResult      # full underlying multi-clock result


# ---------------------------------------------------------------------------
# Return-to-origin functions
# ---------------------------------------------------------------------------

def estimate_return_displacement(
    displacement_result: ClockDisplacementResult,
) -> np.ndarray:
    """Return the negative displacement vector (what correction needs to undo).

    Args:
        displacement_result: Result from evolve_clocks_with_displacement.

    Returns:
        (n_dims,) array: the displacement to apply to return to origin.
    """
    return -displacement_result.displacements


def compute_return_unitary(
    return_displacement: np.ndarray,
    kk_builder,
    gain: float = 1.0,
) -> np.ndarray:
    """Compute the unitary that translates the state back toward origin.

    Uses KK position operators X_d as generators. In the momentum
    eigenbasis (our KK basis), X_d is off-diagonal and satisfies the
    canonical commutation relation [X_d, P_d] ~ i. Therefore:

        U = exp(-i * sum_d c_d * X_d)

    shifts the momentum expectation value: <P_d> -> <P_d> + c_d.

    We set c_d = gain * delta_d / R_d to undo the measured displacement,
    where delta_d is the return displacement (negative of measured drift)
    and R_d is the compactification radius.

    Args:
        return_displacement: (n_dims,) displacement to apply (from
            estimate_return_displacement).
        kk_builder: KKTowerBuilder instance with position operators and radii.
        gain: Scaling factor for the correction strength. gain=1.0 attempts
            full correction; smaller values are more conservative.

    Returns:
        Unitary matrix of shape (dim_q, dim_q).
    """
    X_ops = kk_builder.build_position_operators()
    radii = kk_builder.radii
    n_dims = min(len(return_displacement), len(X_ops))
    dim = X_ops[0].shape[0]

    # Build the generator: G = sum_d c_d * X_d
    # where c_d = gain * delta_d / R_d
    generator = np.zeros((dim, dim), dtype=complex)
    for d in range(n_dims):
        if abs(radii[d]) > 1e-30:
            c_d = gain * return_displacement[d] / radii[d]
            generator += c_d * X_ops[d]

    # U = exp(-i * G)
    U = expm(-1j * generator)
    return U


# ---------------------------------------------------------------------------
# Main analyzer
# ---------------------------------------------------------------------------

class MultiClockAnalyzer:
    """
    Multi-clock trajectory divergence analyzer.

    Evolves N offset copies of the quantum state to empirically measure
    curvature, directional sensitivity, non-Markovianity, and coherence
    width. Complements the analytical Frenet-Serret approach.

    The default configuration uses 5 clocks at offsets [-2, -1, 0, +1, +2]
    times the timestep dt, matching the original FCE compass design.
    """

    def __init__(
        self,
        hamiltonian: np.ndarray,
        lindblad_ops: Optional[List[np.ndarray]] = None,
        decoherence_rates: Optional[List[float]] = None,
        system_params: Optional[Dict[str, float]] = None,
        config: Optional[MultiClockConfig] = None,
    ):
        self.H = hamiltonian
        self.dim = hamiltonian.shape[0]
        self.config = config or MultiClockConfig()

        # Build Lindblad operators and rates (same pattern as engine)
        if lindblad_ops is not None:
            self.lindblad_ops = lindblad_ops
        else:
            self.lindblad_ops, _ = build_lindblad_operators(self.dim)

        if decoherence_rates is not None:
            self.rates = decoherence_rates
        elif system_params is not None:
            self.rates = compute_decoherence_rates(self.dim, system_params)
        else:
            default_params = {'T1': 1e-3, 'T2': 5e-4, 'gate_fidelity': 0.999}
            self.rates = compute_decoherence_rates(self.dim, default_params)

        self.evolver = LindbladEvolver(self.H, self.lindblad_ops, self.rates)
        self.fs = QuantumFrenetSerret()

    @staticmethod
    def compute_offset_multiplier(
        hamiltonian: np.ndarray,
        dt: float,
        target_rotation: float = 0.1,
    ) -> float:
        """Auto-scale offset to produce meaningful phase separation between clocks.

        The offset produces a unitary evolution exp(-i H t_offset). For the
        clocks to meaningfully diverge, the phase rotation between adjacent
        clocks needs to be detectable (~0.1 radians is a good target).

        Adjacent clock offset = 1 * dt * multiplier, so phase rotation =
        E_range * dt * multiplier. We solve for:

            multiplier = target_rotation / (dt * E_range)

        where E_range = max(eigvals(H)) - min(eigvals(H)).

        Args:
            hamiltonian: System Hamiltonian.
            dt: Time step.
            target_rotation: Desired phase rotation in radians between
                adjacent clocks. Default 0.1 radians (~6 degrees).

        Returns:
            Offset multiplier value.
        """
        eigenvalues = np.linalg.eigvalsh(hamiltonian)
        E_range = float(eigenvalues[-1] - eigenvalues[0])
        if E_range < 1e-15 or dt < 1e-30:
            return 1.0
        return target_rotation / (dt * E_range)

    def _generate_offsets(self) -> np.ndarray:
        """Generate integer clock offsets (e.g., [-2, -1, 0, 1, 2])."""
        n = self.config.n_clocks
        half = n // 2
        return np.arange(-half, half + 1, dtype=float)

    def _generate_offset_states(
        self,
        rho_init: np.ndarray,
        dt: float,
    ) -> List[Tuple[float, np.ndarray]]:
        """
        Generate initial states for each clock.

        Clock k with offset t_k starts from the state obtained by
        evolving rho_init unitarily by t_k * dt * offset_multiplier.
        This samples the ideal trajectory at nearby points -- "adjacent timelines."
        """
        offsets = self._generate_offsets()
        states = []
        for offset in offsets:
            t_offset = offset * dt * self.config.offset_multiplier
            if abs(t_offset) < 1e-20:
                rho_offset = rho_init.copy()
            else:
                rho_offset = unitary_evolve(rho_init, self.H, t_offset)
            states.append((offset, rho_offset))
        return states

    def evolve_clocks(
        self,
        rho_init: np.ndarray,
        dt: float,
        n_steps: int,
    ) -> MultiClockResult:
        """
        Run multi-clock evolution and analysis.

        Each clock starts at a different time offset and evolves under
        the same Lindblad decoherence. By tracking how they diverge,
        we measure empirical curvature, directional sensitivity,
        non-Markovianity, and coherence width.

        Args:
            rho_init: Initial density matrix (all clocks derive from this).
            dt: Timestep.
            n_steps: Number of evolution steps.

        Returns:
            MultiClockResult with full analysis.
        """
        checks = validate_density_matrix(rho_init)
        if not checks['valid']:
            raise ValueError(f"Invalid initial density matrix: {checks}")

        # 1. Generate offset initial states
        offset_states = self._generate_offset_states(rho_init, dt)
        offsets = np.array([o for o, _ in offset_states])
        central_idx = len(offsets) // 2  # The t=0 clock
        n_clocks = len(offsets)

        # 2. Initialize storage
        clock_rhos = [rho.copy() for _, rho in offset_states]
        clock_histories = [[rho.copy()] for rho in clock_rhos]
        times = [0.0]

        # Pairwise distances at each step
        dist_history = np.zeros((n_clocks, n_clocks, n_steps + 1))
        dist_history[:, :, 0] = compute_pairwise_distances(
            clock_rhos, self.config.metric
        )

        # Fidelity of each clock to central
        fid_to_central = np.zeros((n_clocks, n_steps + 1))
        for k in range(n_clocks):
            fid_to_central[k, 0] = uhlmann_fidelity(
                clock_rhos[k], clock_rhos[central_idx]
            )

        # Analytical curvature from central clock
        analytical_kappas = [self.fs.compute_curvature(clock_rhos[central_idx], self.H)]

        # Phase tracking (eigenvalues of each clock's state)
        phase_dispersions = [self._compute_phase_dispersion_step(clock_rhos)]

        # 3. Evolve all clocks
        for step in range(n_steps):
            t = (step + 1) * dt

            for k in range(n_clocks):
                clock_rhos[k] = self.evolver.evolve_single_step(clock_rhos[k], dt)
                clock_histories[k].append(clock_rhos[k].copy())

            # Pairwise distances
            dist_history[:, :, step + 1] = compute_pairwise_distances(
                clock_rhos, self.config.metric
            )

            # Fidelity to central
            for k in range(n_clocks):
                fid_to_central[k, step + 1] = uhlmann_fidelity(
                    clock_rhos[k], clock_rhos[central_idx]
                )

            # Analytical curvature
            analytical_kappas.append(
                self.fs.compute_curvature(clock_rhos[central_idx], self.H)
            )

            # Phase dispersion
            phase_dispersions.append(
                self._compute_phase_dispersion_step(clock_rhos)
            )

            times.append(t)

        times = np.array(times)
        analytical_kappas = np.array(analytical_kappas)
        phase_dispersions = np.array(phase_dispersions)

        # 4. Build ClockState objects
        clocks = []
        for k in range(n_clocks):
            clocks.append(ClockState(
                clock_id=k,
                time_offset=float(offsets[k]),
                density_matrices=clock_histories[k],
                times=times,
                fidelity_to_central=fid_to_central[k],
            ))

        # 5. Compute empirical curvature from finite differences
        empirical_kappas = self._compute_empirical_curvature(
            offsets, dist_history, central_idx
        )

        # 6. Compute curvature correlation
        if len(empirical_kappas) > 2 and np.std(empirical_kappas) > 1e-15:
            corr = np.corrcoef(empirical_kappas, analytical_kappas)[0, 1]
            if np.isnan(corr):
                corr = 0.0
        else:
            corr = 0.0

        # 7. Compute divergence tensor
        divergence_tensor = self._compute_divergence_tensor(
            offsets, dist_history, fid_to_central, central_idx, times
        )

        # 8. Compute coherence width
        coherence_width = self._compute_coherence_width(
            offsets, fid_to_central, central_idx
        )

        # 9. Compute non-Markovianity (BLP)
        non_markovianity = self._compute_blp(
            fid_to_central, central_idx, times, dist_history
        )

        return MultiClockResult(
            clocks=clocks,
            central_index=central_idx,
            offsets=offsets,
            times=times,
            empirical_curvatures=empirical_kappas,
            analytical_curvatures=analytical_kappas,
            curvature_correlation=corr,
            divergence_tensor=divergence_tensor,
            coherence_width=coherence_width,
            phase_dispersions=phase_dispersions,
            non_markovianity=non_markovianity,
        )

    def evolve_clocks_with_displacement(
        self,
        rho_init: np.ndarray,
        dt: float,
        n_steps: int,
        kk_builder,
        config: Optional[MultiClockConfig] = None,
        target_rotation: float = 0.1,
    ) -> ClockDisplacementResult:
        """Evolve 5 offset clocks and track per-dimension KK displacement.

        This is the core method for the 5-clock navigation theory. It:
        1. Auto-scales the offset multiplier so clocks produce visible divergence
        2. Records initial KK coordinates x_d(0) = Tr(P_d @ rho) * R_d
        3. Evolves all 5 clocks under Lindblad decoherence
        4. Extracts final KK coordinates for each clock
        5. Computes displacement (drift from origin) and velocity estimates

        Args:
            rho_init: Initial density matrix.
            dt: Time step.
            n_steps: Number of evolution steps.
            kk_builder: KKTowerBuilder instance providing momentum operators
                and compactification radii.
            config: Optional MultiClockConfig override. If None, uses
                self.config with auto-scaled offset_multiplier.
            target_rotation: Target phase rotation in radians between
                adjacent clocks (for auto-scaling).

        Returns:
            ClockDisplacementResult with per-dimension displacement and velocity.
        """
        P_ops = kk_builder.build_momentum_operators()
        radii = kk_builder.radii
        n_dims = len(P_ops)

        # Auto-scale offset multiplier
        multiplier = self.compute_offset_multiplier(
            self.H, dt, target_rotation
        )

        # Override config with auto-scaled multiplier
        cfg = config or MultiClockConfig(
            n_clocks=self.config.n_clocks,
            offset_multiplier=multiplier,
            metric=self.config.metric,
            coherence_threshold=self.config.coherence_threshold,
            blp_tolerance=self.config.blp_tolerance,
            auto_scale=True,
        )
        if cfg.auto_scale:
            cfg = MultiClockConfig(
                n_clocks=cfg.n_clocks,
                offset_multiplier=multiplier,
                metric=cfg.metric,
                coherence_threshold=cfg.coherence_threshold,
                blp_tolerance=cfg.blp_tolerance,
                auto_scale=True,
            )

        # Save original config and swap
        original_config = self.config
        self.config = cfg

        # 1. Record initial KK coordinates
        initial_coords = np.zeros(n_dims)
        for d in range(n_dims):
            initial_coords[d] = np.real(np.trace(P_ops[d] @ rho_init)) * radii[d]

        # 2. Evolve all 5 clocks
        clock_result = self.evolve_clocks(rho_init, dt, n_steps)

        # Restore config
        self.config = original_config

        # 3. Extract final KK coordinates for each clock
        n_clocks = len(clock_result.clocks)
        clock_coords = np.zeros((n_clocks, n_dims))
        for k, clock in enumerate(clock_result.clocks):
            rho_final = clock.density_matrices[-1]
            for d in range(n_dims):
                clock_coords[k, d] = (
                    np.real(np.trace(P_ops[d] @ rho_final)) * radii[d]
                )

        # 4. Central clock coordinates
        central_idx = clock_result.central_index
        central_coords = clock_coords[central_idx]

        # 5. Displacement = central_final - initial
        displacements = central_coords - initial_coords

        # 6. Velocity estimates via central finite differences across clocks
        # v_d ≈ (x_d[clock+1] - x_d[clock-1]) / (2 * offset_dt)
        offset_dt = dt * cfg.offset_multiplier
        velocities = np.zeros(n_dims)
        for d in range(n_dims):
            if n_clocks >= 3:
                # Use clock pairs symmetric around central
                # Central is at index central_idx
                # Use clocks at central-1 and central+1
                if central_idx > 0 and central_idx < n_clocks - 1:
                    dx = clock_coords[central_idx + 1, d] - clock_coords[central_idx - 1, d]
                    velocities[d] = dx / (2.0 * offset_dt)

        # 7. Clock fidelities (final step)
        clock_fidelities = np.zeros(n_clocks)
        rho_central = clock_result.clocks[central_idx].density_matrices[-1]
        for k in range(n_clocks):
            rho_k = clock_result.clocks[k].density_matrices[-1]
            clock_fidelities[k] = uhlmann_fidelity(rho_k, rho_central)

        return ClockDisplacementResult(
            clock_coordinates=clock_coords,
            central_coordinates=central_coords,
            initial_coordinates=initial_coords,
            displacements=displacements,
            clock_velocity_estimates=velocities,
            clock_fidelities=clock_fidelities,
            clock_result=clock_result,
        )

    def _compute_empirical_curvature(
        self,
        offsets: np.ndarray,
        dist_history: np.ndarray,
        central_idx: int,
    ) -> np.ndarray:
        """
        Estimate curvature from parabolic fit of distance vs offset.

        At each timestep, the distance from each clock to the central
        clock is fit to d(offset) = a + b*offset + c*offset^2.
        The coefficient c captures the curvature: how sensitively
        the trajectory responds to small perturbations in initial time.
        """
        n_steps_plus_one = dist_history.shape[2]
        kappas = np.zeros(n_steps_plus_one)

        # Distances from each clock to central at each time
        for t_idx in range(n_steps_plus_one):
            distances_to_central = dist_history[:, central_idx, t_idx]
            _, _, c = fit_divergence_parabola(offsets, distances_to_central)
            kappas[t_idx] = abs(c)  # Curvature is non-negative

        return kappas

    def _compute_divergence_tensor(
        self,
        offsets: np.ndarray,
        dist_history: np.ndarray,
        fid_to_central: np.ndarray,
        central_idx: int,
        times: np.ndarray,
        spectral_projectors: Optional[List[np.ndarray]] = None,
        clock_rhos_central: Optional[List[np.ndarray]] = None,
    ) -> DivergenceTensor:
        """
        Compute pairwise divergence rates and per-dimension sensitivity.

        When spectral_projectors is provided, computes per-dimension
        sensitivity by projecting clock divergence (rho_clock - rho_central)
        onto each spectral subspace and measuring the HS norm. This gives
        genuine per-dimension resolution from the multi-clock analysis.

        Falls back to scalar multiplication by DimensionalScaling factors
        when projectors are not provided (backward compatible).
        """
        n_clocks = len(offsets)
        n_times = len(times)

        # Divergence rates: slope of distance vs time for each pair
        rates = np.zeros((n_clocks, n_clocks))
        if n_times >= 2:
            for i in range(n_clocks):
                for j in range(i + 1, n_clocks):
                    dists = dist_history[i, j, :]
                    # Linear fit: d(t) = a + rate*t
                    if times[-1] > 0:
                        rate = (dists[-1] - dists[0]) / times[-1]
                    else:
                        rate = 0.0
                    rates[i, j] = rate
                    rates[j, i] = rate

        # Per-dimension sensitivity
        if spectral_projectors is not None and clock_rhos_central is not None:
            # Project divergence onto spectral subspaces for genuine
            # per-dimension resolution
            n_bands = len(spectral_projectors)
            dim_sensitivity = np.zeros(max(n_bands, 11))

            # Use the outermost clock's divergence from central
            # at the final timestep
            if len(clock_rhos_central) >= 2:
                rho_central = clock_rhos_central[-1]
                for b, P in enumerate(spectral_projectors):
                    if b >= len(dim_sensitivity):
                        break
                    # Project the central state and compute subspace norm
                    proj_state = P @ rho_central @ P
                    hs_norm = np.sqrt(np.real(np.trace(
                        proj_state.conj().T @ proj_state
                    )))
                    dim_sensitivity[b] = hs_norm

            max_val = np.max(np.abs(dim_sensitivity))
            if max_val > 0:
                dim_sensitivity = dim_sensitivity / max_val

            # Ensure exactly 11 entries
            if len(dim_sensitivity) < 11:
                dim_sensitivity = np.pad(
                    dim_sensitivity, (0, 11 - len(dim_sensitivity))
                )
            dim_sensitivity = dim_sensitivity[:11]
        else:
            # Fallback: scalar multiplication by scaling factors
            try:
                from .navigation import DimensionalScaling
                scaler = DimensionalScaling()
                scaling_factors = np.array([
                    scaler.scaling_factor(d) for d in range(11)
                ])
            except ImportError:
                scaling_factors = np.ones(11)

            central_rates = np.abs(rates[central_idx, :])
            nonzero_rates = central_rates[central_rates > 1e-20]
            dim_sensitivity = map_divergence_to_dimensions(
                nonzero_rates, scaling_factors
            )

        # Max stable offset
        final_fids = fid_to_central[:, -1] if n_times > 0 else fid_to_central[:, 0]
        max_stable = 0.0
        for k in range(n_clocks):
            if final_fids[k] >= self.config.coherence_threshold:
                max_stable = max(max_stable, abs(offsets[k]))

        return DivergenceTensor(
            distance_history=dist_history,
            divergence_rates=rates,
            dimension_sensitivity=dim_sensitivity,
            max_stable_offset=max_stable,
        )

    def _compute_coherence_width(
        self,
        offsets: np.ndarray,
        fid_to_central: np.ndarray,
        central_idx: int,
    ) -> float:
        """
        Maximum time offset where fidelity to central stays above threshold.

        This is the "temporal coherence scale" -- how far you can offset
        your clock and still track the same trajectory.
        """
        threshold = self.config.coherence_threshold
        # Use mean fidelity over time for each clock
        max_offset = 0.0
        for k in range(len(offsets)):
            if k == central_idx:
                continue
            mean_fid = np.mean(fid_to_central[k])
            if mean_fid >= threshold:
                max_offset = max(max_offset, abs(offsets[k]))
        return max_offset

    def _compute_blp(
        self,
        fid_to_central: np.ndarray,
        central_idx: int,
        times: np.ndarray,
        dist_history: Optional[np.ndarray] = None,
    ) -> NonMarkovianityMeasure:
        """
        BLP non-Markovianity measure.

        For Markovian evolution, the distance between any two states
        monotonically *decreases* under CPTP maps (contractivity).
        If distance *increases* (dD/dt > 0), information is flowing
        back from the environment — non-Markovian behavior.

        N_BLP = integral of max(dD/dt, 0) dt

        We use the outermost clock pair for maximum sensitivity.
        The fidelity_flow array stores dD/dt (distance derivative).
        """
        n_clocks = fid_to_central.shape[0]
        n_times = len(times)

        if n_times < 2:
            return NonMarkovianityMeasure(
                blp_measure=0.0,
                reconvergence_times=[],
                is_markovian=True,
                fidelity_flow=np.zeros(0),
            )

        # Use distance between outermost clock and central
        # Distance increasing = non-Markovian (information backflow)
        if dist_history is not None:
            outer_idx = 0  # Most negative offset
            dist_series = dist_history[outer_idx, central_idx, :]
        else:
            # Fallback: convert fidelity to distance
            outer_idx = 0
            fid_series = np.clip(fid_to_central[outer_idx], 0.0, 1.0)
            dist_series = np.arccos(np.sqrt(fid_series))

        # Compute dD/dt via finite differences
        dt_arr = np.diff(times)
        dD = np.diff(dist_series)
        dist_flow = np.zeros(n_times)
        for i in range(len(dD)):
            if dt_arr[i] > 0:
                dist_flow[i + 1] = dD[i] / dt_arr[i]

        # BLP measure: integral of positive dD/dt (distance increasing)
        # Skip first 2 steps (initial transient from offset setup, not dynamics)
        tol = self.config.blp_tolerance
        skip = min(2, n_times - 1)
        blp = 0.0
        reconvergence_times = []
        for i in range(skip + 1, n_times):
            if dist_flow[i] > tol:
                blp += dist_flow[i] * (times[i] - times[i - 1])
                if (len(reconvergence_times) == 0 or
                        times[i] - reconvergence_times[-1] > 2 * (times[1] - times[0])):
                    reconvergence_times.append(float(times[i]))

        # Compare BLP to total distance activity for relative threshold
        total_neg = sum(abs(dist_flow[i]) * (times[i] - times[i-1])
                        for i in range(skip + 1, n_times) if dist_flow[i] < -tol)
        total_activity = blp + total_neg
        is_markovian = blp < max(tol * (times[-1] - times[0]),
                                 0.01 * total_activity if total_activity > 0 else 0)

        return NonMarkovianityMeasure(
            blp_measure=float(blp),
            reconvergence_times=reconvergence_times,
            is_markovian=is_markovian,
            fidelity_flow=dist_flow,  # Now stores dD/dt
        )

    def evolve_kk_probing_clocks(
        self,
        rho_init: np.ndarray,
        dt: float,
        n_steps: int,
        projectors: List[np.ndarray],
    ) -> Dict[str, Any]:
        """
        Probe per-dimension sensitivity using spectral subspace clocks.

        For each spectral projector P_d (corresponding to one energy band /
        dimension), create a projected Hamiltonian H_d = P_d @ H @ P_d and
        evolve the projected state rho_d = P_d @ rho @ P_d / Tr(...) under
        H_d separately. Track divergence between each band evolution and the
        full-H central evolution.

        This answers the question: "Can the multi-clock compass resolve
        individual dimensions?" If band d shows high divergence, dimension d
        is dynamically active and resolvable.

        Args:
            rho_init: Initial density matrix.
            dt: Time step.
            n_steps: Number of evolution steps.
            projectors: List of spectral projectors (one per dimension),
                typically 11 projectors from SpectralDimensionMapper.

        Returns:
            Dict with:
                'divergence_per_band': np.ndarray of divergence rates per band
                'sensitivity_per_dim': np.ndarray of per-dimension sensitivity
                'occupied_bands': list of band indices that had nonzero weight
                'band_fidelities': dict of band_idx -> final fidelity array
        """
        n_bands = len(projectors)

        # Central evolution under full H
        rho_central = rho_init.copy()
        central_history = [rho_central.copy()]

        for _ in range(n_steps):
            rho_central = self.evolver.evolve_single_step(rho_central, dt)
            central_history.append(rho_central.copy())

        # Per-band probing
        divergence_rates = np.zeros(n_bands)
        occupied_bands = []
        band_fidelities = {}

        for b, P in enumerate(projectors):
            # Project initial state into this band
            rho_band = P @ rho_init @ P
            weight = np.real(np.trace(rho_band))

            if weight < 1e-12:
                # Band unoccupied -- no sensitivity measurement possible
                continue

            occupied_bands.append(b)
            rho_band = rho_band / weight
            rho_band = (rho_band + rho_band.conj().T) / 2.0

            # Projected Hamiltonian
            H_band = P @ self.H @ P

            # Build a separate evolver for this band
            band_evolver = LindbladEvolver(
                H_band, self.lindblad_ops, self.rates
            )

            # Evolve and track divergence from central
            rho_b = rho_band.copy()
            fids = np.zeros(n_steps + 1)
            fids[0] = uhlmann_fidelity(rho_b, rho_init)

            for step in range(n_steps):
                rho_b = band_evolver.evolve_single_step(rho_b, dt)
                # Distance from central evolution at same step
                rho_c = central_history[step + 1]
                fids[step + 1] = uhlmann_fidelity(rho_b, rho_c)

            band_fidelities[b] = fids

            # Divergence rate: how fast fidelity drops
            if n_steps > 0 and fids[0] > 1e-15:
                divergence_rates[b] = max(
                    0.0, (fids[0] - fids[-1]) / (n_steps * dt)
                )

        # Per-dimension sensitivity: divergence weighted by scaling
        try:
            from .navigation import DimensionalScaling
            scaler = DimensionalScaling()
            scaling_factors = np.array([
                scaler.scaling_factor(d) for d in range(min(n_bands, 11))
            ])
        except ImportError:
            scaling_factors = np.ones(min(n_bands, 11))

        # Pad or truncate to match bands
        if len(scaling_factors) < n_bands:
            scaling_factors = np.pad(
                scaling_factors, (0, n_bands - len(scaling_factors)),
                constant_values=1.0,
            )

        sensitivity = divergence_rates.copy()
        max_div = np.max(sensitivity)
        if max_div > 1e-20:
            sensitivity = sensitivity / max_div

        return {
            'divergence_per_band': divergence_rates,
            'sensitivity_per_dim': sensitivity,
            'occupied_bands': occupied_bands,
            'band_fidelities': band_fidelities,
        }

    def _compute_phase_dispersion_step(
        self,
        clock_rhos: List[np.ndarray],
    ) -> float:
        """
        Compute phase dispersion across clocks at a single timestep.

        Measures the variance of off-diagonal phases across clock states.
        Unitary evolution preserves eigenvalues but rotates phases, so this
        captures the actual divergence of the clock phases.
        """
        # Extract upper-triangle off-diagonal phases from each clock
        dim = clock_rhos[0].shape[0]
        phases_per_clock = []
        for rho in clock_rhos:
            phases = []
            for i in range(dim):
                for j in range(i + 1, dim):
                    phases.append(np.angle(rho[i, j]))
            phases_per_clock.append(phases)

        phases_arr = np.array(phases_per_clock)  # (n_clocks, n_phases)
        if phases_arr.shape[1] == 0:
            return 0.0
        # Circular variance across clocks for each phase, then mean
        return float(np.mean(np.var(phases_arr, axis=0)))