"""
Spectral analysis for Watkins dynamics.

Provides:
- Spectral gap computation for the Hessian at equilibrium
- Eigenvalue bounds (lambda_min, lambda_max) certification
- Mixing time estimates from spectral gap

Key theorem: At lam* = 1/phi, the constrained Hessian has eigenvalues
[5.934, 20.556], giving spectral gap mu ~ 5.934 and mixing time
O(1/mu) ~ 0.169. The condition number is ~3.46 (moderate anisotropy).

v0.7.1: Fixed Hessian cross-term (was 0, should be T*/eta*).
"""

import torch
import numpy as np
from typing import Tuple, Optional
from .constants import PHI, T_STAR, LAM_STAR, M_PHI

# Certified eigenvalue bounds (from Hessian at golden equilibrium)
from .constants import MU_SLOW, MU_FAST
HESSIAN_LAMBDA_MIN = MU_SLOW   # ~5.934 (spectral gap)
HESSIAN_LAMBDA_MAX = MU_FAST   # ~20.556 (fast mode)
SPECTRAL_GAP = HESSIAN_LAMBDA_MIN
MIXING_TIME = 1 / SPECTRAL_GAP  # ~0.169


def hessian_at_equilibrium(dim: int = 3) -> torch.Tensor:
    """
    Compute Hessian of F at the golden equilibrium.

    For 2-simplex (dim=3): returns 2x2 matrix (projected onto tangent plane)
    with eta = 1 - lam - kap coupling all three variables.

    Eigenvalues: [5.934, 20.556] (positive definite -> stable minimum)

    The Hessian entries are:
        H11 = 1/lam*^2 + T*(1/lam* + 1/eta*)
        H22 = T*(1/kap* + 1/eta*)
        H12 = T*/eta*   (chain rule through eta = 1 - lam - kap)

    Args:
        dim: Simplex dimension (3 or 4)

    Returns:
        Hessian matrix on tangent space
    """
    if dim == 3:
        # At (lam*, kap*, kap*) with kap* = eta* = (1-lam*)/2
        kap_star = (1 - LAM_STAR) / 2

        # d2F/dlam2 = 1/lam^2 + T*(1/lam + 1/eta)
        d2F_dlam2 = 1 / LAM_STAR**2 + T_STAR * (1/LAM_STAR + 1/kap_star)

        # d2F/dkap2 = T*(1/kap + 1/eta)
        d2F_dkap2 = T_STAR * (1/kap_star + 1/kap_star)

        # d2F/dlamdkap = T*/eta  (NOT zero — chain rule through eta = 1-lam-kap)
        d2F_dlamkap = T_STAR / kap_star

        H = torch.tensor([
            [d2F_dlam2, d2F_dlamkap],
            [d2F_dlamkap, d2F_dkap2]
        ], dtype=torch.float64)

        return H
    else:
        raise NotImplementedError(f"dim={dim} not yet supported")


def spectral_gap(H: Optional[torch.Tensor] = None) -> float:
    """
    Compute spectral gap mu = lambda_min(H).

    The spectral gap is the SMALLEST eigenvalue of the Hessian.
    It controls the SLOWEST convergence mode and determines
    the mixing time: tau_mix = 1/mu.

    Args:
        H: Hessian matrix (default: equilibrium Hessian)

    Returns:
        Spectral gap mu
    """
    if H is None:
        H = hessian_at_equilibrium()

    eigenvalues = torch.linalg.eigvalsh(H)
    return float(eigenvalues.min())


def mixing_time_bound(H: Optional[torch.Tensor] = None) -> float:
    """
    Compute mixing time bound tau_mix = 1/mu.

    After time tau_mix, the flow is within 1/e of equilibrium.

    Args:
        H: Hessian matrix (default: equilibrium Hessian)

    Returns:
        Mixing time tau_mix
    """
    return 1 / spectral_gap(H)


def exponential_convergence_rate(
    p: torch.Tensor,
    t: float,
    H: Optional[torch.Tensor] = None
) -> torch.Tensor:
    """
    Bound on distance to equilibrium after time t.

    ||p(t) - p*|| <= ||p(0) - p*|| * exp(-mu*t)

    Args:
        p: Initial state
        t: Evolution time
        H: Hessian (for spectral gap)

    Returns:
        Convergence factor exp(-mu*t)
    """
    mu = spectral_gap(H)
    return torch.exp(torch.tensor(-mu * t))


def verify_positive_definite(H: torch.Tensor) -> Tuple[bool, torch.Tensor]:
    """
    Verify Hessian is positive definite (stable minimum).

    Returns:
        (is_positive_definite, eigenvalues)
    """
    eigenvalues = torch.linalg.eigvalsh(H)
    return bool((eigenvalues > 0).all()), eigenvalues
