#!/usr/bin/env python3
"""Сэмпл из transformer/ (или из .pt чекпойнта обучения).

    python3 src/sample.py --prompt "A photograph of a red fox in fresh snow." --out fox.png
"""
from __future__ import annotations

import argparse
import sys
from pathlib import Path

import torch
from PIL import Image

ROOT = Path(__file__).resolve().parent.parent
sys.path[:0] = [str(ROOT / "src"), str(ROOT / "vae")]

from model import DEV, UDiT          # noqa: E402
from text import EttinText           # noqa: E402
from ae.model import AutoEncoder     # noqa: E402
from export import load_exported     # noqa: E402


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--model", default=str(ROOT / "transformer"), help="папка transformer/ или .pt")
    ap.add_argument("--prompt", required=True)
    ap.add_argument("--out", default=str(ROOT / "tmp" / "sample.png"))
    ap.add_argument("--size", type=int, default=256, help="сторона кадра, кратна 64")
    ap.add_argument("--steps", type=int, default=16, help="шагов тела")
    ap.add_argument("--tail-steps", type=int, default=8, help="шагов декодера (хвоста)")
    ap.add_argument("--cfg", type=float, default=2.5)
    ap.add_argument("--seed", type=int, default=0)
    ap.add_argument("--n", type=int, default=1, help="сколько картинок подряд от --seed")
    ap.add_argument("--text-len", type=int, default=192)
    args = ap.parse_args()

    ae = AutoEncoder.from_pretrained(str(ROOT / "vae"), device=DEV)
    enc = EttinText(str(ROOT / "text_encoder"), DEV, max_len=args.text_len)
    m = UDiT(ae, text_dim=enc.dim).to(DEV)

    src = Path(args.model)
    if src.is_dir():
        step = load_exported(m, enc, ae, src, DEV).get("step")
    else:
        sd = torch.load(src, map_location=DEV, weights_only=False)
        own = m.body.state_dict()
        m.body.load_state_dict({k: v for k, v in sd["body"].items() if k in own}, strict=False)
        enc.mixer.load_state_dict(sd["mixer"])
        if sd.get("decoder") is not None:
            ae.decoder.load_state_dict(sd["decoder"])
        step = sd.get("step")
    m.eval()
    print(f"[model] {src}, шаг {step}", flush=True)

    null_hs, null_mk = enc.slices([""], pad_to_len=args.text_len)
    null_t = enc.mix(null_hs.float())
    hs, mk = enc.slices([args.prompt], pad_to_len=args.text_len)
    txt = enc.mix(hs.float())

    out = Path(args.out)
    out.parent.mkdir(parents=True, exist_ok=True)
    gh = gw = args.size // 8
    for i in range(args.n):
        y = m.sample(txt, mk, gh, gw, steps=args.steps, tail_steps=args.tail_steps, cfg=args.cfg,
                     null_text=null_t, null_mask=null_mk,
                     generator=torch.Generator(DEV).manual_seed(args.seed + i))
        im = Image.fromarray(((y[0].clamp(-1, 1) + 1) * 127.5).byte().permute(1, 2, 0).cpu().numpy())
        path = out if args.n == 1 else out.with_name(f"{out.stem}_{args.seed + i}{out.suffix}")
        im.save(path)
        print(f"[ok] {path} {im.size[0]}x{im.size[1]}, seed {args.seed + i}", flush=True)


if __name__ == "__main__":
    main()