stability.py

"""
AntiAtropos Stability Layer — Phase 3.

This module is the mathematical core of the Lyapunov-inspired reward signal.
It is intentionally a stateless collection of pure functions so it can be
swapped, extended, or unit-tested independently of the environment.

Key concepts implemented
------------------------
1. Lyapunov Energy  V(s) = Σ Q_i²
   The "potential energy" of the cluster.  Zero means all queues are empty.
   A rising V means the cluster is destabilising.

2. Lyapunov Drift  ΔV(t) = V(s_t) − V(s_{t−1})
   The primary stabilising objective.  The reward penalises positive drift.
   Negative drift (energy decreasing) is "good" — the agent drove the system
   toward the equilibrium.

3. Control-Barrier Function  h_i(s) = max(0, Q_i − Q_max)²
   A soft safety constraint.  h_i > 0 only when node i has exceeded the
   hard-queue-depth safety ceiling Q_max.  Summing over all nodes gives the
   total barrier violation penalty.

4. Neely Drift-Plus-Penalty (optional, advanced)
   From Neely's Lyapunov optimisation framework:

       Δ(t) + V_weight · p(t)

   where:
       Δ(t)   = ΔV    — the one-step Lyapunov drift.
       p(t)   = cost  — the per-step penalty (infrastructure cost here).
       V_weight — trade-off parameter:  large V_weight → prioritise cost;
                  small V_weight → prioritise stability.

   Minimising this expression at every step produces a policy that is both
   stable (queue-stable in the mean) and cost-efficient.
"""

from __future__ import annotations

import os
import math
import statistics
from typing import Sequence


# ---------------------------------------------------------------------------
# Safety ceiling used by the barrier function
# ---------------------------------------------------------------------------

Q_BARRIER_MAX: float = 150.0
"""Queue depth above which the barrier function fires (hard safety zone).
Set higher than OVERLOAD_THRESHOLD (80) to allow the agent time to react
before the barrier penalty kicks in."""

MAX_QUEUE_NORM: float = 200.0
"""Normalization divisor shared with environment.py.  Matches FATAL_FAIL_THRESHOLD."""

BARRIER_NORM_SCALE: float = 10000.0
"""Normalization divisor for the barrier term.
The raw barrier H(s) = sum(max(0, Q_i - Q_max)^2) can produce very large values
(e.g. 5 nodes at Q=200, Q_max=150 gives 5*2500=12500). Without normalization,
this dominates the reward. Dividing by this scale keeps barrier in the same
order of magnitude as the other terms when delta=0.005."""

STABILITY_WINDOW: int = 10
"""Number of ticks to look back when judging whether the system is
trend-stable (V is on a decreasing trajectory)."""


# ---------------------------------------------------------------------------
# Reward normalisation defaults (env-overridable)
# ---------------------------------------------------------------------------

REWARD_NORM_MIDPOINT: float = float(os.getenv("ANTIATROPOS_REWARD_MIDPOINT", "-6.0"))
REWARD_NORM_TEMPERATURE: float = float(os.getenv("ANTIATROPOS_REWARD_TEMPERATURE", "2.0"))
REWARD_NORM_EPS: float = float(os.getenv("ANTIATROPOS_REWARD_EPS", "1e-8"))
REWARD_SCALE_VERSION: str = "sigmoid-v2"  # v2: smooth SLA + barrier active


# ---------------------------------------------------------------------------
# Core Lyapunov functions
# ---------------------------------------------------------------------------

def compute_lyapunov(nodes: list[dict]) -> float:
    """
    V(s) = Σ Q_i²

    Sum of squared queue depths across all nodes.  This is the cluster's
    Lyapunov energy.  Lower is more stable; zero means all queues are empty.

    Args:
        nodes: List of per-node state dicts (as returned by simulator.state()).
               Each dict must contain the key ``queue_depth``.

    Returns:
        Scalar Lyapunov energy  ≥ 0.
    """
    return float(
        sum(
            float(n.get("importance_weight", 1.0)) * (n["queue_depth"] ** 2)
            for n in nodes
        )
    )

def compute_lyapunov_graph(
    nodes: list[dict],
    topology: dict[str, list[str]],
    edge_weight: float = 5.0,
) -> float:
    """
    V_graph(s) = Σ w_i·Q_i²  +  edge_weight · Σ_{(i,j)∈edges} |Q_i - Q_j|

    The edge term penalises flow imbalance between connected nodes.
    If node-0 is overloaded but node-1 (its child) is idle, the edge
    term fires even if node-1's individual energy contribution is zero.
    This gives the agent gradient signal to balance load across the graph,
    not just reduce individual node queues.

    With edge_weight=5.0 and max |Q_diff|=200 on 4 edges, the edge term
    adds ~4000 to the Lyapunov energy, or ~2-5 % of the base energy at
    full overload — a meaningful secondary gradient without dominating.
    """
    node_map = {n["node_id"]: n for n in nodes}
    
    # Standard weighted Lyapunov
    base_energy = compute_lyapunov(nodes)
    
    # Edge imbalance penalty (raw queue-depth differences)
    edge_penalty = 0.0
    for parent_id, children in topology.items():
        parent = node_map.get(parent_id)
        if not parent:
            continue
        for child_id in children:
            child = node_map.get(child_id)
            if not child:
                continue
            imbalance = abs(float(parent["queue_depth"]) - float(child["queue_depth"]))
            edge_penalty += imbalance

    return base_energy + edge_weight * edge_penalty


