"""Build superposition instances from a stage's candidate masks.

An *instance* is one concrete assignment: exactly one digit per cell, each drawn
from that cell's candidate set at that stage. Across the instances for a stage,
every (cell, candidate) pair should appear at least once, so the candidate set is
recoverable from the instances without ever being supervised as a set.

Build chronology (per puzzle, per stage):
  1. read the candidate sets from the stage mask
  2. derive the live dependencies from that same mask
  3. propose an instance
  4. drop it if it disobeys any dependency
  5. repeat 3-4 until every (cell, candidate) pair is covered

Dependencies are recomputed from the mask rather than recorded when a technique
fires. The stage-k mask is the cumulative product of every technique applied up
to stage k, so a confinement visible in the mask *is* the disjunction those
techniques established, and it stays visible exactly as long as it is live.

The confinement threshold matters. For a unit U and a digit d not yet placed in
U, S(U,d) is the set of cells of U that can still hold d, and the dependency is
"d must be placed somewhere in S(U,d)". Enforcing this for every unit and every
digit is full Sudoku unit coverage, which forces the unique solution and leaves
no superposition at all. Restricting to small |S(U,d)| keeps only the genuinely
locked disjunctions.
"""
import numpy as np

POPCOUNT = np.array([bin(i).count("1") for i in range(512)], dtype=np.int8)


def build_units():
    """The 27 units, each a list of 9 cell ids (cell = r*9 + c)."""
    units = []
    for r in range(9):
        units.append([r * 9 + c for c in range(9)])
    for c in range(9):
        units.append([r * 9 + c for r in range(9)])
    for br in range(0, 9, 3):
        for bc in range(0, 9, 3):
            units.append([(br + i) * 9 + (bc + j)
                          for i in range(3) for j in range(3)])
    return units


UNITS = build_units()


def digits_of(m):
    m = int(m)
    return [d for d in range(1, 10) if m & (1 << (d - 1))]


def dependencies_from_mask(mask, max_confine=9):
    """Live disjunctions readable from one stage's mask.

    Returns a list of (digit, cells) meaning "digit must be placed in one of
    these cells". Only digits not already pinned in the unit are included, and
    only when the confinement size is at most max_confine.
    """
    deps = []
    for unit in UNITS:
        pinned = 0
        for c in unit:
            m = int(mask[c])
            if m and (m & (m - 1)) == 0:
                pinned |= m
        for d in range(1, 10):
            bit = 1 << (d - 1)
            if pinned & bit:
                continue
            S = [c for c in unit if int(mask[c]) & bit]
            if 1 <= len(S) <= max_confine:
                deps.append((d, S))
    return deps


def split_cells(mask):
    """(forced cell -> digit, list of undetermined cells)."""
    forced, empties = {}, []
    for c in range(81):
        m = int(mask[c])
        if m == 0:
            continue
        if m & (m - 1) == 0:
            forced[c] = m.bit_length()
        else:
            empties.append(c)
    return forced, empties


def _violated(assign, deps):
    return [i for i, (d, S) in enumerate(deps)
            if not any(assign[c] == d for c in S)]


def _violations(assign, deps):
    return len(_violated(assign, deps))


def build_dep_index(deps):
    """cell -> [(dep index, digit that dep wants), ...]"""
    cell_deps = [[] for _ in range(81)]
    for i, (d, S) in enumerate(deps):
        for c in S:
            cell_deps[c].append((i, d))
    return cell_deps


def repair(assign, deps, forced, rng, max_iter=200, cell_deps=None):
    """Min-conflicts repair: repeatedly take a violated disjunction and give its
    digit to whichever of its cells breaks the fewest other disjunctions.

    Purely a proposal-quality step. The accept/reject test still runs afterwards
    and is the only thing that decides whether an instance enters the dataset.
    Violation counts are maintained incrementally, so each move costs only the
    handful of disjunctions that touch the cell being changed.
    """
    if not deps:
        return True
    if cell_deps is None:
        cell_deps = build_dep_index(deps)
    counts = [sum(1 for c in S if assign[c] == d) for d, S in deps]
    nviol = sum(1 for x in counts if x == 0)

    def apply(c, new):
        nonlocal nviol
        old = assign[c]
        if old == new:
            return
        for i, d in cell_deps[c]:
            if d == old:
                counts[i] -= 1
                if counts[i] == 0:
                    nviol += 1
            elif d == new:
                if counts[i] == 0:
                    nviol -= 1
                counts[i] += 1
        assign[c] = new

    for _ in range(max_iter):
        if nviol == 0:
            return True
        bad = [i for i, x in enumerate(counts) if x == 0]
        d, S = deps[bad[rng.integers(len(bad))]]
        free = [c for c in S if c not in forced]
        if not free:
            return False
        best, best_v, prev = None, None, {}
        for c in free:
            prev[c] = assign[c]
            before = nviol
            apply(c, d)
            v = nviol
            apply(c, prev[c])
            assert nviol == before
            if best_v is None or v < best_v:
                best, best_v = c, v
        apply(best, d)
    return nviol == 0


