"""Inference for Qwen-Image-Edit-2511 fine-tuned on Mobile-GUI world-model data.

Loads the last 3 samples from the training metadata and runs them through
either the full checkpoint or the LoRA adapter (or both sequentially).

Usage:
    python infer_mobilegui.py --mode full
    python infer_mobilegui.py --mode lora
    python infer_mobilegui.py --mode both          # full first, then LoRA
    python infer_mobilegui.py --mode both --num-steps 40 --seed 123

Run from any cwd; the script chdirs into DiffSynth-Studio so the base model
cache at ./models/Qwen/... resolves. Outputs go to --output-dir (default
alongside this script: ./infer_outputs/).
"""
import argparse
import gc
import json
import os
import re
import sys
from pathlib import Path

PROJECT = Path("/storage/ljx")
DIFFSYNTH = PROJECT / "repo" / "DiffSynth-Studio"
DATA_BASE = PROJECT / "data" / "Mobile-GUI-Worldmodel-SFT"
DEFAULT_METADATA = DATA_BASE / "metadata_qwen_edit.json"
DEFAULT_FULL_DIR = PROJECT / "models" / "train" / "Qwen-Image-Edit-2511_MobileGUI_full_v2"
DEFAULT_LORA_DIR = PROJECT / "models" / "train" / "Qwen-Image-Edit-2511_MobileGUI_lora_v2"


def pick_free_gpu(min_free_mib: int = 20000) -> int:
    """Return index of the GPU with the most free memory (via nvidia-smi).

    Returns 0 as a last-resort fallback. Intended for pods that requested
    multiple GPUs but only need one, to avoid hitting a device already
    polluted by a non-k8s process.
    """
    import subprocess
    try:
        out = subprocess.check_output(
            ["nvidia-smi",
             "--query-gpu=index,memory.free",
             "--format=csv,noheader,nounits"],
            text=True, timeout=10)
    except Exception as e:
        print(f"[gpu-pick] nvidia-smi failed: {e}; fallback to 0")
        return 0
    rows = []
    for line in out.strip().splitlines():
        idx, free = [x.strip() for x in line.split(",")]
        rows.append((int(idx), int(free)))
    rows.sort(key=lambda r: -r[1])
    print(f"[gpu-pick] nvidia-smi free MiB by index: {rows}")
    if rows and rows[0][1] >= min_free_mib:
        return rows[0][0]
    print(f"[gpu-pick] no GPU has >= {min_free_mib} MiB free, "
          f"using index {rows[0][0] if rows else 0} anyway")
    return rows[0][0] if rows else 0


def latest_step_ckpt(ckpt_dir: Path) -> Path:
    """Pick the step-N.safetensors with the largest N."""
    cands = []
    for p in ckpt_dir.glob("step-*.safetensors"):
        m = re.match(r"step-(\d+)\.safetensors$", p.name)
        if m:
            cands.append((int(m.group(1)), p))
    if not cands:
        raise FileNotFoundError(f"No step-*.safetensors in {ckpt_dir}")
    cands.sort()
    return cands[-1][1]


def target_hw(src_w: int, src_h: int, target_area: int = 1024 * 1024,
              divisor: int = 32) -> tuple[int, int]:
    """Pick (h, w) matching source aspect with area ~= target_area, divisible."""
    import math
    ratio = src_w / src_h
    w = math.sqrt(target_area * ratio)
    h = w / ratio
    w = max(divisor, round(w / divisor) * divisor)
    h = max(divisor, round(h / divisor) * divisor)
    return h, w


def load_samples(n: int = 3, where: str = "tail",
                 metadata: Path = DEFAULT_METADATA) -> list[dict]:
    print(f"[load] reading metadata: {metadata}")
    with open(metadata) as f:
        data = json.load(f)
    if where == "head":
        print(f"[load] total records: {len(data)}; taking first {n}")
        return data[:n]
    print(f"[load] total records: {len(data)}; taking last {n}")
    return data[-n:]


def make_pipe(low_vram: bool = False):
    import torch
    from diffsynth.pipelines.qwen_image import QwenImagePipeline, ModelConfig
    extra = {}
    if low_vram:
        # CPU offload (not disk) so pipe.dit.load_state_dict(assign=True) still
        # finds real tensors to replace. Each transformer block is paged to GPU
        # for computation. ~40GB CPU RAM, ~2GB peak VRAM per block.
        extra = dict(
            offload_dtype=torch.bfloat16,
            offload_device="cpu",
            onload_dtype=torch.bfloat16,
            onload_device="cpu",
            preparing_dtype=torch.bfloat16,
            preparing_device="cuda",
            computation_dtype=torch.bfloat16,
            computation_device="cuda",
        )
        print("[pipe] low-VRAM mode: bf16-on-cpu, bf16-compute-on-gpu")
    print("[pipe] loading base Qwen-Image-Edit-2511 pipeline ...")
    pipe = QwenImagePipeline.from_pretrained(
        torch_dtype=torch.bfloat16,
        device="cuda",
        model_configs=[
            ModelConfig(model_id="Qwen/Qwen-Image-Edit-2511",
                        origin_file_pattern="transformer/diffusion_pytorch_model*.safetensors",
                        **extra),
            ModelConfig(model_id="Qwen/Qwen-Image",
                        origin_file_pattern="text_encoder/model*.safetensors",
                        **extra),
            ModelConfig(model_id="Qwen/Qwen-Image",
                        origin_file_pattern="vae/diffusion_pytorch_model.safetensors",
                        **extra),
        ],
        tokenizer_config=None,
        processor_config=ModelConfig(model_id="Qwen/Qwen-Image-Edit",
                                     origin_file_pattern="processor/"),
    )
    return pipe


