#!/usr/bin/env python3
"""Шаг 2: обучение адаптера на кэше целей (учитель уже не нужен).

Чистая регрессия условия: срезы замороженной 0.8B -> пост-норм hidden states учителя 9B
(см. cache_targets.py). Ни картинок, ни VAE, ни диффузии.

Лосс: MSE с пер-токенным весом 1/max(||y||^2, floor^2), паддинг весом `--pad-weight`.
Пустой промпт и негативы подмешиваются отдельными долями — пустой идёт в негативную ветку CFG
на каждом шаге генерации, его нельзя «выучить случайно».

  python3 src/train_adapter.py --cache cache/teacher --selfcheck
  python3 src/train_adapter.py --cache cache/teacher --epochs 2 --batch 16 --out out/adapter_v1
"""
import argparse
import json
import os
import sys
import time

import numpy as np
import torch
import torch.nn.functional as F

sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
from adapter_lib import build_adapter, save_safetensors                    # noqa: E402
from moss_lib import LAYERS, Student                                       # noqa: E402

DEVICE = "cuda"
ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))


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


def collate(items):
    """list[{ids,y}] -> (ids, mask, y) с правым паддингом до максимума в батче."""
    length = max(it["ids"].numel() for it in items)
    ids = torch.full((len(items), length), 0, dtype=torch.long)
    mask = torch.zeros(len(items), length, dtype=torch.bool)
    y = torch.zeros(len(items), length, items[0]["y"].shape[-1], dtype=torch.bfloat16)
    for i, it in enumerate(items):
        n = it["ids"].numel()
        ids[i, :n] = it["ids"].long()
        mask[i, :n] = True
        y[i, :n] = it["y"]
    return ids.to(DEVICE), mask.to(DEVICE), y.to(DEVICE)


