#!/usr/bin/env python3
"""Deep batch test: generate base image, edit with multiple emotions, track timing."""
import subprocess
import time
import json
from pathlib import Path

OUTPUT_DIR = Path("output/batch_test")
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"
SEED = 42
STEPS = 20
CFG = 4.5

# Emotion variants to edit
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"),
]

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(STEPS),
    "--cfg", str(CFG),
    "--seed", str(SEED),
    "--output", str(base_path),
], check=True)

t_gen = time.perf_counter() - t_start

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

# ── 2. Edit with each emotion ───────────────────────────────────
total_edits_start = time.perf_counter()

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(STEPS),
        "--cfg", str(CFG),
        "--seed", str(SEED),
        "--output", str(edit_path),
    ], check=True)
    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),
        "prompt": prompt,
        "load_s": meta["load_time_s"],
        "inference_s": meta["inference_time_s"],
        "total_wall_s": round(t_edit, 1),
    })
    print(f"  Wall time: {t_edit:.1f}s (load: {meta['load_time_s']}s, inference: {meta['inference_time_s']}s)\n")

total_edits_wall = time.perf_counter() - total_edits_start

# ── 3. Summary ──────────────────────────────────────────────────
print("=" * 60)
print("SUMMARY")
print("=" * 60)
print(f"{'Step':<20} {'Load':>8} {'Infer':>8} {'Wall':>8}")
print("-" * 48)
total_load = 0
total_infer = 0
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")

# Write results JSON
summary = {
    "config": {
        "checkpoint": CHECKPOINT,
        "edit_model": EDIT_MODEL,
        "seed": SEED,
        "steps": STEPS,
        "cfg": 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 saved to {summary_path}")
print(f"Images in {OUTPUT_DIR}/")