#!/usr/bin/env python3
from __future__ import annotations

import importlib.util
import sys
from pathlib import Path


HERE = Path(__file__).resolve().parent


def load_model():
    spec = importlib.util.spec_from_file_location(
        "active_intervention_model", HERE / "model.py"
    )
    module = importlib.util.module_from_spec(spec)
    assert spec.loader is not None
    sys.modules[spec.name] = module
    spec.loader.exec_module(module)
    return module


def example(model, decision: str, action: str, pressure: float, target: float):
    outcome = {
        name: 0.5 for name in model.PREFIX_FEATURES
    }
    telemetry = {name: 0.0 for name in model.TELEMETRY_FEATURES}
    telemetry["queue_waiting_mean"] = pressure
    telemetry["batch_size_mean"] = pressure
    return {
        "decision_id": decision,
        "source": {
            "mns": 16,
            "mbbt": 8192,
            "offered_rate_per_gpu": 2.0,
            "outcome": outcome,
            "telemetry": telemetry,
        },
        "action": {
            "id": action,
            "target_mns": 64 if action == "mns" else 16,
            "target_mbbt": 8192 if action == "mns" else 16384,
        },
        "target_normalized_goodput": target,
        "target_delta_normalized_goodput": target - 0.5,
    }


def main() -> None:
    model = load_model()
    examples = []
    for index, pressure in enumerate((0.2, 0.5, 0.8), 1):
        examples.extend(
            (
                example(model, f"d{index}", "mns", pressure, 0.5 + pressure / 2),
                example(model, f"d{index}", "mbbt", pressure, 0.6 - pressure / 4),
            )
        )
    fitted = model.fit_ridge(
        examples, include_telemetry=True, regularization=1.0
    )
    encoded = fitted.to_json()
    restored = model.RidgeModel.from_json(encoded)
    names, values = model.feature_vector(examples[-2], include_telemetry=True)
    assert tuple(names) == restored.feature_names
    assert abs(fitted.predict(values) - restored.predict(values)) < 1e-12
    ensemble = model.fit_jackknife_ensemble(
        examples, include_telemetry=True, regularization=1.0
    )
    decision = model.select_action(
        ensemble, examples[-2:], include_telemetry=True, minimum_margin=0.0
    )
    assert decision["selected_action"] == "mns"
    assert all(-1.0 <= row["prediction"]["mean"] <= 1.0 for row in decision["candidates"])
    noop = example(model, "noop", "noop", 0.8, 0.5)
    noop["action"]["target_mbbt"] = 8192
    prediction = model.ensemble_predict(
        ensemble, noop, include_telemetry=True
    )
    assert prediction == {
        "mean": 0.0,
        "std": 0.0,
        "min": 0.0,
        "max": 0.0,
        "distinct_n": 1,
    }
    print("active intervention model: PASS")


if __name__ == "__main__":
    main()