import json

import numpy as np

from nayara.data import MixedMemmap
from nayara.prepare_data import make_profile


def test_rank_disjoint_weighted_memmap(tmp_path):
    sources = {}
    for name, weight in (("code", 0.8), ("reasoning", 0.15), ("docs", 0.05)):
        filename = f"{name}.bin"
        np.arange(10_000, dtype=np.uint16).tofile(tmp_path / filename)
        sources[name] = {"file": filename, "tokens": 10_000, "weight": weight}
    (tmp_path / "manifest.json").write_text(json.dumps({"sources": sources}), encoding="utf-8")

    rank0 = MixedMemmap(tmp_path, seq_len=16, batch_size=2, rank=0, world_size=2, seed=7)
    rank1 = MixedMemmap(tmp_path, seq_len=16, batch_size=2, rank=1, world_size=2, seed=7)
    x0, y0, source0 = rank0.next_cpu()
    x1, y1, source1 = rank1.next_cpu()
    assert source0 == source1
    assert x0.shape == y0.shape == (2, 16)
    assert x1.shape == y1.shape == (2, 16)
    assert not np.intersect1d(x0.numpy(), x1.numpy()).size


def test_profile_scales_mix_exactly(tmp_path):
    source = tmp_path / "sources.json"
    output = tmp_path / "scaled.json"
    source.write_text(json.dumps({
        "code": {"weight": 0.80}, "reasoning": {"weight": 0.15}, "docs": {"weight": 0.05},
    }), encoding="utf-8")
    make_profile(source, output, 101)
    scaled = json.loads(output.read_text(encoding="utf-8"))
    assert sum(item["target_tokens"] for item in scaled.values()) == 101
    assert scaled["code"]["target_tokens"] == 81