"""FEP (Free Energy Principle) active inference controller.

Ported from services/cloud-api/fep_controller.py — runs locally as a continuous
background task instead of a Lambda invoked every 5 minutes.
"""

from __future__ import annotations

import asyncio
import json
import logging
from datetime import datetime, timezone

import numpy as np
from scipy.stats import multivariate_normal

from radix_edge.config import get_config
from radix_edge.db import get_db
from radix_edge.meta.generative_model import (
    STATES, GENERATIVE_MODEL,
    obs_dict_to_vector,
)
from radix_edge.observer.gpu_observer import compute_observation_vector

logger = logging.getLogger("radix_edge.meta")


# --- Core FEP math ---

def compute_log_likelihood(obs_vec: np.ndarray, state: str) -> float:
    """Compute log P(obs | state) using diagonal Gaussian likelihood."""
    params = GENERATIVE_MODEL[state]
    mean = params["mean"]
    std = params["std"]
    cov = np.diag(std ** 2)

    try:
        log_prob = float(multivariate_normal.logpdf(obs_vec, mean=mean, cov=cov))
        if np.isnan(log_prob) or np.isinf(log_prob):
            return -1000.0
        return log_prob
    except (ValueError, np.linalg.LinAlgError):
        return -1000.0


def update_beliefs(prior: dict[str, float], obs: dict[str, float]) -> dict[str, float]:
    """Bayesian update: Q(state | obs) ∝ P(obs | state) · P(state)."""
    obs_vec = obs_dict_to_vector(obs)

    log_posterior = {}
    for state in STATES:
        log_lik = compute_log_likelihood(obs_vec, state)
        log_prior = np.log(prior.get(state, 0.25) + 1e-10)
        log_posterior[state] = log_lik + log_prior

    # Log-sum-exp normalization
    max_log = max(log_posterior.values())
    if np.isinf(max_log) or np.isnan(max_log):
        return {s: 0.25 for s in STATES}

    posterior = {}
    for state in STATES:
        val = log_posterior[state] - max_log
        posterior[state] = np.exp(val) if not np.isnan(val) else 0.0

    total = sum(posterior.values())
    if total > 0:
        posterior = {s: p / total for s, p in posterior.items()}
    else:
        posterior = {s: 0.25 for s in STATES}

    return posterior


def compute_variational_free_energy(obs: dict[str, float], posterior: dict[str, float],
                                     prior: dict[str, float]) -> float:
    """F = KL(Q || P) - E_Q[log P(obs | s)]."""
    obs_vec = obs_dict_to_vector(obs)
    free_energy = 0.0

    for state in STATES:
        q_s = posterior[state]
        if q_s < 1e-10:
            continue
        free_energy += q_s * np.log(q_s + 1e-10)
        free_energy -= q_s * compute_log_likelihood(obs_vec, state)
        free_energy -= q_s * np.log(prior.get(state, 0.25) + 1e-10)

    return float(free_energy)


def predict_obs_after_action(obs: dict[str, float], action: str) -> dict[str, float]:
    """Linear approximation of action effects on observations."""
    predicted = obs.copy()

    if action == "balanced":
        predicted["gpu_util_mean"] = obs.get("gpu_util_mean", 50) * 0.9 + 60.0 * 0.1
        predicted["queue_depth"] = obs.get("queue_depth", 2) * 0.8 + 3.0 * 0.2
    elif action == "latency-first":
        predicted["gpu_util_mean"] = obs.get("gpu_util_mean", 50) * 0.85 + 50.0 * 0.15
        predicted["queue_depth"] = obs.get("queue_depth", 2) * 0.7 + 1.5 * 0.3
        predicted["avg_wait_seconds"] = obs.get("avg_wait_seconds", 30) * 0.8
    elif action == "throughput-first":
        predicted["gpu_util_mean"] = obs.get("gpu_util_mean", 50) * 0.9 + 85.0 * 0.1
        predicted["queue_depth"] = obs.get("queue_depth", 2) * 1.2
        predicted["throughput_jobs_per_min"] = obs.get("throughput_jobs_per_min", 5) * 1.15

    return predicted


