"""Экспорт обученного UDM в `transformer/`: safetensors + config.json.

Три части в одном файле, по префиксам:
  body.*        — тело (8 блоков из энкодера AE + 4 новых), поток в латентном пространстве
  text_mixer.*  — обучаемый выбор слоёв Ettin (SliceMixer)
  tail.*        — декодер AE (рендер латента в пиксели), дообученный как обычный flow

  python3 src/export.py --ckpt tmp/run256/last.pt --out transformer
"""
from __future__ import annotations

import argparse
import json
from pathlib import Path

import torch
from safetensors.torch import save_file

ROOT = Path(__file__).resolve().parent.parent


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--ckpt", default=str(ROOT / "tmp" / "run256" / "last.pt"))
    ap.add_argument("--out", default=str(ROOT / "transformer"))
    ap.add_argument("--metrics", default="", help="дописать метрики в config.json (строка)")
    ap.add_argument("--half", action="store_true", help="хранить веса в fp16 (вдвое меньше, для публикации)")
    args = ap.parse_args()

    ck = torch.load(args.ckpt, map_location="cpu")
    tensors = {}
    for name, sd in (("body", ck["body"]), ("text_mixer", ck["mixer"]), ("tail", ck.get("decoder") or {})):
        for k, v in sd.items():
            tensors[f"{name}.{k}"] = (v.half() if args.half else v).contiguous()
    out = Path(args.out)
    out.mkdir(parents=True, exist_ok=True)
    save_file(tensors, str(out / "diffusion_pytorch_model.safetensors"))

    ae_cfg = json.loads((ROOT / "vae" / "config.json").read_text())
    ta = ck.get("args") or {}
    counts = {}                          # сколько блоков в каждом списке — по ключам тензоров
    for k in ck["body"]:
        g = k.split(".")[0]
        if g[0] in "ed" and "." in k:
            idx = k.split(".")[1]
            if idx.isdigit():
                counts.setdefault(g, set()).add(idx)
    nblk = {g: len(v) for g, v in counts.items()}
    # конфиг описывает ТОЛЬКО тело; статистики латента — в vae/config.json, шаг — параметр запуска
    cfg = {
        "_class_name": "UDiT",
        "format": "udm/v1",
        "storage_dtype": "float16" if args.half else "float32",
        "latent_channels": ae_cfg["latent_channels"],
        "text_dim": ta.get("text_dim", 768),
        "body": {
            "kind": "пирамидальный U: ширина растёт к низу",
            "widths": [192, 512, 1024, 1536],
            "heads": [6, 8, 8, 8],
            "blocks": nblk,
            "levels": [
                {"level": 0, "tokens": "÷1", "width": 192,
                 "blocks": "conv 6 enc / 4 dec + 2 self-attn на сторону"},
                {"level": 1, "tokens": "÷2", "width": 512, "blocks": "attn 5 enc / 4 dec"},
                {"level": 2, "tokens": "÷4", "width": 1024, "blocks": "attn 5 enc / 4 dec"},
                {"level": 3, "tokens": "÷8", "width": 1536, "blocks": "attn 4 (боттлнек)"},
            ],
            "skips": "аддитивные (SD-стиль): e0/e1/e2 прибавляются в декодере",
            "down": "mean-pool 2x + Linear",
            "up": "Linear + nearest 2x",
            "positions": {
                "rope": {"kind": "2d", "base": 100, "phase": "в латентных пикселях на каждом уровне"},
                "absolute": {"kind": "2d sincos", "dim": 64,
                             "into": "zero-init Linear(64 -> width) к выходу stem, в паре с RoPE"},
            },
            "text": "кросс-внимание на всех уровнях",
            "parametrization": "x0 латента",
            "params": round(sum(v.numel() for k, v in tensors.items() if k.startswith("body.")) / 1e6, 2),
        },
    }
    (out / "config.json").write_text(json.dumps(cfg, indent=2, ensure_ascii=False))

    print(f"сохранено: {out}/diffusion_pytorch_model.safetensors "
          f"({sum(v.numel() for v in tensors.values())/1e6:.1f}M параметров, "
          f"{(out / 'diffusion_pytorch_model.safetensors').stat().st_size/1e6:.0f} МБ)")
    print(f"тело {cfg['body']['params']}M, блоки {nblk}")


def load_exported(model, enc, ae, path, device=None):
    """Загрузить из `transformer/` (config.json + safetensors с префиксами body./text_mixer./tail.)."""
    import json
    from safetensors.torch import load_file
    path = Path(path)
    sd = load_file(str(path / "diffusion_pytorch_model.safetensors"), device=str(device) if device else None)
    def sub(prefix):
        return {k[len(prefix) + 1:]: v for k, v in sd.items() if k.startswith(prefix + ".")}
    model.body.load_state_dict(sub("body"), strict=False)
    enc.mixer.load_state_dict(sub("text_mixer"))
    if getattr(ae, "decoder", None) is not None and any(k.startswith("tail.") for k in sd):
        ae.decoder.load_state_dict(sub("tail"))      # хвоста может не быть (--vae sdxs)
    cfg_path = path / "config.json"
    return json.loads(cfg_path.read_text()) if cfg_path.exists() else {}


if __name__ == "__main__":
    main()