"""
J-SPACE LEVEL 9c — РЕШАЮЩИЙ ТЕСТ (исправленная версия)
=====================================================
ИСПРАВЛЕНИЯ ОТНОСИТЕЛЬНО L9 (важно для статьи — числа меняются, НЕ переносить старые)
-----------------------------------------------------------------------------------
[FIX-1, блокирующий]  workspace_score раньше брался как max switch% ПО ВСЕМ слоям,
    включая L=N-1. Но интервенция на L=N-1 применяется ПОСЛЕ последнего блока —
    между правкой '=' и ридаутом не остаётся ни одного слоя, только ln_f+head_main
    на позиции '='. patch хирургичен (меняет координату лишь в подпространстве
    {mu_j,mu_i}, сохраняя a,b в ортогональном дополнении), поэтому на L=N-1 ридаут
    почти гарантированно выдаёт op_answer(i,a,b) у ЛЮБОГО загроккавшего варианта.
    Это тавтология — она не отличает A от D.
    Теперь:
        workspace_score = best switch% по слоям 0..N-2  (есть downstream-чтение)
        ws_readout      = switch% на L=N-1              (тавтологический контроль)
    Несущий сигнал живёт на РАННИХ слоях, где правка '=' вынуждает нижележащие
    блоки реально её прочитать, а не переписать заново из a/b.

[FIX-2]  setcoord для режимов 'src'/'all' навязывал позиции-источнику величину
    координаты, посчитанную на '=' (target=ci с последней позиции). Масштабы
    проекций по позициям разные -> цель размерно не та, input_score/all_score
    через setcoord считались мусорно. Теперь target — per-position: ci_pos[s].

[FIX-3]  T3 брал строку с max(op_followed) и на ней мерил разрыв с main — отбор по
    той же величине, которую хотим показать большой. Теперь три числа:
        existence — разрыв на max(op_followed): СУЩЕСТВУЕТ интервенция, где отчёт
                    следует, а вычисление нет (честно как существование);
        at_wsopt  — разрыв на max(switched):    там, где вычислению легче всего
                    последовать — если op всё равно впереди, конфабуляция робастна;
        gap_mean  — средний парный разрыв по всем строкам режима.

[SELF-CONTAINED]  Загрузка весов L8 удалена. Каталог checkpoints_l9c. Первый прогон
    обучает все модели с нуля; далее кэшируются только собственные чекпоинты.
    --retrain принудительно переобучает.

[E-FIX / setclass]  Для многоклассового E (4 класса) парный patch(0,1) тестирует лишь
    ОДНУ под-ось правила и занижает switch (garbage растёт из-за 4 аттракторов). Добавлена
    интервенция setclass: жёстко выставляет классовую идентичность резидуала в целевой
    класс во ВСЁМ между-классовом подпространстве, сохраняя ортогональное дополнение (a,b);
    switched/op усредняются по всем целям. Для бинарных сводится к тому же 1D-направлению,
    что и парный патч -> числа A–H НЕ меняются, setclass идёт сверкой (WSкласс).
    Несущая интервенция выбирается автоматически: setclass для n_ops>2, парная для бинарных.
    E теперь в дефолтном наборе (последним; d=128, 4 класса — МЕДЛЕННО).

ЧТО ПОКАЗАЛ L8 (контекст)
-------------------------
1. ТРИХОТОМИЯ PATCH (грубое вмешательство, все позиции):
       A,B  чётность (читается)      switch  6-8%   garbage 88%   ЛОМАЕТ
       G    a>b (строить, дёшево)    switch  3%     garbage 93%   ЛОМАЕТ
       F    (a+b)%2 (вырожден)       switch  1.5%   garbage 33%   ИГНОРИРУЕТСЯ
       C    QR (строить, дорого)     switch 19%     garbage 73%   частично
       D    QR + 4 потребителя       switch 53%     garbage 43%   ПЕРЕНАПРАВЛЯЕТ
2. КОНФАБУЛЯЦИЯ: голова 'op' СЛЕДУЕТ за патчем в 87-93% даже там, где вычисление
   НЕ следует. Вербальный отчёт диссоциирован от вычисления.
3. Мягкое вмешательство (setcoord) переключает вычисление во ВСЕХ вариантах:
   ось правила каузальна везде — она ХРУПКАЯ в A/B/G и РОБАСТНАЯ в D.

КРИТИЧЕСКИЙ КОНФАУНД (мотивация L9)
----------------------------------
В L8 все вмешательства применялись КО ВСЕМ ПОЗИЦИЯМ, включая токен 'a'. Но в A
правило И ЕСТЬ признак в 'a' — мы редактировали ВХОД, а не рабочее пространство.

РЕШАЮЩИЙ ТЕСТ (T2)
------------------
Вмешательство в три режима:
    'eq'  — только последний токен '='   <- ПРАВИМ РАБОЧЕЕ ПРОСТРАНСТВО
    'src' — только токен-источник правила <- РЕДАКТИРУЕМ ВХОД (контроль)
    'all' — все позиции (как в L8)
    Переключается ли вычисление, когда правится ТОЛЬКО '=' на РАННЕМ слое?
        ДА  -> правило на '=' НЕСУЩЕЕ: настоящее рабочее пространство.
        НЕТ -> сеть читает правило из резидуала источника: маршрутизация.

КОНТРОЛЬ H — «строить, но дёшево»
---------------------------------
H: правило задаётся ТРЕТЬИМ ТОКЕНОМ-КЛЮЧОМ k из K=8 значений (перестановка сидом).
   Не читается из a/b; вычислительно тривиально (таблица 8 строк); но требует
   ПОСТРОЙКИ (k->класс — выученное отображение, не тождество).
       H ≈ D -> достаточно ПОСТРОЙКИ
       H ≈ A -> нужна ДОРОГОВИЗНА пересчёта (амортизация)
   У H четыре позиции [a,b,k,=]; режим 'src' правит токен-ключ (rule_pos=2).

ЗАПУСК
------
    python jspace_level9c.py                 # A,B,C,D,F,H × 5 сидов, с нуля
    python jspace_level9c.py --variants E    # алфавит (медленно)
    python jspace_level9c.py --seeds 7       # больше сидов
    python jspace_level9c.py --retrain       # принудительно переобучить
"""
import os
import json
import time
import argparse
from itertools import permutations
from datetime import datetime
import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
try:
    from sklearn.cluster import KMeans
    from sklearn.metrics import silhouette_score
    SKLEARN = True
except ImportError:
    SKLEARN = False
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
P = 97
EQ = 97
PAD = 98
KEY0 = 99                 # токены-ключи: 99..106 (K=8)
K_KEYS = 8
VOCAB_BASE = 99                    # словарь для правил без ключа
VOCAB_KEY = KEY0 + K_KEYS          # 107 — только для варианта H
N_HEADS, N_LAYERS = 4, 4
LAST = N_LAYERS - 1                 # слой ридаута: интервенция на нём тавтологична для 'eq'


def vocab_for(rule):
    """Пер-вариантный словарь: H (правило от токена-ключа) требует 8 доп. токенов."""
    return VOCAB_KEY if rule == "key" else VOCAB_BASE


