"""Candidate-set size and instance count per stage, level 3 vs level 8.

Two metrics only, measured on the training masks.
"""
import argparse
import multiprocessing as mp

import numpy as np

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):
    out = []
    for s in range(_MASKS.shape[1]):
        mask = np.array(_MASKS[idx, s]).astype(np.uint16)
        _, empties = SI.split_cells(mask)
        sizes = [len(SI.digits_of(mask[c])) for c in empties]
        width = float(np.mean(sizes)) if sizes else 1.0
        wmax = int(np.max(sizes)) if sizes else 1
        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)
        out.append((s, width, wmax, r["n_instances"], len(empties)))
    return out


def run(level, args):
    meta = np.load(args.meta)
    idxs = np.where(meta[:, 1] == level)[0][:args.limit].tolist()
    with mp.Pool(args.workers, initializer=_init,
                 initargs=(args.masks, args)) as pool:
        res = pool.map(_one, idxs, chunksize=4)
    widths, maxes, counts, opens = {}, {}, {}, {}
    for rows in res:
        for s, w, wm, n, ne in rows:
            widths.setdefault(s, []).append(w)
            maxes.setdefault(s, []).append(wm)
            counts.setdefault(s, []).append(n)
            opens.setdefault(s, []).append(ne)
    return len(idxs), widths, maxes, counts, opens


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--masks", default="datasets_multicandidate_s12/train_cand_masks.npy")
    ap.add_argument("--meta", default="datasets_multicandidate_s12/train_meta.npy")
    ap.add_argument("--limit", type=int, default=400)
    ap.add_argument("--max-confine", type=int, default=1)
    ap.add_argument("--max-instances", type=int, default=64)
    ap.add_argument("--max-attempts", type=int, default=300)
    ap.add_argument("--max-repair", type=int, default=120)
    ap.add_argument("--workers", type=int, default=60)
    args = ap.parse_args()

    n3, w3, m3, c3, o3 = run(3, args)
    n8, w8, m8, c8, o8 = run(8, args)

    print(f"training masks, {n3} level-3 puzzles and {n8} level-8 puzzles")
    print(f"confinement threshold |S| <= {args.max_confine}")
    print()
    print(f"{'':>5} {'------------- level 3 -------------':>42} "
          f"{'------------- level 8 -------------':>42}")
    print(f"{'stage':>5} {'mean':>7} {'max':>7} {'worst':>7} {'open':>7} "
          f"{'inst':>7} {'mean':>7} {'max':>7} {'worst':>7} {'open':>7} {'inst':>7}")
    print("-" * 96)
    for s in sorted(w3):
        print(f"{s:>5} "
              f"{np.mean(w3[s]):>7.2f} {np.mean(m3[s]):>7.2f} "
              f"{np.max(m3[s]):>7d} {np.mean(o3[s]):>7.1f} "
              f"{np.mean(c3[s]):>7.1f} "
              f"{np.mean(w8[s]):>7.2f} {np.mean(m8[s]):>7.2f} "
              f"{np.max(m8[s]):>7d} {np.mean(o8[s]):>7.1f} "
              f"{np.mean(c8[s]):>7.1f}")
    print()
    print("mean  = mean candidate set size over undetermined cells")
    print("max   = mean over puzzles of the largest candidate set in that puzzle")
    print("worst = largest candidate set seen in any puzzle at that stage")
    print("open  = undetermined cells remaining;  inst = instances at saturation")


if __name__ == "__main__":
    main()