#!/usr/bin/env python3
"""Run Frontier with pure-decode serving-path MoE and/or collective curves."""

from __future__ import annotations

import atexit
import csv
import json
import os
from collections import Counter
from pathlib import Path


csv.field_size_limit(16 * 1024 * 1024)


COLLECTIVE_PAYLOAD = json.loads(
    Path(os.environ["FRONTIER_COLLECTIVE_CURVE"]).read_text()
)
COLLECTIVE_VARIANT = os.environ.get("FRONTIER_COLLECTIVE_CURVE_VARIANT", "drop_mean")
if "curve_variants_ms_per_step" in COLLECTIVE_PAYLOAD:
    try:
        COLLECTIVE_CURVE = COLLECTIVE_PAYLOAD["curve_variants_ms_per_step"][
            COLLECTIVE_VARIANT
        ]
    except KeyError as error:
        raise ValueError(
            f"collective curve has no variant {COLLECTIVE_VARIANT!r}"
        ) from error
else:
    if COLLECTIVE_VARIANT != "drop_mean":
        raise ValueError("legacy collective curve only supports drop_mean")
    COLLECTIVE_CURVE = COLLECTIVE_PAYLOAD.get(
        "curve_ms_per_step", COLLECTIVE_PAYLOAD
    )
MOE_CURVE_PATH = os.environ.get("FRONTIER_FUSED_MOE_CURVE")
MOE_CURVE = (
    json.loads(Path(MOE_CURVE_PATH).read_text()) if MOE_CURVE_PATH else None
)
USAGE_PATH = Path(os.environ["FRONTIER_CURVE_USAGE"])
USAGE: Counter[str] = Counter()


def write_usage() -> None:
    USAGE_PATH.parent.mkdir(parents=True, exist_ok=True)
    USAGE_PATH.write_text(json.dumps(dict(sorted(USAGE.items())), indent=2) + "\n")


atexit.register(write_usage)

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


_ORIGINAL_ATTN_COLLECTIVE = (
    SklearnMoEExecutionTimePredictor._get_tensor_parallel_communication_time
)
_ORIGINAL_MOE_COLLECTIVE = (
    SklearnMoEExecutionTimePredictor._get_moe_tensor_parallel_allreduce_time
)
_ORIGINAL_GROUPED_GEMM = SklearnMoEExecutionTimePredictor._get_grouped_gemm_time
_ORIGINAL_ATTN_NORM = (
    SklearnMoEExecutionTimePredictor._get_attn_norm_layer_act_execution_time
)
_ORIGINAL_MLP_NORM = (
    SklearnMoEExecutionTimePredictor._get_mlp_norm_layer_act_execution_time
)


def _pure_decode_point(self, batch) -> tuple[str, str] | None:
    if batch is None or not bool(getattr(batch, "is_pure_decode_batch", False)):
        return None
    tp = str(int(self._replica_config.moe_tensor_parallel_size))
    attn_tp = str(int(self._replica_config.attn_tensor_parallel_size))
    if tp != attn_tp:
        raise ValueError(
            "Serving collective curve requires equal attention/MoE TP, got "
            f"attn_tp={attn_tp}, moe_tp={tp}"
        )
    return tp, str(len(batch.requests))