CKPT_DIR, LOG_DIR = "checkpoints_l9c", "logs_l9c"
TRACK_DIR = "track_l9c"
os.makedirs(CKPT_DIR, exist_ok=True)
os.makedirs(LOG_DIR, exist_ok=True)
os.makedirs(TRACK_DIR, exist_ok=True)
# ---- ОНТОГЕНЕЗ ЗНАКА (T10) ----
TRACK_EVERY = 250          # снимать метрики представления каждые N эпох
POST_GROK_EPOCHS = 20000   # продолжать обучение ПОСЛЕ грокинга (кристаллизация?)
VARIANTS = {
    "A": dict(rule="parity", heads="single", n_ops=2, d=64,
              desc="чётность a — ЧИТАЕТСЯ из входа, 1 потребитель"),
    "B": dict(rule="parity", heads="multi",  n_ops=2, d=64,
              desc="чётность a — ЧИТАЕТСЯ из входа, 4 потребителя"),
    "C": dict(rule="qr",     heads="single", n_ops=2, d=64,
              desc="QR(a) — СТРОИТЬ + ДОРОГО, 1 потребитель"),
    "D": dict(rule="qr",     heads="multi",  n_ops=2, d=64,
              desc="QR(a) — СТРОИТЬ + ДОРОГО, 4 потребителя"),
    "F": dict(rule="sumpar", heads="multi",  n_ops=2, d=64,
              desc="(a+b)%2 — ВЫРОЖДЕННЫЙ (правило = функция ответа)"),
    "H": dict(rule="key",    heads="multi",  n_ops=2, d=64,
              desc="токен-ключ -> правило — СТРОИТЬ, но ДЁШЕВО  <<< КОНТРОЛЬ"),
    "E": dict(rule="qrqr",   heads="multi",  n_ops=4, d=128,
              desc="(QR(a),QR(b)) — 4 класса, оба строить  <<< АЛФАВИТ"),
}
DEFAULT_ORDER = ["A", "B", "C", "D", "F", "H", "E"]  # E последним: 4 класса, d=128, МЕДЛЕННО
HEADS_MULTI = ["main", "parity", "big", "op"]
HEADS_SINGLE = ["main"]
CONTROL = "triv"
QR_SET = {(x * x) % P for x in range(1, P)}
OP_NAMES = ["a+b", "a*b", "a-b", "a+2b"]


def op_answer(k, a, b):
    return [(a + b) % P, (a * b) % P, (a - b) % P, (a + 2 * b) % P][k]


