#!/usr/bin/env python3
"""Измерения и проверки для сжатия DiT: эталон, удаление блоков, метрики (LPIPS + CLIP).

    python3 lab/lab.py baseline            # база: картинки + метрики-эталон
    python3 lab/lab.py drop 2 3 8          # удалить блоки (0-based) и измерить


Метрика — против базы с тем же сидом: LPIPS (ниже = ближе), CLIP-скор (выше = точнее промпт).
Всё в 512 px, 20 шагов, guidance 3.0, uncond_skip_layers=[] (иначе индексы блоков съезжают).
"""
import argparse
import glob
import json
import os
import sys
import time

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

MOSS = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))   # папка модели
LAB = os.path.join(os.path.dirname(os.path.abspath(__file__)), "work")  # кэши и прогоны
# diffusers ищет компоненты из model_index.json через importlib по имени модуля
# ("pipeline_moss_image"), поэтому папка модели обязана быть на sys.path.
if MOSS not in sys.path:
    sys.path.insert(0, MOSS)
EVAL = [
    "A red fox walks through a snowy forest at dusk, warm light breaking through the pines, 85mm",
    'A vintage enamel sign bolted to a red brick wall, the sign reads "OPEN 24 HOURS", peeling paint',
    "1girl, solo, long silver hair, blue eyes, black coat, tokyo street at night, cinematic",
    '{"layout":[{"text":"SALE 50%","box":[40,40,420,120]},{"text":"only today","box":[40,600,420,660]}]}',
    "Close-up portrait of an old fisherman, weathered skin, grey beard, sea spray, soft window light",
    "A rustic wooden table with a bowl of ripe pears, a knife and a glass of red wine, morning sun",
    "Two husky puppies playing in deep snow, motion blur, low winter sun, telephoto",
    "A brutalist concrete library at dusk, warm interior glow through narrow windows, wet plaza",
    "A matte black ceramic coffee mug on a linen cloth, product photography, soft shadow, 100mm macro",
    "Abstract flowing ink in water, teal and amber, macro, high contrast, black background",
]
SIZE, STEPS, GUIDANCE, SEED = 512, 20, 3.0, 1234


# ---------------------------------------------------------------- модель и условия
def build():
    from diffusers import DiffusionPipeline
    pipe = DiffusionPipeline.from_pretrained(
        MOSS, custom_pipeline="pipeline_moss_image", trust_remote_code=True,
        dtype={"default": torch.bfloat16, "vae": torch.float32})
    pipe.set_progress_bar_config(disable=True)
    return pipe.to("cuda")


def conditions(pipe):
    """Эмбеддинги промптов и пустого негатива (считает адаптерный ТЕ), затем ТЕ отпускаем."""
    import zlib
    key = zlib.crc32("|".join(EVAL).encode())          # hash() рандомизирован между процессами
    path = os.path.join(LAB, f"eval_cond_{len(EVAL)}_{key}.pt")
    if os.path.exists(path):
        return torch.load(path)
    with torch.no_grad():
        emb, mask = pipe.encode_prompt(EVAL, 4096)
        nemb, nmask = pipe.encode_prompt([""], 4096)
    cond = (emb.cpu(), mask.cpu(), nemb.cpu(), nmask.cpu())
    torch.save(cond, path)
    return cond


@torch.no_grad()
def sample(pipe, cond, size=SIZE, steps=STEPS, seed=SEED, guidance=GUIDANCE, prompts=None):
    emb, mask, nemb, nmask = (t.to("cuda") for t in cond)
    out = []
    n = emb.shape[0] if prompts is None else len(prompts)
    for i in range(n):
        g = torch.Generator("cuda").manual_seed(seed)
        out.append(pipe(prompt_embeds=emb[i:i + 1], prompt_attention_mask=mask[i:i + 1],
                        negative_prompt_embeds=nemb, negative_prompt_attention_mask=nmask,
                        height=size, width=size, num_inference_steps=steps,
                        guidance_scale=guidance, generator=g,
                        uncond_skip_layers=[]).images[0])
    return out


