"""Экспорт обученного 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 {}          # чекпойнт помнит свои аргументы — не сочиняем заново
    training = {
        "dataset": ta.get("hf_dataset") or ta.get("data") or "~/dataset/butterflies_256x64",
        "parts": ta.get("parts"),
        "size": ta.get("size"),
        "steps": int(ck["step"]),
        "batch": ta.get("batch", 8),
        "accum": ta.get("accum", 1),
        "lr": ta.get("lr", 1e-4),
        "lr_tail": ta.get("lr_tail", 3e-5),
        "drop_text": ta.get("drop_text", 0.1),
        "init": ta.get("init"),
        "metrics": args.metrics,
    }
    cfg = {
        "_class_name": "UDiT",
        "name": "Ra",
        "format": "udm/v1",
        "step": int(ck["step"]),
        "latent": {
            "channels": ae_cfg["latent_channels"],
            "spatial_factor": ae_cfg["patch_size"],
            "mean": ae_cfg["latent_mean"],
            "std": ae_cfg["latent_std"],
            "source": "vae/ (Muinez/f8c32-ae)",
        },
        "text": {
            "encoder": "text_encoder/ (jhu-clsp/ettin-encoder-150m, 768-dim, 22 слоя)",
            "max_len": 192,
            "slices": [2, 5, 8, 11, 14, 17, 20],
            "slice_blocks": 2,
            "trainable": "text_mixer.*",
        },
        "body": {
            "kind": "пирамидальный U: ширина растёт к низу",
            "levels": [
                {"level": 0, "tokens": "÷1", "width": 192, "blocks": "conv 6 enc / 4 dec + 1 self-attn на сторону",
                 "text": "кросс-внимание в каждом conv-блоке и в self-attn"},
                {"level": 1, "tokens": "÷2", "width": 512, "blocks": "attn 4 enc / 3 dec", "text": "кросс-внимание"},
                {"level": 2, "tokens": "÷4", "width": 1024, "blocks": "attn 3 enc / 2 dec", "text": "кросс-внимание"},
                {"level": 3, "tokens": "÷8", "width": 1536, "blocks": "attn 2 (боттлнек)", "text": "кросс-внимание"},
            ],
            "skips": "аддитивные (SD-стиль): e0/e1/e2 прибавляются в декодере",
            "down": "mean-pool 2× + Linear",
            "up": "Linear + nearest 2×",
            "rope": "2D RoPE, base 100, фаза в латентных пикселях на каждом уровне",
            "parametrization": "x0 латента",
        },
        "tail": {
            "source": "декодер AE, 6 блоков @512, PixelShuffle ×3",
            "trained_as": "flow по пикселям, условие z0 + (ошибка тела).detach()",
            "sample_steps": 4,
        },
        "params": {
            "body": round(sum(v.numel() for k, v in tensors.items() if k.startswith("body.")) / 1e6, 2),
            "text_mixer": round(sum(v.numel() for k, v in tensors.items()
                                    if k.startswith("text_mixer.")) / 1e6, 2),
            "tail": round(sum(v.numel() for k, v in tensors.items() if k.startswith("tail.")) / 1e6, 2),
        },
        "training": {k: v for k, v in training.items() if v is not None},
        "storage_dtype": "float16" if args.half else "float32",
    }
    cfg["params"]["total"] = round(sum(cfg["params"].values()), 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['params']}")


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"))
    ae.decoder.load_state_dict(sub("tail"))
    cfg_path = path / "config.json"
    return json.loads(cfg_path.read_text()) if cfg_path.exists() else {}


if __name__ == "__main__":
    main()