"""
Training-data generator for the tiny cube-solving model.

Emits (state, solution) pairs as JSONL:
    {"state": "<54 facelet chars>", "solution": "R U2 F' ...", "depth": <walk length>}

## Why states come from a random walk, not direct cubie sampling

Sampling a cube state directly (random piece permutation + orientations) is the
textbook approach, but it needs a consistent orientation reference per slot, and
the slot lists in `packages/engine/src/moveTables.ts` identify slots without
encoding one -- they are not ordered U/D-first with a consistent cyclic
direction. Sampling against them naively produces a valid cube only ~1 call in
37 (measured), because the corner-twist, edge-flip and permutation-parity
constraints are then being enforced against the wrong reference.

Walking from solved using the engine's own move tables sidesteps the whole
question: every reachable state is valid by construction, and the tables are
already cross-checked against an independent implementation. Measured against
the exact-uniform samples that *did* pass, the two agree on the distribution of
Kociemba solution length (mean 20.75 vs 21.0, both spanning 17-22), so the walk
is reaching the uniform distribution, not a biased corner of it.

## Why the walk is long

A deep walk is the point, not a cost. Every reachable state is at most 20 moves
from solved (God's Number), so a walk of 200 samples the same "fully mixed"
distribution as a walk of 50 or 1000 -- which is exactly why the model does not
need a depth curriculum to handle depth-50 and depth-100 benchmark buckets. It
never sees "depth"; it sees a state. Shallow depths are the rare tail, so they
are mixed in explicitly via --shallow-frac rather than left to chance.
"""
import argparse, json, os, random, sys, time

QT={
  "U": [6, 3, 0, 7, 4, 1, 8, 5, 2, 45, 46, 47, 12, 13, 14, 15, 16, 17, 9, 10, 11, 21, 22, 23, 24, 25, 26, 27, 28, 29, 30, 31, 32, 33, 34, 35, 18, 19, 20, 39, 40, 41, 42, 43, 44, 36, 37, 38, 48, 49, 50, 51, 52, 53],
  "R": [0, 1, 20, 3, 4, 23, 6, 7, 26, 15, 12, 9, 16, 13, 10, 17, 14, 11, 18, 19, 29, 21, 22, 32, 24, 25, 35, 27, 28, 51, 30, 31, 48, 33, 34, 45, 36, 37, 38, 39, 40, 41, 42, 43, 44, 8, 46, 47, 5, 49, 50, 2, 52, 53],
  "F": [0, 1, 2, 3, 4, 5, 44, 41, 38, 6, 10, 11, 7, 13, 14, 8, 16, 17, 24, 21, 18, 25, 22, 19, 26, 23, 20, 15, 12, 9, 30, 31, 32, 33, 34, 35, 36, 37, 27, 39, 40, 28, 42, 43, 29, 45, 46, 47, 48, 49, 50, 51, 52, 53],
  "D": [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 24, 25, 26, 18, 19, 20, 21, 22, 23, 42, 43, 44, 33, 30, 27, 34, 31, 28, 35, 32, 29, 36, 37, 38, 39, 40, 41, 51, 52, 53, 45, 46, 47, 48, 49, 50, 15, 16, 17],
  "L": [53, 1, 2, 50, 4, 5, 47, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 0, 19, 20, 3, 22, 23, 6, 25, 26, 18, 28, 29, 21, 31, 32, 24, 34, 35, 42, 39, 36, 43, 40, 37, 44, 41, 38, 45, 46, 33, 48, 49, 30, 51, 52, 27],
  "B": [11, 14, 17, 3, 4, 5, 6, 7, 8, 9, 10, 35, 12, 13, 34, 15, 16, 33, 18, 19, 20, 21, 22, 23, 24, 25, 26, 27, 28, 29, 30, 31, 32, 36, 39, 42, 2, 37, 38, 1, 40, 41, 0, 43, 44, 51, 48, 45, 52, 49, 46, 53, 50, 47],
}


