"""4-state Gaussian generative model for FEP active inference.

Ported from services/cloud-api/fep_controller.py — adapted to use
the local 11-dim observation vector from gpu_observer.compute_observation_vector().
"""

from __future__ import annotations

import numpy as np

# Latent states
STATES = ["UNDERUTILIZED", "BALANCED", "OVERLOADED", "DEGRADED"]

# 11-dimensional observation vector (matches compute_observation_vector):
# 0: gpu_util_mean        1: gpu_util_std         2: queue_depth
# 3: queue_depth_max      4: avg_wait_seconds     5: error_rate
# 6: throughput_jobs/min  7: memory_pressure      8: thermal_headroom
# 9: power_headroom       10: topology_utilization

OBS_DIM = 11

OBS_KEYS = [
    "gpu_util_mean", "gpu_util_std", "queue_depth", "queue_depth_max",
    "avg_wait_seconds", "error_rate", "throughput_jobs_per_min",
    "memory_pressure", "thermal_headroom", "power_headroom", "topology_utilization",
]

# Gaussian parameters per latent state (mean, std for each observation dimension)
GENERATIVE_MODEL = {
    "UNDERUTILIZED": {
        "mean": np.array([15.0, 5.0, 0.5, 1.0, 5.0, 0.01, 2.0, 0.1, 0.3, 0.2, 0.05]),
        "std":  np.array([10.0, 3.0, 0.5, 0.5, 5.0, 0.02, 2.0, 0.1, 0.1, 0.1, 0.05]),
    },
    "BALANCED": {
        "mean": np.array([60.0, 10.0, 3.0, 5.0, 30.0, 0.03, 10.0, 0.5, 0.5, 0.5, 0.3]),
        "std":  np.array([15.0, 5.0, 2.0, 3.0, 20.0, 0.03, 5.0, 0.15, 0.15, 0.15, 0.15]),
    },
    "OVERLOADED": {
        "mean": np.array([90.0, 5.0, 10.0, 15.0, 120.0, 0.05, 20.0, 0.85, 0.8, 0.8, 0.7]),
        "std":  np.array([8.0, 3.0, 5.0, 5.0, 60.0, 0.04, 8.0, 0.1, 0.1, 0.1, 0.15]),
    },
    "DEGRADED": {
        "mean": np.array([70.0, 20.0, 8.0, 12.0, 200.0, 0.15, 8.0, 0.7, 0.9, 0.85, 0.4]),
        "std":  np.array([20.0, 10.0, 4.0, 5.0, 100.0, 0.1, 5.0, 0.2, 0.1, 0.1, 0.2]),
    },
}

# Observation bounds for clamping
OBS_BOUNDS = {
    "gpu_util_mean":          (0.0, 100.0),
    "gpu_util_std":           (0.0, 50.0),
    "queue_depth":            (0.0, 1000.0),
    "queue_depth_max":        (0.0, 5000.0),
    "avg_wait_seconds":       (0.0, 3600.0),
    "error_rate":             (0.0, 1.0),
    "throughput_jobs_per_min": (0.0, 1000.0),
    "memory_pressure":        (0.0, 1.0),
    "thermal_headroom":       (0.0, 2.0),
    "power_headroom":         (0.0, 2.0),
    "topology_utilization":   (0.0, 1.0),
}


def obs_dict_to_vector(obs: dict[str, float]) -> np.ndarray:
    """Convert observation dict to numpy vector with bounds clamping."""
    vec = []
    for key in OBS_KEYS:
        try:
            v = float(obs.get(key, 0.0))
        except (TypeError, ValueError):
            v = 0.0
        if np.isnan(v) or np.isinf(v):
            v = 0.0
        lo, hi = OBS_BOUNDS[key]
        vec.append(max(lo, min(hi, v)))
    return np.array(vec)
