"""
inference.py — single-checkpoint inference for Anime-Elite v2.

Loads EMA weights by default (falls back to live weights if not present).
Mirrors the proven sampling logic from train.py.

CLI:
    python inference.py --prompt "1girl,red hair,smile" --seed 21
    python inference.py --prompt "..." --guidance 2.5 --steps 200 --n 8
"""
import argparse
from pathlib import Path

import torch
import torch.nn as nn
from PIL import Image
from torch.amp import autocast
from diffusers import UNet2DConditionModel, DDIMScheduler


class TagConditioner(nn.Module):
    def __init__(self, vocab_size, dim=256, n_tokens=4):
        super().__init__()
        self.n_tokens, self.dim = n_tokens, dim
        self.net = nn.Sequential(
            nn.Linear(vocab_size, 512), nn.SiLU(),
            nn.Linear(512, n_tokens * dim),
        )

    def forward(self, x):
        return self.net(x).view(-1, self.n_tokens, self.dim)


def build_unet(cross_dim=256, sample_size=96):
    return UNet2DConditionModel(
        sample_size=sample_size,
        in_channels=3, out_channels=3,
        layers_per_block=2,
        block_out_channels=(96, 192, 320, 384),
        down_block_types=("DownBlock2D", "CrossAttnDownBlock2D", "CrossAttnDownBlock2D", "DownBlock2D"),
        up_block_types=("UpBlock2D", "CrossAttnUpBlock2D", "CrossAttnUpBlock2D", "UpBlock2D"),
        cross_attention_dim=cross_dim,
        attention_head_dim=8,
    )


@torch.no_grad()
def generate(unet, tag_cond, vocab, prompt, n=4, seed=42,
             steps=200, guidance=2.0, size=96, device="cuda"):
    unet.eval(); tag_cond.eval()

    tag_to_idx = {t: i for i, t in enumerate(vocab)}
    mh = torch.zeros(len(vocab))
    hits = []
    for t in [s.strip() for s in prompt.split(",") if s.strip()]:
        if t in tag_to_idx:
            mh[tag_to_idx[t]] = 1.0
            hits.append(t)
    mh = mh.unsqueeze(0).repeat(n, 1).to(device)
    null = torch.zeros_like(mh)
    cond = tag_cond(mh)
    uncond = tag_cond(null)

    sched = DDIMScheduler(num_train_timesteps=1000, beta_schedule="squaredcos_cap_v2")
    sched.set_timesteps(steps)

    g = torch.Generator(device=device).manual_seed(seed)
    x = torch.randn(n, 3, size, size, device=device, generator=g)

    for t in sched.timesteps:
        xt = torch.cat([x, x])
        ctx = torch.cat([uncond, cond])
        with autocast("cuda", dtype=torch.bfloat16):
            pred = unet(xt, t, encoder_hidden_states=ctx).sample
        pu, pc = pred.float().chunk(2)
        pred = pu + guidance * (pc - pu)
        x = sched.step(pred, t, x).prev_sample

    arr = ((x.clamp(-1, 1) + 1) * 127.5).byte().permute(0, 2, 3, 1).cpu().numpy()
    return [Image.fromarray(a) for a in arr], hits


def main():
    p = argparse.ArgumentParser()
    p.add_argument("--ckpt", default=r"ckpt_e040_slim.pt")
    p.add_argument("--prompt", default="1girl,red hair,floral background,smile")
    p.add_argument("--n", type=int, default=4)
    p.add_argument("--seed", type=int, default=56)
    p.add_argument("--guidance", type=float, default=2.4)
    p.add_argument("--steps", type=int, default=160)
    p.add_argument("--size", type=int, default=96)
    p.add_argument("--out", default="out")
    p.add_argument("--use_live", action="store_true",
                   help="Use live (non-EMA) weights even if EMA is present")
    args = p.parse_args()

    device = "cuda" if torch.cuda.is_available() else "cpu"
    print(f"Loading {args.ckpt} on {device}...")
    ckpt = torch.load(args.ckpt, map_location=device)
    vocab = ckpt["vocab"]

    sd_key = "unet" if args.use_live or "ema_unet" not in ckpt else "ema_unet"
    print(f"Using weights from key: '{sd_key}'")

    unet = build_unet(sample_size=args.size).to(device)
    tag_cond = TagConditioner(len(vocab)).to(device)
    unet.load_state_dict(ckpt[sd_key])
    tag_cond.load_state_dict(ckpt["tag_cond"])

    print(f"Prompt: {args.prompt!r}")
    print(f"n={args.n}  seed={args.seed}  guidance={args.guidance}  steps={args.steps}")
    imgs, hits = generate(unet, tag_cond, vocab, args.prompt,
                          n=args.n, seed=args.seed, steps=args.steps,
                          guidance=args.guidance, size=args.size, device=device)

    out_dir = Path(args.out)
    out_dir.mkdir(parents=True, exist_ok=True)
    for i, im in enumerate(imgs):
        im.save(out_dir / f"sample_{i:02d}.png")

    print(f"\nMatched tags: {hits}")
    print(f"Saved {len(imgs)} images to {out_dir}/")


if __name__ == "__main__":
    main()