"""Tic-tac-toe as a structured domain.

Type (from the task statement):

    Coor  = {A,C} + {B}          Axis = {X,Y}
    Coor2 = Coor^Axis            Cell = {X,O} + {Empty}
    Field = Coor2 -> Cell        Result = {XWon, OWon, Draw, Illegal}

Structure derived from the type (all canonical constructions):

  - eval maps          F(p) for the 9 positions;
  - the multiplicity space Cell^9 / S_9, coordinatized by counts
    (#X, #O, #Empty), together with its difference relations;
  - the LINE family: fibers of the two projections Coor2 -> Coor
    (rows/columns) plus graphs of the automorphisms of Coor
    (the two diagonals) — 8 lines, not listed ad hoc;
  - per-line profiles: restriction to a line followed by the quotient
    Cell^3 / S_3, i.e. a profile (#X, #O, #Empty) with #=3, 10 values;
  - second-order multiplicities: for each profile p, the count
    N_p(F) = #{lines with profile p} in {0..8}  (e.g. N_{(3,0,0)} = number
    of X winning lines).

Atoms are fibers/regions of these maps; each atom's cost is the sum of
log2(#options) over its derivation's choice points.
"""

from __future__ import annotations

import itertools
import random

import numpy as np

from .core import Atom, Domain, choice, dedupe_atoms

EMPTY, PX, PO = 0, 1, 2
CELL_NAMES = {EMPTY: "Empty", PX: "X", PO: "O"}
RESULTS = ["Illegal", "XWon", "OWon", "Draw"]
ILLEGAL, XWON, OWON, DRAW = 0, 1, 2, 3

N_FIELDS = 3 ** 9
REGION_KINDS = 3  # =c, >=c, <=c


def bits_to_mask(arr: np.ndarray) -> int:
    """bool array -> python int bitmask (bit i = arr[i])."""
    return int.from_bytes(np.packbits(arr, bitorder="little").tobytes(), "little")


def mask_to_bits(mask: int, n: int) -> np.ndarray:
    nbytes = (n + 7) // 8
    b = np.frombuffer(mask.to_bytes(nbytes, "little"), dtype=np.uint8)
    return np.unpackbits(b, bitorder="little")[:n].astype(bool)


