#!/usr/bin/env python3
"""End-to-end portrait pipeline: generate → emotion variants → background removal.

Uses VnAssetsSession to keep models warm across all steps. Lightning LoRA +
flash attention for max edit throughput. Background removal uses isnet-anime.

Usage:
    python test_portrait_pipeline.py
    TORCH_ROCM_AOTRITON_ENABLE_EXPERIMENTAL=1 python test_portrait_pipeline.py  # with flash attention
"""
import json
import os
import time
from pathlib import Path

from vnassets import VnAssetsSession

OUTPUT_DIR = Path("output/portrait_pipeline")
OUTPUT_DIR.mkdir(parents=True, exist_ok=True)
NOBG_DIR = OUTPUT_DIR / "nobg"
NOBG_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
REMOVE_BG_MODEL = "isnet-anime"

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"),
]

flash_attn = os.environ.get("TORCH_ROCM_AOTRITON_ENABLE_EXPERIMENTAL") == "1"
mode = "flash attention" if flash_attn else "matmul fallback"
print(f"Pipeline test — {mode}")
print(f"Seed: {SEED}, SDXL steps: {GEN_STEPS}, Edit steps: {EDIT_STEPS}, CFG: {EDIT_CFG}")
print()

results = []
t_pipeline = time.perf_counter()

# ── 1. Session (load all models once) ───────────────────────────
print("=" * 60)
print("Loading session (SDXL + Qwen Edit + Lightning LoRA)")
print("=" * 60)

t0 = time.perf_counter()
vna = VnAssetsSession(
    sdxl_checkpoint=CHECKPOINT,
    edit_model=EDIT_MODEL,
    edit_lora=LORA,
)
t_session = round(time.perf_counter() - t0, 1)
print(f"Session ready: {t_session}s\n")

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

base_path = str(OUTPUT_DIR / "base.png")
t0 = time.perf_counter()
vna.generate(
    prompt=BASE_PROMPT,
    negative_prompt=NEG_PROMPT,
    steps=GEN_STEPS,
    cfg=4.5,
    seed=SEED,
    output_path=base_path,
)
t_gen = round(time.perf_counter() - t0, 1)

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

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

    edit_path = str(OUTPUT_DIR / f"base_{name}.png")
    t0 = time.perf_counter()
    vna.edit(
        input_path=base_path,
        prompt=prompt,
        steps=EDIT_STEPS,
        cfg=EDIT_CFG,
        seed=SEED,
        output_path=edit_path,
    )
    t_edit = round(time.perf_counter() - t0, 1)

    meta = json.loads((OUTPUT_DIR / f"base_{name}.json").read_text())
    results.append({
        "step": f"edit_{name}",
        "output": edit_path,
        "inference_s": meta["inference_time_s"],
        "lora_load_s": meta.get("lora_load_s"),
        "wall_s": t_edit,
    })
    variant_paths.append(edit_path)
    print(f"  Wall: {t_edit}s (infer: {meta['inference_time_s']}s)\n")

# ── 4. Remove backgrounds (batch) ───────────────────────────────
print("=" * 60)
print("STEP 7: Remove backgrounds (batch, all variants + base)")
print("=" * 60)

all_images = [base_path] + variant_paths
t0 = time.perf_counter()
vna.remove_backgrounds(all_images, str(NOBG_DIR), model=REMOVE_BG_MODEL)
t_rmbg = round(time.perf_counter() - t0, 1)

results.append({
    "step": "remove_backgrounds",
    "files": len(all_images),
    "inference_s": t_rmbg,
    "wall_s": t_rmbg,
})
print(f"  {len(all_images)} images, {t_rmbg}s total\n")

# ── Cleanup ─────────────────────────────────────────────────────
vna.close()

t_total = round(time.perf_counter() - t_pipeline, 1)

# ── Summary ─────────────────────────────────────────────────────
print("=" * 60)
print("SUMMARY")
print("=" * 60)
print(f"{'Step':<22} {'Infer':>8} {'Wall':>8}")
print("-" * 42)
total_infer = 0
total_wall = 0
for r in results:
    s = r["step"]
    infer = r["inference_s"]
    wall = r["wall_s"]
    extra = ""
    if "files" in r:
        extra = f" ({r['files']} files)"
    print(f"{s+extra:<22} {infer:>7.1f}s {wall:>7.1f}s")
    total_infer += infer
    total_wall += wall
print("-" * 42)
print(f"{'TOTAL':<22} {total_infer:>7.1f}s {total_wall:>7.1f}s")
print(f"\nPipeline wall time: {t_total}s (session load: {t_session}s)")

# Write summary
summary = {
    "config": {
        "checkpoint": CHECKPOINT,
        "edit_model": EDIT_MODEL,
        "lora": LORA,
        "flash_attention": flash_attn,
        "seed": SEED,
        "gen_steps": GEN_STEPS,
        "edit_steps": EDIT_STEPS,
        "edit_cfg": EDIT_CFG,
        "remove_bg_model": REMOVE_BG_MODEL,
        "base_prompt": BASE_PROMPT,
    },
    "results": results,
    "session_load_s": t_session,
    "pipeline_wall_s": t_total,
}
summary_path = OUTPUT_DIR / "summary.json"
summary_path.write_text(json.dumps(summary, indent=2))
print(f"\nSummary → {summary_path}")
print(f"Images → {OUTPUT_DIR}/")
print(f"Transparent → {NOBG_DIR}/")