"""Core objects of the structural-certificate formalism.

The formalism has three primitives (see ../../formalism.md):

  1. A structured domain X: a finite set together with a family of
     structurally derived ATOMS.  An atom is a fiber/region of a map
     canonically derived from the type of X (a cell value, a count, a
     line profile, a count comparison, ...).  Its bit cost is the sum of
     log2(#options) over the choice points of its derivation.
  2. A SUBSET DESCRIPTION: a conjunction of atoms (the fiber of a tupled
     derived map).  Cost = sum of atom costs + 1 framing bit per atom.
     The empty conjunction denotes all of X; a conjunction of point-value
     atoms denotes a singleton.  Singleton selection is the degenerate
     case, not a separate primitive.
  3. The OVERWRITE transformation T_{S,y}(g) = y on S, g elsewhere.
     A certificate is a sequence of overwrites applied to a constant
     function (last write wins).

Certificate cost (a prefix-free code):
    L(C) = log2|Y|                    (the initial constant)
         + sum_rules (1 + L(S) + log2|Y|)   (continue-bit, subset, output)
         + 1                          (stop bit)

Everything here works on plain Python ints as bitmasks over X, so all
structural objects are exposed directly as (cost, extension) pairs.
"""

from __future__ import annotations

import heapq
import itertools
import math
import random
from dataclasses import dataclass, field
from typing import Callable, Iterable, Optional


def choice(n: int) -> float:
    """Bit cost of one choice among n options (0 if forced)."""
    return math.log2(n) if n > 1 else 0.0


def popcount(x: int) -> int:
    return x.bit_count()


ATOM_FRAME = 1.0   # framing bit per atom in a conjunction (continue/stop)
RULE_FRAME = 1.0   # framing bit per rule in a certificate


@dataclass(frozen=True)
class Atom:
    """A structurally derived atomic subset of X."""
    name: str
    cost: float
    mask: int


@dataclass
class Rule:
    """One overwrite transformation T_{S,y}: S given as mask, output y."""
    mask: int
    y: int
    scost: float          # cost of the subset description L(S)
    desc: str
    atoms: tuple = ()     # constituent atoms, when the subset is a conjunction

    def cost(self, ny: int) -> float:
        return RULE_FRAME + self.scost + choice(ny)


def make_rule(atoms: Iterable[Atom], y: int, universe: int) -> Rule:
    atoms = tuple(atoms)
    mask = universe
    for a in atoms:
        mask &= a.mask
    scost = sum(a.cost for a in atoms) + ATOM_FRAME * len(atoms)
    desc = " & ".join(a.name for a in atoms) if atoms else "everything"
    return Rule(mask, y, scost, desc, atoms)


class Certificate:
    """A sequence of overwrites applied to a constant function."""

    def __init__(self, default: int, ny: int, rules: Optional[list[Rule]] = None):
        self.default = default
        self.ny = ny
        self.rules: list[Rule] = rules if rules is not None else []

    def cost(self) -> float:
        return (choice(self.ny) + 1.0
                + sum(r.cost(self.ny) for r in self.rules))

    def predict(self, x: int) -> int:
        bit = 1 << x
        for r in reversed(self.rules):
            if r.mask & bit:
                return r.y
        return self.default

    def predict_all(self, n: int) -> list[int]:
        out = [self.default] * n
        for r in self.rules:
            m = r.mask
            while m:
                low = m & -m
                out[low.bit_length() - 1] = r.y
                m ^= low
        return out

    def consistent(self, obs_ymask: list[int], universe: int) -> bool:
        """Check consistency with observations given as per-class bitmasks."""
        claimed = 0
        for r in reversed(self.rules):
            pts = r.mask & ~claimed
            if pts:
                for y in range(self.ny):
                    if y != r.y and pts & obs_ymask[y]:
                        return False
            claimed |= r.mask
        rest = universe & ~claimed
        for y in range(self.ny):
            if y != self.default and rest & obs_ymask[y]:
                return False
        return True

    def explain(self, ylabels: list[str]) -> str:
        lines = [f"START: everything -> {ylabels[self.default]}   "
                 f"[{choice(self.ny) + 1.0:.2f} bits]"]
        for r in self.rules:
            lines.append(f"  OVERWRITE {r.desc} -> {ylabels[r.y]}   "
                         f"[{r.cost(self.ny):.2f} bits, |S|={popcount(r.mask)}]")
        lines.append(f"TOTAL: {self.cost():.2f} bits, {len(self.rules)} rules")
        return "\n".join(lines)

    def copy(self) -> "Certificate":
        return Certificate(self.default, self.ny, list(self.rules))