def build_fields() -> np.ndarray:
    """(19683, 9) uint8 array; cell j of field i is (i // 3**j) % 3."""
    idx = np.arange(N_FIELDS)
    return np.stack([(idx // 3 ** j) % 3 for j in range(9)], axis=1).astype(np.uint8)


# cells are indexed row-major: cell = 3*row + col, rows/cols in {0,1,2}
LINES = ([[3 * r + c for c in range(3)] for r in range(3)]      # rows: fibers of proj_Y
         + [[3 * r + c for r in range(3)] for c in range(3)]    # cols: fibers of proj_X
         + [[0, 4, 8], [2, 4, 6]])                              # graphs of Aut(Coor)
LINE_NAMES = ["row0", "row1", "row2", "col0", "col1", "col2", "diag", "anti"]

# the 10 profiles (#X, #O, #Empty) with sum 3
PROFILES = [(x, o, 3 - x - o) for x in range(4) for o in range(4 - x)]
PROFILE_IDX = {p: i for i, p in enumerate(PROFILES)}
WIN_X_PROFILE = PROFILE_IDX[(3, 0, 0)]
WIN_O_PROFILE = PROFILE_IDX[(0, 3, 0)]


class TTTData:
    def __init__(self):
        A = build_fields()
        self.fields = A
        self.counts = np.stack([(A == v).sum(axis=1) for v in (PX, PO, EMPTY)],
                               axis=1)  # columns: nX, nO, nEmpty
        lut = np.zeros((4, 4), dtype=np.int8)  # (nX, nO) -> profile id
        for (x, o, _e), i in PROFILE_IDX.items():
            lut[x, o] = i
        lp = np.zeros((N_FIELDS, 8), dtype=np.int8)
        for li, line in enumerate(LINES):
            sub = A[:, line]
            lp[:, li] = lut[(sub == PX).sum(axis=1), (sub == PO).sum(axis=1)]
        self.line_profiles = lp
        pc = np.zeros((N_FIELDS, len(PROFILES)), dtype=np.int8)
        for pi in range(len(PROFILES)):
            pc[:, pi] = (lp == pi).sum(axis=1)
        self.profile_counts = pc
        self.wx = pc[:, WIN_X_PROFILE].astype(int)
        self.wo = pc[:, WIN_O_PROFILE].astype(int)


def labels_simple(d: TTTData) -> np.ndarray:
    """The extension of the hand-written certificate in the task statement:
    Illegal by default; Draw on full boards with no winning line; OWon on
    any O-line; XWon on any X-line (applied last, so it wins ties)."""
    nx, no, ne = d.counts[:, 0], d.counts[:, 1], d.counts[:, 2]
    lab = np.full(N_FIELDS, ILLEGAL, dtype=np.int8)
    lab[(ne == 0) & (d.wx == 0) & (d.wo == 0)] = DRAW
    lab[d.wo > 0] = OWON
    lab[d.wx > 0] = XWON
    return lab


def labels_legal(d: TTTData) -> np.ndarray:
    """Reachable-terminal classification: XWon / OWon / Draw for positions
    reachable in play whose game has ended; everything else Illegal
    (including non-terminal positions, per the task statement)."""
    nx, no, ne = d.counts[:, 0], d.counts[:, 1], d.counts[:, 2]
    diff = nx.astype(int) - no.astype(int)
    lab = np.full(N_FIELDS, ILLEGAL, dtype=np.int8)
    lab[(d.wx > 0) & (d.wo == 0) & (diff == 1)] = XWON
    lab[(d.wo > 0) & (d.wx == 0) & (diff == 0)] = OWON
    lab[(ne == 0) & (d.wx == 0) & (d.wo == 0) & (diff == 1)] = DRAW
    return lab


def build_domain(d: TTTData) -> Domain:
    """Atom library derived from the type structure."""
    atoms: list[Atom] = []
    nfam = 5  # cell, count, lineprofile, profilecount, countcompare
    tag = choice(nfam)
    A = d.fields

    # 1. cell values: family, position (9), value (3)
    cell_cost = tag + choice(9) + choice(3)
    cell_atoms: dict[tuple[int, int], Atom] = {}
    for p in range(9):
        r, c = divmod(p, 3)
        for v in (EMPTY, PX, PO):
            a = Atom(f"cell({r},{c})={CELL_NAMES[v]}", cell_cost,
                     bits_to_mask(A[:, p] == v))
            atoms.append(a)
            cell_atoms[(p, v)] = a

    # 2. global counts: family, value (3), region kind (3), threshold (10)
    for vi, vname in enumerate(("X", "O", "Empty")):
        base = tag + choice(3)
        col = d.counts[:, vi]
        for cth in range(10):
            cost = base + choice(REGION_KINDS) + choice(10)
            atoms.append(Atom(f"#{vname}={cth}", cost, bits_to_mask(col == cth)))
            atoms.append(Atom(f"#{vname}>={cth}", cost, bits_to_mask(col >= cth)))
            atoms.append(Atom(f"#{vname}<={cth}", cost, bits_to_mask(col <= cth)))

    # 3. per-line profiles: family, line (8), profile (10)
    lp_cost = tag + choice(8) + choice(10)
    for li in range(8):
        for pi, prof in enumerate(PROFILES):
            atoms.append(Atom(f"prof[{LINE_NAMES[li]}]={prof}", lp_cost,
                              bits_to_mask(d.line_profiles[:, li] == pi)))

    # 4. profile counts N_p: family, profile (10), region kind (3), threshold (9)
    for pi, prof in enumerate(PROFILES):
        base = tag + choice(10)
        col = d.profile_counts[:, pi]
        for cth in range(9):
            cost = base + choice(REGION_KINDS) + choice(9)
            atoms.append(Atom(f"N{prof}={cth}", cost, bits_to_mask(col == cth)))
            atoms.append(Atom(f"N{prof}>={cth}", cost, bits_to_mask(col >= cth)))
            atoms.append(Atom(f"N{prof}<={cth}", cost, bits_to_mask(col <= cth)))

    # 5. count comparisons: family, ordered value pair (6), kind (2: =, >=),
    #    offset in -9..9 (19)
    names = ("X", "O", "Empty")
    for ui, vi in itertools.permutations(range(3), 2):
        base = tag + choice(6)
        du = d.counts[:, ui].astype(int) - d.counts[:, vi].astype(int)
        for off in range(-9, 10):
            cost = base + choice(2) + choice(19)
            atoms.append(Atom(f"#{names[ui]}-#{names[vi]}={off}", cost,
                              bits_to_mask(du == off)))
            atoms.append(Atom(f"#{names[ui]}-#{names[vi]}>={off}", cost,
                              bits_to_mask(du >= off)))

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

    def singleton_atoms(x: int) -> list[Atom]:
        return [cell_atoms[(p, int(A[x, p]))] for p in range(9)]

    return Domain(N_FIELDS, RESULTS, atoms, singleton_atoms, "tictactoe")


# ---------------------------------------------------------------------------
# Symmetry group (D4 from the type: Aut(Coor)^Axis ⋊ Aut(Axis)) and sampling
# ---------------------------------------------------------------------------

def d4_perms() -> list[list[int]]:
    def transpose(p):  # swap the two axes
        return [(i % 3) * 3 + i // 3 for i in p]

    def hflip(p):      # A<->C on one axis
        return [(i // 3) * 3 + (2 - i % 3) for i in p]

    perms = {tuple(range(9))}
    frontier = [list(range(9))]
    while frontier:
        nxt = []
        for p in frontier:
            for q in (transpose(p), hflip(p)):
                if tuple(q) not in perms:
                    perms.add(tuple(q))
                    nxt.append(q)
        frontier = nxt
    return [list(p) for p in perms]


def orbits(d: TTTData) -> list[list[int]]:
    perms = d4_perms()
    pow3 = 3 ** np.arange(9)
    A = d.fields.astype(np.int64)
    images = np.stack([(A[:, perm] * pow3).sum(axis=1) for perm in perms], axis=1)
    canon = images.min(axis=1)
    reps: dict[int, list[int]] = {}
    for i, c in enumerate(canon):
        reps.setdefault(int(c), []).append(i)
    return list(reps.values())


def orbit_stratified_sample(d: TTTData, labels: np.ndarray, frac: float,
                            rng: random.Random) -> list[int]:
    """Take whole D4-orbits, stratified by class so every class contributes
    ~frac of its orbits (orbits are label-pure since labels are invariant)."""
    by_class: dict[int, list[list[int]]] = {}
    for orb in orbits(d):
        by_class.setdefault(int(labels[orb[0]]), []).append(orb)
    out: list[int] = []
    for y, orbs in by_class.items():
        rng.shuffle(orbs)
        k = max(1, round(frac * len(orbs)))
        for orb in orbs[:k]:
            out.extend(orb)
    rng.shuffle(out)
    return out


def certificate_predictions(cert, n: int = N_FIELDS) -> np.ndarray:
    pred = np.full(n, cert.default, dtype=np.int8)
    for r in cert.rules:
        pred[mask_to_bits(r.mask, n)] = r.y
    return pred