def apply_full(pipe, ckpt_path: Path):
    import torch
    from diffsynth import load_state_dict
    print(f"[full] loading checkpoint: {ckpt_path}")
    sd = load_state_dict(str(ckpt_path))
    # assign=True is needed when the base model has meta parameters (low_vram
    # disk offload). strict=False tolerates wrapper prefixes. After assign,
    # parameters live wherever the checkpoint lived — move back to the
    # pipeline's compute device.
    missing, unexpected = pipe.dit.load_state_dict(sd, strict=False, assign=True)
    if missing:
        print(f"[full] WARN missing keys: {len(missing)} (sample: {missing[:3]})")
    if unexpected:
        print(f"[full] WARN unexpected keys: {len(unexpected)} (sample: {unexpected[:3]})")
    dev = getattr(pipe, "device", "cuda")
    dtype = getattr(pipe, "torch_dtype", torch.bfloat16)
    pipe.dit.to(device=dev, dtype=dtype)
    print(f"[full] moved dit to device={dev} dtype={dtype}")


def apply_lora(pipe, ckpt_path: Path):
    print(f"[lora] loading adapter: {ckpt_path}")
    pipe.load_lora(pipe.dit, str(ckpt_path))


def run_mode(mode: str, samples: list[dict], out_dir: Path,
             num_steps: int, seed: int, ckpt_override: Path | None,
             low_vram: bool = False,
             full_dir: Path = DEFAULT_FULL_DIR,
             lora_dir: Path = DEFAULT_LORA_DIR):
    from PIL import Image

    if mode == "full":
        ckpt = ckpt_override or latest_step_ckpt(full_dir)
    elif mode == "lora":
        ckpt = ckpt_override or latest_step_ckpt(lora_dir)
    else:
        raise ValueError(mode)

    pipe = make_pipe(low_vram=low_vram)
    if mode == "full":
        apply_full(pipe, ckpt)
    else:
        apply_lora(pipe, ckpt)

    mode_dir = out_dir / mode
    mode_dir.mkdir(parents=True, exist_ok=True)
    (out_dir / "ckpt_used.txt").open("a").write(f"{mode}: {ckpt}\n")

    for i, s in enumerate(samples):
        prompt = s["prompt"]
        in_path = DATA_BASE / s["edit_image"]
        gt_path = DATA_BASE / s["image"]
        print(f"\n[{mode}] sample {i}: in={in_path.name}  gt={gt_path.name}")
        print(f"[{mode}]   prompt: {prompt[:120].replace(chr(10), ' ')}...")

        src = Image.open(in_path).convert("RGB")
        h, w = target_hw(src.size[0], src.size[1])
        print(f"[{mode}]   src {src.size} -> gen {w}x{h}")

        out = pipe(
            prompt=prompt,
            edit_image=src,
            seed=seed,
            num_inference_steps=num_steps,
            height=h, width=w,
            zero_cond_t=True,
        )
        out_path = mode_dir / f"sample{i}_pred.png"
        out.save(out_path)
        print(f"[{mode}]   saved {out_path}")

        # Copy input/gt/prompt once (they're mode-independent).
        shared_in = out_dir / f"sample{i}_input.png"
        shared_gt = out_dir / f"sample{i}_gt.png"
        if not shared_in.exists():
            src.save(shared_in)
        if not shared_gt.exists():
            Image.open(gt_path).convert("RGB").save(shared_gt)
        (out_dir / f"sample{i}_prompt.txt").write_text(prompt)

    # Free GPU before next mode.
    import torch
    del pipe
    gc.collect()
    torch.cuda.empty_cache()