def _collective_path_time(self, batch, path: str) -> float:
    point = _pure_decode_point(self, batch)
    if point is None:
        original = (
            _ORIGINAL_ATTN_COLLECTIVE
            if path == "attention"
            else _ORIGINAL_MOE_COLLECTIVE
        )
        return original(self, batch)
    tp, decode_batch = point
    if tp == "1":
        original = (
            _ORIGINAL_ATTN_COLLECTIVE
            if path == "attention"
            else _ORIGINAL_MOE_COLLECTIVE
        )
        value = original(self, batch)
        if value != 0.0:
            raise ValueError(f"TP1 {path} collective must be zero, got {value}")
        USAGE[f"collective:{path}:tp1-b{decode_batch}:structural-zero"] += 1
        return value
    if tp not in COLLECTIVE_CURVE or decode_batch not in COLLECTIVE_CURVE[tp]:
        raise ValueError(
            "Collective curve has no exact pure-decode point for "
            f"TP={tp}, batch={decode_batch}; refusing to extrapolate"
        )
    layers = int(self._num_layers_per_pipeline_stage)
    if layers <= 0:
        raise ValueError(f"invalid layers per pipeline stage: {layers}")
    USAGE[f"collective:{path}:tp{tp}-b{decode_batch}"] += 1
    # The serving curve is the sum of attention and MoE output reductions over
    # the full model step. ExecutionTime later multiplies each per-layer path.
    return float(COLLECTIVE_CURVE[tp][decode_batch]) / (2 * layers)


def _attention_collective_from_curve(self, batch) -> float:
    return _collective_path_time(self, batch, "attention")


def _moe_collective_from_curve(self, batch) -> float:
    return _collective_path_time(self, batch, "moe")


def _grouped_gemm_from_curve(self, num_tokens_or_allocation, batch=None) -> float:
    if MOE_CURVE is None:
        return _ORIGINAL_GROUPED_GEMM(
            self, num_tokens_or_allocation, batch=batch
        )
    point = _pure_decode_point(self, batch)
    if point is None:
        return _ORIGINAL_GROUPED_GEMM(
            self, num_tokens_or_allocation, batch=batch
        )
    tp, decode_batch = point
    if tp not in MOE_CURVE or decode_batch not in MOE_CURVE[tp]:
        raise ValueError(
            "Fused MoE curve has no exact pure-decode point for "
            f"TP={tp}, batch={decode_batch}; refusing to extrapolate"
        )
    layers = int(self._num_layers_per_pipeline_stage)
    if layers <= 0:
        raise ValueError(f"invalid layers per pipeline stage: {layers}")
    USAGE[f"moe:tp{tp}-b{decode_batch}"] += 1
    return float(MOE_CURVE[tp][decode_batch]) / layers


def _norm_without_fused_duplicate(self, batch, *, name: str) -> float:
    original = _ORIGINAL_ATTN_NORM if name == "attn" else _ORIGINAL_MLP_NORM
    value = original(self, batch)
    point = _pure_decode_point(self, batch)
    if point is None or point[0] == "1":
        return value
    if value < 0:
        raise ValueError(f"negative {name} norm prediction cannot be deducted: {value}")
    tp, decode_batch = point
    USAGE[f"fused_norm_deduction:{name}:tp{tp}-b{decode_batch}"] += 1
    # AllReduceFusionPattern 1 already contains residual add + RMSNorm. Returning
    # zero removes the exact predictor value that would otherwise be emitted in
    # the frozen Frontier attn_norm_time/mlp_norm_time ledger row.
    return 0.0


def _attn_norm_without_fused_duplicate(self, batch) -> float:
    return _norm_without_fused_duplicate(self, batch, name="attn")


def _mlp_norm_without_fused_duplicate(self, batch) -> float:
    return _norm_without_fused_duplicate(self, batch, name="mlp")


SklearnMoEExecutionTimePredictor._get_tensor_parallel_communication_time = (
    _attention_collective_from_curve
)
SklearnMoEExecutionTimePredictor._get_moe_tensor_parallel_allreduce_time = (
    _moe_collective_from_curve
)
SklearnMoEExecutionTimePredictor._get_grouped_gemm_time = _grouped_gemm_from_curve
SklearnMoEExecutionTimePredictor._get_attn_norm_layer_act_execution_time = (
    _attn_norm_without_fused_duplicate
)
SklearnMoEExecutionTimePredictor._get_mlp_norm_layer_act_execution_time = (
    _mlp_norm_without_fused_duplicate
)

from frontier.main import main  # noqa: E402


if __name__ == "__main__":
    main()