#!/usr/bin/env python3
"""Шаг 1: цели учителя (Qwen3.5-9B) на диск.

Зачем отдельным шагом: 15.63 ГиБ весов учителя при 16 ГиБ карты не живут вместе ни с чем, а цель
(пост-норм `last_hidden_state` слоя 32) считается один раз на промпт и дальше не меняется.
Кэш позволяет перебирать гиперпараметры адаптера часами без единого прогона учителя.

Формат кэша (shard_XXXXX.pt, torch.save, грузится с weights_only=True):
    {"items": [{"ids": int32[L], "y": bf16[L, 4096]}, ...], "texts": [...], "kinds": [...]}
Пробы и val — те же поля отдельными файлами, плюс meta.json с параметрами и статистикой норм.

  python3 src/cache_targets.py --prompts-dir ~/dataset --out cache/teacher --len 256
  # продолжить прерванный прогон: те же аргументы, готовые шарды пропускаются
"""
import argparse
import json
import os
import sys
import time

import numpy as np
import torch

sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
from moss_lib import MOSS_DIR, encode, load_teacher, load_tokenizer                    # noqa: E402
from prompts import NEGATIVES, PROBE, load_prompts                                     # noqa: E402

ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))


def save_items(path, items, texts, kinds):
    torch.save({"items": items, "texts": texts, "kinds": kinds}, path)


def load_items(path):
    return torch.load(path, map_location="cpu", weights_only=True)