def grid(images, path, cols=5):
    w = images[0].width
    rows = (len(images) + cols - 1) // cols
    canvas = Image.new("RGB", (cols * (w + 6), rows * (w + 6)), "white")
    for i, im in enumerate(images):
        canvas.paste(im, ((i % cols) * (w + 6), (i // cols) * (w + 6)))
    canvas.save(path)


# ---------------------------------------------------------------- метрики
class Metrics:
    """LPIPS к базе + CLIP-скор абсолютный. Модели грузятся один раз."""
    def __init__(self, device="cuda"):
        import lpips
        from transformers import CLIPModel, CLIPImageProcessor, CLIPTokenizer
        self.lp = lpips.LPIPS(net="alex", verbose=False).to(device).eval()
        name = "openai/clip-vit-large-patch14"
        self.clip = CLIPModel.from_pretrained(name, local_files_only=True).to(device).eval()
        self.ip = CLIPImageProcessor.from_pretrained(name, local_files_only=True)
        self.tk = CLIPTokenizer.from_pretrained(name, local_files_only=True)
        self.device = device

    @torch.no_grad()
    def lpips_vs(self, images, base):
        if images[0].size != base[0].size:
            return None
        vals = []
        for a, b in zip(images, base):
            ta = torch.tensor(list(a.convert("RGB").getdata()), dtype=torch.float32).reshape(a.height, a.width, 3)
            tb = torch.tensor(list(b.convert("RGB").getdata()), dtype=torch.float32).reshape(b.height, b.width, 3)
            ta = (ta / 127.5 - 1).permute(2, 0, 1)[None].to(self.device)
            tb = (tb / 127.5 - 1).permute(2, 0, 1)[None].to(self.device)
            vals.append(float(self.lp(ta, tb)))
        return sum(vals) / len(vals)

    @torch.no_grad()
    def clip_score(self, images, prompts):
        pix = self.ip(images=[im.convert("RGB") for im in images], return_tensors="pt")["pixel_values"].to(self.device)
        tok = self.tk(prompts, padding=True, truncation=True, return_tensors="pt").to(self.device)
        out = self.clip(pixel_values=pix, **tok)
        ie = F.normalize(out.image_embeds, dim=-1)
        te = F.normalize(out.text_embeds, dim=-1)
        return float((ie * te).sum(-1).mean() * 100)


# ---------------------------------------------------------------- хирургия
def drop_blocks(tr, idxs):
    drop = set(idxs)
    tr.layers = torch.nn.ModuleList([l for i, l in enumerate(tr.layers) if i not in drop])
    return tr


# ---------------------------------------------------------------- CLI
def cmd_metrics(a):
    """Пересчитать метрики по уже сгенерированным lab/runs/<tag>/ без прогона модели."""
    m = Metrics()
    base_dir = os.path.join(LAB, "runs", "base")
    base = [Image.open(os.path.join(base_dir, f"{i:02d}.png")) for i in range(len(EVAL))]
    rows = []
    for tag in sorted(os.listdir(os.path.join(LAB, "runs"))):
        d = os.path.join(LAB, "runs", tag)
        if not os.path.isdir(d) or tag == "base":
            continue
        files = sorted(glob.glob(os.path.join(d, "*.png")))
        if len(files) != len(EVAL):
            continue
        imgs = [Image.open(f) for f in files]
        lp = m.lpips_vs(imgs, base)
        rows.append({"tag": tag, "lpips_vs_base": None if lp is None else round(lp, 4),
                     "clip_score": round(m.clip_score(imgs, EVAL), 2)})
        print(json.dumps(rows[-1]), flush=True)
    print(json.dumps({"base_clip": round(m.clip_score(base, EVAL), 2)}), flush=True)


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("cmd", choices=["baseline", "drop", "metrics"])
    ap.add_argument("args", nargs="*")
    ap.add_argument("--out", default=None, help="имя прогона в lab/runs/")
    ap.add_argument("--size", type=int, default=SIZE)
    ap.add_argument("--steps", type=int, default=STEPS)
    ap.add_argument("--seed", type=int, default=SEED)
    ap.add_argument("--probe", action="store_true", help="только скорость/VRAM, без картинок")
    ap.add_argument("--tags", default="", help="сначала собрать модель из слияний lab/lora/<tag>.pt")
    a = ap.parse_args()

    if a.cmd == "metrics":
        cmd_metrics(a)
        return

    pipe = build()
    tr = pipe.transformer.eval()
    cond = conditions(pipe)
    if os.environ.get("DROP_TE", "1") == "1" and getattr(pipe, "text_encoder", None) is not None:
        pipe.text_encoder = None
        torch.cuda.empty_cache()

    if a.tags:
        from window import apply_spec
        print("сборка из слияний:", apply_spec(tr, [t for t in a.tags.split(",") if t]), flush=True)
    tag = a.out or (a.cmd + "_" + "_".join(a.args))
    if a.cmd == "baseline":
        tag = a.out or "base"
    if a.cmd == "drop":
        drop_blocks(tr, [int(x) for x in a.args])

    n_par = sum(p.numel() for p in tr.parameters())
    t0 = time.time()
    images = sample(pipe, cond, size=a.size, steps=a.steps, seed=a.seed)
    dt = time.time() - t0
    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(n_par / 1e9, 3),
           "size": a.size, "steps": a.sample_steps, "seed": a.seed,
           "gen_seconds": round(dt / len(images), 1),
           "peak_vram_gib": round(torch.cuda.max_memory_allocated() / 2 ** 30, 2)}
    base_dir = os.path.join(LAB, "runs", "base")
    if tag != "base" and os.path.isdir(base_dir):
        base_imgs = [Image.open(os.path.join(base_dir, f"{i:02d}.png")) for i in range(len(images))]
        if (base_imgs[0].size == images[0].size):
            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), 2)
    path = os.path.join(LAB, "results.jsonl")
    with open(path, "a") as f:
        f.write(json.dumps(res, ensure_ascii=False) + "\n")
    print(json.dumps(res, ensure_ascii=False), flush=True)


if __name__ == "__main__":
    main()