#!/usr/bin/env python3
"""Слияние смежных блоков: учим LoRA блока s отдавать то же, что отдавали блоки s..s+k.

Идея: блоки в середине меняют поток мало и частично друг друга отменяют. Значит из пары-тройки
смежных блоков можно оставить один, добавив ему низкоранговую поправку, обученную на состояниях
учителя (исходной модели). Так глубина падает без переобучения сети.

    python3 src/window.py cache --start 3 --drop 2 --samples 120 --size 1024
    python3 src/window.py train --start 3 --drop 2 --samples 120 --size 1024 --rank 32 --steps 3000
    python3 src/window.py eval  --start 3 --drop 2

Кэш: для каждого сэмпла (промпт, латент, t) пишем вход блока s и выход блока s+k учителя.
Память тренировки — один блок (226 МБ) + LoRA, поэтому шаг дешёвый, а учитель нужен только на кэше.
"""
import argparse
import json
import time
import os
import sys

import numpy as np
import torch
import torch.nn.functional as F
from PIL import Image

sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
from lab import (EVAL, LAB, Metrics, build, conditions, grid, sample,   # noqa: E402
                 drop_blocks)

CACHE_DIR = os.path.join(LAB, "cache")
TRAIN_PROMPTS = None


def prompts_for_training(n, seed=0):
    """Строки из пула ~/dataset (по одной на файл), как в обучении адаптера ТЕ."""
    global TRAIN_PROMPTS
    if TRAIN_PROMPTS is not None:
        return TRAIN_PROMPTS
    files = sorted(f for f in __import__("glob").glob(os.path.expanduser("~/dataset/civit/*.txt"))
                   if os.path.getsize(f) > 40)
    lines = []
    for f in files:
        with open(f, errors="ignore") as fh:
            for line in fh:
                line = line.strip()
                if 20 < len(line) < 400:
                    lines.append(line)
    rng = np.random.default_rng(seed)
    TRAIN_PROMPTS = [lines[i] for i in rng.permutation(len(lines))[:n]]
    return TRAIN_PROMPTS


# ---------------------------------------------------------------- LoRA
class LoRALinear(torch.nn.Module):
    """y = W x + (B A x) * scale; A и B — единственные обучаемые тензоры."""

    def __init__(self, base, rank, alpha=None):
        super().__init__()
        self.base = base
        for p in self.base.parameters():
            p.requires_grad_(False)
        self.a = torch.nn.Parameter(torch.randn(rank, base.in_features) / base.in_features ** 0.5)
        self.b = torch.nn.Parameter(torch.zeros(base.out_features, rank))
        self.scale = (alpha or rank) / rank

    # _Block читает .weight.dtype у проекций
    @property
    def weight(self):
        return self.base.weight

    @property
    def bias(self):
        return self.base.bias

    def forward(self, x):
        # a/b держим в fp32 (точность мелких поправок), считаем в dtype входа
        a, b = self.a.to(x.dtype), self.b.to(x.dtype)
        return self.base(x) + F.linear(F.linear(x, a), b) * self.scale


def inject_lora(block, rank, targets=("q", "k", "v", "o", "ffn_up_gate", "ffn_down")):
    for name in targets:
        lin = block.proj[name]
        if not isinstance(lin, LoRALinear):
            block.proj[name] = LoRALinear(lin, rank)


def merge_lora(block):
    """Влить LoRA в базовые веса: W' = W + B A * scale. На инференсе поправка ничего не стоит."""
    for name, lin in list(block.proj.items()):
        if isinstance(lin, LoRALinear):
            delta = (lin.b.float() @ lin.a.float()) * lin.scale
            base = lin.base
            base.weight.data = (base.weight.data.float() + delta).to(base.weight.dtype)
            block.proj[name] = base
    return block


def lora_state(block):
    return {k: v.detach().cpu() for k, v in block.state_dict().items() if k.endswith((".a", ".b"))}