class Domain:
    """A structured finite domain: |X|, output labels, atom library and a
    way to build the degenerate singleton description for any point."""

    def __init__(self, n: int, ylabels: list[str], atoms: list[Atom],
                 singleton_atoms: Callable[[int], list[Atom]], name: str = ""):
        self.n = n
        self.ylabels = list(ylabels)
        self.ny = len(ylabels)
        self.atoms = atoms
        self.singleton_atoms = singleton_atoms
        self.name = name
        self.universe = (1 << n) - 1
        # atoms containing x, sorted by cost — computed lazily per point
        self._atoms_at: dict[int, list[Atom]] = {}

    def atoms_at(self, x: int) -> list[Atom]:
        if x not in self._atoms_at:
            bit = 1 << x
            lst = [a for a in self.atoms if a.mask & bit]
            lst.sort(key=lambda a: a.cost)
            self._atoms_at[x] = lst
        return self._atoms_at[x]

    def singleton_rule_parts(self, x: int) -> tuple[float, int, tuple]:
        ats = tuple(self.singleton_atoms(x))
        cost = sum(a.cost for a in ats) + ATOM_FRAME * len(ats)
        mask = self.universe
        for a in ats:
            mask &= a.mask
        return cost, mask, ats

    def broader_atoms(self, atom: Atom) -> list[Atom]:
        """Atoms whose extension strictly contains atom's (cached).  These are
        the structural generalizations: e.g. the existential lift of a
        line-indexed profile atom to the profile-count atom N_p >= 1."""
        if not hasattr(self, "_broader"):
            self._broader: dict[int, list[Atom]] = {}
            for a in self.atoms:
                self._broader[a.mask] = [b for b in self.atoms
                                         if b.mask != a.mask
                                         and a.mask & ~b.mask == 0]
        return self._broader.get(atom.mask, [])


def dedupe_atoms(atoms: Iterable[Atom], universe: int) -> list[Atom]:
    """Keep, per extension, the cheapest atom; drop empty and full masks."""
    best: dict[int, Atom] = {}
    for a in atoms:
        if a.mask == 0 or a.mask == universe:
            continue
        b = best.get(a.mask)
        if b is None or a.cost < b.cost:
            best[a.mask] = a
    return sorted(best.values(), key=lambda a: a.cost)


# ---------------------------------------------------------------------------
# Subset library (explicit, for exact search on tiny domains)
# ---------------------------------------------------------------------------

def build_library(domain: Domain, max_atoms: int = 3) -> list[tuple[int, float, str]]:
    """All conjunctions of <= max_atoms atoms, deduped semantically:
    one entry per distinct extension, at its minimum description cost."""
    best: dict[int, tuple[float, str]] = {}

    def consider(mask: int, cost: float, desc: str):
        if mask == 0:
            return
        cur = best.get(mask)
        if cur is None or cost < cur[0]:
            best[mask] = (cost, desc)

    atoms = domain.atoms
    for k in range(1, max_atoms + 1):
        for combo in itertools.combinations(atoms, k):
            mask = domain.universe
            cost = ATOM_FRAME * k
            for a in combo:
                mask &= a.mask
                cost += a.cost
            consider(mask, cost, " & ".join(a.name for a in combo))
    return [(m, c, d) for m, (c, d) in best.items()]


# ---------------------------------------------------------------------------
# Exact minimum certificate: backward Dijkstra over "settled" sets.
#
# Reduction: because later overwrites win, only the LAST write at each point
# matters.  Reading a certificate backwards, each rule (S, y) settles the
# points of S not settled by later rules, and f must equal y there.  So a
# minimum certificate is a minimum-weight chain 0 = D_k ⊂ ... ⊂ D_0 in the
# lattice of settled sets, each step adding one library subset that is
# f-pure on its unsettled part, until the rest is monochromatic.
# ---------------------------------------------------------------------------

