"""
Quantum Frenet-Serret apparatus for state trajectories in projective Hilbert space.

Computes curvature, torsion, and the Frenet-Serret frame along quantum state
trajectories. These geometric quantities characterize how the state deviates
from geodesic (error-free) evolution.

Mathematical framework (Alsing & Cafaro, 2024):
    - Evolution speed: v(t) = Delta_E / hbar
    - Curvature: kappa^2 = alpha_4 - 1
      where alpha_4 = <(Delta_H)^4> / <(Delta_H)^2>^2 (energy kurtosis)
    - Torsion: tau^2 = alpha_4 - 1 - alpha_3^2
      where alpha_3 = <(Delta_H)^3> / <(Delta_H)^2>^(3/2) (energy skewness)
    - Zero curvature = geodesic evolution (error-free)
    - Nonzero curvature = deviation from ideal path (error signature)

Pi enters naturally through the Fubini-Study geometry of CP^{n-1}:
a qubit on the Bloch sphere traverses arc length pi*R for a half rotation.

References:
    [1] Alsing & Cafaro, "The quantum mechanical Frenet-Serret frame" (2024)
        Part I: arXiv:2311.18458, Part II: arXiv:2311.18463
    [2] Brody & Hughston, "Geometric quantum mechanics" (2001)
"""

import numpy as np
from scipy import linalg
from dataclasses import dataclass
from typing import List, Optional


@dataclass
class FrenetSerretFrame:
    """Quantum Frenet-Serret frame at a point along the trajectory."""
    time: float                      # Physical time
    arc_length: float                # Arc-length parameter s
    curvature: float                 # kappa(s) >= 0
    torsion: float                   # tau(s), can be any real value
    speed: float                     # v(t) = Delta_E / hbar
    tangent: Optional[np.ndarray]    # Tangent direction in density matrix space
    normal: Optional[np.ndarray]     # Normal direction


