"""Materialize superposition instances for every puzzle and every stage.

Input:  datasets_multicandidate_s12/{split}_cand_masks.npy   (N, 12, 81)
Output: datasets_superposition/{split}_assignments.npy       (M, 81) uint8
        datasets_superposition/{split}_index.npy             (M, 3)  int32
            columns: [puzzle_idx, stage, instance_within_stage]

M is about 48 x N (one puzzle expands into ~6 instances at stage 0 down to 1
at stage 11). Generation is embarrassingly parallel over puzzles.
"""
import argparse
import os
import time

import numpy as np
import multiprocessing as mp

import superposition_instances as SI

_MASKS = None
_ARGS = None


def _init(path, args):
    global _MASKS, _ARGS
    _MASKS = np.load(path, mmap_mode="r")
    _ARGS = args


def _one(idx):
    assigns, index = [], []
    for s in range(_MASKS.shape[1]):
        mask = np.array(_MASKS[idx, s]).astype(np.uint16)
        r = SI.instances_for_stage(
            mask,
            max_confine=_ARGS.max_confine,
            max_instances=_ARGS.max_instances,
            max_attempts=_ARGS.max_attempts,
            seed=idx * 100 + s,
            max_repair=_ARGS.max_repair,
            require_new=True,
            patience=_ARGS.patience,
        )
        inst = r["instances"]
        if len(inst) == 0:
            continue
        assigns.append(inst.astype(np.uint8))
        for k in range(len(inst)):
            index.append((idx, s, k))
    if not assigns:
        return (np.zeros((0, 81), dtype=np.uint8),
                np.zeros((0, 3), dtype=np.int32))
    return (np.concatenate(assigns, axis=0),
            np.array(index, dtype=np.int32))


def process_split(split, args):
    mask_path = os.path.join(args.mask_dir, f"{split}_cand_masks.npy")
    masks = np.load(mask_path, mmap_mode="r")
    n = len(masks) if args.limit is None else min(args.limit, len(masks))
    print(f"[{split}] {n:,} puzzles from {mask_path}", flush=True)

    t0 = time.time()
    chunks_a, chunks_i = [], []
    done = 0
    with mp.Pool(args.workers, initializer=_init,
                 initargs=(mask_path, args)) as pool:
        for a, i in pool.imap(_one, range(n), chunksize=8):
            if len(a):
                chunks_a.append(a)
                chunks_i.append(i)
            done += 1
            if done % 2000 == 0 or done == n:
                rate = done / max(time.time() - t0, 1e-6)
                kept = sum(len(x) for x in chunks_i)
                print(f"  {done:,}/{n:,}  {rate:.0f} puzzles/s  "
                      f"{kept:,} instances  "
                      f"({kept / done:.1f} per puzzle)", flush=True)

    assignments = (np.concatenate(chunks_a, axis=0) if chunks_a
                   else np.zeros((0, 81), dtype=np.uint8))
    index = (np.concatenate(chunks_i, axis=0) if chunks_i
             else np.zeros((0, 3), dtype=np.int32))

    os.makedirs(args.out_dir, exist_ok=True)
    ap = os.path.join(args.out_dir, f"{split}_assignments.npy")
    ip = os.path.join(args.out_dir, f"{split}_index.npy")
    np.save(ap, assignments)
    np.save(ip, index)

    elapsed = time.time() - t0
    print(f"[{split}] saved {len(assignments):,} instances  "
          f"({len(assignments) / n:.2f} per puzzle) in {elapsed / 60:.1f} min",
          flush=True)
    print(f"         {ap}  {os.path.getsize(ap) / 1e9:.2f} GB", flush=True)
    print(f"         {ip}  {os.path.getsize(ip) / 1e9:.2f} GB", flush=True)

    # per-stage instance counts
    print(f"[{split}] instances per stage:", flush=True)
    for s in range(masks.shape[1]):
        c = int((index[:, 1] == s).sum())
        print(f"    stage {s:>2}: {c:>12,}  ({c / n:.2f} per puzzle)",
              flush=True)
    return len(assignments)


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--mask_dir", default="datasets_multicandidate_s12")
    ap.add_argument("--out_dir", default="datasets_superposition")
    ap.add_argument("--splits", default="train,test")
    ap.add_argument("--limit", type=int, default=None)
    ap.add_argument("--max-confine", type=int, default=1)
    ap.add_argument("--max-instances", type=int, default=32)
    ap.add_argument("--max-attempts", type=int, default=400)
    ap.add_argument("--max-repair", type=int, default=80)
    ap.add_argument("--patience", type=int, default=40)
    ap.add_argument("--workers", type=int, default=64)
    args = ap.parse_args()

    print(f"confine |S|<={args.max_confine}  workers={args.workers}  "
          f"out={args.out_dir}", flush=True)
    for split in args.splits.split(","):
        process_split(split.strip(), args)


if __name__ == "__main__":
    main()