"""Vectorized cube moves over numpy arrays of facelet indices.

`gen_data.apply_move` rebuilds a 54-character Python string per move, which is
fine when the bottleneck is a Kociemba call. Value iteration has no solver in
the loop: every training step expands all 18 children of every state in the
batch, so at batch 1024 that is ~18k move applications per step and the string
path becomes the whole cost.

A move is a fixed permutation of the 54 facelets -- `new[i] = old[perm[i]]` --
so a batch of states is one fancy-index away: `states[:, perm]`. Quarter turns
compose by indexing a permutation with itself (`p2 = p[p]`), which is how the
half- and counter-turns are built rather than by applying the quarter turn
repeatedly.

States are (N, 54) uint8 of colour indices into COLORS, not strings. Use
`encode`/`decode` at the boundaries only.
"""
import numpy as np

from gen_data import QT, MOVES, SOLVED

COLORS = "URFDLB"
_C2I = {c: i for i, c in enumerate(COLORS)}

# MOVES order is (face, suffix) with suffix in ("", "'", "2") -- 1, 3 and 2
# quarter turns respectively. PERMS is indexed by position in MOVES so a move
# index from the model maps straight to a permutation.
_TURNS = {"": 1, "'": 3, "2": 2}


def _compose(perm: np.ndarray, times: int) -> np.ndarray:
    out = np.arange(54, dtype=np.int64)
    for _ in range(times):
        out = out[perm]
    return out


PERMS = np.stack([_compose(np.array(QT[m[0]], dtype=np.int64), _TURNS[m[1:]])
                  for m in MOVES])
NUM_MOVES = len(MOVES)
assert PERMS.shape == (NUM_MOVES, 54)


def encode(states) -> np.ndarray:
    """Facelet string(s) -> (N, 54) uint8."""
    if isinstance(states, str):
        states = [states]
    return np.array([[_C2I[c] for c in s] for s in states], dtype=np.uint8)


def decode(arr: np.ndarray):
    """(N, 54) uint8 -> list of facelet strings."""
    return ["".join(COLORS[i] for i in row) for row in np.asarray(arr)]


SOLVED_ARR = encode(SOLVED)[0]


def apply_moves(states: np.ndarray, move_ids) -> np.ndarray:
    """Applies one move per state. `move_ids` is an int array of length N."""
    rows = np.arange(len(states))[:, None]
    return states[rows, PERMS[np.asarray(move_ids)]]


def all_children(states: np.ndarray) -> np.ndarray:
    """(N, 54) -> (N, 18, 54): every state's 18 successors, in MOVES order."""
    return states[:, PERMS]


def is_solved(states: np.ndarray) -> np.ndarray:
    """(N, 54) -> (N,) bool."""
    return (states == SOLVED_ARR).all(axis=1)


def random_walk(n: int, max_depth: int, rng: np.random.Generator):
    """n states from independent walks off solved, depths uniform in 1..max_depth.

    Returns (states, depths). The walk never turns the same face twice in a row,
    matching gen_data.random_state -- consecutive turns of one face are a single
    different move, so allowing them biases the effective depth downward.

    Depths are returned for reporting only. They are an upper bound on the true
    distance, never the label: value iteration derives its own targets, and a
    walk of length 12 routinely lands somewhere far nearer to solved.
    """
    depths = rng.integers(1, max_depth + 1, size=n)
    states = np.repeat(SOLVED_ARR[None, :], n, axis=0)
    prev_face = np.full(n, -1)
    for step in range(max_depth):
        live = np.where(depths > step)[0]
        if not len(live):
            break
        mv = rng.integers(0, NUM_MOVES, size=len(live))
        clash = (mv // 3) == prev_face[live]
        while clash.any():
            mv[clash] = rng.integers(0, NUM_MOVES, size=int(clash.sum()))
            clash = (mv // 3) == prev_face[live]
        states[live] = states[live][np.arange(len(live))[:, None], PERMS[mv]]
        prev_face[live] = mv // 3
    return states, depths