"""Matched pilot work units, including masked tails and end-to-end estimates."""
from __future__ import annotations

import math
from typing import Protocol, TypedDict

from .loop import BatchPlan, Run, lr_multiplier

QDQ_SEMANTICS = "official act_qdq before dense/expert, output dtype rounding, identity STE backward; additive FP32 LoRA before output rounding; raw uncached taps"


class TokenCounts(Protocol):
    def __len__(self) -> int: ...
    def supervised_tokens(self, index: int) -> int: ...


class Units(TypedDict):
    blocks: int
    valid_supervised_tokens: int
    optimizer_updates: int
    accumulation: int
    batch_ceiling: int
    lr_schedule: str
    learning_rates: list[float]
    qdq_semantics: str


def work_units(data: TokenCounts, plan: BatchPlan, run: Run, step: int | None = None) -> Units:
    stop = len(plan.batches) if step is None else step
    if not 0 < stop <= len(plan.batches):
        raise ValueError("invalid matched-unit checkpoint")
    counts = [data.supervised_tokens(i) for batch in plan.batches[:stop] for i in batch]
    if any(type(n) is not int or not 1 <= n <= 5 for n in counts):
        raise ValueError("invalid supervised position count")
    return {"blocks": len(counts), "valid_supervised_tokens": sum(counts), "optimizer_updates": stop,
            "accumulation": 1, "batch_ceiling": run.batch_size,
            "lr_schedule": ("linear-warmup-10pct-then-cosine; " if run.warmup_fraction else "")+"cosine-per-optimizer-update; initial lr then scheduler.step after update; total="+str(len(plan.batches)),
            "learning_rates": [run.lr*lr_multiplier(run, s, len(plan.batches)) for s in range(stop)],
            "qdq_semantics": QDQ_SEMANTICS}


def require_matched(arms: dict[str, dict[str, Units]]) -> None:
    if set(arms) != {"REF", "FAST", "FAST-B"} or not arms["REF"]:
        raise ValueError("matched units require all arms/checkpoints")
    if arms["REF"] != arms["FAST"] or arms["REF"] != arms["FAST-B"]:
        raise ValueError("pilot blocks/tokens/updates/accumulation/LR/QDQ units differ")


class Estimate(TypedDict):
    label: str
    training_gpu_hours: float
    evaluation_gpu_hours: float
    replay_gpu_hours: float
    serving_gpu_hours: float
    total_gpu_hours: float
    ideal_two_solo_wall_hours: float
    training_blocks: int
    optimizer_updates: int
    valid_supervised_tokens: int


def estimate(units: Units, blocks_per_second: dict[str, float], evaluation_rounds: int,
             eval_seconds_per_round: float, replay_gpu_seconds: float, serving_gpu_seconds: float) -> Estimate:
    costs = [*blocks_per_second.values(), eval_seconds_per_round, replay_gpu_seconds, serving_gpu_seconds]
    if (set(blocks_per_second) != {"REF", "FAST", "FAST-B"} or any(not math.isfinite(x) or x <= 0 for x in costs)
            or evaluation_rounds < 1):
        raise ValueError("positive measured/declared costs required including replay and serving qualification")
    training = sum(units["blocks"]/rate for rate in blocks_per_second.values())/3600
    evaluation = 13*evaluation_rounds*eval_seconds_per_round/3600
    replay, serving = replay_gpu_seconds/3600, serving_gpu_seconds/3600
    total = training+evaluation+replay+serving
    return {"label": "ESTIMATED; additive costs, perfect two-solo balance lower bound; excludes load/I/O/queue overhead",
            "training_gpu_hours": training, "evaluation_gpu_hours": evaluation, "replay_gpu_hours": replay,
            "serving_gpu_hours": serving, "total_gpu_hours": total, "ideal_two_solo_wall_hours": total/2,
            "training_blocks": 3*units["blocks"], "optimizer_updates": 3*units["optimizer_updates"],
            "valid_supervised_tokens": 3*units["valid_supervised_tokens"]}