"""Worked example of superposition-instance generation for one puzzle/stage.

Prints the candidate sets, the dependencies read off the mask, the instances that
survived, the ones that were ruled out, and the resulting per-cell coverage.
Display is restricted to one unit so the table is readable; generation always
runs on the full grid.
"""
import argparse

import numpy as np

import superposition_instances as SI

NAMES = ([f"row {r}" for r in range(9)] + [f"col {c}" for c in range(9)]
         + [f"box {b}" for b in range(9)])


def fmt(cells):
    return "{" + ",".join(str(d) for d in cells) + "}"


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--masks", default="datasets_multicandidate_s12/test_cand_masks.npy")
    ap.add_argument("--puzzle", type=int, default=2)
    ap.add_argument("--stage", type=int, default=0)
    ap.add_argument("--unit", type=int, default=14, help="unit index to display (14 = col 5)")
    ap.add_argument("--max-confine", type=int, default=3)
    ap.add_argument("--max-instances", type=int, default=16)
    args = ap.parse_args()

    masks = np.load(args.masks, mmap_mode="r")
    mask = np.array(masks[args.puzzle, args.stage]).astype(np.uint16)

    unit = SI.UNITS[args.unit]
    forced, empties = SI.split_cells(mask)

    print("=" * 78)
    print(f"puzzle {args.puzzle}, stage {args.stage}, displaying {NAMES[args.unit]}")
    print("=" * 78)
    print(f"grid: {len(forced)} determined cells, {len(empties)} undetermined, "
          f"mean candidate count {np.mean([len(SI.digits_of(mask[c])) for c in empties]):.2f}")

    print(f"\ncandidate sets in {NAMES[args.unit]}:")
    for c in unit:
        r, cc = divmod(c, 9)
        ds = SI.digits_of(mask[c])
        tag = "determined" if len(ds) == 1 else ""
        print(f"    r{r}c{cc}  {fmt(ds):<22} {tag}")

    all_deps = SI.dependencies_from_mask(mask, max_confine=9)
    conf_deps = SI.dependencies_from_mask(mask, max_confine=args.max_confine)
    print(f"\ndependencies over the whole grid:")
    print(f"    {len(all_deps)} total (every unit, every unplaced digit)")
    print(f"    {len(conf_deps)} with confinement size <= {args.max_confine} "
          f"(these are the ones enforced)")
    sizes = {}
    for _, S in all_deps:
        sizes[len(S)] = sizes.get(len(S), 0) + 1
    print("    by confinement size: "
          + ", ".join(f"|S|={k}: {v}" for k, v in sorted(sizes.items())))

    print(f"\ndependencies enforced inside {NAMES[args.unit]}:")
    shown = 0
    for d, S in conf_deps:
        if not set(S) <= set(unit):
            continue
        locs = " or ".join(f"r{c//9}c{c%9}" for c in S)
        print(f"    digit {d} must be placed at {locs}")
        shown += 1
    if shown == 0:
        print("    (none at this confinement threshold)")

    res = SI.instances_for_stage(mask, max_confine=args.max_confine,
                                 max_instances=args.max_instances, seed=0)

    print(f"\ngeneration: {res['n_instances']} instances kept, "
          f"{res['n_rejected']} ruled out over {res['n_attempts']} attempts")
    print(f"coverage: {res['n_covered']}/{res['n_pairs']} "
          f"(cell, candidate) pairs = {100*res['coverage']:.1f}%")

    inst = res["instances"]
    if len(inst):
        print(f"\nwhat each instance assigned in {NAMES[args.unit]}:")
        hdr = "    inst  " + "  ".join(f"r{c//9}c{c%9}" for c in unit)
        print(hdr)
        for i, a in enumerate(inst):
            print(f"    {i:>4}  " + "  ".join(f"{a[c]:>4}" for c in unit))

        print(f"\nper-cell coverage in {NAMES[args.unit]}:")
        for c in unit:
            ds = set(SI.digits_of(mask[c]))
            seen = set(int(a[c]) for a in inst)
            miss = sorted(ds - seen)
            r, cc = divmod(c, 9)
            status = "complete" if not miss else f"missing {fmt(miss)}"
            print(f"    r{r}c{cc}  candidates {fmt(sorted(ds)):<22} "
                  f"seen {fmt(sorted(seen)):<22} {status}")


if __name__ == "__main__":
    main()