def make_grid(out_dir: Path, num: int, modes: list[str]):
    from PIL import Image
    print(f"\n[grid] building side-by-side comparison")
    rows = []
    for i in range(num):
        cols = []
        labels = []
        for name in ["input", "gt"] + [f"{m}" for m in modes]:
            if name in ("input", "gt"):
                p = out_dir / f"sample{i}_{name}.png"
            else:
                p = out_dir / name / f"sample{i}_pred.png"
            if p.exists():
                cols.append(Image.open(p).convert("RGB"))
                labels.append(name)
        if not cols:
            continue
        target_h = 768
        resized = []
        for img in cols:
            ratio = target_h / img.size[1]
            resized.append(img.resize(
                (int(img.size[0] * ratio), target_h), Image.LANCZOS))
        row_w = sum(im.size[0] for im in resized) + 8 * (len(resized) - 1)
        row = Image.new("RGB", (row_w, target_h + 30), (20, 20, 20))
        x = 0
        from PIL import ImageDraw, ImageFont
        draw = ImageDraw.Draw(row)
        try:
            font = ImageFont.truetype(
                "/usr/share/fonts/truetype/dejavu/DejaVuSans-Bold.ttf", 20)
        except Exception:
            font = ImageFont.load_default()
        for img, lab in zip(resized, labels):
            row.paste(img, (x, 30))
            draw.text((x + 4, 4), f"[{lab}]", fill=(255, 255, 255), font=font)
            x += img.size[0] + 8
        rows.append(row)
    if not rows:
        print("[grid] nothing to stitch")
        return
    max_w = max(r.size[0] for r in rows)
    total_h = sum(r.size[1] for r in rows) + 12 * (len(rows) - 1)
    grid = Image.new("RGB", (max_w, total_h), (10, 10, 10))
    y = 0
    for r in rows:
        grid.paste(r, (0, y))
        y += r.size[1] + 12
    grid_path = out_dir / "comparison.jpg"
    grid.save(grid_path, quality=92)
    print(f"[grid] saved {grid_path}")


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--mode", choices=["full", "lora", "both", "grid"],
                    default="both")
    ap.add_argument("--num-samples", type=int, default=3)
    ap.add_argument("--sample-from", choices=["head", "tail", "both"],
                    default="tail",
                    help="which slice of metadata: head / tail / both")
    ap.add_argument("--num-steps", type=int, default=40)
    ap.add_argument("--seed", type=int, default=123)
    ap.add_argument("--output-dir", type=Path,
                    default=PROJECT / "infer_outputs")
    ap.add_argument("--metadata", type=Path, default=DEFAULT_METADATA,
                    help="path to metadata_qwen_edit.json to sample prompts from")
    ap.add_argument("--full-dir", type=Path, default=DEFAULT_FULL_DIR,
                    help="directory of full-finetune step-*.safetensors")
    ap.add_argument("--lora-dir", type=Path, default=DEFAULT_LORA_DIR,
                    help="directory of LoRA step-*.safetensors")
    ap.add_argument("--ckpt-full", type=Path, default=None,
                    help="override full checkpoint path")
    ap.add_argument("--ckpt-lora", type=Path, default=None,
                    help="override LoRA adapter path")
    ap.add_argument("--low-vram", action="store_true",
                    help="offload weights to CPU in fp8; ~20GB VRAM, slower")
    args = ap.parse_args()

    # Pin to a single least-loaded physical GPU before torch initialises.
    if "CUDA_VISIBLE_DEVICES" not in os.environ or \
            "," in os.environ.get("CUDA_VISIBLE_DEVICES", ""):
        gpu = pick_free_gpu()
        os.environ["CUDA_VISIBLE_DEVICES"] = str(gpu)
        print(f"[init] CUDA_VISIBLE_DEVICES={gpu}")

    os.chdir(DIFFSYNTH)
    print(f"[init] cwd -> {os.getcwd()}")

    args.output_dir.mkdir(parents=True, exist_ok=True)

    modes_run = []
    if args.mode in ("full", "both"):
        modes_run.append("full")
    if args.mode in ("lora", "both"):
        modes_run.append("lora")

    if args.sample_from == "both":
        slices = [("head", args.output_dir / "head"),
                  ("tail", args.output_dir / "tail")]
    else:
        slices = [(args.sample_from, args.output_dir)]

    if modes_run:
        # Load base pipe once per (slice, mode). We let run_mode handle pipe
        # creation/teardown so checkpoint swapping stays isolated.
        for slice_name, slice_dir in slices:
            slice_dir.mkdir(parents=True, exist_ok=True)
            samples = load_samples(args.num_samples, slice_name, args.metadata)
            (slice_dir / "samples.json").write_text(
                json.dumps(samples, indent=2, ensure_ascii=False))
            for m in modes_run:
                ckpt = args.ckpt_full if m == "full" else args.ckpt_lora
                run_mode(m, samples, slice_dir, args.num_steps, args.seed,
                         ckpt, low_vram=args.low_vram,
                         full_dir=args.full_dir, lora_dir=args.lora_dir)

    for _, slice_dir in slices:
        make_grid(slice_dir, args.num_samples, ["full", "lora"])


if __name__ == "__main__":
    main()