def propose(mask, forced, empties, deps, uncovered, rng, aware=True):
    """One candidate instance. `aware` first satisfies the confined
    disjunctions, then fills the rest preferring not-yet-covered pairs.
    Without it, every cell is an independent draw from its candidate set."""
    assign = np.zeros(81, dtype=np.int8)
    for c, d in forced.items():
        assign[c] = d
    taken = set(forced)

    if aware and deps:
        order = sorted(range(len(deps)), key=lambda i: len(deps[i][1]))
        for i in order:
            d, S = deps[i]
            if any(assign[c] == d for c in S):
                continue
            free = [c for c in S if c not in taken]
            if not free:
                continue          # unsatisfiable in this proposal; test catches it
            c = free[rng.integers(len(free))]
            assign[c] = d
            taken.add(c)

    for c in empties:
        if c in taken:
            continue
        cands = digits_of(mask[c])
        unc = [d for d in cands if (c, d) in uncovered]
        pool = unc if unc else cands
        assign[c] = pool[rng.integers(len(pool))]
    return assign


def instances_for_stage(mask, max_confine=3, max_instances=64,
                        max_attempts=400, seed=0, aware=True, max_repair=200,
                        require_new=True, patience=60):
    """Run the generate/test/loop for one stage.

    Returns a dict with the instances and the statistics the caller wants:
    how many were produced, how many were ruled out, and how much of each
    candidate set the survivors cover.
    """
    rng = np.random.default_rng(seed)
    forced, empties = split_cells(mask)
    deps = dependencies_from_mask(mask, max_confine)

    all_pairs = {(c, d) for c in empties for d in digits_of(mask[c])}
    uncovered = set(all_pairs)

    cell_deps = build_dep_index(deps)
    accepted, rejected, attempts = [], 0, 0
    seen = set()
    if not empties:
        # Fully determined stage (the last one): the single instance is the
        # solution itself, which the loop below would never emit.
        only = np.zeros(81, dtype=np.int8)
        for c, d in forced.items():
            only[c] = d
        accepted.append(only)
    since_new = 0
    while (uncovered and attempts < max_attempts
           and len(accepted) < max_instances and since_new < patience):
        attempts += 1
        inst = propose(mask, forced, empties, deps, uncovered, rng, aware)
        if max_repair:
            repair(inst, deps, forced, rng, max_repair, cell_deps)
        key = inst.tobytes()
        if _violations(inst, deps) != 0 or key in seen:
            rejected += 1
            since_new += 1
            continue
        seen.add(key)
        gained = [(c, int(inst[c])) for c in empties
                  if (c, int(inst[c])) in uncovered]
        if require_new and not gained:
            # Valid but redundant: adds no candidate the set already lacks.
            since_new += 1
            continue
        since_new = 0
        accepted.append(inst)
        for pair in gained:
            uncovered.discard(pair)

    covered = len(all_pairs) - len(uncovered)
    widths = [len(digits_of(mask[c])) for c in empties]
    # A disjunction says "digit d appears in S", never "at most once", so an
    # instance may place the same digit twice in a unit. Count how often.
    dups = []
    for a in accepted:
        n = 0
        for unit in UNITS:
            seen_d = {}
            for c in unit:
                seen_d[a[c]] = seen_d.get(a[c], 0) + 1
            n += sum(v - 1 for v in seen_d.values() if v > 1)
        dups.append(n)
    return {
        "instances": np.array(accepted, dtype=np.int8) if accepted
                     else np.zeros((0, 81), dtype=np.int8),
        "n_instances": len(accepted),
        "n_rejected": rejected,
        "n_attempts": attempts,
        "n_deps": len(deps),
        "dep_sizes": [len(S) for _, S in deps],
        "n_empty": len(empties),
        "mean_width": float(np.mean(widths)) if widths else 0.0,
        "n_pairs": len(all_pairs),
        "n_covered": covered,
        "coverage": covered / len(all_pairs) if all_pairs else 1.0,
        "hit_cap": bool(uncovered),
        "mean_dups": float(np.mean(dups)) if dups else 0.0,
    }