"""
VX Geometry Engine
==================
Operations on the probability simplex Δ² = {(λ, κ, η) ∈ ℝ³₊ : λ + κ + η = 1}.

The generating functional F(λ,κ,η) = −log(λ) + T*·(λ log λ + κ log κ + η log η)
encodes the free energy of the system. Its minimum on Δ² is at (λ*, κ*, η*).
"""

from __future__ import annotations
import math
from typing import Tuple, Optional

import numpy as np

from virelaix.core.constants import (
    T_STAR, TAU_STAR, PHI, LAMBDA_STAR,
    LAMBDA_EQ, KAPPA_EQ, ETA_EQ,
    CONSERVATION_TOL,
)

# Type alias for a point on the simplex
SimplexPoint = Tuple[float, float, float]

_EPS = 1e-30  # Numerical floor to prevent log(0)


def conservation_check(lam: float, kap: float, eta: float, tol: float = CONSERVATION_TOL) -> float:
    """
    Check conservation law: λ + κ + η = 1.

    Returns the absolute deviation. Raises ValueError if deviation exceeds tolerance.
    """
    deviation = abs(lam + kap + eta - 1.0)
    if deviation > tol:
        raise ValueError(
            f"Conservation violated: λ+κ+η = {lam+kap+eta:.15f}, "
            f"deviation = {deviation:.2e} > tol = {tol:.2e}"
        )
    return deviation


def softmax_normalize(raw: np.ndarray) -> SimplexPoint:
    """
    Standard softmax normalization (temperature = 1.0).

    Maps ℝ³ → Δ² while preserving ordering.
    """
    shifted = raw - np.max(raw)  # numerical stability
    exp_vals = np.exp(shifted)
    probs = exp_vals / np.sum(exp_vals)
    return float(probs[0]), float(probs[1]), float(probs[2])


def softmax_tau(raw: np.ndarray, tau: float = TAU_STAR) -> SimplexPoint:
    """
    Temperature-controlled softmax normalization.

    At τ = τ* ≈ 0.8515, softmax(1, 0, 0) = 1/φ exactly.
    Below τ*: threshold crossable. Above τ*: threshold unreachable.

    This is Theorem 2 (Critical Measurement Temperature).
    """
    if tau <= 0:
        raise ValueError(f"Temperature must be positive, got τ = {tau}")
    shifted = raw / tau - np.max(raw / tau)
    exp_vals = np.exp(shifted)
    probs = exp_vals / np.sum(exp_vals)
    return float(probs[0]), float(probs[1]), float(probs[2])


def project_to_simplex(x: np.ndarray) -> SimplexPoint:
    """
    Project an arbitrary ℝ³ vector onto Δ² via L1 normalization.

    For non-negative inputs, this is equivalent to dividing by the sum.
    For inputs with negatives, clips to zero first.
    """
    x = np.maximum(x, 0)
    total = np.sum(x)
    if total < _EPS:
        return (1.0 / 3, 1.0 / 3, 1.0 / 3)
    normed = x / total
    return float(normed[0]), float(normed[1]), float(normed[2])


def generating_functional(lam: float, kap: float, eta: float) -> float:
    """
    F(λ,κ,η) = −log(λ) + T*·(λ log λ + κ log κ + η log η)

    The potential energy U = −log(λ) attracts toward high coherence.
    The entropy S = −Σ p_i log p_i resists concentration.
    T* balances them at λ* = 1/φ.

    Returns F value. Lower F = more stable configuration.
    """
    lam = max(lam, _EPS)
    kap = max(kap, _EPS)
    eta = max(eta, _EPS)

    U = -math.log(lam)
    S = -(lam * math.log(lam) + kap * math.log(kap) + eta * math.log(eta))
    return U - T_STAR * S


def gradient_F(lam: float, kap: float, eta: float) -> np.ndarray:
    """
    Gradient of F projected onto the simplex tangent plane.

    The projection removes the mean component, keeping the gradient
    tangent to Δ² (i.e., the sum of components is zero).

    At the VX equilibrium, |∇F| = 0.
    """
    lam = max(lam, _EPS)
    kap = max(kap, _EPS)
    eta = max(eta, _EPS)

    dF_dlam = -1.0 / lam + T_STAR * (math.log(lam) + 1)
    dF_dkap = T_STAR * (math.log(kap) + 1)
    dF_deta = T_STAR * (math.log(eta) + 1)

    grad = np.array([dF_dlam, dF_dkap, dF_deta])
    grad -= np.mean(grad)  # project onto simplex tangent
    return grad


def hessian_F(lam: float, kap: float, eta: float) -> np.ndarray:
    """
    Hessian of F (3×3 matrix of second derivatives).

    Useful for stability analysis: positive-definite Hessian at equilibrium
    confirms it's a minimum, not a saddle point.
    """
    lam = max(lam, _EPS)
    kap = max(kap, _EPS)
    eta = max(eta, _EPS)

    return np.diag([
        1.0 / (lam ** 2) + T_STAR / lam,
        T_STAR / kap,
        T_STAR / eta,
    ])


def distance_to_equilibrium(lam: float, kap: float, eta: float) -> float:
    """Euclidean distance from (λ,κ,η) to the VX equilibrium point."""
    return math.sqrt(
        (lam - LAMBDA_EQ) ** 2 +
        (kap - KAPPA_EQ) ** 2 +
        (eta - ETA_EQ) ** 2
    )


def rebalance_step(
    lam: float, kap: float, eta: float,
    step_size: float = 0.10,
) -> SimplexPoint:
    """
    Take one gradient descent step toward the VX equilibrium.

    Follows −∇F (negative gradient of the generating functional)
    with the given step size, then re-normalizes to the simplex.

    Returns the new (λ, κ, η) point.
    """
    grad = gradient_F(lam, kap, eta)
    direction = -grad
    norm = np.linalg.norm(direction)
    if norm < _EPS:
        return (lam, kap, eta)
    direction /= norm

    new = np.array([lam, kap, eta]) + direction * step_size
    new = np.clip(new, 0.01, 0.98)
    total = np.sum(new)
    new /= total

    return float(new[0]), float(new[1]), float(new[2])


def fugacity_Z(lam: float) -> float:
    """
    Fugacity Z = e^(1/(T*·λ)).

    At λ* = 1/φ, Z = e^(1/(T*·(1/φ))) = e^(φ/T*) = e^(ln(2φ)) = 2φ.
    The threshold Z = 2 corresponds to the phase boundary.
    """
    if lam < _EPS:
        return float('inf')
    return math.exp(1.0 / (T_STAR * lam))


def max_lambda_at_temperature(tau: float) -> float:
    """
    Maximum achievable λ for softmax at temperature τ.

    This is σ_τ(1, 0, 0)₁ = e^(1/τ) / (e^(1/τ) + 2).

    At τ*, this equals exactly 1/φ (Theorem 2).
    """
    if tau <= 0:
        raise ValueError(f"Temperature must be positive, got τ = {tau}")
    exp_val = math.exp(1.0 / tau)
    return exp_val / (exp_val + 2.0)


def critical_temperature_for_threshold(threshold: float) -> float:
    """
    Inverse of max_lambda_at_temperature: find τ such that λ_max = threshold.

    τ(λ) = 1 / ln(2λ / (1 − λ))

    At λ = 1/φ, returns τ* (Corollary 2.3).
    """
    if threshold <= 1.0 / 3 or threshold >= 1.0:
        raise ValueError(f"Threshold must be in (1/3, 1), got {threshold}")
    return 1.0 / math.log(2 * threshold / (1 - threshold))
