"""Exact fixed 30B-token data mixture for NanoCoder and MiniCoder."""

from __future__ import annotations


TOTAL_TOKENS = 30_000_000_000
SOURCE_KEYS = ("stack_v3", "fineweb_edu")
TARGET_SHARES = {"stack_v3": 0.75, "fineweb_edu": 0.25}
TARGET_SOURCE_TOKENS = {
    "stack_v3": 22_500_000_000,
    "fineweb_edu": 7_500_000_000,
}
SOURCES = {
    "stack_v3": {
        "label": "The Stack v3 train",
        "dataset_id": "HuggingFaceCode/stack-v3-train",
        "config_name": None,
        "kind": "repository",
    },
    "fineweb_edu": {
        "label": "FineWeb-Edu",
        "dataset_id": "HuggingFaceFW/fineweb-edu",
        "config_name": "sample-100BT",
        "kind": "text",
    },
}


class TokenCreditScheduler:
    """Deterministic weighted-fair scheduling in supervised tokens."""

    def __init__(self):
        self.credits = {key: 0.0 for key in SOURCE_KEYS}
        self.consumed = {key: 0 for key in SOURCE_KEYS}

    def choose(self, tokens_this_step: int) -> str:
        for key in SOURCE_KEYS:
            self.credits[key] += TARGET_SHARES[key] * tokens_this_step
        selected = max(SOURCE_KEYS, key=lambda key: self.credits[key])
        self.credits[selected] -= tokens_this_step
        self.consumed[selected] += tokens_this_step
        return selected

    def state_dict(self) -> dict[str, dict[str, float | int]]:
        return {
            "credits": dict(self.credits),
            "consumed": dict(self.consumed),
        }

    def load_state_dict(self, state: dict) -> None:
        if set(state["credits"]) != set(SOURCE_KEYS):
            raise ValueError("Checkpoint curriculum keys do not match")
        self.credits = {
            key: float(state["credits"][key]) for key in SOURCE_KEYS
        }
        self.consumed = {
            key: int(state["consumed"][key]) for key in SOURCE_KEYS
        }


__all__ = [
    "SOURCES",
    "SOURCE_KEYS",
    "TARGET_SHARES",
    "TARGET_SOURCE_TOKENS",
    "TOTAL_TOKENS",
    "TokenCreditScheduler",
]