class QuantumFrenetSerret:
    """
    Compute the Frenet-Serret frame along quantum state trajectories.

    Given a Hamiltonian H and a sequence of density matrices rho(t),
    computes curvature and torsion at each point using the energy
    distribution moments (Alsing-Cafaro formalism).
    """

    def __init__(self, hbar: float = 1.0):
        self.hbar = hbar

    def energy_moment(self, rho: np.ndarray, H: np.ndarray, order: int) -> float:
        """
        Compute the n-th central moment of the energy distribution.

        <(Delta_H)^n> = Tr((H - <H>*I)^n * rho)

        where <H> = Tr(H * rho).

        Args:
            rho: Density matrix
            H: Hamiltonian
            order: Moment order (2, 3, 4, ...)

        Returns:
            n-th central moment (real for even n, may be complex for odd n)
        """
        # Mean energy
        E_mean = np.real(np.trace(H @ rho))

        # Centered Hamiltonian
        dim = H.shape[0]
        Delta_H = H - E_mean * np.eye(dim)

        # Compute (Delta_H)^n
        Delta_H_n = np.linalg.matrix_power(Delta_H, order)

        # Trace with density matrix
        moment = np.trace(Delta_H_n @ rho)
        return np.real(moment)

    def evolution_speed(self, rho: np.ndarray, H: np.ndarray) -> float:
        """
        Quantum evolution speed (Fubini-Study metric velocity).

        v(t) = Delta_E / hbar = sqrt(<(Delta_H)^2>) / hbar

        This is the rate of change of arc length: ds/dt = v(t).
        """
        variance = self.energy_moment(rho, H, 2)
        # Variance must be non-negative; clamp for numerical safety
        variance = max(variance, 0.0)
        return np.sqrt(variance) / self.hbar

    def compute_curvature(self, rho: np.ndarray, H: np.ndarray) -> float:
        """
        Quantum curvature from energy distribution kurtosis.

        kappa^2 = alpha_4 - 1
        where alpha_4 = <(Delta_H)^4> / <(Delta_H)^2>^2

        For a Gaussian energy distribution: alpha_4 = 3, kappa^2 = 2
        For an eigenstate (delta distribution): alpha_4 = 1, kappa = 0
        For geodesic evolution: kappa = 0

        Returns:
            kappa >= 0 (curvature magnitude)
        """
        mu2 = self.energy_moment(rho, H, 2)

        if mu2 < 1e-20:
            # Pure eigenstate or zero Hamiltonian: zero curvature
            return 0.0

        mu4 = self.energy_moment(rho, H, 4)
        alpha_4 = mu4 / (mu2 ** 2)

        # kappa^2 = alpha_4 - 1
        kappa_sq = max(alpha_4 - 1.0, 0.0)
        return np.sqrt(kappa_sq)

    def compute_torsion(self, rho: np.ndarray, H: np.ndarray) -> float:
        """
        Quantum torsion from energy distribution skewness.

        tau^2 = alpha_4 - 1 - alpha_3^2
        where alpha_3 = <(Delta_H)^3> / <(Delta_H)^2>^(3/2)

        Returns:
            tau (torsion), can be positive, negative, or zero
        """
        mu2 = self.energy_moment(rho, H, 2)

        if mu2 < 1e-20:
            return 0.0

        mu3 = self.energy_moment(rho, H, 3)
        mu4 = self.energy_moment(rho, H, 4)

        alpha_3 = mu3 / (mu2 ** 1.5)
        alpha_4 = mu4 / (mu2 ** 2)

        tau_sq = max(alpha_4 - 1.0 - alpha_3 ** 2, 0.0)
        # Sign convention: use sign of alpha_3
        sign = np.sign(alpha_3) if abs(alpha_3) > 1e-15 else 1.0
        return sign * np.sqrt(tau_sq)

    def compute_tangent(self, rho: np.ndarray, H: np.ndarray) -> np.ndarray:
        """
        Tangent vector to the trajectory in density matrix space.

        For the Lindblad equation without dissipation (unitary part):
            d(rho)/ds = -i[H_centered, rho] / (hbar * v)

        This is normalized by the evolution speed to give a unit tangent
        in the Hilbert-Schmidt inner product.

        Returns:
            Tangent matrix (same shape as rho), or zero matrix if speed=0
        """
        v = self.evolution_speed(rho, H)
        if v < 1e-20:
            return np.zeros_like(rho)

        E_mean = np.real(np.trace(H @ rho))
        dim = H.shape[0]
        H_centered = H - E_mean * np.eye(dim)

        # drho/dt (unitary part) = -i[H, rho]
        commutator = H_centered @ rho - rho @ H_centered
        drho_dt = -1j * commutator

        # Tangent = drho/ds = (drho/dt) / v
        tangent = drho_dt / (self.hbar * v)

        # Normalize in Hilbert-Schmidt norm
        hs_norm = np.sqrt(np.real(np.trace(tangent.conj().T @ tangent)))
        if hs_norm > 1e-15:
            tangent = tangent / hs_norm

        return tangent

    def compute_normal(
        self,
        rho: np.ndarray,
        H: np.ndarray,
        tangent: np.ndarray,
    ) -> np.ndarray:
        """
        Normal vector via covariant derivative of tangent.

        The normal direction is obtained by:
        1. Computing the derivative of the tangent (d|T>/ds)
        2. Projecting out the component along the tangent (Gram-Schmidt)
        3. Normalizing

        Returns:
            Normal matrix, or zero if curvature is zero
        """
        v = self.evolution_speed(rho, H)
        if v < 1e-20:
            return np.zeros_like(rho)

        dim = H.shape[0]
        E_mean = np.real(np.trace(H @ rho))
        H_c = H - E_mean * np.eye(dim)

        # Second derivative: d^2(rho)/dt^2 (unitary part)
        comm1 = H_c @ rho - rho @ H_c
        comm2 = H_c @ comm1 - comm1 @ H_c
        d2rho = -comm2  # (-i)^2 = -1

        # Convert to arc-length parameterization
        d2rho_ds2 = d2rho / (self.hbar * v) ** 2

        # Gram-Schmidt: remove tangent component
        proj = np.real(np.trace(d2rho_ds2.conj().T @ tangent))
        normal = d2rho_ds2 - proj * tangent

        # Normalize
        hs_norm = np.sqrt(np.real(np.trace(normal.conj().T @ normal)))
        if hs_norm > 1e-15:
            normal = normal / hs_norm
        else:
            normal = np.zeros_like(rho)

        return normal

    def compute_frame(
        self,
        rho_history: List[np.ndarray],
        H: np.ndarray,
        dt: float,
    ) -> List[FrenetSerretFrame]:
        """
        Compute the full Frenet-Serret frame at each point in a trajectory.

        Args:
            rho_history: List of density matrices at successive times
            H: System Hamiltonian
            dt: Time step between density matrices

        Returns:
            List of FrenetSerretFrame, one per density matrix
        """
        frames = []
        cumulative_arc = 0.0

        for i, rho in enumerate(rho_history):
            t = i * dt
            speed = self.evolution_speed(rho, H)
            kappa = self.compute_curvature(rho, H)
            tau = self.compute_torsion(rho, H)
            tangent = self.compute_tangent(rho, H)
            normal = self.compute_normal(rho, H, tangent)

            # Arc-length accumulation (trapezoidal)
            if i > 0:
                prev_speed = frames[-1].speed
                cumulative_arc += 0.5 * (prev_speed + speed) * dt

            frame = FrenetSerretFrame(
                time=t,
                arc_length=cumulative_arc,
                curvature=kappa,
                torsion=tau,
                speed=speed,
                tangent=tangent,
                normal=normal,
            )
            frames.append(frame)

        return frames

    def curvature_from_trajectory(
        self,
        rho_history: List[np.ndarray],
        dt: float,
    ) -> np.ndarray:
        """
        Compute curvature numerically from finite differences of the trajectory.

        This is an alternative to the energy-moment formula that works
        even when the Hamiltonian is unknown. Uses the discrete Frenet-Serret
        approach: kappa = |d^2r/ds^2| where r is the trajectory in DM space.

        Args:
            rho_history: Sequence of density matrices
            dt: Time step

        Returns:
            Array of curvature values (length = len(rho_history) - 2)
        """
        if len(rho_history) < 3:
            return np.array([])

        curvatures = []
        for i in range(1, len(rho_history) - 1):
            # Finite differences in Hilbert-Schmidt space
            drho_forward = rho_history[i + 1] - rho_history[i]
            drho_backward = rho_history[i] - rho_history[i - 1]

            # Second derivative
            d2rho = (drho_forward - drho_backward) / dt

            # First derivative (central)
            drho = (rho_history[i + 1] - rho_history[i - 1]) / (2 * dt)

            # Speed = |drho/dt| in HS norm
            speed = np.sqrt(np.real(np.trace(drho.conj().T @ drho)))

            if speed < 1e-15:
                curvatures.append(0.0)
                continue

            # Curvature = |d^2r/ds^2| = |d^2rho/dt^2 * (1/v^2) - (a.T/v^3)*drho/dt|
            # Simplified: |second_deriv - (second_deriv . tangent) * tangent| / v^2
            tangent = drho / speed
            d2rho_normalized = d2rho / speed ** 2
            proj = np.real(np.trace(d2rho_normalized.conj().T @ tangent))
            normal_component = d2rho_normalized - proj * tangent
            kappa = np.sqrt(np.real(np.trace(
                normal_component.conj().T @ normal_component
            )))

            curvatures.append(kappa)

        return np.array(curvatures)

    # ------------------------------------------------------------------
    # Spectral subspace decomposition for per-dimension curvature
    # ------------------------------------------------------------------

    @staticmethod
    def spectral_decompose_hamiltonian(
        H: np.ndarray,
        energy_bands: list,
    ) -> list:
        """
        Decompose Hamiltonian into spectral projectors for energy bands.

        For each band [E_low, E_high], builds a projector P = sum |i><i|
        over eigenvectors whose eigenvalues fall in that band.

        Args:
            H: Hamiltonian matrix (dim x dim).
            energy_bands: List of (E_low, E_high) tuples, one per
                target dimension.

        Returns:
            List of projector matrices (one per band), same shape as H.
            Empty bands get a zero projector.
        """
        dim = H.shape[0]
        eigenvalues, eigenvectors = np.linalg.eigh(H)

        projectors = []
        for E_low, E_high in energy_bands:
            P = np.zeros((dim, dim), dtype=complex)
            for i, E_i in enumerate(eigenvalues):
                if E_low <= E_i <= E_high:
                    v = eigenvectors[:, i].reshape(-1, 1)
                    P += v @ v.conj().T
            projectors.append(P)

        return projectors

    def compute_subspace_curvatures(
        self,
        rho: np.ndarray,
        H: np.ndarray,
        projectors: list,
    ) -> tuple:
        """
        Compute curvature and torsion in each spectral subspace.

        For each projector P_d:
          1. Project: rho_d = P_d @ rho @ P_d
          2. If Tr(rho_d) > threshold, normalize and compute curvature
          3. Otherwise, kappa=0, tau=0 (subspace unoccupied)

        This gives a per-dimension curvature that reflects how the
        state's energy distribution deviates from geodesic within
        each energy band -- the geometric signature of each dimension.

        Args:
            rho: Full density matrix.
            H: Full Hamiltonian.
            projectors: List of spectral projectors (from
                spectral_decompose_hamiltonian).

        Returns:
            (curvatures, torsions, weights) -- each np.ndarray of
            length len(projectors).
              curvatures[d] = kappa in subspace d
              torsions[d]   = tau in subspace d
              weights[d]    = Tr(P_d @ rho) (occupation of subspace d)
        """
        n_bands = len(projectors)
        curvatures = np.zeros(n_bands)
        torsions = np.zeros(n_bands)
        weights = np.zeros(n_bands)

        threshold = 1e-12

        for d, P in enumerate(projectors):
            # Subspace weight: how much of rho lives in this band
            rho_sub = P @ rho @ P
            w = np.real(np.trace(rho_sub))
            weights[d] = w

            if w < threshold:
                # Subspace is empty or negligible
                continue

            # Check if projector spans at least 2 eigenvectors
            # (single eigenvector → eigenstate → kappa=0 exactly)
            rank = np.real(np.trace(P))
            if rank < 1.5:
                # Rank-1 subspace: geodesic by definition
                continue

            # Normalize projected state
            rho_sub_norm = rho_sub / w

            # Enforce Hermiticity for numerical safety
            rho_sub_norm = (rho_sub_norm + rho_sub_norm.conj().T) / 2.0

            # Projected Hamiltonian
            H_sub = P @ H @ P

            # Compute curvature and torsion on the subspace
            curvatures[d] = self.compute_curvature(rho_sub_norm, H_sub)
            torsions[d] = self.compute_torsion(rho_sub_norm, H_sub)

        return curvatures, torsions, weights