# ---------------------------------------------------------------- cache
def cmd_cache(a):
    os.makedirs(CACHE_DIR, exist_ok=True)
    out = cache_path(a)
    if os.path.exists(out):
        print("cache exists:", out)
        return
    pipe = build()
    tr = pipe.transformer.eval()
    prompts = prompts_for_training(a.samples)
    # по одному-восемь промптов за раз: батч из 240 строк паддится до 4096 токенов и съедает 8 ГБ
    conds = []
    with torch.no_grad():
        for i in range(0, len(prompts), 8):
            e, m = pipe.encode_prompt(prompts[i:i + 8], 4096)
            conds += [(e[j:j + 1].cpu(), m[j:j + 1].cpu()) for j in range(e.shape[0])]
        nemb, nmask = pipe.encode_prompt([""], 4096)
    nemb, nmask = nemb.cpu(), nmask.cpu()
    pipe.text_encoder = None
    torch.cuda.empty_cache()

    start, last = a.start, a.start + a.drop
    rec = {}
    def make_hook(i):
        def hook(m, inp, out):
            if i in (start - 1, last):
                rec.setdefault("states", {})[i] = out.to(torch.bfloat16).cpu()
            if i == start - 1:
                rec["rope"] = inp[1]
            if i == start:
                rec["mod"] = inp[2]
        return hook
    handles = [l.register_forward_hook(make_hook(i)) for i, l in enumerate(tr.layers)]

    g = a.size // 16
    xs, mods, tgts, ropes, meta = [], [], [], [], []
    rng = np.random.default_rng(a.seed)

    @torch.no_grad()
    def collect(si, ei_emb, ei_mask):
        rec.clear()
        lat = torch.randn(1, 128, g, g,
                          generator=torch.Generator("cpu").manual_seed(int(rng.integers(1 << 30)))).cuda()
        lat = lat.bfloat16()
        t0 = float(rng.uniform(0.05, 0.95))
        # пара шагов Эйлера, чтобы латент был «как на реальной траектории», а не чистый шум
        for _ in range(a.warm):
            v = tr(lat, torch.tensor([1000.0 * t0]), ei_emb.cuda(), ei_mask.cuda()).sample.float()
            lat = (lat.float() - 0.08 * v).bfloat16()
        t0 = max(0.02, t0 - 0.08 * a.warm)
        tr(lat, torch.tensor([1000.0 * t0]), ei_emb.cuda(), ei_mask.cuda())
        return (rec["states"][start - 1], rec["states"][last],
                rec["mod"][0].cpu(), rec["rope"].cpu(), t0)

    for si in range(a.samples):
        x_s, t_s, mod_s, rope_s, t0 = collect(si, *conds[si])
        xs.append(x_s)
        tgts.append(t_s)
        mods.append(mod_s)
        ropes.append(rope_s)
        meta.append({"prompt": si, "t": round(t0, 3)})
        if si == 0:
            print(f"  state norm {x_s.float().norm():.1f} -> target norm {t_s.float().norm():.1f}, "
                  f"rel diff {((t_s.float() - x_s.float()).norm() / x_s.float().norm()):.3f}", flush=True)
        if (si + 1) % 25 == 0:
            print(f"  cached {si + 1}/{a.samples}", flush=True)
    for h in handles:
        h.remove()
    torch.save({"x": xs, "target": tgts,
                "mod": torch.stack(mods), "rope": ropes, "meta": meta,
                "start": start, "drop": a.drop, "size": a.size, "warm": a.warm}, out)
    print("saved", out, f"{os.path.getsize(out) / 2 ** 30:.1f} GiB")


def cache_path(a):
    return os.path.join(CACHE_DIR, f"win{a.start}_{a.drop}_{a.samples}_{a.size}.pt")


