#!/usr/bin/env python3
"""Small-data action-response model for the active intervention pilot.

The model predicts the paired normalized SLO-goodput treatment effect from a
source measurement and a full MNS/MBBT action.  Telemetry features are direct,
continuous engine measurements; there is no diagnosis-to-action rule here.
"""

from __future__ import annotations

import math
from dataclasses import dataclass
from typing import Any, Iterable, Mapping, Sequence

import numpy as np


PREFIX_FEATURES = (
    "normalized_slo_goodput",
    "admitted_fraction",
    "completed_over_admitted",
    "completed_pass_rate",
    "completed_fail_fraction_of_total",
    "outstanding_over_admitted",
    "ttft_max_over_slo_max",
    "ttft_mean_over_slo_max",
    "tpot_max_over_slo",
    "tpot_mean_over_slo",
    "admitted_input_tokens_mean_over_limit",
)

TELEMETRY_FEATURES = (
    "scheduler_steps_per_s",
    "batch_size_mean",
    "batch_size_cv",
    "batch_tokens_mean",
    "batch_tokens_cv",
    "decode_batch_size_mean",
    "decode_batch_size_cv",
    "prefill_token_fraction",
    "queue_waiting_mean",
    "queue_running_mean",
    "preemptions_per_step",
    "kv_usage_mean",
    "kv_usage_max",
    "kv_usage_end_minus_start",
    "graph_none_share",
    "graph_full_share",
    "graph_padding_fraction",
)


def _finite(value: Any, name: str) -> float:
    result = float(value)
    if not math.isfinite(result):
        raise ValueError(f"{name} must be finite")
    return result


def feature_vector(
    example: Mapping[str, Any], *, include_telemetry: bool
) -> tuple[list[str], np.ndarray]:
    source = example["source"]
    action = example["action"]
    source_log_mns = math.log2(_finite(source["mns"], "source MNS"))
    source_log_mbbt = math.log2(_finite(source["mbbt"], "source MBBT"))
    target_log_mns = math.log2(_finite(action["target_mns"], "target MNS"))
    target_log_mbbt = math.log2(_finite(action["target_mbbt"], "target MBBT"))
    delta_mns = target_log_mns - source_log_mns
    delta_mbbt = target_log_mbbt - source_log_mbbt
    names = [
        "source_log2_mns",
        "source_log2_mbbt",
        "target_log2_mns",
        "target_log2_mbbt",
        "delta_log2_mns",
        "delta_log2_mbbt",
        "delta_product",
        "offered_rate_per_gpu",
    ]
    values = [
        source_log_mns,
        source_log_mbbt,
        target_log_mns,
        target_log_mbbt,
        delta_mns,
        delta_mbbt,
        delta_mns * delta_mbbt,
        _finite(source["offered_rate_per_gpu"], "offered rate"),
    ]
    for name in PREFIX_FEATURES:
        names.append(f"outcome.{name}")
        values.append(_finite(source["outcome"][name], name))
    if include_telemetry:
        for name in TELEMETRY_FEATURES:
            value = _finite(source["telemetry"][name], name)
            names.extend(
                (
                    f"telemetry.{name}",
                    f"telemetry.{name}*delta_mns",
                    f"telemetry.{name}*delta_mbbt",
                )
            )
            values.extend((value, value * delta_mns, value * delta_mbbt))
    vector = np.asarray(values, dtype=np.float64)
    if not np.all(np.isfinite(vector)):
        raise ValueError("feature vector contains a non-finite value")
    return names, vector


@dataclass(frozen=True)
class RidgeModel:
    feature_names: tuple[str, ...]
    mean: np.ndarray
    scale: np.ndarray
    weights: np.ndarray
    intercept: float
    regularization: float

    def predict(self, values: np.ndarray) -> float:
        if values.shape != self.mean.shape:
            raise ValueError("ridge prediction feature shape mismatch")
        normalized = (values - self.mean) / self.scale
        return float(self.intercept + normalized @ self.weights)

    def to_json(self) -> dict[str, Any]:
        return {
            "feature_names": list(self.feature_names),
            "mean": self.mean.tolist(),
            "scale": self.scale.tolist(),
            "weights": self.weights.tolist(),
            "intercept": self.intercept,
            "regularization": self.regularization,
        }

    @classmethod
    def from_json(cls, payload: Mapping[str, Any]) -> "RidgeModel":
        return cls(
            feature_names=tuple(str(value) for value in payload["feature_names"]),
            mean=np.asarray(payload["mean"], dtype=np.float64),
            scale=np.asarray(payload["scale"], dtype=np.float64),
            weights=np.asarray(payload["weights"], dtype=np.float64),
            intercept=float(payload["intercept"]),
            regularization=float(payload["regularization"]),
        )