def compute_drift(v_prev: float, v_curr: float) -> float:
    """
    ΔV(t) = V(s_t) − V(s_{t−1})

    One-step drift in Lyapunov energy.

    Negative drift → the agent moved the system toward a lower-energy state.
    Positive drift → the cluster is destabilising.

    Args:
        v_prev: Lyapunov energy at the *previous* tick.
        v_curr: Lyapunov energy at the *current* tick.

    Returns:
        Signed scalar drift value.
    """
    return v_curr - v_prev


def compute_barrier(nodes: list[dict], q_max: float = Q_BARRIER_MAX) -> float:
    """
    Control-Barrier Function (CBF) violation penalty.

        H(s) = Σ_i  max(0, Q_i − Q_max)²

    This is zero when no node exceeds the ceiling and grows quadratically as
    queues enter the "hard danger zone" above Q_max.  It can be added to the
    reward as an extra penalty for unsafe states.

    Args:
        nodes:  Per-node state dicts (must contain ``queue_depth``).
        q_max:  Safety ceiling for queue depth.  Default: Q_BARRIER_MAX.

    Returns:
        Scalar barrier violation energy  ≥ 0.
    """
    violation = 0.0
    for n in nodes:
        excess = n["queue_depth"] - q_max
        if excess > 0:
            violation += excess ** 2
    return violation


# ---------------------------------------------------------------------------
# Trend analysis
# ---------------------------------------------------------------------------

def is_lyapunov_stable(
    v_history: Sequence[float],
    window: int = STABILITY_WINDOW,
) -> bool:
    """
    Return True if the Lyapunov energy has been on a non-increasing trend
    over the last ``window`` ticks.

    Uses a simple linear regression slope: if slope ≤ 0 the system is
    considered trend-stable.

    Args:
        v_history: Ordered sequence of Lyapunov energy values (oldest first).
        window:    How many recent values to consider.

    Returns:
        True if the system is trend-stable, False otherwise.
    """
    recent = list(v_history[-window:])
    if len(recent) < 2:
        return True   # not enough data — assume stable at episode start

    n = len(recent)
    xs = list(range(n))
    mean_x = (n - 1) / 2.0
    mean_y = statistics.mean(recent)

    num = sum((xs[i] - mean_x) * (recent[i] - mean_y) for i in range(n))
    den = sum((xs[i] - mean_x) ** 2 for i in range(n))

    if den == 0:
        return True

    slope = num / den
    return slope <= 0.0


def lyapunov_variance(v_history: Sequence[float]) -> float:
    """
    Variance of the Lyapunov energy trajectory over an episode.

    Used by the grader as the primary stability metric: a lower variance
    means the agent kept the cluster in a consistently stable state, rather
    than allowing wild oscillations.

    Args:
        v_history: All per-tick V(s) values for the episode.

    Returns:
        Population variance of the energy trajectory.
    """
    if len(v_history) < 2:
        return 0.0
    return statistics.variance(v_history)


# ---------------------------------------------------------------------------
# Neely Drift-Plus-Penalty (advanced reward signal)
# ---------------------------------------------------------------------------

def drift_plus_penalty(
    v_prev: float,
    v_curr: float,
    penalty_cost: float,
    V_weight: float = 1.0,
) -> float:
    """
    Neely's Drift-Plus-Penalty objective:

        DPP(t) = ΔV(t) + V_weight · p(t)

    where:
        ΔV(t)        = v_curr − v_prev  (Lyapunov drift)
        p(t)         = penalty_cost     (infrastructure cost this tick)
        V_weight     = trade-off coefficient:
                         large  → agent optimises cost more aggressively,
                         small  → agent focuses on stability.

    Minimising this at each step produces a queue-stable policy with bounded
    average cost — the theoretical guarantee from Neely's framework.

    This function can substitute for the simpler ΔV term in the reward
    when you want to make the cost trade-off explicit and theoretically
    grounded (rather than the ad-hoc β·Cost term).

    Args:
        v_prev:       Lyapunov energy at previous tick.
        v_curr:       Lyapunov energy at current tick.
        penalty_cost: Per-step cost to penalise (e.g. current_cost_per_hour).
        V_weight:     Trade-off weight  V  in Neely's framework.

    Returns:
        Scalar DPP value.  The reward should negate this:
        R_t = −DPP(t) − γ·SLA_violation_step
    """
    delta_v = compute_drift(v_prev, v_curr)
    return delta_v + V_weight * penalty_cost