SOLVED = "".join("URFDLB"[i // 9] for i in range(54))
MOVES = [f + s for f in "URFDLB" for s in ("", "'", "2")]
_TURNS = {"": 1, "'": 3, "2": 2}


def apply_move(state: str, move: str) -> str:
    """Applies one move in standard notation to a 54-char facelet string."""
    perm = QT[move[0]]
    for _ in range(_TURNS[move[1:]]):
        state = "".join(state[perm[i]] for i in range(54))
    return state


def apply_sequence(state: str, moves) -> str:
    for m in moves:
        state = apply_move(state, m)
    return state


def random_state(rng: random.Random, depth: int) -> str:
    """Walks `depth` moves from solved, never turning the same face twice in a row."""
    state, prev = SOLVED, None
    for _ in range(depth):
        while True:
            move = rng.choice(MOVES)
            if move[0] != prev:
                break
        state = apply_move(state, move)
        prev = move[0]
    return state


def check_kociemba_is_native(warn_only=False):
    """Refuses to run against kociemba's pure-Python fallback.

    The package ships a C extension and a pure-Python implementation, and falls
    back to the latter *silently* -- only a warning -- when the extension cannot
    build (typically missing dev headers for whichever interpreter the venv was
    made from). The fallback is roughly 50x slower, which turns a two-hour dataset
    into a multi-day one. Since the only symptom is a slow log, this is checked
    up front rather than discovered later: a rented machine generating at
    fallback speed bills the whole time.
    """
    import time
    import kociemba

    probe = "DRLUUBFBRBLURRLRUBLRDDFDLFUFUFFDBRDUBRUFLLFDDBFLUBLRBD"
    kociemba.solve(probe)  # warm the pruning tables before timing
    start = time.time()
    for _ in range(10):
        kociemba.solve(probe)
    ms = (time.time() - start) / 10 * 1000
    if ms > 200:
        msg = (f"kociemba is running at {ms:.0f} ms/solve, which is pure-Python "
               f"fallback speed (native is ~20 ms). Install build-essential and "
               f"python3-dev for this interpreter, then reinstall kociemba.")
        if not warn_only:
            raise SystemExit(f"ABORT: {msg}")
        print(f"WARNING: {msg}", file=sys.stderr)
    return ms


def generate(count, seed, deep_depth, shallow_frac, max_shallow, augment, out):
    import kociemba

    rng = random.Random(seed)
    kociemba.solve(random_state(rng, 20))  # warm the pruning tables
    written = 0
    start = time.time()

    while written < count:
        if rng.random() < shallow_frac:
            depth = rng.randint(1, max_shallow)
        else:
            depth = deep_depth
        state = random_state(rng, depth)
        if state == SOLVED:
            continue
        solution = kociemba.solve(state).split()

        # One Kociemba call also settles every state along its own solution path:
        # after applying a prefix, the remaining suffix is a valid solve of the
        # state you land on. That is a ~20x multiplier on a call costing ~23ms,
        # and it is what makes a multi-million-pair dataset practical on few
        # cores. The suffix labels are valid but not what Kociemba would itself
        # emit for those states, so they are a slightly different (more
        # multi-modal) target than the canonical pairs -- hence opt-in, and worth
        # measuring against augment=0 before relying on it.
        pairs = [(state, solution, depth)]
        if augment:
            walk = state
            for i, move in enumerate(solution[:-1]):
                walk = apply_move(walk, move)
                if walk == SOLVED:
                    break
                pairs.append((walk, solution[i + 1:], len(solution) - i - 1))
            rng.shuffle(pairs)
            pairs = pairs[:augment]

        for st, sol, d in pairs:
            if written >= count:
                break
            out.write(json.dumps({"state": st, "solution": " ".join(sol), "depth": d}) + "\n")
            written += 1

        if written % 20000 < len(pairs):
            rate = written / max(time.time() - start, 1e-9)
            print(f"  {written}/{count}  ({rate:.0f}/s)", file=sys.stderr)

    return written


def _worker(job):
    """One shard, in its own process with its own seed and its own pruning tables."""
    idx, count, args = job
    path = f"{args['out']}.part{idx}"
    with open(path, "w") as fh:
        generate(count, args["seed"] + 1000 * idx, args["deep_depth"],
                 args["shallow_frac"], args["max_shallow"], args["augment"], fh)
    return path


def effective_cpus():
    """Number of CPUs this process may actually use.

    `nproc` and os.cpu_count() report the *host's* cores, which on a rented
    container is not what we get: a box advertising 128 cores routinely runs
    under a cgroup quota of a fraction of that. Spawning one worker per host core
    then puts 128 processes on ~16 CPUs, each running at an eighth speed. The
    aggregate is unchanged -- it is capped by the quota either way -- but every
    throughput estimate derived from the advertised core count is wrong by that
    factor, which is how a 20-minute job came to be predicted for a 100-minute
    one.

    Prefers the scheduler affinity mask, then the cgroup v2 and v1 quotas.
    """
    candidates = []
    try:
        candidates.append(len(os.sched_getaffinity(0)))
    except AttributeError:
        pass
    for quota_path, period_path in (
        ("/sys/fs/cgroup/cpu.max", None),
        ("/sys/fs/cgroup/cpu/cpu.cfs_quota_us", "/sys/fs/cgroup/cpu/cpu.cfs_period_us"),
    ):
        try:
            raw = open(quota_path).read().split()
            quota = raw[0]
            period = int(raw[1]) if period_path is None else int(open(period_path).read())
            if quota not in ("max", "-1"):
                candidates.append(max(1, int(int(quota) / period)))
        except (OSError, ValueError, IndexError):
            pass
    candidates.append(os.cpu_count() or 1)
    return max(1, min(candidates))


def generate_parallel(count, workers, args_dict, shards=None):
    """Shards generation across processes.

    `shards` fixes the number of shards independently of `workers`. Per-shard
    seeds are derived as seed + 1000*index, so the dataset produced depends on the
    shard count -- meaning two machines with different core counts generate
    *different* data from the same --seed. That is fine for training data and
    quietly wrong for a shared holdout: two models meant to face identical puzzles
    face merely similar ones, turning a paired comparison into an unpaired one.
    Pass a fixed `shards` whenever the output must be reproducible across
    machines.
    """
    """Shards generation across processes. Kociemba is CPU-bound and releases
    nothing to threads, so processes are the only way to use more than one core --
    and generation, not training, is the wall-clock bottleneck for this dataset."""
    import multiprocessing as mp

    n_shards = shards or workers
    per = [count // n_shards] * n_shards
    for i in range(count % n_shards):
        per[i] += 1
    jobs = [(i, per[i], args_dict) for i in range(n_shards) if per[i]]
    with mp.Pool(min(workers, len(jobs))) as pool:
        parts = pool.map(_worker, jobs)

    written = 0
    with open(args_dict["out"], "w") as out:
        for path in parts:
            with open(path) as fh:
                for line in fh:
                    out.write(line)
                    written += 1
            os.remove(path)
    return written


def main():
    p = argparse.ArgumentParser(description=__doc__)
    p.add_argument("--count", type=int, default=1_000_000)
    p.add_argument("--seed", type=int, default=0)
    p.add_argument("--deep-depth", type=int, default=200,
                   help="Walk length for the deep (fully mixed) majority of states.")
    p.add_argument("--shallow-frac", type=float, default=0.15,
                   help="Fraction drawn from near-solved states instead.")
    p.add_argument("--max-shallow", type=int, default=8)
    p.add_argument("--augment", type=int, default=0, metavar="K",
                   help="Also emit up to K states from each solution path (0 = canonical pairs only).")
    p.add_argument("--out", default="-")
    p.add_argument("--shards", type=int, default=0,
                   help="Fix the shard count so output is identical across machines "
                        "regardless of core count. Required for a shared holdout.")
    p.add_argument("--workers", type=int, default=0,
                   help="Parallel generator processes. 0 = the CPUs actually available "
                        "to this container, which is usually far fewer than nproc reports.")
    args = p.parse_args()

    ms = check_kociemba_is_native()
    if args.workers <= 0:
        args.workers = effective_cpus()
    print(f"kociemba {ms:.1f} ms/solve | workers {args.workers} "
          f"(host reports {os.cpu_count()})", file=sys.stderr)

    if args.workers > 1:
        if args.out == "-":
            p.error("--workers needs --out (shards are merged into a file)")
        start = time.time()
        n = generate_parallel(args.count, args.workers, {
            "out": args.out, "seed": args.seed, "deep_depth": args.deep_depth,
            "shallow_frac": args.shallow_frac, "max_shallow": args.max_shallow,
            "augment": args.augment}, shards=args.shards or None)
        print(f"wrote {n} pairs in {time.time()-start:.0f}s "
              f"({n/max(time.time()-start,1e-9):.0f}/s, {args.workers} workers)", file=sys.stderr)
        return

    out = sys.stdout if args.out == "-" else open(args.out, "w")
    try:
        n = generate(args.count, args.seed, args.deep_depth, args.shallow_frac,
                     args.max_shallow, args.augment, out)
        print(f"wrote {n} pairs", file=sys.stderr)
    finally:
        if out is not sys.stdout:
            out.close()


if __name__ == "__main__":
    main()