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

import importlib.util
import math
from pathlib import Path
from types import SimpleNamespace


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


def load_module():
    spec = importlib.util.spec_from_file_location(
        "intervention_response_phase_aware_v2", HERE / "analyze_existing.py"
    )
    module = importlib.util.module_from_spec(spec)
    assert spec.loader is not None
    spec.loader.exec_module(module)
    return module


def load_prepare_module():
    spec = importlib.util.spec_from_file_location(
        "intervention_response_phase_aware_prepare", HERE / "prepare_pilot.py"
    )
    module = importlib.util.module_from_spec(spec)
    assert spec.loader is not None
    spec.loader.exec_module(module)
    return module


def load_controller_module():
    spec = importlib.util.spec_from_file_location(
        "intervention_response_phase_aware_controller", HERE / "pilot_controller.py"
    )
    module = importlib.util.module_from_spec(spec)
    assert spec.loader is not None
    spec.loader.exec_module(module)
    return module


def load_pilot_analysis_module():
    spec = importlib.util.spec_from_file_location(
        "intervention_response_phase_aware_pilot_analysis", HERE / "analyze_pilot.py"
    )
    module = importlib.util.module_from_spec(spec)
    assert spec.loader is not None
    spec.loader.exec_module(module)
    return module


def main() -> None:
    module = load_module()
    assert module.common_decile_fractions(
        trace_duration_s=60.0, minimum_elapsed_s=19.448
    ) == (0.1, 0.2, 0.3)
    assert module.common_decile_fractions(
        trace_duration_s=60.0, minimum_elapsed_s=60.0
    )[-1] == 1.0
    stats = module.numeric([0.0, 1.0, 2.0])
    assert stats == {
        "n": 3,
        "min": 0.0,
        "max": 2.0,
        "distinct_n": 3,
        "median": 1.0,
    }
    assert math.isclose(module._pearson([1.0, 2.0], [2.0, 4.0]), 1.0)
    assert module._pearson([1.0, 1.0], [2.0, 3.0]) is None
    prepare = load_prepare_module()
    requests = [
        SimpleNamespace(
            sampling_u=index / 10.0,
            row_id=f"r{index}",
            arrival_s=float(index),
            prompt_tokens_hint=100 + index,
        )
        for index in range(1, 6)
    ]
    _anchor, selected = prepare.attainable_anchor(requests, 3)
    assert len(selected) == 3
    record = prepare.selection_record(selected, duration_s=3.0)
    assert record["selected_count"] == 3
    assert record["offered_req_s_per_gpu"] == 0.25
    assert len(prepare.SESSION_ORDER) == 6
    assert {mns for _replicate, mns in prepare.SESSION_ORDER} == {16, 64}
    first_id = prepare.private_request_id(
        source_sha256="a" * 64, line_number=1, original_id="1"
    )
    assert first_id == prepare.private_request_id(
        source_sha256="a" * 64, line_number=1, original_id="1"
    )
    assert first_id != prepare.private_request_id(
        source_sha256="b" * 64, line_number=1, original_id="1"
    )
    controller = load_controller_module()
    assert math.isclose(controller.remaining_projection(6, 0), 7.7)
    assert math.isclose(controller.remaining_projection(6, 5), 1.45)
    parsed = controller.parser().parse_args(
        [
            "--manifest",
            "/tmp/manifest.json",
            "--run-root",
            "/tmp/run",
            "--aituner-root",
            "/tmp/aituner",
            "--vllm-source",
            "/tmp/vllm",
            "--venv",
            "/tmp/venv",
            "--model",
            "/tmp/model",
            "--client",
            "/tmp/client.py",
            "--dry-run",
        ]
    )
    assert parsed.dry_run is True
    pilot_analysis = load_pilot_analysis_module()
    stable = pilot_analysis.stable_adjacent_features(
        [
            {"end_fraction": 0.1, "qualifying_response_features": ["queue"]},
            {
                "end_fraction": 0.25,
                "qualifying_response_features": ["kv", "queue"],
            },
            {"end_fraction": 0.5, "qualifying_response_features": ["queue"]},
        ]
    )
    assert stable == {"0.10->0.25": ["queue"], "0.25->0.50": ["queue"]}
    load_consistency = {
        "0.10->0.25:queue": {"passes_two_regimes": True},
        "0.25->0.50:queue": {"passes_two_regimes": True},
    }
    mechanism = pilot_analysis.mechanism_gate(stable, load_consistency)
    assert mechanism["passes"] is False
    stable["0.25->0.50"].append("kv")
    load_consistency["0.25->0.50:kv"] = {"passes_two_regimes": True}
    mechanism = pilot_analysis.mechanism_gate(stable, load_consistency)
    assert mechanism["passes"] is True
    assert mechanism["passing_transitions"] == ["0.25->0.50"]
    efficacy = pilot_analysis.stable_adjacent_efficacy_features(
        [
            {
                "end_fraction": 0.1,
                "efficacy": {"telemetry_qualifying_features": ["early"]},
            },
            {
                "end_fraction": 0.25,
                "efficacy": {"telemetry_qualifying_features": ["queue"]},
            },
            {
                "end_fraction": 0.5,
                "efficacy": {"telemetry_qualifying_features": ["kv", "queue"]},
            },
        ]
    )
    assert efficacy == {"0.25->0.50": ["queue"]}
    coverage = pilot_analysis.telemetry_coverage(
        [
            {"step_index": 1, "submit_mono_ns": 100_000_000},
            {"step_index": 2, "submit_mono_ns": 200_000_000},
        ],
        start_ns=0,
        end_ns=300_000_000,
    )
    assert coverage == {
        "start_gap_s": 0.1,
        "end_gap_s": 0.1,
        "max_internal_gap_s": 0.1,
    }
    coverage_gate = pilot_analysis.cumulative_coverage_gate(
        [
            {
                "trial_sanity": [
                    {
                        "trial_id": "a",
                        "admitted_fraction": 0.25,
                        "completed_fraction": 0.2,
                    }
                ]
            },
            {
                "trial_sanity": [
                    {
                        "trial_id": "a",
                        "admitted_fraction": 0.5,
                        "completed_fraction": 0.4,
                    }
                ]
            },
        ]
    )
    assert coverage_gate["red_flags"] == []
    print("phase-aware intervention response v2 analysis: PASS")


if __name__ == "__main__":
    main()