"""Multi-source training data (``--records path[:weight]``, repeatable): parsing, per-source
split handling and reloads, and the per-micro-batch source shares (CPU, no model)."""

from __future__ import annotations

import json
import random
from collections import Counter
from pathlib import Path

import pytest

pytest.importorskip("torch")

from deg.model import EncodedContext  # noqa: E402
from deg.train import (  # noqa: E402
    DEG_RECORDS,
    Example,
    Source,
    active_sources,
    batches,
    load_sources,
    make_sources,
    parse_records_arg,
)


def _ex(uid: str, n: int = 50) -> Example:
    return Example("deg", uid, EncodedContext(tuple(range(n)), 0, ()), [])


def _write(path: Path, rows: list[dict]) -> None:
    path.parent.mkdir(parents=True, exist_ok=True)
    with path.open("a") as fh:
        for r in rows:
            fh.write(json.dumps(r) + "\n")


def test_parse_records_weights_and_names():
    s = parse_records_arg("/data/deg/datagen/v3-long/records.jsonl:0.2")
    assert s.path == Path("/data/deg/datagen/v3-long/records.jsonl") and s.weight == 0.2
    assert s.name == "v3-long"
    s = parse_records_arg("/x/y/records.jsonl")
    assert s.weight == 1.0 and s.name == "y"
    assert parse_records_arg("/x/a:b/records.jsonl").path == Path("/x/a:b/records.jsonl")
    assert parse_records_arg("other.jsonl:3").name == "other"
    with pytest.raises(ValueError):
        parse_records_arg("/x/records.jsonl:0")
    assert [s.path for s in make_sources(None)] == [DEG_RECORDS]
    a, b = make_sources(["/a/records.jsonl", "/b/a/records.jsonl:2"])
    assert (a.name, b.name, b.weight) == ("a", "a#2", 2.0)


def test_load_sources_keeps_each_split_and_reloads_all(tmp_path):
    v2, v3 = tmp_path / "v2" / "records.jsonl", tmp_path / "v3" / "records.jsonl"
    _write(v2, [{"uid": f"a{i}", "split": s} for i, s in enumerate(
        ["train", "val", "test", "train"])])
    sources = make_sources([str(v2), f"{v3}:0.2"])  # v3 not generated yet
    built = []

    def build(rec):
        built.append(rec["uid"])
        return None if rec["uid"] == "a3" else _ex(rec["uid"])

    assert load_sources(sources, build, 0, 1) == {"v2": 1, "v3": 0}
    assert built == ["a0", "a3"]  # val/test never built; a3 rejected by build, not retried
    _write(v3, [{"uid": "v3-d1", "split": "train"}, {"uid": "v3-d2", "split": "val"}])
    _write(v2, [{"uid": "a9", "split": "train"}])
    assert load_sources(sources, build, 0, 1) == {"v2": 1, "v3": 1}
    assert built == ["a0", "a3", "a9", "v3-d1"]
    assert [e.uid for e in sources[1].pool] == ["v3-d1"]

    def broken(rec):
        raise KeyError("schema")

    logged = []
    _write(v3, [{"uid": "v3-d3", "split": "train"}])
    assert load_sources(sources, broken, 0, 1, logged.append) == {"v2": 0, "v3": 0}
    assert logged[0]["skip"] == "v3-d3" and logged[0]["source"] == "v3"


def test_load_sources_shards_by_uid(tmp_path):
    path = tmp_path / "r" / "records.jsonl"
    _write(path, [{"uid": f"u{i}", "split": "train"} for i in range(40)])
    pools = []
    for rank in range(2):
        srcs = make_sources([str(path)])
        load_sources(srcs, lambda r: _ex(r["uid"]), rank, 2)
        pools.append({e.uid for e in srcs[0].pool})
    assert pools[0].isdisjoint(pools[1]) and len(pools[0] | pools[1]) == 40


def _shares(sources, n=4000, min_train=50):
    rng = random.Random(0)
    gen = batches(sources, None, rng, budget=10_000, jev_frac=0.0, max_records=8,
                  min_train=min_train)
    got = Counter()
    for _ in range(n):
        batch = next(gen)
        srcs = {e.uid.split("-")[0] for e in batch}
        assert len(srcs) == 1  # a micro-batch comes from one source
        got[srcs.pop()] += 1
    return {k: v / n for k, v in got.items()}


def test_small_source_gets_fixed_share_once_it_has_min_train():
    big = Source(Path("/v2/records.jsonl"), 0.8)
    big.pool = [_ex(f"v2-{i}", 20 + i % 300) for i in range(5000)]
    small = Source(Path("/v3/records.jsonl"), 0.2)
    small.pool = [_ex(f"v3-{i}", 3000 + i) for i in range(49)]
    assert active_sources([big, small], 50) == [big]
    assert _shares([big, small]) == {"v2": 1.0}
    small.pool.append(_ex("v3-49", 4000))  # 50th train record: share switches on
    assert active_sources([big, small], 50) == [big, small]
    share = _shares([big, small]).get("v3", 0.0)
    assert abs(share - 0.2) < 0.03, share


def test_no_source_over_threshold_still_trains():
    a = Source(Path("/a/records.jsonl"))
    a.pool = [_ex("a-1")]
    assert active_sources([a, Source(Path("/b/records.jsonl"))], 50) == [a]
    assert _shares([a], n=10) == {"a": 1.0}
    with pytest.raises(RuntimeError):
        next(batches([Source(Path("/c/records.jsonl"))], None, random.Random(0), 100, 0.0, 4))


def test_lr_scale_restart_continues_the_schedule():
    from deg.train import lr_scale

    full = [lr_scale(s, 20000, 30) for s in range(1, 20001)]
    assert full[0] == pytest.approx(1 / 30 * (0.5 * (1 + 1) * 0.9 + 0.1), rel=1e-6)
    assert full[-1] == pytest.approx(0.1)
    # a restart at step 1000 re-warms for 30 steps, then matches the uninterrupted run
    assert lr_scale(1001, 20000, 30, 1000) == pytest.approx(full[1000] / 30)
    for s in (1030, 5000, 20000):
        assert lr_scale(s, 20000, 30, 1000) == pytest.approx(full[s - 1])


def test_heads_only_steps_then_backbone_warms_up():
    from deg.train import backbone_live, lr_scale

    assert not backbone_live(1001, 1000, 500)
    assert not backbone_live(1500, 1000, 500)
    assert backbone_live(1501, 1000, 500)
    assert backbone_live(1, 0, 0)
    # the backbone LR re-warms from the end of the heads-only phase
    assert lr_scale(1501, 20000, 30, 1500) == pytest.approx(lr_scale(1501, 20000, 30) / 30,
                                                            rel=1e-6)