"""Compose: turn one prepared corpus into one or more training mixes that fit a token budget.

Defaults come from our own 1B-token study (huggingface.co/blog/codelion/optimal-dataset-mixing): a static mix of
50% textbook-style text, 30% filtered web and 20% educational web beat every curriculum we tried at 70M parameters
and 1B tokens. That is one model size and one budget, so treat it as a starting point and fork it.

When a source is too small for its share, it is repeated up to `max_epochs` (about four epochs of repetition cost
little in Muennighoff et al., 2023). Past that, the gap is filled by rephrasing that source's best documents,
up to `max_synthetic_share` of the mix, since a mix that leans on synthetic text generalized worst in our study.
"""
from __future__ import annotations

GROUP_OF_KIND = {"textbook": "textbook", "synthetic": "textbook", "web": "web", "multilingual": "web",
                 "educational": "educational", "math": "math", "code": "code", "other": "web"}
BASE_SHARES = {"textbook": 0.50, "web": 0.30, "educational": 0.20}


def default_weights(sources: list[dict]) -> dict[str, float]:
    """Share per source name. Math and code, when present, take 10% each off the top; the rest follows 50/30/20
    over the groups that exist, split evenly inside each group."""
    groups: dict[str, list[str]] = {}
    for s in sources:
        groups.setdefault(GROUP_OF_KIND.get(s.get("kind", "web"), "web"), []).append(s["name"])
    shares: dict[str, float] = {}
    fixed = {g: 0.10 for g in ("math", "code") if g in groups}
    rest = 1.0 - sum(fixed.values())
    base = {g: v for g, v in BASE_SHARES.items() if g in groups}
    if not base:
        # only math and/or code sources: split evenly
        return {s["name"]: round(1 / len(sources), 4) for s in sources}
    tot = sum(base.values())
    for g, v in base.items():
        shares[g] = rest * v / tot
    shares.update(fixed)
    out = {}
    for g, names in groups.items():
        for n in names:
            out[n] = round(shares[g] / len(names), 4)
    return out


def plan_mix(budget_tokens: int, sources: list[dict], weights: dict[str, float], max_epochs: float = 4.0,
             max_synthetic_share: float = 0.25, rephrase_yield: float = 4.0, avg_out_tokens: int = 700) -> dict:
    """sources: [{name, kind, available_tokens (after cleaning and dedup), high_tokens (top tier)}].
    Returns per-source rows and totals, including how much synthetic text to generate and roughly how many calls."""
    tot_w = sum(weights.get(s["name"], 0) for s in sources) or 1.0
    rows, synth_total, short_total = [], 0, 0
    for s in sources:
        share = weights.get(s["name"], 0) / tot_w
        want = int(budget_tokens * share)
        avail = max(0, int(s.get("available_tokens", 0)))
        fresh = min(want, avail)
        repeat_cap = int(avail * max_epochs)
        from_repeat = min(want, repeat_cap)
        gap = want - from_repeat
        synth = min(gap, int(s.get("high_tokens", avail) * rephrase_yield)) if gap > 0 else 0
        short = gap - synth
        epochs = round(from_repeat / avail, 2) if avail else 0.0
        rows.append({"name": s["name"], "kind": s.get("kind", "web"), "share": round(share, 4), "want_tokens": want,
                     "available_tokens": avail, "fresh_tokens": fresh, "epochs": epochs, "synthetic_tokens": synth,
                     "short_tokens": short})
        synth_total += synth
        short_total += short
    cap = int(budget_tokens * max_synthetic_share)
    warnings = []
    if synth_total > cap:
        warnings.append(f"The plan needs {fmt(synth_total)} synthetic tokens, over the {int(max_synthetic_share * 100)}% cap "
                        f"({fmt(cap)}). Add sources, lower the budget, or raise the cap knowingly.")
    if short_total:
        warnings.append(f"{fmt(short_total)} tokens cannot be filled even with repetition and rephrasing.")
    for r in rows:
        if r["epochs"] > 1.0:
            warnings.append(f"{r['name']} repeats {r['epochs']}x to reach its share.")
    return {"budget_tokens": budget_tokens, "rows": rows, "synthetic_tokens": synth_total,
            "synthetic_calls": synth_total // max(1, avg_out_tokens), "short_tokens": short_total, "warnings": warnings}


def fmt(n: float) -> str:
    n = float(n)
    for unit, div in (("T", 1e12), ("B", 1e9), ("M", 1e6), ("K", 1e3)):
        if abs(n) >= div:
            return f"{n / div:.3g}{unit}"
    return f"{int(n)}"