def exact_min_certificate(labels: list[int], domain: Domain,
                          library: list[tuple[int, float, str]]
                          ) -> Certificate:
    n, ny = domain.n, domain.ny
    full = domain.universe
    ymasks = [0] * ny
    for i, l in enumerate(labels):
        ymasks[l] |= 1 << i

    ycost = choice(ny)
    dist: dict[int, float] = {0: 0.0}
    parent: dict[int, tuple[int, int, float, str, int]] = {}
    heap: list[tuple[float, int]] = [(0.0, 0)]
    best_total = math.inf
    best_state = None
    best_default = 0

    while heap:
        g, D = heapq.heappop(heap)
        if g > dist.get(D, math.inf):
            continue
        if g >= best_total:
            break
        rem = full & ~D
        # terminal test: remainder monochromatic (or empty)
        if rem == 0:
            total = g + ycost + 1.0
            if total < best_total:
                best_total, best_state, best_default = total, D, labels[0]
        else:
            for y in range(ny):
                if rem & ~ymasks[y] == 0:
                    total = g + ycost + 1.0
                    if total < best_total:
                        best_total, best_state, best_default = total, D, y
                    break
        for mask, scost, desc in library:
            new = mask & ~D
            if new == 0:
                continue
            ry = -1
            for y in range(ny):
                if new & ~ymasks[y] == 0:
                    ry = y
                    break
            if ry < 0:
                continue
            nD = D | mask
            ng = g + RULE_FRAME + scost + ycost
            if ng < dist.get(nD, math.inf):
                dist[nD] = ng
                parent[nD] = (D, mask, scost, desc, ry)
                heapq.heappush(heap, (ng, nD))

    # reconstruct: path from 0 to best_state picks rules LAST-first
    rules_rev = []
    D = best_state
    while D != 0:
        pD, mask, scost, desc, ry = parent[D]
        rules_rev.append(Rule(mask, ry, scost, desc))
        D = pD
    cert = Certificate(best_default, ny, rules_rev)  # reversed order = application order
    return cert


# ---------------------------------------------------------------------------
# Cheapest admissible structural subset containing a point (repair search).
#
# Branch-and-bound DFS over conjunctions of atoms containing x, atoms sorted
# by cost.  A conjunction is admissible when it avoids `forbidden`
# (observed points whose label differs from the intended output).  Once a
# prefix is admissible, extending it only raises cost, so we stop there.
# ---------------------------------------------------------------------------

def cheapest_subset_containing(domain: Domain, must: int, forbidden: int,
                               max_atoms: int = 4, node_budget: int = 200_000,
                               fallback: Optional[tuple] = None
                               ) -> Optional[tuple[float, int, tuple]]:
    """Cheapest describable subset S with must ⊆ S and S ∩ forbidden = ∅.
    With must a single point this is the repair search; with must the union
    of several rules' claimed points it is the merge search.  Returns
    (cost, mask, atoms) or None (or the fallback) if nothing admissible is
    found within the search bounds."""
    if forbidden == 0:
        return 0.0, domain.universe, ()
    if popcount(must) == 1:
        atoms_m = domain.atoms_at(must.bit_length() - 1)
    else:
        atoms_m = [a for a in domain.atoms if must & ~a.mask == 0]
    return _dfs_cheapest(domain, atoms_m, forbidden, max_atoms,
                         node_budget, fallback)


def cheapest_subset(domain: Domain, x: int, forbidden: int,
                    max_atoms: int = 4, node_budget: int = 2_000_000
                    ) -> tuple[float, int, tuple]:
    """Repair search: cheapest admissible subset containing the point x.
    Always succeeds — the singleton conjunction is the fallback."""
    return cheapest_subset_containing(
        domain, 1 << x, forbidden, max_atoms, node_budget,
        fallback=domain.singleton_rule_parts(x))