def main():
    ap = argparse.ArgumentParser(description="адаптер Qwen3.5-0.8B -> текстовый энкодер MOSS (Qwen3.5-9B)")
    ap.add_argument("--cache", default=os.path.join(ROOT, "cache", "teacher"))
    ap.add_argument("--student", default=None,
                    help="по умолчанию локальная копия в <moss-dir>/text_encoder")
    ap.add_argument("--layers", default=",".join(map(str, LAYERS)))
    ap.add_argument("--hidden", type=int, default=6144)
    ap.add_argument("--proj-layers", type=int, default=3)
    ap.add_argument("--norm", default="ln-per-layer")
    ap.add_argument("--attention", type=int, default=2, help="блоков ветки внимания по токенам")
    ap.add_argument("--attn-dim", type=int, default=1024)
    ap.add_argument("--attn-heads", type=int, default=8)
    ap.add_argument("--attn-max-len", type=int, default=1024, help="таблица позиций ветки внимания")
    ap.add_argument("--mixer", type=int, default=2, help="блоков внимания ПО ОСИ СРЕЗОВ")
    ap.add_argument("--mixer-heads", type=int, default=8)
    ap.add_argument("--mixer-ffn", type=int, default=2)
    ap.add_argument("--epochs", type=int, default=2)
    ap.add_argument("--batch", type=int, default=16)
    ap.add_argument("--lr", type=float, default=3e-4)
    ap.add_argument("--min-lr", type=float, default=1e-5)
    ap.add_argument("--attn-lr", type=float, default=None)
    ap.add_argument("--norm-floor", type=float, default=0.0, help="0 — взять медиану нормы из meta.json")
    ap.add_argument("--pad-weight", type=float, default=0.1)
    ap.add_argument("--empty-ratio", type=float, default=0.03, help="доля пустого промпта (негатив CFG)")
    ap.add_argument("--neg-ratio", type=float, default=0.02)
    ap.add_argument("--probe-every", type=int, default=200)
    ap.add_argument("--save-every", type=int, default=500)
    ap.add_argument("--max-steps", type=int, default=0)
    ap.add_argument("--seed", type=int, default=0)
    ap.add_argument("--init-from", default=None, help="старт с готового адаптера (.safetensors)")
    ap.add_argument("--out", default=os.path.join(ROOT, "out", "adapter_v1"))
    ap.add_argument("--selfcheck", action="store_true", help="только проверить контракт и выйти")
    args = ap.parse_args()

    torch.manual_seed(args.seed)
    rng = np.random.default_rng(args.seed)
    layers = [int(x) for x in args.layers.split(",")]

    meta = json.load(open(os.path.join(args.cache, "meta.json")))
    probe = load_pt(os.path.join(args.cache, "probe.pt"))
    val = load_pt(os.path.join(args.cache, "val.pt"))
    special = load_pt(os.path.join(args.cache, "special.pt"))
    floor = args.norm_floor or float(meta["norm_floor"])
    print(f"кэш: {meta['n_train']} промптов, val {meta['n_val']}, probe {meta['n_probe']}, "
          f"учитель {meta['teacher_mode']}, len {meta['len']}, norm_floor {floor:.0f}", flush=True)

    student = Student(args.student, layers, device=DEVICE)
    in_dim = student.in_dim
    out_dim = int(probe["items"][0]["y"].shape[-1])
    adapter = build_adapter(in_dim, out_dim, args.hidden, args.proj_layers, args.norm,
                            n_slices=len(layers), attention=args.attention, attn_dim=args.attn_dim,
                            attn_heads=args.attn_heads, max_len=args.attn_max_len,
                            mixer=args.mixer, mixer_heads=args.mixer_heads,
                            mixer_ffn=args.mixer_ffn).to(DEVICE)
    n_par = sum(p.numel() for p in adapter.parameters())
    n_branch = sum(p.numel() for n, p in adapter.named_parameters() if n.startswith(("attn.", "mixer.")))
    print(f"адаптер: {n_par / 1e6:.1f}M ({in_dim} -> {args.hidden} x{args.proj_layers - 1} -> {out_dim}), "
          f"нормировка {args.norm}, срезов {len(layers)}, ветки {n_branch / 1e6:.1f}M", flush=True)
    if args.init_from:
        from safetensors.torch import load_file
        sd = load_file(args.init_from)
        missing, unexpected = adapter.load_state_dict(sd, strict=False)
        print(f"старт с {args.init_from}: missing {list(missing)[:3]}, unexpected {list(unexpected)[:3]}",
              flush=True)

    def token_weights(y, mask):
        w = 1.0 / y.pow(2).sum(-1).clamp_min(floor ** 2)
        return w * torch.where(mask, 1.0, args.pad_weight)

    def loss_of(ids, mask, y):
        pred = adapter(student.hidden(ids, mask), mask).float()
        yf = y.float()
        w = token_weights(yf, mask)[..., None].expand_as(yf)
        return (w * (pred - yf).pow(2)).sum() / w.sum()

    @torch.no_grad()
    def measure(items):
        """Относительная ошибка и косинус по реальным токенам — по одному примеру, без паддинга."""
        adapter.eval()
        rel, cos = [], []
        for it in items:
            ids = it["ids"].long().to(DEVICE)[None]
            mask = torch.ones_like(ids, dtype=torch.bool)
            pred = adapter(student.hidden(ids, mask), mask)[0].float()
            y = it["y"].to(DEVICE).float()
            rel.append(((pred - y).norm(dim=-1) / y.norm(dim=-1).clamp_min(1e-6)).mean().item())
            cos.append(F.cosine_similarity(pred, y, dim=-1).mean().item())
        adapter.train()
        return float(np.mean(rel)), float(np.mean(cos))

    def report(step, tag):
        names = probe["kinds"]
        rel, cos = measure(probe["items"])
        vrel, vcos = measure(val["items"])
        print(f"    {tag}@{step}: probe cos {cos:.3f} / rel {rel * 100:.1f}% | "
              f"val cos {vcos:.3f} / rel {vrel * 100:.1f}%", flush=True)
        return {"step": step, "probe_cos": cos, "probe_rel": rel, "val_cos": vcos, "val_rel": vrel,
                "probe_each": {n: measure([it])[1] for n, it in zip(names, probe["items"])}}

    if args.selfcheck:
        hist = report(0, "probe")
        print(f"  probe по кейсам: " + ", ".join(f"{k}={v:.3f}" for k, v in hist["probe_each"].items()))
        print(f"  VRAM {torch.cuda.max_memory_allocated() / 2**30:.1f} ГиБ, самопроверка пройдена")
        return

    def save(step):
        save_safetensors(adapter, args.out + ".safetensors", {
            "student_layers": ",".join(map(str, layers)), "norm": args.norm, "hidden": args.hidden,
            "proj_layers": args.proj_layers, "attention": args.attention, "attn_dim": args.attn_dim,
            "attn_heads": args.attn_heads, "attn_max_len": args.attn_max_len, "mixer": args.mixer,
            "mixer_heads": args.mixer_heads, "n_slices": len(layers), "in_dim": in_dim,
            "out_dim": out_dim, "norm_floor": floor, "pad_weight": args.pad_weight,
            "student": student.path, "teacher_mode": meta["teacher_mode"], "len": meta["len"],
            "prompts_dir": meta["prompts_dir"], "step": step})

    mlp_params = [p for n, p in adapter.named_parameters() if not n.startswith(("attn.", "mixer."))]
    branch_params = [p for n, p in adapter.named_parameters() if n.startswith(("attn.", "mixer."))]
    groups = [{"params": mlp_params, "base": args.lr, "floor": args.min_lr}]
    if branch_params:
        branch_lr = args.attn_lr if args.attn_lr is not None else args.lr
        groups.append({"params": branch_params, "base": branch_lr, "floor": args.min_lr})
    opt = torch.optim.AdamW(groups, betas=(0.9, 0.95), weight_decay=0.0)

    n_train = int(meta["n_train"])
    steps_per_epoch = n_train // args.batch
    total_steps = steps_per_epoch * args.epochs
    if args.max_steps:
        total_steps = min(total_steps, args.max_steps)
    print(f"{steps_per_epoch} шагов/эпоху, {total_steps} всего", flush=True)

    adapter.train()
    history = [report(0, "probe")]
    step, seen, t_epoch = 0, 0, time.time()
    for epoch in range(args.epochs):
        order = rng.permutation(len(meta["shards"]))
        for si in order:
            shard = load_pt(os.path.join(args.cache, meta["shards"][si]))
            items = shard["items"]
            idx = rng.permutation(len(items))
            for b in range(0, len(idx) - args.batch + 1, args.batch):
                batch = [items[i] for i in idx[b:b + args.batch]]
                mixed = []
                for it in batch:                      # подмешиваем пустой/негативы к пулу
                    r = rng.random()
                    if r < args.empty_ratio:
                        mixed.append(special["items"][0])
                    elif r < args.empty_ratio + args.neg_ratio:
                        mixed.append(special["items"][1 + int(rng.integers(len(special["items"]) - 1))])
                    else:
                        mixed.append(it)
                ids, mask, y = collate(mixed)
                frac = 0.5 * (1 + np.cos(np.pi * step / max(1, total_steps)))
                for g in opt.param_groups:
                    g["lr"] = g["floor"] + (g["base"] - g["floor"]) * frac
                opt.zero_grad()
                loss = loss_of(ids, mask, y)
                loss.backward()
                torch.nn.utils.clip_grad_norm_(adapter.parameters(), 1.0)
                opt.step()
                step += 1
                seen += ids.numel()
                if step % 20 == 0:
                    el = time.time() - t_epoch
                    print(f"  эпоха {epoch} шаг {step}/{total_steps}: loss {loss.item():.4f}, "
                          f"lr {opt.param_groups[0]['lr']:.1e}, {el / 20:.2f} с/шаг, "
                          f"{seen / el / 1000:.0f}k ток/с", flush=True)
                    t_epoch = time.time()
                if args.probe_every and step % args.probe_every == 0:
                    history.append(report(step, "probe"))
                if args.save_every and step % args.save_every == 0:
                    save(step)
                if args.max_steps and step >= args.max_steps:
                    break
            if args.max_steps and step >= args.max_steps:
                print(f"достигнут --max-steps {args.max_steps}, стоп", flush=True)
                break
        if args.max_steps and step >= args.max_steps:
            break

    history.append(report(step, "final"))
    save(step)
    with open(args.out + ".history.json", "w") as f:
        json.dump(history, f, indent=1)
    print(f"готово: {args.out}.safetensors ({n_par / 1e6:.1f}M, шагов {step}, "
          f"пик VRAM {torch.cuda.max_memory_allocated() / 2**30:.1f} ГиБ)", flush=True)


if __name__ == "__main__":
    main()