def fit_ridge(
    examples: Sequence[Mapping[str, Any]],
    *,
    include_telemetry: bool,
    regularization: float,
) -> RidgeModel:
    if not examples:
        raise ValueError("ridge fit requires examples")
    if regularization <= 0:
        raise ValueError("ridge regularization must be positive")
    encoded = [
        feature_vector(example, include_telemetry=include_telemetry)
        for example in examples
    ]
    names = encoded[0][0]
    if any(item[0] != names for item in encoded):
        raise ValueError("feature names changed across examples")
    x = np.stack([item[1] for item in encoded])
    y = np.asarray(
        [
            _finite(example["target_delta_normalized_goodput"], "target effect")
            for example in examples
        ],
        dtype=np.float64,
    )
    mean = x.mean(axis=0)
    scale = x.std(axis=0)
    scale[scale < 1e-12] = 1.0
    normalized = (x - mean) / scale
    intercept = float(y.mean())
    centered = y - intercept
    system = normalized.T @ normalized + regularization * np.eye(x.shape[1])
    weights = np.linalg.solve(system, normalized.T @ centered)
    return RidgeModel(
        feature_names=tuple(names),
        mean=mean,
        scale=scale,
        weights=weights,
        intercept=intercept,
        regularization=regularization,
    )


def fit_jackknife_ensemble(
    examples: Sequence[Mapping[str, Any]],
    *,
    include_telemetry: bool,
    regularization: float,
    group_key: str = "decision_id",
) -> list[RidgeModel]:
    groups = sorted({str(example[group_key]) for example in examples})
    if len(groups) < 3:
        raise ValueError("jackknife ensemble requires at least three groups")
    models = []
    for held_out in groups:
        training = [
            example for example in examples if str(example[group_key]) != held_out
        ]
        models.append(
            fit_ridge(
                training,
                include_telemetry=include_telemetry,
                regularization=regularization,
            )
        )
    return models


def ensemble_predict(
    models: Sequence[RidgeModel],
    example: Mapping[str, Any],
    *,
    include_telemetry: bool,
) -> dict[str, float]:
    if not models:
        raise ValueError("ensemble prediction requires models")
    source = example["source"]
    action = example["action"]
    if (
        int(action["target_mns"]) == int(source["mns"])
        and int(action["target_mbbt"]) == int(source["mbbt"])
    ):
        return {"mean": 0.0, "std": 0.0, "min": 0.0, "max": 0.0, "distinct_n": 1}
    names, values = feature_vector(example, include_telemetry=include_telemetry)
    if any(model.feature_names != tuple(names) for model in models):
        raise ValueError("ensemble feature schema mismatch")
    raw = np.asarray([model.predict(values) for model in models], dtype=np.float64)
    clipped = np.clip(raw, -1.0, 1.0)
    return {
        "mean": float(clipped.mean()),
        "std": float(clipped.std(ddof=0)),
        "min": float(clipped.min()),
        "max": float(clipped.max()),
        "distinct_n": len(set(float(value) for value in clipped)),
    }


def select_action(
    models: Sequence[RidgeModel],
    candidates: Sequence[Mapping[str, Any]],
    *,
    include_telemetry: bool,
    confidence_z: float = 1.0,
    minimum_margin: float = 0.02,
) -> dict[str, Any]:
    if len(candidates) < 2:
        raise ValueError("action selection requires at least two candidates")
    rows = []
    for example in candidates:
        prediction = ensemble_predict(
            models, example, include_telemetry=include_telemetry
        )
        rows.append(
            {
                "action_id": str(example["action"]["id"]),
                "prediction": prediction,
                "lower": prediction["mean"] - confidence_z * prediction["std"],
                "upper": prediction["mean"] + confidence_z * prediction["std"],
            }
        )
    rows.sort(key=lambda row: (-row["prediction"]["mean"], row["action_id"]))
    best, second = rows[:2]
    margin = float(best["prediction"]["mean"] - second["prediction"]["mean"])
    confident = bool(
        margin >= minimum_margin and best["lower"] > second["upper"]
    )
    return {
        "selected_action": best["action_id"],
        "confident": confident,
        "predicted_margin": margin,
        "candidates": rows,
    }


def models_to_json(models: Iterable[RidgeModel]) -> list[dict[str, Any]]:
    return [model.to_json() for model in models]


def models_from_json(payload: Iterable[Mapping[str, Any]]) -> list[RidgeModel]:
    return [RidgeModel.from_json(item) for item in payload]