def _dfs_cheapest(domain: Domain, atoms_list: list[Atom], forbidden: int,
                  max_atoms: int, node_budget: int,
                  fallback: Optional[tuple]) -> Optional[tuple[float, int, tuple]]:
    if fallback is not None:
        best_cost, best_mask, best_atoms = fallback
        best_pop = popcount(best_mask)
    else:
        best_cost, best_mask, best_atoms, best_pop = math.inf, 0, None, 0
    atoms_x = atoms_list
    nodes = 0

    def dfs(start: int, mask: int, cost: float, picked: list[Atom], depth: int):
        nonlocal best_cost, best_mask, best_atoms, best_pop, nodes
        for j in range(start, len(atoms_x)):
            a = atoms_x[j]
            c2 = cost + a.cost + ATOM_FRAME
            if c2 > best_cost + 1e-12:
                return  # atoms sorted by cost: all later ones are worse
            nodes += 1
            if nodes > node_budget:
                return
            m2 = mask & a.mask
            if m2 == mask:
                continue  # redundant atom
            if m2 & forbidden == 0:
                p2 = popcount(m2)
                if c2 < best_cost - 1e-12 or p2 > best_pop:
                    best_cost, best_mask, best_pop = c2, m2, p2
                    best_atoms = tuple(picked + [a])
            elif depth + 1 < max_atoms:
                dfs(j + 1, m2, c2, picked + [a], depth + 1)

    dfs(0, domain.universe, 0.0, [], 0)
    if best_atoms is None and fallback is None:
        return None
    return best_cost, best_mask, best_atoms


# ---------------------------------------------------------------------------
# Incremental learner: observe examples one at a time; when the current
# extension is wrong at x, append the cheapest admissible overwrite whose
# subset contains x.  Correctly predicted points change nothing in the
# certificate — they only extend the observation masks, which shrink the
# admissible set of all FUTURE overwrites.
# ---------------------------------------------------------------------------

class IncrementalLearner:
    def __init__(self, domain: Domain, max_atoms: int = 4,
                 node_budget: int = 2_000_000):
        self.domain = domain
        self.max_atoms = max_atoms
        self.node_budget = node_budget
        self.cert: Optional[Certificate] = None
        self.obs: dict[int, int] = {}
        self.obs_ymask = [0] * domain.ny
        self.n_repairs = 0
        self.repair_costs: list[float] = []

    def observe(self, x: int, y: int) -> bool:
        """Returns True if a repair was needed."""
        self.obs[x] = y
        self.obs_ymask[y] |= 1 << x
        if self.cert is None:
            self.cert = Certificate(y, self.domain.ny)
            return False
        if self.cert.predict(x) == y:
            return False
        forbidden = 0
        for yy in range(self.domain.ny):
            if yy != y:
                forbidden |= self.obs_ymask[yy]
        scost, mask, atoms = cheapest_subset(
            self.domain, x, forbidden, self.max_atoms, self.node_budget)
        self.cert.rules.append(make_rule(atoms, y, self.domain.universe))
        self.n_repairs += 1
        self.repair_costs.append(RULE_FRAME + scost + choice(self.domain.ny))
        return True

    # -- recompression -----------------------------------------------------

    def prune(self) -> int:
        """Remove rules whose removal keeps the certificate consistent with
        all observations.  Returns number of rules removed."""
        assert self.cert is not None
        removed = 0
        changed = True
        while changed:
            changed = False
            # try removing the most expensive rules first
            order = sorted(range(len(self.cert.rules)),
                           key=lambda i: -self.cert.rules[i].scost)
            for i in order:
                trial = Certificate(self.cert.default, self.cert.ny,
                                    self.cert.rules[:i] + self.cert.rules[i + 1:])
                if trial.consistent(self.obs_ymask, self.domain.universe):
                    self.cert = trial
                    removed += 1
                    changed = True
                    break
        return removed

    def recompress(self, orders: int = 2, rng: Optional[random.Random] = None) -> None:
        """Prune; replay the observations in other orders (majority class
        first, plus random shuffles); then apply structural generalization
        (broaden) to the best candidate.  Keeps the cheapest consistent
        certificate found."""
        assert self.cert is not None
        rng = rng or random.Random(0)
        self.prune()
        best = self.cert
        items = list(self.obs.items())
        counts = {y: 0 for y in range(self.domain.ny)}
        for _, y in items:
            counts[y] += 1
        candidates = [sorted(items, key=lambda t: (-counts[t[1]], t[0]))]
        for _ in range(orders - 1):
            s = list(items)
            rng.shuffle(s)
            candidates.append(s)
        for order in candidates:
            trial = IncrementalLearner(self.domain, self.max_atoms, self.node_budget)
            for x, y in order:
                trial.observe(x, y)
            trial.prune()
            if (trial.cert is not None
                    and trial.cert.consistent(self.obs_ymask, self.domain.universe)
                    and trial.cert.cost() < best.cost()):
                best = trial.cert
        self.cert = best
        broaden_certificate(self.cert, self.domain, self.obs_ymask)
        merge_pass(self.cert, self.domain, self.obs_ymask, self.max_atoms)
        broaden_certificate(self.cert, self.domain, self.obs_ymask)
        self.prune()


