#!/usr/bin/env python3
"""A/B картинками: родной энкодер MOSS (Qwen3.5-9B) против 0.8B + адаптер.

Один промпт, один сид, один размер — картинки встают рядом, слева родной, справа адаптер.
Условие родного энкодера должно быть посчитано заранее (`native_embeds.py`): 15.6 ГБ учителя и
DiT в 16 ГБ не влезают одновременно, а условие от промпта не зависит.

  python3 src/native_embeds.py --out out/native_embeds.pt
  python3 src/sample_ab.py --adapter out/adapter_v1.safetensors --out-dir out/samples
"""
import argparse
import os
import sys

import torch
import torch.nn.functional as F
from PIL import Image, ImageDraw, ImageFont

sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
from adapter_lib import load_adapter                                             # noqa: E402
from moss_lib import (MOSS_DIR, STUDENT, Student, build_pipeline, condition,   # noqa: E402
                      empty_cache, generate, load_tokenizer)
from prompts import AB                                                           # noqa: E402

ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
FONT = "/usr/share/fonts/truetype/dejavu/DejaVuSans.ttf"


def grid(rows, out_path, gap=8, header=40, caption=28, bg=(24, 24, 24), fg=(235, 235, 235)):
    """rows: [(подпись, prompt, [PIL, PIL])]; колонки = родной / адаптер."""
    try:
        font_h, font_c = ImageFont.truetype(FONT, 20), ImageFont.truetype(FONT, 14)
    except OSError:
        font_h = font_c = ImageFont.load_default()
    width = max(sum(im.width for im in r[2]) + gap * (len(r[2]) - 1) for r in rows)
    height = header + sum(caption + gap + max(im.height for im in r[2]) for r in rows)
    canvas = Image.new("RGB", (width, height), bg)
    draw = ImageDraw.Draw(canvas)
    draw.text((8, 10), "родной Qwen3.5-9B (слева)  /  0.8B + адаптер (справа)", font=font_h, fill=fg)
    y = header
    for name, prompt, images in rows:
        draw.text((8, y), f"{name}: {prompt[:110]}", font=font_c, fill=fg)
        y += caption
        x = 0
        for im in images:
            canvas.paste(im, (x, y))
            x += im.width + gap
        y += max(im.height for im in images) + gap
    canvas.save(out_path)
    return out_path


def cosine(a, b):
    n = min(a.shape[0], b.shape[0])
    return F.cosine_similarity(a[:n].float(), b[:n].float(), dim=-1).mean().item()


def main():
    ap = argparse.ArgumentParser(description="A/B: родной энкодер MOSS против адаптера")
    ap.add_argument("--adapter", default=os.path.join(ROOT, "adapter.safetensors"))
    ap.add_argument("--native", default=os.path.join(ROOT, "out", "native_embeds.pt"))
    ap.add_argument("--moss-dir", default=MOSS_DIR)
    ap.add_argument("--student", default=STUDENT)
    ap.add_argument("--out-dir", default=os.path.join(ROOT, "out", "samples"))
    ap.add_argument("--size", type=int, default=1024)
    ap.add_argument("--steps", type=int, default=30)
    ap.add_argument("--guidance", type=float, default=3.0)
    ap.add_argument("--seed", type=int, default=1234)
    args = ap.parse_args()
    if not os.path.exists(args.native):
        raise SystemExit(f"нет {args.native} — сначала запустите src/native_embeds.py")

    native = torch.load(args.native, map_location="cpu", weights_only=True)
    os.makedirs(args.out_dir, exist_ok=True)
    tok = load_tokenizer(args.moss_dir)
    adapter, meta = load_adapter(args.adapter, device="cuda")
    student = Student(args.student, [int(x) for x in str(meta["student_layers"]).split(",")])
    pos = dict(zip([n for n, _ in AB], condition(adapter, student, tok, [t for _, t in AB])))
    neg = condition(adapter, student, tok, [""])[0]
    for name, emb in pos.items():
        print(f"  {name:10s} cos(условие адаптера, родное) = {cosine(emb[0], native['items'][name]):.3f}",
              flush=True)
    del student, adapter
    empty_cache()

    pipe = build_pipeline(args.moss_dir)
    pipe.to("cuda")
    rows = []
    for name, prompt in AB:
        adapter_img = generate(pipe, pos[name], neg, size=args.size, steps=args.steps,
                               seed=args.seed, guidance=args.guidance)
        nat_emb = (native["items"][name], torch.ones(native["items"][name].shape[0], dtype=torch.bool))
        nat_img = generate(pipe, nat_emb, (native["items"]["negative"],
                                          torch.ones(native["items"]["negative"].shape[0], dtype=torch.bool)),
                           size=args.size, steps=args.steps, seed=args.seed, guidance=args.guidance)
        adapter_img.save(os.path.join(args.out_dir, f"{name}_adapter.png"))
        nat_img.save(os.path.join(args.out_dir, f"{name}_native.png"))
        rows.append((name, prompt, [nat_img, adapter_img]))
        print(f"  {name}: готово", flush=True)

    out = grid(rows, os.path.join(args.out_dir, "ab_native_vs_adapter.png"))
    print(f"готово: {out}, пик VRAM {torch.cuda.max_memory_allocated() / 2**30:.1f} ГиБ", flush=True)


if __name__ == "__main__":
    main()