"""Build the curriculum streams for the multi-part score.

A curriculum is three kinds of text one after another (Python source, C headers,
Wikipedia), then a second, different helping of each. A .json file beside each stream
marks the segments, and the scorer uses it to reward several things at once: each kind
on its own, fast adaptation to a new kind, and retention when a kind returns.

Needs data/search.bin and data/holdout.bin (python make_data.py) and the Wikipedia
sample at data/dl/enwik8 (http://mattmahoney.net/dc/enwik8.zip, unzipped).

  python make_curriculum.py      # writes data/multi.bin and data/multi_holdout.bin (+ .json)
"""
from __future__ import annotations

import json
import random
import sysconfig
from pathlib import Path

KB = 1024
FIRST, SECOND, ADAPT = 224 * KB, 112 * KB, 16 * KB
MB = 1024 * KB


def c_headers(wanted: int) -> bytes:
    root = Path(sysconfig.get_paths()["stdlib"]).parent
    files = sorted(root.rglob("*.h"))
    random.Random(0).shuffle(files)
    chunks, total = [], 0
    for path in files:
        try:
            chunk = path.read_bytes()
        except OSError:
            continue
        chunks.append(chunk)
        total += len(chunk)
        if total >= wanted:
            break
    joined = b"".join(chunks)
    if len(joined) < wanted:
        raise SystemExit(f"only {len(joined)} bytes of C headers found; need {wanted}")
    return joined


def wikipedia(offset: int, length: int) -> bytes:
    with open("data/dl/enwik8", "rb") as handle:
        handle.seek(offset)
        return handle.read(length)


def write(name: str, helpings: list[tuple[str, int, bytes]]) -> None:
    stream, segments = b"", []
    for kind, which_pass, text in helpings:
        segments.append({"kind": kind, "pass": which_pass,
                         "start": len(stream), "end": len(stream) + len(text)})
        stream += text
    Path(f"data/{name}.bin").write_bytes(stream)
    Path(f"data/{name}.json").write_text(
        json.dumps({"adapt_bytes": ADAPT, "segments": segments}, indent=1), encoding="utf-8")
    print(f"wrote data/{name}.bin ({len(stream)} bytes, {len(segments)} segments)")


def main() -> int:
    search = Path("data/search.bin").read_bytes()
    holdout = Path("data/holdout.bin").read_bytes()
    headers = c_headers(2 * (FIRST + SECOND))
    half = FIRST + SECOND
    write("multi", [
        ("python", 1, search[:FIRST]),
        ("c-header", 1, headers[:FIRST]),
        ("wikipedia", 1, wikipedia(0, FIRST)),
        ("python", 2, search[600 * KB:600 * KB + SECOND]),
        ("c-header", 2, headers[FIRST:half]),
        ("wikipedia", 2, wikipedia(5 * MB, SECOND)),
    ])
    write("multi_holdout", [
        ("python", 1, holdout[:FIRST]),
        ("c-header", 1, headers[half:half + FIRST]),
        ("wikipedia", 1, wikipedia(50 * MB, FIRST)),
        ("python", 2, holdout[600 * KB:600 * KB + SECOND]),
        ("c-header", 2, headers[half + FIRST:2 * half]),
        ("wikipedia", 2, wikipedia(60 * MB, SECOND)),
    ])
    return 0


if __name__ == "__main__":
    raise SystemExit(main())