def merge_pass(cert: Certificate, domain: Domain, obs_ymask: list[int],
               max_atoms: int = 4, node_budget: int = 100_000) -> Certificate:
    """Merge move: for each pair of same-output rules, search for ONE
    describable subset containing the observed points both rules are
    responsible for, avoiding observed points of other classes; replace the
    pair if the whole certificate gets cheaper and stays consistent.  This is
    the same search as repair, with "must contain a point" generalized to
    "must contain a set"."""
    universe = domain.universe
    improved = True
    while improved:
        improved = False
        R = cert.rules
        n = len(R)
        suffix = [0] * (n + 1)  # suffix[k] = union of masks of rules k..n-1
        for k in range(n - 1, -1, -1):
            suffix[k] = suffix[k + 1] | R[k].mask
        for j in range(n - 1, 0, -1):
            for i in range(j - 1, -1, -1):
                if R[i].y != R[j].y:
                    continue
                y = R[j].y
                claimed = ((R[i].mask & ~suffix[i + 1])
                           | (R[j].mask & ~suffix[j + 1]))
                keep = claimed & obs_ymask[y]
                if keep == 0:
                    continue
                forbid = 0
                for yy in range(cert.ny):
                    if yy != y:
                        forbid |= obs_ymask[yy]
                forbid &= ~suffix[j + 1]  # later rules re-fix those anyway
                res = cheapest_subset_containing(domain, keep, forbid,
                                                 max_atoms, node_budget)
                if res is None:
                    continue
                _, _, atoms = res
                nr = make_rule(atoms, y, universe)
                trial = Certificate(cert.default, cert.ny,
                                    R[:i] + R[i + 1:j] + [nr] + R[j + 1:])
                if (trial.cost() < cert.cost() - 1e-9
                        and trial.consistent(obs_ymask, universe)):
                    cert.rules = trial.rules
                    improved = True
                    break
            if improved:
                break
    return cert


