"""Pointer-chasing dataset generation, shortcut checks, and the deterministic
lookup reference.

Task: a fresh uniform-random mapping f: S -> S per example, a start state x,
and a hop count k. The answer is f^k(x). Training uses k ~ Uniform{1..8};
evaluation uses k up to 64 and smaller state spaces.

Sequence serialization (see models.py): the model consumes
[f(0), f(1), ..., f(S-1), x, k]; the position of token i in the mapping part
encodes state i, so learned positional embeddings play the role of "state id".
Sequence length is S+2 regardless of k, so length generalisation here means
generalising over the hop count k, whose token embedding is OOD for k > 8.

Shortcut controls (per Liu et al. 2210.10749, Anil et al. 2207.04901,
Li et al. 2604.07822):
  - f is uniform random per example -> no counter-free / aperiodic structure,
    so no constant-depth lookup shortcut exists in principle.
  - answer marginals f^k(x) are approximately uniform for every k; we check
    this empirically (check_answer_marginals).
  - k is the ONLY length-correlated token; it sits at a fixed position.
  - fresh mappings every batch -> nothing to memorise; a held-out probe with
    never-seen f is generated on the fly at eval time.
"""

import numpy as np


def sample_mappings(rng, batch, S):
    """Uniform random f: S -> S for each example. Shape [batch, S]."""
    return rng.integers(0, S, size=(batch, S))


def iterate_mapping(f, x, k_max):
    """Apply f repeatedly to x. Returns outs of shape [k_max, batch] where
    outs[t-1] = f^t(x), for t = 1..k_max. Vectorised lookup."""
    f = np.asarray(f)
    y = np.asarray(x).copy()
    outs = np.empty((k_max, f.shape[0]), dtype=np.int64)
    for t in range(1, k_max + 1):
        y = np.take_along_axis(f, y[:, None], axis=1).squeeze(1)
        outs[t - 1] = y
    return outs


def sample_train_batch(rng, batch, S, k_min=1, k_max=8):
    """A training batch: fresh f per example, x uniform, k ~ U{k_min..k_max}.
    Returns dict with numpy arrays f [B,S], x [B], k [B], answer [B]."""
    f = sample_mappings(rng, batch, S)
    x = rng.integers(0, S, size=batch)
    k = rng.integers(k_min, k_max + 1, size=batch)
    all_ans = iterate_mapping(f, x, k_max)          # [k_max, B]
    ans = all_ans[k - 1, np.arange(batch)]
    return {"f": f, "x": x, "k": k, "answer": ans}


def sample_eval_batch(rng, batch, S, k):
    """An evaluation batch at a FIXED hop count k (fresh mappings, fresh x).
    answer = f^k(x)."""
    f = sample_mappings(rng, batch, S)
    x = rng.integers(0, S, size=batch)
    all_ans = iterate_mapping(f, x, k)
    ans = all_ans[k - 1, np.arange(batch)]
    return {"f": f, "x": x, "k": np.full(batch, k, dtype=np.int64), "answer": ans}


class LookupSolver:
    """Deterministic exact-answer reference: materialise f as a table and
    iterate k times. Must score 100% by construction; used as a pipeline
    integrity check in evaluation, not as a learned baseline."""

    def solve(self, f, x, k):
        y = np.asarray(x).copy()
        for _ in range(int(k)):
            y = np.take_along_axis(np.asarray(f), y[:, None], axis=1).squeeze(1)
        return y


def check_answer_marginals(rng, S=64, k_max=8, n=200_000, tol=0.02):
    """Empirical shortcut check: the marginal distribution of the answer
    f^k(x) must be ~uniform over S for every training k. Returns dict of
    max deviation from uniform (in total-variation distance) per k."""
    worst = {}
    for k in range(1, k_max + 1):
        b = sample_eval_batch(rng, n // 1000 * 1000, S, k)
        counts = np.bincount(b["answer"], minlength=S).astype(np.float64)
        p = counts / counts.sum()
        tv = 0.5 * np.abs(p - 1.0 / S).sum()
        worst[k] = float(tv)
    ok = all(v < tol for v in worst.values())
    return {"tv_per_k": worst, "tol": tol, "pass": ok}


def sequence_length(S):
    return S + 2