"""Tiny structured domains for exact experiments.

X = V^P where P is a disjoint union of "bare" blocks (sets with full
symmetric-group symmetry) and V is a bare value set.  The type induces:

  - evaluation maps  eval_p : X -> V        (one per position)
  - block counts     c_{b,v} : X -> {0..|b|} (the multiplicity space
                     V^b / S_b, coordinatized by counts)

Atoms are fibers/regions of these maps.  Costs follow the single
principle: sum of log2(#options) at each choice point of the derivation.
A derivation of an atom chooses: family, then the family's parameters,
then a region of the value space (kind + threshold).
"""

from __future__ import annotations

import itertools

from .core import Atom, Domain, choice, dedupe_atoms

REGION_KINDS = 3  # =c, >=c, <=c


def tiny_domain(blocks: list[int], nvals: int = 2,
                ylabels: tuple[str, ...] = ("0", "1"),
                vlabels: str = "abcdef") -> Domain:
    npos = sum(blocks)
    n = nvals ** npos
    nfam = 2  # eval, count
    tag = choice(nfam)

    def digit(x: int, p: int) -> int:
        return (x // nvals ** p) % nvals

    # block membership
    block_of = []
    for bi, sz in enumerate(blocks):
        block_of += [bi] * sz
    positions_of = [[p for p in range(npos) if block_of[p] == bi]
                    for bi in range(len(blocks))]

    atoms = []
    # eval atoms: choose family, position (flat among npos), value
    eval_cost = tag + choice(npos) + choice(nvals)
    eval_atom_cache: dict[tuple[int, int], Atom] = {}
    for p in range(npos):
        for v in range(nvals):
            mask = 0
            for x in range(n):
                if digit(x, p) == v:
                    mask |= 1 << x
            a = Atom(f"x[{p}]={vlabels[v]}", eval_cost, mask)
            atoms.append(a)
            eval_atom_cache[(p, v)] = a

    # count atoms: choose family, block, value, region kind, threshold
    for bi, sz in enumerate(blocks):
        for v in range(nvals):
            base = tag + choice(len(blocks)) + choice(nvals)
            for c in range(sz + 1):
                cost = base + choice(REGION_KINDS) + choice(sz + 1)
                m_eq = m_ge = m_le = 0
                for x in range(n):
                    cnt = sum(1 for p in positions_of[bi] if digit(x, p) == v)
                    if cnt == c:
                        m_eq |= 1 << x
                    if cnt >= c:
                        m_ge |= 1 << x
                    if cnt <= c:
                        m_le |= 1 << x
                bname = f"b{bi}" if len(blocks) > 1 else ""
                atoms.append(Atom(f"#{bname}{vlabels[v]}={c}", cost, m_eq))
                atoms.append(Atom(f"#{bname}{vlabels[v]}>={c}", cost, m_ge))
                atoms.append(Atom(f"#{bname}{vlabels[v]}<={c}", cost, m_le))

    universe = (1 << n) - 1
    atoms = dedupe_atoms(atoms, universe)

    def singleton_atoms(x: int) -> list[Atom]:
        return [eval_atom_cache[(p, digit(x, p))] for p in range(npos)]

    name = f"V{nvals}^P{'+'.join(map(str, blocks))}"
    return Domain(n, list(ylabels), atoms, singleton_atoms, name)


def all_functions(n: int, ny: int):
    """Iterate over every extensional function X -> Y as a tuple of labels."""
    return itertools.product(range(ny), repeat=n)