# ---------------------------------------------------------------- train
def cmd_train(a):
    data = cache_path(a)
    if not os.path.exists(data):
        raise SystemExit(f"нет кэша {data}; сначала запустите cache")
    d = torch.load(data, map_location="cpu", weights_only=False)
    xs, tgts = d["x"], d["target"]
    mod = d["mod"].to(torch.bfloat16)
    rope_map = {r.shape[1]: r.to("cuda") for r in d["rope"]}      # rope зависит от числа токенов
    n = len(xs)
    # длины последовательностей разные (у промптов разное число токенов) — батчи собираем
    # по одинаковой длине, чтобы не паддить и не менять поведение блока
    buckets = {}
    for i, x in enumerate(xs):
        buckets.setdefault(x.shape[1], []).append(i)
    rng = np.random.default_rng(0)
    for k in buckets:
        buckets[k] = rng.permutation(buckets[k]).tolist()

    pipe = build()
    tr = pipe.transformer.eval()
    idx = d["start"]
    if a.base_tags:                       # refit учится ПОВЕРХ первых слияний, а не на голом блоке
        drops = apply_spec(tr, [t for t in a.base_tags.split(",") if t])
        # после удаления блоков нумерация съезжает: переводим исходный номер в индекс живой модели
        idx = d["start"] - sum(1 for x in drops if x < d["start"])
        print(f"база: {a.base_tags}, блок {d['start']} -> индекс {idx}", flush=True)
    block = tr.layers[idx].to("cuda")
    pipe.transformer = None
    pipe.vae = None
    pipe.text_encoder = None
    torch.cuda.empty_cache()
    if not a.full:
        inject_lora(block, a.rank)
        block.to("cuda")
    params = [p for p in block.parameters() if p.requires_grad]
    print(f"trainable {sum(p.numel() for p in params) / 1e6:.2f}M"
          f" (base {sum(p.numel() for p in block.parameters()) / 1e6:.1f}M)", flush=True)
    import bitsandbytes as bnb
    opt = bnb.optim.AdamW8bit(params, lr=a.lr, betas=(0.9, 0.95), weight_decay=0.0)

    def pack(idx):
        """индексы одной длины -> (x, target, mod) на GPU"""
        x = torch.cat([xs[i] for i in idx]).cuda().bfloat16()
        t = torch.cat([tgts[i] for i in idx]).cuda().bfloat16()
        return x, t, mod[list(idx)].cuda(), rope_map[x.shape[1]]

    @torch.no_grad()
    def quality(idx):
        block.eval()
        rels, coss = [], []
        for i in idx:
            x, t, m, rope = pack([i])
            pred = block(x, rope, m).float()
            rels.append(float((((pred - t.float()) ** 2).mean(-1)
                               / t.float().pow(2).mean(-1).clamp_min(1e-6)).mean()))
            coss.append(float(F.cosine_similarity(pred.flatten(1), t.float().flatten(1), dim=1)))
        return sum(rels) / len(rels), sum(coss) / len(coss)

    held = [int(v) for k in buckets for v in buckets[k][:max(1, len(buckets[k]) // 10)]][:12]
    rel0, cos0 = quality(held)
    print(f"до обучения (просто удалить блоки): rel {rel0:.4f} cos {cos0:.5f}", flush=True)

    tag = a.tag or f"win{d['start']}_{d['start'] + d['drop']}"
    out = os.path.join(LAB, "lora", f"{tag}.pt")
    os.makedirs(os.path.dirname(out), exist_ok=True)
    keys = [k for k in buckets if len(buckets[k]) > 1] or list(buckets)   # бывает и по одному
    cursor = {k: 0 for k in keys}

    def next_batch():
        k = keys[int(rng.integers(len(keys)))]
        b = buckets[k]
        c = cursor[k]
        idx = [b[(c + j) % len(b)] for j in range(min(a.batch, len(b)))]
        cursor[k] = c + a.batch
        return idx

    step, t_start = 0, time.time()
    while step < a.steps:
        idx = next_batch()
        block.train()
        x, t, m, rope = pack(idx)
        pred = block(x, rope, m)
        tf = t.float()
        rel_tok = ((pred.float() - tf) ** 2).mean(-1) / tf.pow(2).mean(-1).clamp_min(1e-6)
        loss = rel_tok.mean() + a.cos_w * (1 - F.cosine_similarity(
            pred.float().flatten(1), tf.flatten(1), dim=1).mean())
        opt.zero_grad(set_to_none=True)
        loss.backward()
        torch.nn.utils.clip_grad_norm_(params, 1.0)
        opt.step()
        step += 1
        if step % 500 == 0 or step == a.steps:
            rel, cos = quality(held)
            print(f"step {step:5d} loss {float(loss):.4f} | held rel {rel:.4f} cos {cos:.5f} "
                  f"| {step / (time.time() - t_start):.1f} шаг/с", flush=True)
            save(out, block, a, d)
    rel, cos = quality(held)
    save(out, block, a, d)
    row = {"tag": tag, "start": d["start"], "drop": d["drop"], "rank": a.rank, "steps": step,
           "held_rel_before": round(rel0, 4), "held_cos_before": round(cos0, 5),
           "held_rel_after": round(rel, 4), "held_cos_after": round(cos, 5),
           "trainable_m": round(sum(p.numel() for p in params) / 1e6, 2)}
    with open(os.path.join(LAB, "window_results.jsonl"), "a") as f:
        f.write(json.dumps(row) + "\n")
    print(json.dumps(row), flush=True)


def save(path, block, a, d):
    torch.save({"state": lora_state(block) if not a.full else
                {k: v.cpu() for k, v in block.state_dict().items()},
                "full": a.full, "rank": a.rank, "start": d["start"], "drop": d["drop"]}, path)


# ---------------------------------------------------------------- eval
def cmd_eval(a):
    tag = a.tag or f"win{a.start}_{a.start + a.drop}"
    path = os.path.join(LAB, "lora", f"{tag}.pt")
    ck = torch.load(path, map_location="cpu", weights_only=False)
    pipe = build()
    tr = pipe.transformer.eval()
    cond = conditions(pipe)
    pipe.text_encoder = None
    torch.cuda.empty_cache()
    block = tr.layers[ck["start"]]
    if ck["full"]:
        block.load_state_dict(ck["state"])
    else:
        inject_lora(block, ck["rank"])
        block.to("cuda")
        block.load_state_dict(ck["state"], strict=False)
    if not ck["full"] and not a.no_merge:
        merge_lora(block)
    for l in block.proj.values():
        for p in l.parameters():
            p.requires_grad_(False)
    drop_blocks(tr, list(range(ck["start"] + 1, ck["start"] + ck["drop"] + 1)))
    images = sample(pipe, cond, size=a.size, steps=a.sample_steps, seed=a.seed)
    run_dir = os.path.join(LAB, "runs", tag)
    os.makedirs(run_dir, exist_ok=True)
    for i, im in enumerate(images):
        im.save(os.path.join(run_dir, f"{i:02d}.png"))
    grid(images, os.path.join(LAB, "runs", f"{tag}.png"))
    res = {"tag": tag, "blocks": len(tr.layers),
           "params_b": round(sum(p.numel() for p in tr.parameters()) / 1e9, 3),
           "size": a.size, "steps": a.steps, "seed": a.seed,
           "peak_vram_gib": round(torch.cuda.max_memory_allocated() / 2 ** 30, 2)}
    base_dir = os.path.join(LAB, "runs", "base")
    if os.path.isdir(base_dir):
        base_imgs = [Image.open(os.path.join(base_dir, f"{i:02d}.png")) for i in range(len(images))]
        m = Metrics()
        lp = m.lpips_vs(images, base_imgs)
        res["lpips_vs_base"] = None if lp is None else round(lp, 4)
        res["clip_score"] = round(m.clip_score(images, EVAL[:len(images)]), 2)
    with open(os.path.join(LAB, "results.jsonl"), "a") as f:
        f.write(json.dumps(res, ensure_ascii=False) + "\n")
    print(json.dumps(res, ensure_ascii=False), flush=True)


def apply_spec(tr, tags):
    """Влить LoRA всех слияний в веса и удалить поглощённые блоки. Возвращает список удалённых."""
    drops = set()
    for tag in tags:
        ck = torch.load(os.path.join(LAB, "lora", f"{tag}.pt"), map_location="cpu", weights_only=False)
        block = tr.layers[ck["start"]]
        if ck["full"]:
            block.load_state_dict(ck["state"])
        else:
            inject_lora(block, ck["rank"])
            block.to("cuda")
            block.load_state_dict(ck["state"], strict=False)
            merge_lora(block)
        drops |= set(range(ck["start"] + 1, ck["start"] + ck["drop"] + 1))
    drop_blocks(tr, sorted(drops))
    return sorted(drops)


def cmd_combo(a):
    """Собрать все слияния в одну модель: каждый LoRA влит в свои веса, удалённые блоки — вон."""
    tags = [t for t in a.tags.split(",") if t]
    pipe = build()
    tr = pipe.transformer.eval()
    cond = conditions(pipe)
    pipe.text_encoder = None
    torch.cuda.empty_cache()
    drops = set(apply_spec(tr, tags))
    tag = a.tag or "combo"
    images = sample(pipe, cond, size=a.size, steps=a.sample_steps, seed=a.seed)
    run_dir = os.path.join(LAB, "runs", tag)
    os.makedirs(run_dir, exist_ok=True)
    for i, im in enumerate(images):
        im.save(os.path.join(run_dir, f"{i:02d}.png"))
    grid(images, os.path.join(LAB, "runs", f"{tag}.png"))
    res = {"tag": tag, "blocks": len(tr.layers), "dropped": sorted(drops),
           "params_b": round(sum(p.numel() for p in tr.parameters()) / 1e9, 3),
           "size": a.size, "steps": a.sample_steps, "seed": a.seed,
           "peak_vram_gib": round(torch.cuda.max_memory_allocated() / 2 ** 30, 2)}
    base_dir = os.path.join(LAB, "runs", "base")
    if os.path.isdir(base_dir):
        base_imgs = [Image.open(os.path.join(base_dir, f"{i:02d}.png")) for i in range(len(images))]
        m = Metrics()
        lp = m.lpips_vs(images, base_imgs)
        res["lpips_vs_base"] = None if lp is None else round(lp, 4)
        res["clip_score"] = round(m.clip_score(images, EVAL[:len(images)]), 2)
    with open(os.path.join(LAB, "results.jsonl"), "a") as f:
        f.write(json.dumps(res, ensure_ascii=False) + "\n")
    print(json.dumps(res, ensure_ascii=False), flush=True)


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("cmd", choices=["cache", "train", "eval", "combo"])
    ap.add_argument("--tags", default="", help="для combo: через запятую")
    ap.add_argument("--start", type=int, default=3, help="блок, который остаётся и получает LoRA")
    ap.add_argument("--drop", type=int, default=2, help="сколько блоков после него удалить")
    ap.add_argument("--samples", type=int, default=300)
    ap.add_argument("--warm", type=int, default=2, help="шагов Эйлера перед записью состояния")
    ap.add_argument("--size", type=int, default=512)
    ap.add_argument("--steps", type=int, default=3000, help="шаги обучения")
    ap.add_argument("--sample-steps", type=int, default=20, help="шаги сэмплинга на eval")
    ap.add_argument("--seed", type=int, default=1234)
    ap.add_argument("--rank", type=int, default=32)
    ap.add_argument("--lr", type=float, default=1e-3)
    ap.add_argument("--batch", type=int, default=4)
    ap.add_argument("--cos-w", type=float, default=0.1)
    ap.add_argument("--full", action="store_true", help="учить весь блок, не только LoRA")
    ap.add_argument("--tag", default=None)
    ap.add_argument("--base-tags", default="", help="сначала влить эти слияния в веса модели")
    ap.add_argument("--no-merge", action="store_true", help="не вливать LoRA в веса на eval")
    a = ap.parse_args()
    {"cache": cmd_cache, "train": cmd_train, "eval": cmd_eval, "combo": cmd_combo}[a.cmd](a)


if __name__ == "__main__":
    main()