#!/usr/bin/env python3
"""Batch test with Lightning LoRA + flash attention for max speed."""
import subprocess
import time
import json
import os
from pathlib import Path

OUTPUT_DIR = Path("output/batch_fast")
OUTPUT_DIR.mkdir(parents=True, exist_ok=True)

BASE_PROMPT = "1girl, solo, red hair, glasses, blue eyes, white crop top, standing, portrait"
NEG_PROMPT = "deformed, ugly, bad quality, lowres"
CHECKPOINT = "models/novaAnimeXL_ilV190.safetensors"
EDIT_MODEL = "models/qwen_image_edit_2509_fp8_e4m3fn.safetensors"
LORA = "models/Qwen-Image-Edit-2509-Lightning-4steps-V1.0-bf16.safetensors"
SEED = 42
GEN_STEPS = 20
EDIT_STEPS = 4
EDIT_CFG = 1.0

EMOTIONS = [
    ("smile", "make her smile happily with a warm genuine smile"),
    ("angry", "make her look angry and furious, furrowed brow"),
    ("sad", "make her look sad and crying, tears in her eyes"),
    ("surprised", "make her look surprised, wide eyes, mouth slightly open"),
    ("blushing", "make her blush intensely, embarrassed expression, pink cheeks"),
]

env = os.environ.copy()
env["TORCH_ROCM_AOTRITON_ENABLE_EXPERIMENTAL"] = "1"

results = []

# ── 1. Generate base image ─────────────────────────────────────
print("=" * 60)
print("STEP 1: Generate base image")
print("=" * 60)

base_path = OUTPUT_DIR / "base.png"
t_start = time.perf_counter()
subprocess.run([
    "vnasset", "generate",
    "--checkpoint", CHECKPOINT,
    "--prompt", BASE_PROMPT,
    "--negative-prompt", NEG_PROMPT,
    "--steps", str(GEN_STEPS),
    "--cfg", "4.5",
    "--seed", str(SEED),
    "--output", str(base_path),
], check=True, env=env)
t_gen = time.perf_counter() - t_start

meta = json.loads((OUTPUT_DIR / "base.json").read_text())
results.append({
    "step": "generate_base",
    "output": str(base_path),
    "load_s": meta["load_time_s"],
    "inference_s": meta["inference_time_s"],
    "total_wall_s": round(t_gen, 1),
})
print(f"  Wall: {t_gen:.1f}s (load: {meta['load_time_s']}s, infer: {meta['inference_time_s']}s)\n")

# ── 2. Edit with each emotion ───────────────────────────────────
for i, (name, prompt) in enumerate(EMOTIONS):
    print("=" * 60)
    print(f"STEP {i+2}: Edit → {name}")
    print("=" * 60)

    edit_path = OUTPUT_DIR / f"base_{name}.png"
    t_start = time.perf_counter()
    subprocess.run([
        "vnasset", "edit",
        "--model", EDIT_MODEL,
        "--input", str(base_path),
        "--prompt", prompt,
        "--steps", str(EDIT_STEPS),
        "--cfg", str(EDIT_CFG),
        "--seed", str(SEED),
        "--lora", LORA,
        "--output", str(edit_path),
    ], check=True, env=env)
    t_edit = time.perf_counter() - t_start

    meta = json.loads((OUTPUT_DIR / f"base_{name}.json").read_text())
    results.append({
        "step": f"edit_{name}",
        "output": str(edit_path),
        "load_s": meta["load_time_s"],
        "inference_s": meta["inference_time_s"],
        "lora_load_s": meta["lora_load_s"],
        "total_wall_s": round(t_edit, 1),
    })
    print(f"  Wall: {t_edit:.1f}s (load: {meta['load_time_s']}s, infer: {meta['inference_time_s']}s)\n")

# ── 3. Summary ──────────────────────────────────────────────────
print("=" * 60)
print("SUMMARY (Lightning LoRA + Flash Attention)")
print("=" * 60)
print(f"{'Step':<20} {'Load':>8} {'Infer':>8} {'Wall':>8}")
print("-" * 48)
total_load = total_infer = total_wall = 0
for r in results:
    print(f"{r['step']:<20} {r['load_s']:>7.1f}s {r['inference_s']:>7.1f}s {r['total_wall_s']:>7.1f}s")
    total_load += r["load_s"]
    total_infer += r["inference_s"]
    total_wall += r["total_wall_s"]
print("-" * 48)
print(f"{'TOTAL':<20} {total_load:>7.1f}s {total_infer:>7.1f}s {total_wall:>7.1f}s")

summary = {
    "config": {
        "checkpoint": CHECKPOINT,
        "edit_model": EDIT_MODEL,
        "lora": LORA,
        "flash_attention": True,
        "seed": SEED,
        "gen_steps": GEN_STEPS,
        "edit_steps": EDIT_STEPS,
        "edit_cfg": EDIT_CFG,
        "base_prompt": BASE_PROMPT,
    },
    "results": results,
    "totals": {
        "load_s": round(total_load, 1),
        "inference_s": round(total_infer, 1),
        "total_wall_s": round(total_wall, 1),
    },
}
summary_path = OUTPUT_DIR / "summary.json"
summary_path.write_text(json.dumps(summary, indent=2))
print(f"\nSummary → {summary_path}")