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

import json
import tempfile
from pathlib import Path

import numpy as np

from analyze_prefixes import PrefixExample
from prepare_pilot_simulator import load_module as load_prepare_module
from run_pilot_simulator import load_module as load_run_module
from analyze_strong_pilot import (
    covariate_shift,
    fit_model,
    load_pilot_simulator,
    predict_model,
)


def example(index: int) -> PrefixExample:
    label = int(index >= 4)
    return PrefixExample(
        cell=f"cell-{index // 2}",
        anchor=float(index),
        cutoff_s=5.0,
        tp=1,
        full_elapsed_s=10.0,
        feasible=label,
        primary_feasible=label,
        outcome=tuple(float(index + offset) for offset in range(13)),
        instrumentation=tuple(float(index * offset + 1) for offset in range(17)),
        completion_time_source="exact_monotonic",
    )


def main() -> None:
    examples = [example(index) for index in range(8)]
    simulator = [(float(index), index / 10.0, float(index >= 4)) for index in range(8)]
    for instrumentation_aware in (False, True):
        model = fit_model(
            examples,
            simulator,
            instrumentation_aware=instrumentation_aware,
            regularization=1.0,
        )
        probability = predict_model(model, examples, simulator)
        assert probability.shape == (8,)
        assert np.all((probability >= 0.0) & (probability <= 1.0))
        shift = covariate_shift(
            examples,
            simulator,
            examples,
            simulator,
            instrumentation_aware=instrumentation_aware,
        )
        assert shift["values"]["min"] >= 0.0
        assert shift["count_gt_3"] == 0

    payload = {
        "status": "PASS",
        "results": [
            {
                "cell": f"cell-{index // 2}",
                "role": "low1" if index % 2 == 0 else "high1",
                "scorer": {
                    "throughput_requests_per_second_per_gpu": 1.0 + index,
                    "slo": {
                        "pass_rate": index / 12.0,
                        "feasible": index % 2 == 0,
                    },
                },
            }
            for index in range(12)
        ],
    }
    with tempfile.TemporaryDirectory() as temporary:
        path = Path(temporary) / "metrics.json"
        path.write_text(json.dumps(payload), encoding="utf-8")
        features, red_flags = load_pilot_simulator(path)
    assert len(features) == 12
    assert red_flags == []

    with tempfile.TemporaryDirectory() as temporary:
        root = Path(temporary)
        (root / "prepare_dependency.py").write_text("VALUE = 17\n", encoding="utf-8")
        (root / "prepare_target.py").write_text(
            "from prepare_dependency import VALUE\n", encoding="utf-8"
        )
        assert load_prepare_module(root / "prepare_target.py").VALUE == 17
        (root / "run_dependency.py").write_text("VALUE = 23\n", encoding="utf-8")
        (root / "run_target.py").write_text(
            "from run_dependency import VALUE\n", encoding="utf-8"
        )
        assert load_run_module("run_target", root / "run_target.py").VALUE == 23
    print("fidelity strong pilot: PASS")


if __name__ == "__main__":
    main()