def main():
    ap = argparse.ArgumentParser(description="цели учителя MOSS (Qwen3.5-9B) -> кэш на диске")
    ap.add_argument("--moss-dir", default=MOSS_DIR)
    ap.add_argument("--teacher-dir", default=None, help="папка учителя (по умолчанию <moss-dir>/text_encoder)")
    ap.add_argument("--prompts-dir", default=os.path.expanduser("~/dataset"))
    ap.add_argument("--out", default=os.path.join(ROOT, "cache", "teacher"))
    ap.add_argument("--limit", type=int, default=0, help="0 — весь пул")
    ap.add_argument("--len", type=int, default=256, help="макс. длина последовательности (обрезка хвоста)")
    ap.add_argument("--shard", type=int, default=500, help="промптов в шарде")
    ap.add_argument("--val-size", type=int, default=128, help="сколько промптов отложить на val")
    ap.add_argument("--batch", type=int, default=1,
                    help="промптов за один forward; 1 — бит в бит как в пайплайне, >1 меняет "
                         "цель на ~0.5%% (разные блоки ядер, bf16)")
    ap.add_argument("--prefer-8bit", action="store_true", help="сразу 8 бит, без попытки bf16")
    ap.add_argument("--seed", type=int, default=0)
    args = ap.parse_args()

    tok = load_tokenizer(args.moss_dir)
    pool = load_prompts(args.prompts_dir, limit=args.limit, seed=args.seed)
    if len(pool) < args.val_size + 1:
        raise SystemExit(f"пул слишком мал: {len(pool)}")

    # спец-промпты (пустой + негативы) отдельным файлом, val — хвост пула
    special_texts = [""] + list(NEGATIVES)
    special_kinds = ["empty"] + ["neg"] * len(NEGATIVES)
    val_texts = pool[:args.val_size]
    train_texts = pool[args.val_size:]
    probe_texts = [t for _, t in PROBE]
    print(f"кэш: {len(train_texts)} обучающих + special {len(special_texts)} + val {len(val_texts)} "
          f"+ probe {len(probe_texts)}, len={args.len}", flush=True)

    t0 = time.time()
    ids = encode(tok, special_texts + train_texts + val_texts + probe_texts, args.len)
    n_special, n_train, n_val = len(special_texts), len(train_texts), len(val_texts)
    print(f"токенизация: {time.time() - t0:.1f} c, медиана длины "
          f"{int(np.median([len(r) for r in ids[n_special:n_special + n_train]]))} из {args.len}",
          flush=True)

    teacher, mode = load_teacher(prefer_8bit=args.prefer_8bit, path=args.teacher_dir)

    def targets(rows):
        """list[list[int]] -> list[(ids, y)]: батч с правым паддингом, цель обрезается по длине.

        Паддинг справа + маска: реальные токены от него не зависят (внимание причинное), это
        проверяется selfcheck-ом ниже.
        """
        out = []
        for s in range(0, len(rows), args.batch):
            chunk = rows[s:s + args.batch]
            length = max(len(r) for r in chunk)
            idb = torch.full((len(chunk), length), tok.pad_token_id, dtype=torch.long)
            mb = torch.zeros((len(chunk), length), dtype=torch.long)
            for i, r in enumerate(chunk):
                idb[i, :len(r)] = torch.tensor(r)
                mb[i, :len(r)] = 1
            y = teacher.last_hidden(idb.to("cuda"), mb.to("cuda"))
            if not torch.isfinite(y).all():          # 8-бит считает в fp16: ловим переполнение сразу
                raise SystemExit("в цели NaN/inf — учитель посчитал в fp16 и переполнился")
            for i, r in enumerate(chunk):
                out.append((torch.tensor(r, dtype=torch.int32),
                            y[i, :len(r)].to(torch.bfloat16).cpu()))
        return out

    if args.batch > 1:                       # паддинг в батче меняет только округление bf16
        check_rows = ids[:4]
        single = torch.cat([y for r in check_rows for _, y in targets([r])]).float()
        batched = torch.cat([y for _, y in targets(check_rows)]).float()
        rel = (single - batched).abs().max().item() / single.norm(dim=-1).mean().item()
        print(f"  selfcheck: батч меняет цель на {rel * 100:.2f}% (округление bf16, ждём <1%)", flush=True)
        if rel > 0.01:
            raise SystemExit("батч с паддингом меняет цель сильнее округления — считай --batch 1")

    os.makedirs(args.out, exist_ok=True)
    special = [{"ids": r, "y": y} for r, y in targets(ids[:n_special])]
    save_items(os.path.join(args.out, "special.pt"), special, special_texts, special_kinds)
    probe = [{"ids": r, "y": y} for r, y in targets(ids[n_special + n_train + n_val:])]
    save_items(os.path.join(args.out, "probe.pt"), probe, probe_texts, [n for n, _ in PROBE])
    val = [{"ids": r, "y": y} for r, y in targets(ids[n_special + n_train:n_special + n_train + n_val])]
    save_items(os.path.join(args.out, "val.pt"), val, val_texts, ["val"] * n_val)

    norms, shards, tok_total, done, t0 = [], [], 0, 0, time.time()
    for s in range(0, n_train, args.shard):
        path = os.path.join(args.out, f"shard_{s // args.shard:05d}.pt")
        shards.append(os.path.basename(path))
        if os.path.exists(path):                      # возобновление прерванного прогона
            data = load_items(path)
            items = data["items"]
            done += len(items)
        else:
            chunk = list(range(s, min(s + args.shard, n_train)))
            items = [{"ids": r, "y": y} for r, y in targets([ids[n_special + i] for i in chunk])]
            save_items(path, items, [train_texts[i] for i in chunk], ["pool"] * len(chunk))
            done += len(items)
            el = time.time() - t0
            t0 = time.time()
            print(f"  {done}/{n_train} промптов, {el:.0f} c на шард, "
                  f"осталось ~{el * (n_train - done) / max(1, len(items)) / 60:.0f} мин", flush=True)
        norms += [it["y"].float().norm(dim=-1).numpy() for it in items]
        tok_total += int(sum(it["ids"].numel() for it in items))

    norm = np.concatenate(norms)
    meta = {"moss_dir": args.moss_dir, "prompts_dir": args.prompts_dir, "teacher_mode": mode,
            "len": args.len, "batch": args.batch,
            "n_train": n_train, "n_val": n_val, "n_probe": len(probe_texts),
            "n_special": n_special, "shards": shards, "tokens": tok_total,
            "norms": {f"p{q}": float(np.percentile(norm, q)) for q in (1, 10, 50, 90, 99)},
            "norm_mean": float(norm.mean()), "norm_floor": float(np.median(norm))}
    with open(os.path.join(args.out, "meta.json"), "w") as f:
        json.dump(meta, f, indent=1)
    raw = sum(os.path.getsize(os.path.join(args.out, s)) for s in shards)
    print(f"готово: {args.out}, {n_train} промптов, {tok_total / 1e6:.1f}M токенов, "
          f"{raw / 2**30:.1f} ГиБ; норма цели медиана {meta['norm_floor']:.0f} "
          f"(p1 {meta['norms']['p1']:.0f}, p99 {meta['norms']['p99']:.0f}); учитель {mode}", flush=True)


if __name__ == "__main__":
    main()