import math
import os
import random

import gradio as gr
import numpy as np
import spaces
import torch

torch.backends.cuda.matmul.allow_tf32 = True
torch.backends.cudnn.allow_tf32 = True

from lakonlab.models.architectures import OklabColorEncoder
from lakonlab.models.diffusions.schedulers import FlowAdapterScheduler
from lakonlab.pipelines.pipeline_pixelflux2_klein import PixelFlux2KleinPipeline
from lakonlab.pipelines.prompt_rewriters.qwen3_vl import Qwen3VLPromptRewriter
from lakonlab.ui.gradio.create_text_to_img import create_interface_text_to_img

from huggingface_hub import login
login(token=os.getenv('HF_TOKEN'))


DEFAULT_PROMPT = (
    'Restored color photo from the 1900s. A middle-aged man with cybernetic metal hands is sitting on an old wooden '
    'chair and reading the newspaper. The newspaper has the prominent headline "AsymFLOW RELEASED" in large bold font. '
    'Close-up shot focusing on the newspaper.'
)
DEFAULT_NEG_PROMPT = 'Low quality, worst quality, blurry, deformed, bad anatomy, unclear text'

SYSTEM_PROMPT_TEXT_ONLY_PATH = 'lakonlab/pipelines/prompt_rewriters/system_prompts/default_text_only.txt'


def set_random_seed(seed: int, deterministic: bool = True) -> None:
    random.seed(seed)
    np.random.seed(seed)
    torch.manual_seed(seed)
    torch.cuda.manual_seed(seed)
    torch.cuda.manual_seed_all(seed)
    os.environ['PYTHONHASHSEED'] = str(seed)
    if deterministic:
        torch.backends.cudnn.deterministic = True
        torch.backends.cudnn.benchmark = False


pipe = PixelFlux2KleinPipeline.from_pretrained(
    'black-forest-labs/FLUX.2-klein-base-9B',
    vae=OklabColorEncoder(
        use_affine_norm=True,
        mean=(0.56, 0.0, 0.01),
        std=0.16),
    scheduler=FlowAdapterScheduler(
        shift=17.0,
        use_dynamic_shifting=True,
        base_seq_len=1024 ** 2,
        max_seq_len=2048 ** 2,
        base_logshift=math.log(17.0),
        max_logshift=math.log(34.0),
        dynamic_shifting_type='sqrt',
        base_scheduler='UniPCMultistep'),
    torch_dtype=torch.bfloat16)
pipe.load_lakonlab_adapter(
    'Lakonik/AsymFLUX.2-klein-9B',
    target_module_name='transformer')
pipe = pipe.to('cuda')

prompt_rewriter = Qwen3VLPromptRewriter(
    device_map='cuda',
    system_prompt_text_only=open(SYSTEM_PROMPT_TEXT_ONLY_PATH, 'r').read(),
    max_new_tokens_default=512,
)


@spaces.GPU
def run_rewrite_prompt_gpu(seed, prompt, progress):
    set_random_seed(seed)
    progress(0.05, desc='Rewriting prompt...')
    return prompt_rewriter.rewrite_text_batch([prompt])[0]


def run_rewrite_prompt(seed, prompt, rewrite_prompt, progress=gr.Progress(track_tqdm=True)):
    if rewrite_prompt:
        return run_rewrite_prompt_gpu(seed, prompt, progress), None
    return '', None


@spaces.GPU
def generate(
        seed, prompt, negative_prompt, rewrite_prompt, rewritten_prompt, width, height, steps, guidance_scale,
        progress=gr.Progress(track_tqdm=True)):
    return pipe(
        prompt=rewritten_prompt if rewrite_prompt else prompt,
        negative_prompt=negative_prompt,
        width=width,
        height=height,
        num_inference_steps=steps,
        guidance_scale=guidance_scale,
        generator=torch.Generator().manual_seed(seed),
    ).images[0]


with gr.Blocks(
        analytics_enabled=False,
        title='AsymFLUX.2-klein Demo',
        css_paths='lakonlab/ui/gradio/style.css') as demo:
    gr.Markdown(
        '# AsymFLUX.2-klein Demo\n\n'
        'Pixel-space text-to-image generation demo of the paper '
        '[Asymmetric Flow Models](https://arxiv.org/abs/2605.12964). '
        '**Base model:** [FLUX.2 klein Base 9B](https://huggingface.co/black-forest-labs/FLUX.2-klein-base-9B). '
        '**Code:** [https://github.com/Lakonik/LakonLab](https://github.com/Lakonik/LakonLab).\n'
        '<br> Use and distribution of this app are governed by the '
        '[FLUX Non-Commercial License](https://huggingface.co/black-forest-labs/FLUX.2-klein-base-9B/blob/main/LICENSE.md).'
    )

    create_interface_text_to_img(
        generate,
        prompt=DEFAULT_PROMPT,
        negative_prompt=DEFAULT_NEG_PROMPT,
        steps=38,
        min_steps=4,
        max_steps=50,
        guidance_scale=4.0,
        height=1280,
        width=960,
        create_negative_prompt=True,
        create_prompt_rewrite=True,
        args=[
            'last_seed',
            'prompt',
            'negative_prompt',
            'rewrite_prompt',
            'rewritten_prompt',
            'width',
            'height',
            'steps',
            'guidance_scale',
        ],
        rewrite_prompt_api=run_rewrite_prompt,
        rewrite_prompt_args=['last_seed', 'prompt', 'rewrite_prompt'])

demo.queue().launch()