def evaluate_action(action: str, obs: dict[str, float], beliefs: dict[str, float],
                    prior: dict[str, float]) -> tuple[float, dict]:
    """Compute expected free energy G(a) = E_Q[F(obs' | a)]."""
    predicted_obs = predict_obs_after_action(obs, action)
    predicted_posterior = update_beliefs(prior, predicted_obs)
    efe = compute_variational_free_energy(predicted_obs, predicted_posterior, prior)
    return efe, {"predicted_gpu_util": predicted_obs.get("gpu_util_mean", 0)}


def choose_action(obs: dict[str, float], beliefs: dict[str, float],
                  prior: dict[str, float]) -> dict:
    """Evaluate all actions and choose minimum expected free energy."""
    actions = ["balanced", "latency-first", "throughput-first"]

    evaluations = []
    for action in actions:
        efe, components = evaluate_action(action, obs, beliefs, prior)
        evaluations.append({"action": action, "efe": efe, "components": components})

    best = min(evaluations, key=lambda x: x["efe"])
    return {
        "chosen_action": best["action"],
        "expected_free_energy": best["efe"],
        "all_evaluations": evaluations,
    }


def apply_action(action: str, db_conn):
    """Apply a chosen action by adjusting brain_params weights."""
    adjustments = {
        "balanced": {"w_queue_depth": 0.5, "w_gpu_saturation": 0.3, "w_node_pressure": 0.2},
        "latency-first": {"w_queue_depth": 0.65, "w_gpu_saturation": 0.2, "w_node_pressure": 0.15},
        "throughput-first": {"w_queue_depth": 0.3, "w_gpu_saturation": 0.55, "w_node_pressure": 0.15},
    }

    weights = adjustments.get(action)
    if not weights:
        return

    now = datetime.now(timezone.utc).strftime("%Y-%m-%dT%H:%M:%S.%fZ")
    db_conn.execute(
        """UPDATE brain_params SET
            w_queue_depth = ?, w_gpu_saturation = ?, w_node_pressure = ?, updated_at = ?
        WHERE key = 'current'""",
        (weights["w_queue_depth"], weights["w_gpu_saturation"], weights["w_node_pressure"], now),
    )
    db_conn.commit()


# --- Background loop ---

async def fep_loop():
    """Background task: run FEP active inference at configured interval."""
    config = get_config()
    if not config.fep_enabled:
        logger.info("FEP controller disabled")
        return

    # Uniform prior
    prior = {s: 0.25 for s in STATES}
    beliefs = {s: 0.25 for s in STATES}

    logger.info("FEP controller started (interval=%.0fs)", config.fep_interval)

    while True:
        try:
            with get_db() as db:
                # Compute local observations
                obs = compute_observation_vector(db)

                # Bayesian update
                beliefs = update_beliefs(prior, obs)

                # Choose action
                decision = choose_action(obs, beliefs, prior)
                fe = compute_variational_free_energy(obs, beliefs, prior)

                # Apply action
                apply_action(decision["chosen_action"], db)

                # Record decision
                latent_state = max(beliefs, key=beliefs.get)
                now = datetime.now(timezone.utc).strftime("%Y-%m-%dT%H:%M:%S.%fZ")
                db.execute(
                    """INSERT INTO fep_decisions (latent_state, state_probs, observations, chosen_action, free_energy, decided_at)
                    VALUES (?, ?, ?, ?, ?, ?)""",
                    (
                        latent_state,
                        json.dumps({s: round(p, 4) for s, p in beliefs.items()}),
                        json.dumps({k: round(v, 4) for k, v in obs.items()}),
                        decision["chosen_action"],
                        round(fe, 4),
                        now,
                    ),
                )

                # Prune old decisions (keep 7 days)
                db.execute(
                    "DELETE FROM fep_decisions WHERE decided_at < datetime('now', '-7 days')"
                )

                logger.debug(
                    "FEP: state=%s (%.2f), action=%s, F=%.2f",
                    latent_state, beliefs[latent_state],
                    decision["chosen_action"], fe,
                )

                # Update prior for next cycle (posterior becomes prior)
                prior = beliefs.copy()

        except Exception:
            logger.exception("FEP controller error")

        await asyncio.sleep(config.fep_interval)
