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

import math
import sys
from pathlib import Path


HERE = Path(__file__).resolve().parent
sys.path.insert(0, str(HERE))

import analyze_prefixes as analysis  # noqa: E402


def main() -> None:
    exact, exact_source = analysis.completion_elapsed_s(
        {"completed_elapsed_s": 7.25}
    )
    assert exact == 7.25 and exact_source == "exact_monotonic"

    reconstructed, reconstructed_source = analysis.completion_elapsed_s(
        {
            "success": True,
            "arrival_s": 2.0,
            "ttft_ms": 100.0,
            "tpot_ms": 10.0,
            "completion_tokens": 11,
        }
    )
    assert math.isclose(reconstructed or 0.0, 2.2)
    assert reconstructed_source == "reconstructed_from_latency"
    missing, missing_source = analysis.completion_elapsed_s({"success": False})
    assert missing is None and missing_source == "unobserved_failure"

    examples = [
        analysis.PrefixExample(
            cell=f"c{index}",
            anchor=float(index),
            cutoff_s=5.0,
            tp=1,
            full_elapsed_s=65.0,
            feasible=label,
            primary_feasible=label,
            outcome=(float(index),),
            instrumentation=(float(index % 2),),
            completion_time_source="exact_monotonic",
        )
        for index, label in enumerate((0, 1, 1))
    ]
    labels = analysis.np.asarray([0, 1, 1])
    probabilities = analysis.np.asarray([0.01, 0.99, 0.60])
    policy = analysis.policy_metrics(examples, labels, probabilities, 0.95)
    assert policy["early_accept"] == 1
    assert policy["early_reject"] == 1
    assert policy["abstain_continue_full"] == 1
    assert policy["false_accept"] == 0 and policy["false_reject"] == 0
    assert policy["valid_zero_error_policy"]
    assert policy["valid_cost_reduction_fraction"] is not None
    model = analysis.fit_frozen_model(
        examples,
        instrumentation_aware=True,
        regularization=1.0,
    )
    frozen_probability = analysis.predict_frozen_model(model, examples)
    assert len(frozen_probability) == len(examples)
    assert analysis.np.all(frozen_probability >= 0.0)
    assert analysis.np.all(frozen_probability <= 1.0)
    print("fidelity prefix analysis: PASS")


if __name__ == "__main__":
    main()