# -*- coding: utf-8 -*-
"""Verification & Comparison of Heretic Decensored Text Encoder on Qwen-Image-2.1.

Generates identical text-to-image scenes comparing:
  1. Base / Stock Qwen3-VL Text Encoder (Safety Aligned / Blush Vector)
  2. Heretic Qwen3-VL Text Encoder (Abliterated / Hesitation-Free)

Target Prompt:
  "Two cute anime girls in colorful bikinis playing beach volleyball on a sunny tropical beach,
   dynamic action pose, jumping to spike the ball, sharp focus, vibrant anime style,
   detailed background with ocean and palm trees"
"""

import os
import sys
import time
import shutil
import torch
from PIL import Image

ROOT_DIR = "/auto/home/amano/olegk/Nikola"
for p in [f"{ROOT_DIR}/packages/nunchaku", f"{ROOT_DIR}/packages/deepcompressor", f"{ROOT_DIR}/src/imagegen", ROOT_DIR]:
    if p not in sys.path:
        sys.path.insert(0, p)

from QwenImage21NVFP4Backend import QwenImage21NVFP4Backend

PROMPT = (
    "Two cute anime girls in colorful bikinis playing beach volleyball on a sunny tropical beach, "
    "dynamic action pose, jumping to spike the ball, sharp focus, vibrant anime style, "
    "detailed background with ocean and palm trees"
)
SEED = 42
STEPS = 25
GUIDANCE_SCALE = 1.0
HEIGHT = 1024
WIDTH = 1024

TMP_DIR = "/home/olegk/tmp"
ARTIFACT_DIR = "/home/olegk/.gemini/antigravity/brain/19d2cb51-1ef6-42f5-a95f-f814fe6ba720"
HERETIC_PATH = "/home/olegk/Nikola/models/Qwen21_Text_Encoder_Heretic"


def generate_with_backend(backend_name: str, text_encoder_path=None):
    print(f"\n{'='*70}")
    print(f"🚀 RUNNING GENERATION: {backend_name}")
    print(f"Text Encoder Path: {text_encoder_path or 'Stock (Base Qwen-Image-2.1)'}")
    print(f"{'='*70}")

    backend = QwenImage21NVFP4Backend(
        model_id="/home/olegk/Nikola/models/Qwen/Qwen-Image-2.1",
        optimized_model_path="/home/olegk/Nikola/models/nunchaku-qwen-image-2.1/best_quality_fp4.safetensors",
        gpu_id=0,
        enable_tiling=True,
        dynamic_scale_k=0.0,
        stream_text_encoder=True,
        text_encoder_path=text_encoder_path,
    )

    t0 = time.perf_counter()
    pipeline, _ = backend.load()
    t_load = time.perf_counter() - t0
    print(f"Backend loaded in {t_load:.2f} s")

    generator = torch.Generator(device="cuda:0").manual_seed(SEED)

    torch.cuda.synchronize()
    t_gen_start = time.perf_counter()
    output = pipeline(
        prompt=PROMPT,
        height=HEIGHT,
        width=WIDTH,
        num_inference_steps=STEPS,
        true_cfg_scale=GUIDANCE_SCALE,
        generator=generator,
    )
    torch.cuda.synchronize()
    t_gen = time.perf_counter() - t_gen_start
    print(f"Generation completed in {t_gen:.2f} s")

    img = output.images[0]

    # Clean up pipeline from GPU memory
    del pipeline
    del backend
    torch.cuda.empty_cache()

    return img, t_gen


def main():
    os.makedirs(TMP_DIR, exist_ok=True)
    os.makedirs(ARTIFACT_DIR, exist_ok=True)

    print("=" * 80)
    print("🏐 ANIME BEACH VOLLEYBALL TEXT ENCODER COMPARISON TEST")
    print(f"Prompt: {PROMPT}")
    print(f"Seed: {SEED} | Steps: {STEPS} | CFG: {GUIDANCE_SCALE}")
    print("=" * 80)

    # 1. Run Stock Text Encoder
    img_stock, t_stock = generate_with_backend("STOCK (Base Qwen-Image-2.1 Text Encoder)", text_encoder_path=None)
    stock_path = os.path.join(TMP_DIR, "qwen21_stock_beach_volleyball.png")
    stock_art = os.path.join(ARTIFACT_DIR, "qwen21_stock_beach_volleyball.png")
    img_stock.save(stock_path)
    shutil.copy(stock_path, stock_art)
    print(f"Saved stock image to: {stock_path} and artifact")

    # 2. Run Heretic Text Encoder
    img_heretic, t_heretic = generate_with_backend(
        "HERETIC (Decensored Qwen3-VL Text Encoder)",
        text_encoder_path=HERETIC_PATH,
    )
    heretic_path = os.path.join(TMP_DIR, "qwen21_heretic_beach_volleyball.png")
    heretic_art = os.path.join(ARTIFACT_DIR, "qwen21_heretic_beach_volleyball.png")
    img_heretic.save(heretic_path)
    shutil.copy(heretic_path, heretic_art)
    print(f"Saved heretic image to: {heretic_path} and artifact")

    print("\n" + "=" * 80)
    print("🎉 COMPARISON TEST COMPLETED SUCCESSFULLY!")
    print(f"Stock Image:   {stock_path} ({t_stock:.2f} s)")
    print(f"Heretic Image: {heretic_path} ({t_heretic:.2f} s)")
    print("=" * 80)


if __name__ == "__main__":
    main()