#!/usr/bin/env python3
"""Run Frontier with the existing joint repair plus a whole-layer decode curve."""

from __future__ import annotations

import json
import os
import runpy
from pathlib import Path


ROOT = Path(__file__).resolve().parent
REPO = ROOT.parents[1]
GRID = json.loads(Path(os.environ["FRONTIER_WHOLE_DECODE_GRID"]).read_text())
WHOLE_CURVE = {
    str(tp): {
        str(cell["batch"]): float(cell["median_execute_ms"])
        for cell in GRID["cells"]
        if int(cell["tp"]) == tp
    }
    for tp in (2, 4)
}

# Apply the existing serving-path collective/MoE correction first.
joint = runpy.run_path(
    str(REPO / "runs/frontier-collective-joint-v0/run_frontier_with_curves.py")
)
USAGE = joint["USAGE"]
_pure_decode_point = joint["_pure_decode_point"]

from frontier.entities import ExecutionTime  # noqa: E402
from frontier.execution_time_predictor.sklearn_moe_execution_time_predictor import (  # noqa: E402
    SklearnMoEExecutionTimePredictor,
)


_JOINT_STAGE_PREDICT = SklearnMoEExecutionTimePredictor.predict_stage_execution_time


def _whole_decode_stage_time(
    self,
    batch,
    stage_id,
    cluster_type,
    num_layers=1,
    layer_id=0,
):
    base = _JOINT_STAGE_PREDICT(
        self,
        batch,
        stage_id,
        cluster_type,
        num_layers=num_layers,
        layer_id=layer_id,
    )
    point = _pure_decode_point(self, batch)
    if point is None or point[0] == "1" or point[1] == "1":
        return base
    tp, decode_batch = point
    if tp not in WHOLE_CURVE or decode_batch not in WHOLE_CURVE[tp]:
        raise ValueError(
            "Whole-layer curve has no exact pure-decode point for "
            f"TP={tp}, batch={decode_batch}; refusing to extrapolate"
        )
    layers = int(base.num_layers)
    if layers <= 0:
        raise ValueError(f"invalid stage layer count: {layers}")
    target_ms = WHOLE_CURVE[tp][decode_batch]
    USAGE[f"whole_decode:tp{tp}-b{decode_batch}:target_ms={target_ms:.9f}"] += 1
    return ExecutionTime(
        num_layers_per_pipeline_stage=layers,
        attention_rope_execution_time=0.0,
        attention_kv_cache_save_execution_time=0.0,
        attention_decode_execution_time=0.0,
        attention_prefill_execution_time=0.0,
        attention_layer_pre_proj_execution_time=0.0,
        attention_layer_post_proj_execution_time=0.0,
        attn_norm_time=0.0,
        mlp_norm_time=0.0,
        add_time=0.0,
        tensor_parallel_communication_time=0.0,
        pipeline_parallel_communication_time=0.0,
        expert_parallel_communication_time=0.0,
        moe_gating_time=0.0,
        moe_shuffling_time=0.0,
        schedule_time=base._schedule_time,
        sampler_e2e_time=base._sampler_e2e_time,
        prepare_inputs_e2e_time=base._prepare_inputs_e2e_time,
        process_model_outputs_time=base._process_model_outputs_time,
        ray_comm_time=base._ray_comm_time,
        is_moe=True,
        moe_grouped_gemm_time=target_ms / layers,
        pp_producer_send_path_runtime_time=(
            base._pp_producer_send_path_runtime_time
        ),
        pp_receiver_head_runtime_time=base._pp_receiver_head_runtime_time,
        pp_prefill_consumer_active_runtime_time=(
            base._pp_prefill_consumer_active_runtime_time
        ),
        pp_stage_boundary_residual_runtime_time=(
            base._pp_stage_boundary_residual_runtime_time
        ),
    )


SklearnMoEExecutionTimePredictor.predict_stage_execution_time = (
    _whole_decode_stage_time
)

from frontier.main import main  # noqa: E402


if __name__ == "__main__":
    main()