def broaden_certificate(cert: Certificate, domain: Domain,
                        obs_ymask: list[int], max_sweeps: int = 5) -> Certificate:
    """Structural generalization pass ("recompress if the modification creates
    a more general structural pattern").  For each rule, try (a) dropping an
    atom and (b) replacing an atom by a strictly broader atom — e.g. lifting a
    line-indexed profile atom prof[l]=p to the existential profile-count atom
    N_p >= 1, which undoes one symmetry-breaking choice.  Same-output rules
    subsumed by the broadened rule are dropped in the same move.  A move is
    kept iff the certificate stays consistent with the observations and its
    total cost strictly decreases."""
    universe = domain.universe
    sweeps = 0
    improved = True
    while improved and sweeps < max_sweeps:
        improved = False
        sweeps += 1
        for i in range(len(cert.rules)):
            if i >= len(cert.rules):
                break
            r = cert.rules[i]
            if not r.atoms:
                continue
            base_cost = cert.cost()
            best_trial = None
            variants: list[tuple] = []
            for k in range(len(r.atoms)):
                variants.append(tuple(a for j, a in enumerate(r.atoms) if j != k))
                for b in domain.broader_atoms(r.atoms[k]):
                    if b in r.atoms:
                        continue
                    variants.append(tuple(b if j == k else a
                                          for j, a in enumerate(r.atoms)))
            forbidden_y = 0
            for yy in range(cert.ny):
                if yy != r.y:
                    forbidden_y |= obs_ymask[yy]
            later = 0
            for rr in cert.rules[i + 1:]:
                later |= rr.mask

            def try_rule(nr: Rule):
                nonlocal best_trial
                rules2 = list(cert.rules)
                rules2[i] = nr
                rules2 = [rr for j, rr in enumerate(rules2)
                          if j == i or not (rr.y == nr.y
                                            and rr.mask & ~nr.mask == 0)]
                trial = Certificate(cert.default, cert.ny, rules2)
                if trial.cost() >= base_cost - 1e-9:
                    return False
                if not trial.consistent(obs_ymask, universe):
                    return False
                if best_trial is None or trial.cost() < best_trial.cost():
                    best_trial = trial
                return True

            for atoms2 in variants:
                nr = make_rule(atoms2, r.y, universe)
                if try_rule(nr):
                    continue
                # generalize-then-respecialize: the broadened rule may claim
                # observed points of other classes; look for ONE additional
                # atom that excludes those while keeping the points this rule
                # correctly claims (e.g. N_X>=1 & diff=1 needs N_O=0).
                claimed = nr.mask & ~later
                viol = claimed & forbidden_y
                keep = claimed & obs_ymask[r.y]
                if viol == 0 or keep == 0:
                    continue
                for c in domain.atoms:
                    if c in atoms2:
                        continue
                    if keep & ~c.mask == 0 and c.mask & viol == 0:
                        try_rule(make_rule(atoms2 + (c,), r.y, universe))
                        break  # atoms are cost-sorted: first hit is cheapest
            if best_trial is not None:
                cert.rules = best_trial.rules
                cert.default = best_trial.default
                improved = True
    return cert


# ---------------------------------------------------------------------------
# Global greedy learner (backward construction), used as the "global
# restructuring" baseline: repeatedly commit the most cost-effective pure
# overwrite on the not-yet-settled observations, then reverse the order.
# ---------------------------------------------------------------------------

def greedy_learner(obs: dict[int, int], domain: Domain, max_atoms: int = 4,
                   seeds_per_class: int = 4, node_budget: int = 500_000,
                   rng: Optional[random.Random] = None) -> Certificate:
    rng = rng or random.Random(0)
    ny = domain.ny
    ymask = [0] * ny
    for x, y in obs.items():
        ymask[y] |= 1 << x
    obs_all = 0
    for m in ymask:
        obs_all |= m

    done = 0
    picks: list[Rule] = []
    while True:
        rem = obs_all & ~done
        mono = None
        if rem == 0:
            mono = max(range(ny), key=lambda y: popcount(ymask[y]))
        else:
            for y in range(ny):
                if rem & ~ymask[y] == 0:
                    mono = y
                    break
        if mono is not None:
            cert = Certificate(mono, ny, list(reversed(picks)))
            return cert
        best = None  # (score, scost, mask, y, desc)
        for y in range(ny):
            live = ymask[y] & ~done
            if live == 0:
                continue
            forbidden = (obs_all & ~ymask[y]) & ~done
            pts = []
            m = live
            while m and len(pts) < seeds_per_class * 8:
                low = m & -m
                pts.append(low.bit_length() - 1)
                m ^= low
            rng.shuffle(pts)
            for x in pts[:seeds_per_class]:
                scost, mask, atoms = cheapest_subset(
                    domain, x, forbidden, max_atoms, node_budget)
                gain = popcount(mask & obs_all & ~done)
                score = (RULE_FRAME + scost + choice(ny)) / max(gain, 1)
                if best is None or score < best[0]:
                    best = (score, atoms, y)
        assert best is not None
        _, atoms, y = best
        rule = make_rule(atoms, y, domain.universe)
        picks.append(rule)
        done |= rule.mask