# ---------------------------------------------------------------------------
# Convenience: full reward computation (matches environment.py formula)
# ---------------------------------------------------------------------------

def smooth_sla_penalty(
    avg_latency_norm: float,
    error_rate: float,
    latency_threshold: float = 0.20,
    error_threshold: float = 0.05,
    latency_temperature: float = 0.03,
    error_temperature: float = 0.01,
) -> float:
    """
    Smooth SLA penalty in [0, 1] that ramps up as metrics approach thresholds.

    Unlike the binary cliff (0 or 1), this gives the agent gradient signal
    BEFORE the SLA is actually violated, enabling preventive learning.

    Uses two sigmoids (one for latency, one for errors) and takes the max
    so whichever dimension is worse dominates.

    Args:
        avg_latency_norm:     Normalized average latency [0, 1].
        error_rate:           Cluster-wide error rate [0, 1].
        latency_threshold:    Normalized latency SLA boundary.
        error_threshold:      Error rate SLA boundary.
        latency_temperature:  Sigmoid temperature for latency (lower = sharper).
        error_temperature:    Sigmoid temperature for errors (lower = sharper).

    Returns:
        Smooth penalty in [0, 1]. Near 0 when safe, near 1 when violating.

    Raises:
        ValueError: If inputs are outside [0, 1], indicating raw (non-normalized)
            values were passed by mistake. This is a common bug: passing latency
            in raw ms (e.g. 200.0) instead of normalized [0,1] (e.g. 0.20).
    """
    if avg_latency_norm < -0.01 or avg_latency_norm > 1.5:
        raise ValueError(
            f"smooth_sla_penalty: avg_latency_norm={avg_latency_norm:.4f} is outside "
            f"expected [0, 1] range. Did you pass raw ms instead of normalized? "
            f"Divide by MAX_LATENCY_NORM before calling."
        )
    if error_rate < -0.01 or error_rate > 1.5:
        raise ValueError(
            f"smooth_sla_penalty: error_rate={error_rate:.4f} is outside "
            f"expected [0, 1] range."
        )
    lat_z = (avg_latency_norm - latency_threshold) / max(1e-8, latency_temperature)
    err_z = (error_rate - error_threshold) / max(1e-8, error_temperature)
    lat_penalty = 1.0 / (1.0 + math.exp(-lat_z))
    err_penalty = 1.0 / (1.0 + math.exp(-err_z))
    return max(lat_penalty, err_penalty)


def compute_reward(
    v_prev: float,
    v_curr: float,
    cost: float,
    sla_violation_step: float = 0.0,
    alpha: float = 1.0,
    beta: float = 0.05,
    gamma: float = 2.0,
    barrier: float = 0.0,
    delta: float = 0.005,
) -> float:
    """
    R_t = -(alpha * DeltaV(s) + beta * Cost + gamma * SLA_smooth + delta * Barrier)

    Convenience wrapper that mirrors the reward formula in environment.py.

    Args:
        v_prev:         Lyapunov energy at previous tick.
        v_curr:         Lyapunov energy at current tick.
        cost:           Infrastructure cost this tick (USD/hr).
        sla_violation_step: Smooth SLA penalty in [0, 1] (was binary 0/1).
        alpha:          Weight on Lyapunov drift.
        beta:           Weight on cost.
        gamma:          Weight on SLA violations.
        barrier:        Control-barrier function violation energy.
        delta:          Weight on barrier penalty.

    Returns:
        Scalar reward (higher is better, always <= 0 in a stable episode).
    """
    delta_v = compute_drift(v_prev, v_curr)
    # Normalize barrier to prevent reward domination: raw barrier can be ~12500,
    # after dividing by BARRIER_NORM_SCALE it's ~1.25, then scaled by delta=0.005
    # gives ~0.006 which is comparable to other terms.
    barrier_normalized = barrier / BARRIER_NORM_SCALE if BARRIER_NORM_SCALE > 0 else barrier
    return -(alpha * delta_v + beta * cost + gamma * sla_violation_step + delta * barrier_normalized)


def normalize_reward(
    raw_reward: float,
    midpoint: float = REWARD_NORM_MIDPOINT,
    temperature: float = REWARD_NORM_TEMPERATURE,
    eps: float = REWARD_NORM_EPS,
) -> float:
    """
    Deterministically map raw reward to [0, 1] using a smooth sigmoid.

    reward_01 = 1 / (1 + exp(-(raw_reward - midpoint) / temperature))
    """
    temp = max(float(eps), abs(float(temperature)))
    z = (float(raw_reward) - float(midpoint)) / temp
    if z >= 0:
        exp_neg = math.exp(-z)
        out = 1.0 / (1.0 + exp_neg)
    else:
        exp_pos = math.exp(z)
        out = exp_pos / (1.0 + exp_pos)
    return min(1.0, max(0.0, float(out)))