# ==========================================================
# ДАТАСЕТ
# ==========================================================
def make_dataset(rule, seed=0, train_frac=0.8):
    """
    Возвращает (TR, TE, meta). Для rule='key' последовательность [a,b,k,=],
    для остальных [a,b,=]. meta['rule_pos'] — позиция токена-источника правила.
    """
    a, b = torch.meshgrid(torch.arange(P), torch.arange(P), indexing="ij")
    a, b = a.flatten(), b.flatten()
    is_qr = lambda t: torch.tensor([0 if int(x) in QR_SET else 1 for x in t]).long()
    key = None
    if rule == "parity":
        cls = (a % 2).long(); rule_pos = 0
    elif rule == "qr":
        cls = is_qr(a); rule_pos = 0
    elif rule == "sumpar":
        cls = ((a + b) % 2).long(); rule_pos = 0
    elif rule == "qrqr":
        cls = (is_qr(a) * 2 + is_qr(b)).long(); rule_pos = 0
    elif rule == "key":
        # ключ k из 8; отображение k -> класс фиксировано сидом (не тождество)
        g = torch.Generator().manual_seed(1000 + seed)
        key = torch.randint(0, K_KEYS, (a.numel(),), generator=g)
        perm = torch.randperm(K_KEYS, generator=g)
        key2cls = torch.zeros(K_KEYS, dtype=torch.long)
        key2cls[perm[:K_KEYS // 2]] = 0
        key2cls[perm[K_KEYS // 2:]] = 1
        cls = key2cls[key]
        rule_pos = 2                      # токен ключа
    else:
        raise ValueError(rule)
    n_ops = 4 if rule == "qrqr" else 2
    y_main = torch.zeros_like(a)
    for k in range(n_ops):
        y_main = torch.where(cls == k, op_answer(k, a, b), y_main)
    if rule == "key":
        X = torch.stack([a, b, KEY0 + key, torch.full_like(a, EQ)], dim=1)
    else:
        X = torch.stack([a, b, torch.full_like(a, EQ)], dim=1)
    data = dict(X=X, main=y_main, cls=cls, a=a, b=b,
                parity=(y_main % 2 == 0).long(),
                big=(y_main > P // 2).long(),
                op=cls,
                triv=(b % 2 == 0).long())
    g = torch.Generator().manual_seed(seed)
    idx = torch.randperm(X.size(0), generator=g)
    split = int(len(idx) * train_frac)
    pack = lambda ii: {k: v[ii].to(device) for k, v in data.items()}
    # насколько правило читается из токена a (доля a с однозначным классом)
    pure_a = float(np.mean([len(set(cls[a == v].tolist())) == 1 for v in range(P)]))
    meta = dict(rule=rule, n_ops=n_ops, seq_len=X.size(1), rule_pos=rule_pos,
                readable_from_a=round(pure_a, 3),
                class_balance=[round(float((cls == k).float().mean()), 3) for k in range(n_ops)])
    if rule == "key":
        meta["key2cls"] = key2cls.tolist()
    return pack(idx[:split]), pack(idx[split:]), meta


# ==========================================================
# МОДЕЛЬ
# ==========================================================
class Block(nn.Module):
    def __init__(self, d, h):
        super().__init__()
        self.ln1, self.ln2 = nn.LayerNorm(d), nn.LayerNorm(d)
        self.attn = nn.MultiheadAttention(d, h, batch_first=True)
        self.mlp = nn.Sequential(nn.Linear(d, 4 * d), nn.GELU(), nn.Linear(4 * d, d))

    def forward(self, x):
        n = self.ln1(x)
        o, _ = self.attn(n, n, n)
        x = x + o
        return x + self.mlp(self.ln2(x)), o, None


class Net(nn.Module):
    def __init__(self, heads_mode, n_ops, d, seq_len, vocab=VOCAB_BASE):
        super().__init__()
        self.heads_mode, self.n_ops, self.d, self.vocab = heads_mode, n_ops, d, vocab
        self.embed = nn.Embedding(vocab, d)
        self.pos = nn.Parameter(torch.randn(1, seq_len, d) * 0.02)
        self.layers = nn.ModuleList([Block(d, N_HEADS) for _ in range(N_LAYERS)])
        self.ln_f = nn.LayerNorm(d)
        self.head_main = nn.Linear(d, vocab)
        self.head_triv = nn.Linear(d, 2)
        if heads_mode == "multi":
            self.head_parity = nn.Linear(d, 2)
            self.head_big = nn.Linear(d, 2)
            self.head_op = nn.Linear(d, n_ops)

    def active_heads(self):
        return HEADS_MULTI if self.heads_mode == "multi" else HEADS_SINGLE

    def forward(self, x, intervention=None, cache=False):
        res = self.embed(x) + self.pos[:, :x.size(1), :]
        C = {"res": [res.clone()]} if cache else None
        for i, layer in enumerate(self.layers):
            res, _, _ = layer(res)
            if intervention is not None and intervention["layer"] == i:
                res = intervene(res, intervention)
            if cache:
                C["res"].append(res.clone())
        h = self.ln_f(res[:, -1, :])
        out = {"main": self.head_main(h), "triv": self.head_triv(h)}
        if self.heads_mode == "multi":
            out["parity"] = self.head_parity(h)
            out["big"] = self.head_big(h)
            out["op"] = self.head_op(h)
        return (out, C) if cache else out


def positions_for(mode, seq_len, rule_pos):
    """Какие позиции правим: 'eq' — только последняя, 'src' — токен-источник, 'all' — все."""
    if mode == "eq":
        return [seq_len - 1]
    if mode == "src":
        return [rule_pos]
    return list(range(seq_len))


def intervene(res, iv):
    """
    patch    — обмен координат в базисе {v_src, v_tgt} на ЗАДАННЫХ позициях
    setcoord — вырезать проекцию на ось и подставить target на ЗАДАННЫХ позициях
               (target может быть скаляром ИЛИ тензором [S] — per-position)  [FIX-2]
    ablate   — занулить проекцию
    add      — грубое прибавление (случайный контроль)
    """
    res = res.clone()
    t = iv["type"]
    pos = iv.get("positions", list(range(res.size(1))))
    if t == "patch":
        vs, vt, al = iv["v_src"], iv["v_tgt"], iv.get("alpha", 1.0)
        for s in pos:
            V = torch.stack([vs[s], vt[s]], dim=1)
            c = res[:, s, :] @ torch.linalg.pinv(V).T
            res[:, s, :] = res[:, s, :] + al * ((c.flip(-1) - c) @ V.T)
        return res
    if t == "setcoord":
        u, tgt = iv["axis"], iv["target"]
        for s in pos:
            ts = tgt[s] if torch.is_tensor(tgt) else tgt   # [FIX-2] per-position target
            cur = res[:, s, :] @ u[s]
            res[:, s, :] = res[:, s, :] + (ts - cur).unsqueeze(-1) * u[s].unsqueeze(0)
        return res
    if t == "setclass":
        # [E-FIX] жёстко выставить КЛАССОВУЮ идентичность в целевой класс i во ВСЁМ
        # между-классовом подпространстве, сохранив ортогональное дополнение (там a,b).
        # Обобщение парного патча на n классов. Для n_ops=2 сводится к 1D-оси.
        cmeans, g, basis, i = iv["cmeans"], iv["grand"], iv["basis"], iv["target"]
        for s in pos:
            Q = basis[s]                                   # [d, r] ортонормир. столбцы
            x = res[:, s, :] - g[s]                         # центрируем
            x_between = (x @ Q) @ Q.T                       # проекция на между-классовое
            x_within = x - x_between                        # ортог. дополнение (a,b)
            tgt_between = cmeans[i][s] - g[s]               # идентичность класса i
            res[:, s, :] = g[s] + x_within + tgt_between
        return res
    if t == "ablate":
        v = iv["vec"]
        for s in pos:
            vs_ = v[s]
            pr = (res[:, s, :] @ vs_) / ((vs_ * vs_).sum() + 1e-8)
            res[:, s, :] = res[:, s, :] - pr.unsqueeze(-1) * vs_.unsqueeze(0)
        return res
    if t == "add":
        v, al = iv["vec"], iv.get("alpha", 1.0)
        for s in pos:
            vn = v[s] / (v[s].norm() + 1e-8)
            sc = res[:, s, :].norm(dim=-1, keepdim=True).mean()
            res[:, s, :] = res[:, s, :] + al * vn.unsqueeze(0) * sc
        return res
    raise ValueError(t)


# ==========================================================
# T10. ОНТОГЕНЕЗ ЗНАКА — когда рождается ось правила?
# ==========================================================
def track_snapshot(model, TR, TE, n_ops, meta, pair=(0, 1)):
    """
    Быстрый срез представления. ws_switch здесь считается по РАННИМ слоям (0..N-2),
    как и headline-метрика в T2 [FIX-1] — трекинг тавтологического L=N-1 бессмыслен.
    """
    i, j = pair
    model.eval()
    S, rp = meta["seq_len"], meta["rule_pos"]
    with torch.no_grad():
        acc = float((model(TE["X"])["main"].argmax(1) == TE["main"]).float().mean())
    with torch.no_grad():
        _, Ctr = model(TR["X"], cache=True)
    emb = Ctr["res"][0]
    e_i = emb[TR["cls"] == i].mean(0)[rp]
    e_j = emb[TR["cls"] == j].mean(0)[rp]
    u_emb = (e_i - e_j)
    u_emb = u_emb / (u_emb.norm() + 1e-8)
    m = TE["cls"] == j
    X, a_, b_ = TE["X"][m], TE["a"][m], TE["b"][m]
    ans_tgt = op_answer(i, a_, b_)
    pos_eq = [S - 1]
    best = dict(ws_switch=0.0, layer=-1, axis_norm=0.0, cos_to_input=0.0, silhouette=0.0)
    for L in range(LAST):                         # [FIX-1] только 0..N-2
        r = Ctr["res"][L + 1]
        mu_i = r[TR["cls"] == i].mean(0)
        mu_j = r[TR["cls"] == j].mean(0)
        delta = mu_i - mu_j                                # [S, d]
        with torch.no_grad():
            o = model(X, intervention=dict(layer=L, type="patch", positions=pos_eq,
                                           v_src=mu_j, v_tgt=mu_i, alpha=1.0))
            sw = float((o["main"].argmax(1) == ans_tgt).float().mean()) * 100
        if sw > best["ws_switch"]:
            d_eq = delta[-1]                               # ось правила на токене '='
            cos = float(torch.abs(torch.dot(d_eq / (d_eq.norm() + 1e-8), u_emb)))
            sil = 0.0
            if SKLEARN:
                with torch.no_grad():
                    _, Cte = model(TE["X"], cache=True)
                v = Cte["res"][L + 1][:, -1, :].cpu().numpy()
                try:
                    km = KMeans(n_ops, random_state=0, n_init=5).fit_predict(v)
                    sil = float(silhouette_score(v, km))
                except Exception:
                    sil = 0.0
            best = dict(ws_switch=round(sw, 2), layer=L,
                        axis_norm=round(float(d_eq.norm()), 4),
                        cos_to_input=round(cos, 4), silhouette=round(sil, 4))
    return dict(acc=round(acc, 4), **best)


# ==========================================================
# ОБУЧЕНИЕ
# ==========================================================
def accs(out, D, keys):
    return {k: (out[k].argmax(1) == D[k]).float().mean().item() for k in keys if k in out}


def train_variant(vid, seed, retrain, target=0.95, track=False):
    """
    track=True -> каждые TRACK_EVERY эпох снимается срез представления (T10),
    и обучение продолжается ещё POST_GROK_EPOCHS эпох после грокинга.
    SELF-CONTAINED: загрузка весов L8 удалена; используются только собственные
    чекпоинты L9c. Первый прогон обучает с нуля.
    """
    cfg = VARIANTS[vid]
    max_epochs = 60000 if cfg["n_ops"] == 4 else 40000
    TR, TE, meta = make_dataset(cfg["rule"], seed=seed)
    torch.manual_seed(seed)
    vocab = vocab_for(cfg["rule"])
    model = Net(cfg["heads"], cfg["n_ops"], cfg["d"], meta["seq_len"], vocab).to(device)
    tag = f"{vid}_s{seed}" + ("_trk" if track else "")
    ckpt = os.path.join(CKPT_DIR, tag + ".pth")
    trk_path = os.path.join(TRACK_DIR, f"{vid}_s{seed}.json")
    # ПЕРЕИСПОЛЬЗОВАНИЕ ТОЛЬКО СОБСТВЕННЫХ ВЕСОВ L9c (никаких L8).
    # При --track переобучение обязательно: нужна траектория по ходу обучения.
    if not retrain and not track and os.path.exists(ckpt):
        try:
            sd = torch.load(ckpt, map_location=device, weights_only=True)
            model.load_state_dict(sd)
            print(f"[+] {tag}: веса из собственного чекпоинта L9c ({ckpt})")
            with torch.no_grad():
                g = float((model(TE["X"])["main"].argmax(1) == TE["main"]).float().mean()) >= target
            tl = dict(loaded=True, source="L9c", grokked=bool(g))
            if not tl["grokked"]:
                print(f"  [!] загруженная модель НЕ ЗАГРОККАЛА — из агрегата исключается")
            return model, TR, TE, meta, tl
        except Exception as ex:
            print(f"[!] чекпоинт {ckpt} не подошёл ({type(ex).__name__}) — обучаю с нуля")
    print(f"\n[*] {tag}: {cfg['desc']}")
    print(f"    правило читается из a: {meta['readable_from_a']*100:.0f}% | "
          f"позиция источника: {meta['rule_pos']} | seq_len {meta['seq_len']}")
    if track:
        print(f"    [T10] трекинг каждые {TRACK_EVERY} эпох; после грокинга "
              f"ещё {POST_GROK_EPOCHS} эпох (кристаллизация)")
    opt = torch.optim.AdamW(model.parameters(), lr=1e-3, weight_decay=1.0, betas=(0.9, 0.98))
    keys = model.active_heads() + [CONTROL]
    curve, trace, t0, grok = [], [], time.time(), None
    stop_at = max_epochs
    for e in range(max_epochs):
        model.train(); opt.zero_grad()
        out = model(TR["X"])
        loss = sum(F.cross_entropy(out[k], TR[k]) for k in keys)
        loss.backward(); opt.step()
        if track and e % TRACK_EVERY == 0:
            snap = track_snapshot(model, TR, TE, cfg["n_ops"], meta)
            snap["epoch"] = e
            snap["grokked"] = grok is not None
            trace.append(snap)
            if e % (TRACK_EVERY * 4) == 0:
                print(f"  [T10] E{e:6d} acc {snap['acc']:.3f} | ws_switch {snap['ws_switch']:5.1f}% "
                      f"| |ось| {snap['axis_norm']:.3f} | cos→вход {snap['cos_to_input']:.3f} "
                      f"| sil {snap['silhouette']:.3f} | L{snap['layer']}")
        if e % 500 == 0:
            model.eval()
            with torch.no_grad():
                a_te = accs(model(TE["X"]), TE, keys)
            curve.append(dict(epoch=e, loss=round(loss.item(), 4),
                              test={k: round(v, 4) for k, v in a_te.items()}))
            if not track or e % 2000 == 0:
                print(f"  E{e:6d} loss {loss.item():6.3f} | " +
                      " ".join(f"{k} {v:.3f}" for k, v in a_te.items()))
            if grok is None and a_te["main"] >= target:
                grok = e
                print(f"  [!] ГРОККИНГ @ {e} ({time.time()-t0:.0f}s)")
                if not track:
                    break
                stop_at = min(max_epochs, e + POST_GROK_EPOCHS)
                print(f"  [T10] продолжаю до {stop_at} — смотрим, дозревает ли ось")
            model.train()
        if e >= stop_at:
            print(f"  [T10] пост-грокинг завершён на {e}")
            break
    if grok is None:
        print(f"  [!!!] НЕ ЗАГРОККАЛ — результаты аудита НЕВАЛИДНЫ, из агрегата исключён")
    torch.save(model.state_dict(), ckpt)
    tl = dict(epochs_to_grok=grok, grokked=grok is not None,
              wall_sec=round(time.time() - t0, 1), curve=curve)
    if track and trace:
        tl["track"] = trace
        with open(trk_path, "w", encoding="utf-8") as f:
            json.dump(dict(variant=vid, seed=seed, grok=grok, trace=trace),
                      f, ensure_ascii=False, indent=2)
        analyze_ontogeny(trace, grok, vid, seed)
    return model, TR, TE, meta, tl


def analyze_ontogeny(trace, grok, vid, seed):
    """Разбор кривой онтогенеза: КОГДА рождается ось относительно грокинга?"""
    if not trace or grok is None:
        return None
    ws = np.array([t["ws_switch"] for t in trace])
    ep = np.array([t["epoch"] for t in trace])
    acc = np.array([t["acc"] for t in trace])
    cos = np.array([t["cos_to_input"] for t in trace])
    peak = float(ws.max())
    if peak < 5:
        verdict = "оси нет ни на одном этапе"
        e_axis = None
    else:
        half = peak / 2
        idx = int(np.argmax(ws >= half))
        e_axis = int(ep[idx])
        d = e_axis - grok
        if d < -500:
            verdict = f"ось ДО грокинга (на {-d} эпох раньше) -> знак = ПРЕДПОСЫЛКА"
        elif d > 500:
            verdict = f"ось ПОСЛЕ грокинга (на {d} эпох позже) -> КРИСТАЛЛИЗАЦИЯ"
        else:
            verdict = "ось ВМЕСТЕ с грокингом -> представление = обобщение"
    ws_at_grok = float(ws[np.argmin(np.abs(ep - grok))])
    ws_final = float(ws[-1])
    growth = ws_final - ws_at_grok
    print(f"\n  [T10] ОНТОГЕНЕЗ ЗНАКА ({vid}/s{seed}):")
    print(f"        грокинг @ {grok} | ось (полумакс) @ {e_axis} | ВЕРДИКТ: {verdict}")
    print(f"        ws_switch: в момент грокинга {ws_at_grok:.1f}% -> в конце {ws_final:.1f}% "
          f"(рост {growth:+.1f} п.п.)")
    print(f"        cos→вход:  в момент грокинга "
          f"{float(cos[np.argmin(np.abs(ep - grok))]):.3f} -> в конце {float(cos[-1]):.3f}")
    if growth > 10:
        print(f"        >>> КРИСТАЛЛИЗАЦИЯ ПОДТВЕРЖДЕНА: ось дозревает при неизменной точности "
              f"(acc {float(acc[np.argmin(np.abs(ep-grok))]):.3f} -> {float(acc[-1]):.3f})")
    return dict(grok=grok, axis_epoch=e_axis, verdict=verdict,
                ws_at_grok=round(ws_at_grok, 2), ws_final=round(ws_final, 2),
                growth=round(growth, 2),
                cos_at_grok=round(float(cos[np.argmin(np.abs(ep - grok))]), 4),
                cos_final=round(float(cos[-1]), 4))


# ==========================================================
# ИНСТРУМЕНТЫ
# ==========================================================
def class_means(model, X, cls, L, n_ops):
    model.eval()
    with torch.no_grad():
        _, C = model(X, cache=True)
    r = C["res"][L + 1]
    return [r[cls == k].mean(0) for k in range(n_ops)]


def pair_axis(mu, i, j):
    d = mu[i] - mu[j]
    return d / (d.norm(dim=-1, keepdim=True) + 1e-8)


def class_subspace(mu):
    """
    Для setclass: из классовых средних mu (list из [S,d]) строим
      cmeans [n_ops,S,d], grand [S,d], basis — список по позициям ортонормированных
      базисов Q_s [d,r] между-классового подпространства (r = ранг центрированных средних,
      обычно n_ops-1). Для n_ops=2 r=1 -> одно дискриминантное направление, и setclass
      совпадает с парным патчем по механике.
    """
    M = torch.stack(mu)                    # [n_ops, S, d]
    g = M.mean(0)                          # [S, d]
    Mc = M - g                             # [n_ops, S, d]
    S = M.size(1)
    r_cap = max(1, M.size(0) - 1)          # центрир. средние: ранг <= n_ops-1
    basis = []
    for s in range(S):
        A = Mc[:, s, :]                    # [n_ops, d]
        _, Sg, Vh = torch.linalg.svd(A, full_matrices=False)
        thr = 1e-5 * float(Sg[0]) if Sg.numel() and float(Sg[0]) > 0 else 1e-8
        r = min(r_cap, max(1, int((Sg > thr).sum())))
        basis.append(Vh[:r].T.contiguous())   # [d, r]
    return M, g, basis


# ==========================================================
# T0. ПОТОКИ
# ==========================================================
def T0_flows(model, TE, n_ops, meta):
    model.eval()
    with torch.no_grad():
        _, C = model(TE["X"], cache=True)
    c, S = TE["cls"], TE["X"].size(1)

    def spread(t):
        ms = torch.stack([t[c == k].mean(0) for k in range(n_ops)])
        g = ms.mean(0)
        return [round(float(torch.stack([torch.norm(ms[k, s] - g[s])
                                         for k in range(n_ops)]).mean()), 4) for s in range(S)]
    out = dict(embedding=spread(C["res"][0]),
               residuals=[dict(layer=i, spread=spread(C["res"][i + 1])) for i in range(N_LAYERS)])
    rp = meta["rule_pos"]
    eq = [r["spread"][-1] for r in out["residuals"]]
    src = [r["spread"][rp] for r in out["residuals"]]
    out["summary"] = dict(emb_at_src=out["embedding"][rp], emb_at_eq=out["embedding"][-1],
                          peak_at_eq=round(max(eq), 4), peak_layer=int(np.argmax(eq)),
                          eq_traj=eq, src_traj=src)
    print(f"  T0: эмб@источник {out['summary']['emb_at_src']:.4f} | "
          f"пик@'=' {out['summary']['peak_at_eq']:.4f} (L{out['summary']['peak_layer']})")
    print(f"      траектория '=': {eq}")
    print(f"      траектория источника: {src}")
    return out


# ==========================================================
# T2. РЕШАЮЩИЙ ТЕСТ — ВМЕШАТЕЛЬСТВО ПО ПОЗИЦИЯМ
# ==========================================================
def T2_positions(model, TR, TE, n_ops, meta, pair=(0, 1)):
    """
    Три режима вмешательства:
        'eq'  — только последний токен '='  -> ПРАВИМ РАБОЧЕЕ ПРОСТРАНСТВО
        'src' — только токен-источник       -> РЕДАКТИРУЕМ ВХОД (контроль)
        'all' — все позиции                 -> как в L8
    Для каждого: patch (грубо) и setcoord (мягко, per-position target [FIX-2]).

    [FIX-1] headline workspace_score берётся по РАННИМ слоям 0..N-2; интервенция
    на L=N-1 (позиция ридаута, слой ридаута) тавтологична и логируется отдельно
    как ws_readout.
    """
    i, j = pair
    m = TE["cls"] == j
    X, a_, b_ = TE["X"][m], TE["a"][m], TE["b"][m]
    ans_src, ans_tgt = op_answer(j, a_, b_), op_answer(i, a_, b_)
    S, rp = meta["seq_len"], meta["rule_pos"]
    multi = model.heads_mode == "multi"
    tgt_op = torch.full_like(TE["op"][m], i)
    print("  T2: РЕШАЮЩИЙ ТЕСТ — вмешательство по позициям")
    print(f"      {'режим':<5} {'L':<2} {'α':>4} | {'orig%':>7} {'SWITCH%':>8} {'garb%':>7} | "
          f"{'op→%':>6}")
    rows = []
    for mode in ["eq", "src", "all"]:
        pos = positions_for(mode, S, rp)
        for L in range(N_LAYERS):
            mu = class_means(model, TR["X"], TR["cls"], L, n_ops)
            u = pair_axis(mu, i, j)                 # [S, d]
            ci_pos = (mu[i] * u).sum(-1)            # [FIX-2] per-position target [S]
            for al in [0.5, 1.0, 2.0]:
                with torch.no_grad():
                    o = model(X, intervention=dict(layer=L, type="patch", positions=pos,
                                                   v_src=mu[j], v_tgt=mu[i], alpha=al))
                    pr = o["main"].argmax(1)
                    so = float((pr == ans_src).float().mean()) * 100
                    sp = float((pr == ans_tgt).float().mean()) * 100
                    r = dict(mode=mode, kind="patch", layer=L, alpha=al,
                             orig=round(so, 2), switched=round(sp, 2),
                             garbage=round(100 - so - sp, 2))
                    if multi:
                        r["op_followed"] = round(
                            float((o["op"].argmax(1) == tgt_op).float().mean()) * 100, 2)
                    rows.append(r)
            # мягкое: выставить координату в значение целевого класса ПО КАЖДОЙ позиции
            with torch.no_grad():
                o = model(X, intervention=dict(layer=L, type="setcoord", positions=pos,
                                               axis=u, target=ci_pos))
                pr = o["main"].argmax(1)
                so = float((pr == ans_src).float().mean()) * 100
                sp = float((pr == ans_tgt).float().mean()) * 100
                r = dict(mode=mode, kind="setcoord", layer=L, alpha=None,
                         orig=round(so, 2), switched=round(sp, 2),
                         garbage=round(100 - so - sp, 2))
                if multi:
                    r["op_followed"] = round(
                        float((o["op"].argmax(1) == tgt_op).float().mean()) * 100, 2)
                rows.append(r)

    # ---- setclass: КОРРЕКТНАЯ многоклассовая интервенция (обязательна для E) ----
    # Жёстко выставляем классовую идентичность в целевой класс i во всём между-классовом
    # подпространстве. switched/orig/op усредняются по ВСЕМ целям i (для бинарных это i∈{0,1},
    # для E — i∈{0,1,2,3}). Для бинарных совпадает по механике с парным патчем -> сверка.
    Xall, clsall, mainall = TE["X"], TE["cls"], TE["main"]
    aall, ball = TE["a"], TE["b"]
    for mode in ["eq", "src", "all"]:
        pos = positions_for(mode, S, rp)
        for L in range(N_LAYERS):
            mu = class_means(model, TR["X"], TR["cls"], L, n_ops)
            M, gmean, basis = class_subspace(mu)
            sw_t, or_t, op_t = [], [], []
            for it in range(n_ops):
                msk = clsall != it
                Xi = Xall[msk]
                ans_i = op_answer(it, aall[msk], ball[msk])
                with torch.no_grad():
                    o = model(Xi, intervention=dict(layer=L, type="setclass", positions=pos,
                                                    cmeans=M, grand=gmean, basis=basis, target=it))
                    pr = o["main"].argmax(1)
                sw_t.append(float((pr == ans_i).float().mean()) * 100)
                or_t.append(float((pr == mainall[msk]).float().mean()) * 100)
                if multi:
                    op_t.append(float((o["op"].argmax(1) == it).float().mean()) * 100)
            sw, orr = float(np.mean(sw_t)), float(np.mean(or_t))
            r = dict(mode=mode, kind="setclass", layer=L, alpha=None,
                     orig=round(orr, 2), switched=round(sw, 2),
                     garbage=round(100 - orr - sw, 2),
                     per_target_switch=[round(x, 2) for x in sw_t])
            if multi:
                r["op_followed"] = round(float(np.mean(op_t)), 2)
            rows.append(r)

    PAIR_KINDS = ("patch", "setcoord")

    def _best(mode, kinds, early=False):
        sub = [r for r in rows if r["mode"] == mode and r["kind"] in kinds]
        if early:                                   # [FIX-1] исключить тавтологический L=N-1
            sub2 = [r for r in sub if r["layer"] < LAST]
            sub = sub2 if sub2 else sub
        return max(sub, key=lambda r: r["switched"]) if sub else None

    best_pair       = {m: _best(m, PAIR_KINDS)             for m in ["eq", "src", "all"]}
    best_pair_early = {m: _best(m, PAIR_KINDS, early=True) for m in ["eq", "src", "all"]}
    best_cls        = {m: _best(m, ("setclass",))             for m in ["eq", "src", "all"]}
    best_cls_early  = {m: _best(m, ("setclass",), early=True) for m in ["eq", "src", "all"]}

    # авто: несущая интервенция — КЛАССОВАЯ для многоклассовых (E), ПАРНАЯ для бинарных
    use_class = n_ops > 2
    head_early = best_cls_early if use_class else best_pair_early
    head_full  = best_cls if use_class else best_pair
    hk = ("setclass",) if use_class else PAIR_KINDS
    eq_last = [r for r in rows if r["mode"] == "eq" and r["layer"] == LAST and r["kind"] in hk]
    readout_eq = max(eq_last, key=lambda r: r["switched"]) if eq_last else head_full["eq"]

    tag_iv = "setclass" if use_class else "pair"
    print(f"      [несущая интервенция: {tag_iv}; n_ops={n_ops}]")
    for mode in ["eq", "src", "all"]:
        b = head_early[mode] if mode == "eq" else head_full[mode]
        print(f"      {mode:<5} {b['layer']:<2} {str(b['alpha'] or 'set'):>4} | "
              f"{b['orig']:7.1f} {b['switched']:8.1f} {b['garbage']:7.1f} | "
              f"{b.get('op_followed', float('nan')):6.1f}   ({b['kind']})")
    print(f"      eq@L{LAST} (ТАВТОЛОГ. КОНТРОЛЬ ридаута): switch {readout_eq['switched']:.1f}% "
          f"(op {readout_eq.get('op_followed', float('nan')):.1f}%)")
    # сверка второй интервенции на eq (для бинарных — setclass; для E — парная под-ось)
    other_eq = best_cls_early["eq"] if not use_class else best_pair_early["eq"]
    if other_eq is not None:
        lbl = "setclass" if not use_class else "парная(0,1)"
        print(f"      сверка eq [{lbl}]: switch {other_eq['switched']:.1f}% "
              f"(L{other_eq['layer']}, {other_eq['kind']})")
    print("\n      [ЧТЕНИЕ] 'eq'(ранние) высок  -> правило на '=' НЕСУЩЕЕ = рабочее пространство")
    print("               'eq' низок, 'src' высок -> сеть читает правило из входа = маршрутизация")
    print("               eq@Lласт высок ПРИ низком eq(ранние) -> это лишь правка входа ридаута,")
    print("                                                       НЕ рабочее пространство")
    return dict(pair=[i, j], rows=rows,
                best=head_full, best_early=head_early,
                best_pair_early=best_pair_early, best_cls_early=best_cls_early,
                intervention=tag_iv,
                workspace_score=head_early["eq"]["switched"],    # авто, [FIX-1] без L=N-1
                ws_switch_pair=(best_pair_early["eq"]["switched"] if best_pair_early["eq"] else None),
                ws_switch_class=(best_cls_early["eq"]["switched"] if best_cls_early["eq"] else None),
                ws_readout=readout_eq["switched"],               # тавтологический контроль
                ws_garbage=head_early["eq"]["garbage"],
                ws_orig=head_early["eq"]["orig"],
                input_score=head_full["src"]["switched"],
                all_score=head_full["all"]["switched"])


# ==========================================================
# T3. КОНФАБУЛЯЦИЯ  [FIX-3]
# ==========================================================
def T3_confabulation(t2, model):
    """
    Расхождение между «модель говорит, что применяет правило X» (голова op)
    и «модель действительно применяет X» (main). Три честных числа на режим:
      existence — разрыв на max(op_followed): СУЩЕСТВУЕТ интервенция, где отчёт
                  следует ~100%, а вычисление нет;
      at_wsopt  — разрыв на max(switched): даже там, где вычислению легче всего
                  последовать, op впереди -> конфабуляция робастна;
      gap_mean  — средний парный разрыв по всем строкам режима.
    """
    if model.heads_mode != "multi":
        return None
    out = {}
    for mode in ["eq", "all"]:
        sub = [r for r in t2["rows"] if r["mode"] == mode and "op_followed" in r]
        if not sub:
            continue
        b_op = max(sub, key=lambda r: r["op_followed"])       # существование
        b_ws = max(sub, key=lambda r: r["switched"])          # ws-оптимум
        gap_exist = round(b_op["op_followed"] - b_op["switched"], 2)
        gap_wsopt = round(b_ws["op_followed"] - b_ws["switched"], 2)
        gap_mean = round(float(np.mean([r["op_followed"] - r["switched"] for r in sub])), 2)
        out[mode] = dict(
            existence=dict(layer=b_op["layer"], kind=b_op["kind"], alpha=b_op["alpha"],
                           op_followed=b_op["op_followed"], main_switched=b_op["switched"],
                           gap=gap_exist),
            at_wsopt=dict(layer=b_ws["layer"], kind=b_ws["kind"], alpha=b_ws["alpha"],
                          op_followed=b_ws["op_followed"], main_switched=b_ws["switched"],
                          gap=gap_wsopt),
            gap_mean=gap_mean)
        print(f"  T3: конфабуляция [{mode}]: "
              f"существ. op {b_op['op_followed']:.1f}%/main {b_op['switched']:.1f}% "
              f"-> разрыв {gap_exist:.1f} | "
              f"ws-опт. op {b_ws['op_followed']:.1f}%/main {b_ws['switched']:.1f}% "
              f"-> разрыв {gap_wsopt:.1f} | средний разрыв {gap_mean:.1f}")
    return out


# ==========================================================
# T5. ИГНИШН (на режиме 'eq' — правим только рабочее пространство)
# ==========================================================
def T5_ignition(model, TR, TE, n_ops, meta, pair=(0, 1), mode="eq"):
    i, j = pair
    m01 = (TE["cls"] == i) | (TE["cls"] == j)
    X, a_, b_ = TE["X"][m01], TE["a"][m01], TE["b"][m01]
    ans_i, ans_j = op_answer(i, a_, b_), op_answer(j, a_, b_)
    pos = positions_for(mode, meta["seq_len"], meta["rule_pos"])
    out = dict(pair=[i, j], mode=mode, layers={})
    print(f"  T5: игнишн (режим '{mode}') — доля цели среди ЧИСТЫХ ответов")
    for L in range(N_LAYERS):
        mu = class_means(model, TR["X"], TR["cls"], L, n_ops)
        u = pair_axis(mu, i, j)
        ci = float((mu[i][-1] * u[-1]).sum())
        cj = float((mu[j][-1] * u[-1]).sum())
        lo, hi = min(ci, cj), max(ci, cj)
        span = (hi - lo) if (hi - lo) > 1e-6 else 1.0
        ts = [lo - 0.5 * span + k * (2.0 * span) / 24 for k in range(25)]
        curve = []
        for t in ts:
            with torch.no_grad():
                # свип по СКАЛЯРНОЙ цели на позиции '=' (ось u[-1]); setcoord с pos=[-1]
                pr = model(X, intervention=dict(layer=L, type="setcoord", positions=pos,
                                                axis=u, target=t))["main"].argmax(1)
            fi = float((pr == ans_i).float().mean()) * 100
            fj = float((pr == ans_j).float().mean()) * 100
            clean = fi + fj
            curve.append(dict(t=round(t, 4), cls_i=round(fi, 2), cls_j=round(fj, 2),
                              clean=round(clean, 2),
                              p=(round(fi / clean, 4) if clean > 1e-6 else None)))
        ps = [c["p"] for c in curve if c["p"] is not None]
        if len(ps) < 5:
            eff = sharp = 0.0
        else:
            arr = np.array(ps)
            eff = float(abs(arr[-1] - arr[0]))
            sharp = float(np.abs(np.diff(arr)).max() / (eff / len(arr) + 1e-6)) if eff > 0.05 else 0.0
        mc = float(np.mean([c["clean"] for c in curve[8:17]]))
        out["layers"][f"L{L}"] = dict(curve=curve, effect=round(eff, 3),
                                      sharpness=round(sharp, 2), mid_clean=round(mc, 1))
        print(f"      L{L}: эффект {eff:.2f} | резкость {sharp:5.1f} | чистых {mc:5.1f}%")
    bl, b = max(out["layers"].items(), key=lambda kv: kv[1]["effect"])
    out["summary"] = dict(best_layer=bl, effect=b["effect"], sharpness=b["sharpness"],
                          mid_clean=b["mid_clean"])
    return out


# ==========================================================
# T8. АБЛЯЦИЯ (по режимам позиций)
# ==========================================================
def T8_ablation(model, TR, TE, n_ops, meta):
    keys = model.active_heads() + [CONTROL]
    model.eval()
    with torch.no_grad():
        base = accs(model(TE["X"]), TE, keys)
    S, rp = meta["seq_len"], meta["rule_pos"]
    print("  T8: абляция оси правила по режимам позиций")
    rows = []
    for mode in ["eq", "src", "all"]:
        pos = positions_for(mode, S, rp)
        for L in range(N_LAYERS):
            mu = class_means(model, TR["X"], TR["cls"], L, n_ops)
            M = torch.stack(mu); Mc = M - M.mean(0, keepdim=True)
            vec = torch.stack([torch.linalg.svd(Mc[:, s, :], full_matrices=False)[2][0]
                               for s in range(S)])
            rnd = torch.randn_like(vec); rnd = rnd / rnd.norm(dim=-1, keepdim=True)
            with torch.no_grad():
                a = accs(model(TE["X"], intervention=dict(layer=L, type="ablate",
                                                          positions=pos, vec=vec)), TE, keys)
                ar = accs(model(TE["X"], intervention=dict(layer=L, type="ablate",
                                                           positions=pos, vec=rnd)), TE, keys)
            cons = model.active_heads()
            dc = float(np.mean([base[k] - a[k] for k in cons]))
            dr = float(np.mean([base[k] - ar[k] for k in cons]))
            dt = float(base[CONTROL] - a[CONTROL])
            rows.append(dict(mode=mode, layer=L, consumer_drop=round(dc, 4),
                             random_drop=round(dr, 4), control_drop=round(dt, 4)))
    for mode in ["eq", "src", "all"]:
        w = max([r for r in rows if r["mode"] == mode], key=lambda r: r["consumer_drop"])
        print(f"      {mode:<4} L{w['layer']}: потребители −{w['consumer_drop']*100:5.1f}% | "
              f"случайно −{w['random_drop']*100:5.1f}% | контроль −{w['control_drop']*100:5.1f}%")
    # headline по 'eq' — по РАННИМ слоям (0..N-2), симметрично T2 [FIX-1]
    eq_early = [r for r in rows if r["mode"] == "eq" and r["layer"] < LAST] or \
               [r for r in rows if r["mode"] == "eq"]
    worst_eq = max(eq_early, key=lambda r: r["consumer_drop"])
    return dict(baseline={k: round(v, 4) for k, v in base.items()}, rows=rows,
                summary=dict(eq_consumer_drop=worst_eq["consumer_drop"],
                             eq_random_drop=worst_eq["random_drop"],
                             eq_control_drop=worst_eq["control_drop"],
                             eq_layer=worst_eq["layer"],
                             eq_margin=round(worst_eq["consumer_drop"] -
                                             worst_eq["random_drop"], 4)))


def T9_clusters(model, TE, n_ops):
    if not SKLEARN:
        return None
    model.eval()
    with torch.no_grad():
        _, C = model(TE["X"], cache=True)
    true = TE["cls"].cpu().numpy()
    rows = []
    for L in range(N_LAYERS):
        v = C["res"][L + 1][:, -1, :].cpu().numpy()
        km = KMeans(n_ops, random_state=0, n_init=10).fit_predict(v)
        s = float(silhouette_score(v, km))
        acc = max(float((np.array([pm[k] for k in km]) == true).mean())
                  for pm in permutations(range(n_ops)))
        rows.append(dict(layer=L, silhouette=round(s, 4), match=round(acc, 4)))
    print(f"  T9: sil {[r['silhouette'] for r in rows]} | match {[r['match'] for r in rows]}")
    best = max(rows, key=lambda r: r["match"])
    return dict(rows=rows, best=best)


# ==========================================================
# ПРОГОН
# ==========================================================
def run(vid, seed, retrain, track=False):
    cfg = VARIANTS[vid]
    n_ops = cfg["n_ops"]
    print("\n" + "=" * 86)
    print(f" {vid} / seed {seed}: {cfg['desc']}")
    print("=" * 86)
    model, TR, TE, meta, tr = train_variant(vid, seed, retrain, track=track)
    keys = model.active_heads() + [CONTROL]
    with torch.no_grad():
        final = {k: round(v, 4) for k, v in accs(model(TE["X"]), TE, keys).items()}
    grokked = tr.get("grokked", final["main"] >= 0.95)
    print(f"  точность: {final} | ГРОККИНГ: {grokked}")
    log = dict(variant=vid, seed=seed, config=cfg, dataset=meta,
               hyper=dict(p=P, d_model=cfg["d"], n_layers=N_LAYERS),
               timestamp=datetime.now().isoformat(), training=tr,
               final_test_acc=final, grokked=bool(grokked))
    log["T0_flows"] = T0_flows(model, TE, n_ops, meta)
    log["T2_positions"] = T2_positions(model, TR, TE, n_ops, meta)
    log["T3_confab"] = T3_confabulation(log["T2_positions"], model)
    log["T5_ignition_eq"] = T5_ignition(model, TR, TE, n_ops, meta, mode="eq")
    log["T8_ablation"] = T8_ablation(model, TR, TE, n_ops, meta)
    log["T9_clusters"] = T9_clusters(model, TE, n_ops)
    if track and tr.get("track"):
        log["T10_ontogeny"] = dict(
            trace=tr["track"],
            analysis=analyze_ontogeny(tr["track"], tr.get("epochs_to_grok"), vid, seed))
    t2, s5, s8, s0 = (log["T2_positions"], log["T5_ignition_eq"]["summary"],
                      log["T8_ablation"]["summary"], log["T0_flows"]["summary"])
    cf = log["T3_confab"] or {}

    def _cf(mode, field):
        d = cf.get(mode) or {}
        if field == "exist":
            return (d.get("existence") or {}).get("gap")
        if field == "wsopt":
            return (d.get("at_wsopt") or {}).get("gap")
        if field == "mean":
            return d.get("gap_mean")
        return None
    log["metrics"] = dict(
        grokked=bool(grokked),
        emb_at_src=s0["emb_at_src"], peak_at_eq=s0["peak_at_eq"],
        ws_intervention=t2["intervention"],   # 'pair' (бинарные) | 'setclass' (E)
        ws_switch=t2["workspace_score"],      # <<< ГЛАВНАЯ (авто): правка только '=' на РАННИХ слоях
        ws_switch_pair=t2["ws_switch_pair"],  # парная интервенция (для бинарных = ws_switch)
        ws_switch_class=t2["ws_switch_class"],# setclass (для E = ws_switch; для бинарных — сверка)
        ws_readout=t2["ws_readout"],          # <<< тавтологический контроль (L=N-1)
        input_switch=t2["input_score"],       # правка только входа
        all_switch=t2["all_score"],
        ws_garbage=t2["ws_garbage"],
        ws_orig=t2["ws_orig"],
        confab_gap_eq=_cf("eq", "exist"),
        confab_gap_all=_cf("all", "exist"),
        confab_gap_all_wsopt=_cf("all", "wsopt"),
        confab_gap_all_mean=_cf("all", "mean"),
        ign_effect_eq=s5["effect"], ign_sharp_eq=s5["sharpness"], ign_clean_eq=s5["mid_clean"],
        eq_consumer_drop=s8["eq_consumer_drop"], eq_control_drop=s8["eq_control_drop"],
        eq_margin=s8["eq_margin"],
    )
    if log.get("T10_ontogeny") and log["T10_ontogeny"].get("analysis"):
        an = log["T10_ontogeny"]["analysis"]
        log["metrics"].update(
            onto_axis_epoch=an["axis_epoch"], onto_grok=an["grok"],
            onto_ws_at_grok=an["ws_at_grok"], onto_ws_final=an["ws_final"],
            onto_growth=an["growth"], onto_cos_at_grok=an["cos_at_grok"],
            onto_cos_final=an["cos_final"])
    print(f"\n  МЕТРИКИ: {json.dumps(log['metrics'], ensure_ascii=False)}")
    with open(os.path.join(LOG_DIR, f"{vid}_s{seed}.json"), "w", encoding="utf-8") as f:
        json.dump(log, f, ensure_ascii=False, indent=2)
    return log


def aggregate(logs):
    print("\n" + "=" * 122)
    print(" СВОДКА (mean ± std). Не загроккавшие прогоны исключены из агрегата.")
    print(" WS(ранн)  = несущая правка только '=' на слоях 0..N-2 (авто: парная|setclass).")
    print(" WSкласс   = setclass-версия (для E — она же несущая; для бинарных — сверка с парной).")
    print(" WS(ридаут)= та же правка на L=N-1: тавтологический контроль (высок почти всегда).")
    print("=" * 122)
    print(f"{'вар':<4}{'grok':<6}| {'WS(ранн)%':>10} {'WSкласс%':>10} {'WS(ридаут)%':>12} "
          f"{'input%':>8} {'all%':>7} | {'конфаб(all)':>12} {'сред':>6} | {'игн.эфф':>8}")
    print("-" * 122)
    agg = []
    for vid in list(VARIANTS):
        ls = [l for l in logs if l["variant"] == vid]
        if not ls:
            continue
        ok = [l for l in ls if l["metrics"]["grokked"]]
        use = ok if ok else ls

        def ms(k):
            v = [l["metrics"][k] for l in use if l["metrics"].get(k) is not None]
            return (round(float(np.mean(v)), 2), round(float(np.std(v)), 2)) if v else (None, None)
        row = dict(variant=vid, n=len(ls), n_grokked=len(ok), desc=VARIANTS[vid]["desc"])
        for k in ["emb_at_src", "peak_at_eq", "ws_switch", "ws_switch_pair", "ws_switch_class",
                  "ws_readout", "input_switch", "all_switch", "ws_garbage", "ws_orig",
                  "confab_gap_eq", "confab_gap_all", "confab_gap_all_wsopt", "confab_gap_all_mean",
                  "ign_effect_eq", "ign_sharp_eq", "ign_clean_eq",
                  "eq_consumer_drop", "eq_control_drop", "eq_margin"]:
            m, s = ms(k)
            row[k] = dict(mean=m, std=s) if m is not None else None
        agg.append(row)
        g = lambda k: (f"{row[k]['mean']:.1f}±{row[k]['std']:.1f}" if row[k] else "—")
        print(f"{vid:<4}{len(ok)}/{len(ls):<4}| {g('ws_switch'):>10} {g('ws_switch_class'):>10} "
              f"{g('ws_readout'):>12} {g('input_switch'):>8} {g('all_switch'):>7} | "
              f"{g('confab_gap_all'):>12} {g('confab_gap_all_mean'):>6} | {g('ign_effect_eq'):>8}")
    print("\n[РЕШАЮЩИЙ ТЕСТ — колонка «WS(ранн)»]")
    print("  Правим ТОЛЬКО токен '=' на слоях 0..N-2 (есть downstream-чтение), вход не трогаем.")
    print("  Высокий  -> правило на '=' НЕСУЩЕЕ: настоящее рабочее пространство.")
    print("  Низкий при высоком «input» -> сеть читает правило из входа: маршрутизация.")
    print("  ВАЖНО: если WS(ридаут) высок, а WS(ранн) низок — это лишь правка входа ридаута,")
    print("         рабочего пространства НЕТ. Разрыв двух колонок — прямое доказательство.")
    print("\n[КЛЮЧЕВЫЕ СРАВНЕНИЯ]")
    print("  A,B (читается из a)        -> ожидаем: WS(ранн) низкий, input высокий")
    print("  C,D (QR: строить + дорого) -> ожидаем: WS(ранн) высокий")
    print("  H (ключ: строить, но ДЁШЕВО):")
    print("      H ≈ D -> достаточно ПОСТРОЙКИ")
    print("      H ≈ A -> нужна ДОРОГОВИЗНА пересчёта (амортизация)")
    print("  F (вырожденный) -> ожидаем игнорирование")
    print("  E (4 класса, оба строить) -> несущая = setclass (WSкласс). Проверка, что картина")
    print("      рабочего пространства масштабируется за пределы бинарного правила.")
    print("  Сверка WS(ранн) vs WSкласс: для бинарных должны совпадать (парная≈setclass);")
    print("      расхождение -> сигнал, что одна из интервенций артефактна.")
    print("\n[КОНФАБУЛЯЦИЯ]  колонка «конфаб(all)» = существование, «сред» = средний разрыв")
    print("  Разрыв = (голова 'op' следует за патчем) − (вычисление следует).")
    print("  «существование» велико -> ЕСТЬ интервенция, где отчёт следует ~100%, а вычисление нет.")
    print("  «средний» велик -> так ведёт себя не одна подобранная точка, а режим в целом:")
    print("  представление проходит тест на самоотчёт, не будучи несущим.")
    path = os.path.join(LOG_DIR, f"summary_{datetime.now():%Y%m%d_%H%M%S}.json")
    with open(path, "w", encoding="utf-8") as f:
        json.dump(agg, f, ensure_ascii=False, indent=2)
    print(f"\n[+] сводка: {path}")


def ontogeny_summary(logs):
    rows = [l for l in logs if l.get("metrics", {}).get("onto_grok") is not None]
    if not rows:
        return
    print("\n" + "=" * 100)
    print(" T10. ОНТОГЕНЕЗ ЗНАКА: когда рождается ось относительно грокинга?")
    print("=" * 100)
    print(f"{'вар':<4}{'сид':<4}| {'грокинг':>8} {'ось@полумакс':>13} {'Δ':>7} | "
          f"{'ws@грок':>8} {'ws финал':>9} {'рост':>7} | {'cos@грок':>9} {'cos финал':>10}")
    print("-" * 100)
    for l in rows:
        m = l["metrics"]
        d = (m["onto_axis_epoch"] - m["onto_grok"]) if m["onto_axis_epoch"] is not None else None
        print(f"{l['variant']:<4}{l['seed']:<4}| {m['onto_grok']:8d} "
              f"{str(m['onto_axis_epoch']):>13} {str(d):>7} | "
              f"{m['onto_ws_at_grok']:8.1f} {m['onto_ws_final']:9.1f} {m['onto_growth']:+7.1f} | "
              f"{m['onto_cos_at_grok']:9.3f} {m['onto_cos_final']:10.3f}")
    print("\n[ТРИ СЦЕНАРИЯ]")
    print("  Δ < 0  -> ось ДО грокинга:      знак есть ПРЕДПОСЫЛКА обобщения")
    print("  Δ ≈ 0  -> ось ВМЕСТЕ:           представление и обобщение — одно событие")
    print("  Δ > 0  -> ось ПОСЛЕ:            КРИСТАЛЛИЗАЦИЯ под weight decay")
    print("\n[КРИСТАЛЛИЗАЦИЯ]")
    print("  рост ws_switch после грокинга при НЕИЗМЕННОЙ точности -> ось ДОЗРЕВАЕТ,")
    print("  когда задача уже решена. Знак — продукт сжатия, а не вычислительной нужды.")
    print("\n[ОТРЫВ ОТ ВХОДА]")
    print("  cos→вход ~1  -> ось на '=' есть просто тень входного признака")
    print("  cos→вход →0  -> представление ПЕРЕСТРОИЛОСЬ во что-то своё")
    print("  падение cos при росте ws_switch = знак отрывается от входа и становится несущим")
    path = os.path.join(LOG_DIR, f"ontogeny_{datetime.now():%Y%m%d_%H%M%S}.json")
    with open(path, "w", encoding="utf-8") as f:
        json.dump([dict(variant=l["variant"], seed=l["seed"],
                        analysis=l["T10_ontogeny"]["analysis"],
                        trace=l["T10_ontogeny"]["trace"]) for l in rows],
                  f, ensure_ascii=False, indent=2)
    print(f"\n[+] онтогенез: {path}")


if __name__ == "__main__":
    ap = argparse.ArgumentParser()
    ap.add_argument("--variants", nargs="+", default=DEFAULT_ORDER, choices=list(VARIANTS))
    ap.add_argument("--seeds", type=int, default=5)
    ap.add_argument("--retrain", action="store_true")
    ap.add_argument("--track", action="store_true",
                    help="T10: онтогенез знака — трекинг каждые 250 эпох + "
                         "продолжение обучения после грокинга (ДОЛГО)")
    args = ap.parse_args()
    print(f"[*] J-Space Level 9c (исправленный) | {device} | sklearn={SKLEARN}")
    print(f"[*] варианты {args.variants} × {args.seeds} сидов | самостоятельный (без весов L8)")
    logs = [run(v, s, args.retrain, track=args.track)
            for v in args.variants for s in range(args.seeds)]
    aggregate(logs)
    if args.track:
        ontogeny_summary(logs)