import os
import uuid
import gradio as gr
import spaces
from clip_slider_pipeline import CLIPSliderFlux
from diffusers import FluxPipeline, AutoencoderTiny
import torch
import numpy as np
import cv2
from PIL import Image
from diffusers.utils import load_image
from diffusers.utils import export_to_video
import random
from transformers import pipeline

# Translation model load
translator = pipeline("translation", model="Helsinki-NLP/opus-mt-ko-en")

# English menu labels
english_labels = {
    "Prompt": "Prompt",
    "1st direction to steer": "1st Direction",
    "2nd direction to steer": "2nd Direction",
    "Strength": "Strength",
    "Generate directions": "Generate Directions",
    "Generated Images": "Generated Images",
    "From 1st to 2nd direction": "From 1st to 2nd Direction",
    "Strip": "Image Strip",
    "Looping video": "Looping Video",
    "Advanced options": "Advanced Options",
    "Num of intermediate images": "Number of Intermediate Images",
    "Num iterations for clip directions": "Number of CLIP Direction Iterations",
    "Num inference steps": "Number of Inference Steps",
    "Guidance scale": "Guidance Scale",
    "Randomize seed": "Randomize Seed",
    "Seed": "Seed"
}

# Load pipelines
base_model = "black-forest-labs/FLUX.1-schnell"
taef1 = AutoencoderTiny.from_pretrained("madebyollin/taef1", torch_dtype=torch.bfloat16).to("cuda")
pipe = FluxPipeline.from_pretrained(
    base_model,
    vae=taef1,
    torch_dtype=torch.bfloat16
)
pipe.transformer.to(memory_format=torch.channels_last)
clip_slider = CLIPSliderFlux(pipe, device=torch.device("cuda"))
MAX_SEED = 2**32 - 1

def save_images_with_unique_filenames(image_list, save_directory):
    if not os.path.exists(save_directory):
        os.makedirs(save_directory)
    paths = []
    for image in image_list:
        unique_filename = f"{uuid.uuid4()}.png"
        file_path = os.path.join(save_directory, unique_filename)
        image.save(file_path)
        paths.append(file_path)
    return paths

def convert_to_centered_scale(num):
    if num % 2 == 0:  # even
        start = -(num // 2 - 1)
        end = num // 2
    else:  # odd
        start = -(num // 2)
        end = num // 2
    return tuple(range(start, end + 1))

def translate_if_korean(text):
    if any('\u3131' <= char <= '\u3163' or '\uac00' <= char <= '\ud7a3' for char in text):
        return translator(text)[0]['translation_text']
    return text

@spaces.GPU(duration=85)
def generate(prompt,
             concept_1,
             concept_2,
             scale,
             randomize_seed=True,
             seed=42,
             recalc_directions=True,
             iterations=200,
             steps=3,
             interm_steps=33,
             guidance_scale=3.5,
             x_concept_1="", x_concept_2="",
             avg_diff_x=None,
             total_images=[],
             gradio_progress=gr.Progress()):
    # Translate prompt and concepts if Korean
    prompt = translate_if_korean(prompt)
    concept_1 = translate_if_korean(concept_1)
    concept_2 = translate_if_korean(concept_2)
    print(f"Prompt: {prompt}, ← {concept_2}, {concept_1} ➡️ . scale {scale}, interm steps {interm_steps}")
    slider_x = [concept_2, concept_1]
    if randomize_seed:
        seed = random.randint(0, MAX_SEED)
    if not sorted(slider_x) == sorted([x_concept_1, x_concept_2]) or recalc_directions:
        gradio_progress(0, desc="Calculating directions...")
        avg_diff = clip_slider.find_latent_direction(slider_x[0], slider_x[1], num_iterations=iterations)
        x_concept_1, x_concept_2 = slider_x[0], slider_x[1]
    else:
        avg_diff = avg_diff_x
    images = []
    high_scale = scale
    low_scale = -1 * scale
    for i in gradio_progress.tqdm(range(interm_steps), desc="Generating images"):
        cur_scale = low_scale + (high_scale - low_scale) * i / (interm_steps - 1)
        image = clip_slider.generate(
            prompt,
            width=768,
            height=768,
            guidance_scale=guidance_scale,
            scale=cur_scale,
            seed=seed,
            num_inference_steps=steps,
            avg_diff=avg_diff
        )
        images.append(image)
    canvas = Image.new('RGB', (256 * interm_steps, 256))
    for i, im in enumerate(images):
        canvas.paste(im.resize((256, 256)), (256 * i, 0))
    comma_concepts_x = f"{slider_x[1]}, {slider_x[0]}"
    scale_total = convert_to_centered_scale(interm_steps)
    scale_min = scale_total[0]
    scale_max = scale_total[-1]
    scale_middle = scale_total.index(0)
    post_generation_slider_update = gr.update(label=comma_concepts_x, value=0, minimum=scale_min, maximum=scale_max, interactive=True)
    avg_diff_x = avg_diff.cpu()
    video_path = f"{uuid.uuid4()}.mp4"
    print(video_path)
    return x_concept_1, x_concept_2, avg_diff_x, export_to_video(images, video_path, fps=5), canvas, images, images[scale_middle], post_generation_slider_update, seed

def update_pre_generated_images(slider_value, total_images):
    number_images = len(total_images) if total_images else 0
    if number_images > 0:
        scale_tuple = convert_to_centered_scale(number_images)
        return total_images[scale_tuple.index(slider_value)][0]
    else:
        return None

def reset_recalc_directions():
    return True

# Five "Time Stream" themed examples (one Korean example included)
examples = [
    ["신선한 토마토가 부패한 토마토로 변해가는 과정", "Fresh", "Rotten", 2.0],
    ["A blooming flower gradually withers into decay", "Bloom", "Wither", 1.5],
    ["A vibrant cityscape transforms into a derelict ruin over time", "Modern", "Ruined", 2.5],
    ["A lively forest slowly changes into an autumnal landscape", "Spring", "Autumn", 2.0],
    ["A calm ocean evolves into a stormy seascape as time passes", "Calm", "Stormy", 3.0]
]

# CSS for a bright and modern UI with a background image
css = """
/* Bright and modern UI with background image */
body {
    background: #ffffff url('https://images.unsplash.com/photo-1506748686214-e9df14d4d9d0?ixlib=rb-1.2.1&auto=format&fit=crop&w=1600&q=80') no-repeat center center fixed;
    background-size: cover;
    font-family: "Helvetica Neue", Helvetica, Arial, sans-serif;
    color: #333;
}
footer {
    visibility: hidden;
}
.container {
    max-width: 1200px;
    margin: 20px auto;
    padding: 0 10px;
}
.main-panel {
    background-color: rgba(255, 255, 255, 0.9);
    border-radius: 12px;
    padding: 20px;
    margin-bottom: 20px;
    box-shadow: 0 4px 8px rgba(0, 0, 0, 0.1);
}
.controls-panel {
    background-color: rgba(255, 255, 255, 0.85);
    border-radius: 8px;
    padding: 16px;
    box-shadow: inset 0 2px 4px rgba(0, 0, 0, 0.05);
}
.image-display {
    min-height: 400px;
    display: flex;
    flex-direction: column;
    justify-content: center;
}
.slider-container {
    padding: 10px 0;
}
.advanced-panel {
    margin-top: 20px;
    border-top: 1px solid #eaeaea;
    padding-top: 20px;
}
"""

with gr.Blocks(css=css, title="Time Stream") as demo:
    gr.Markdown("# Time Stream")
    
    x_concept_1 = gr.State("")
    x_concept_2 = gr.State("")
    total_images = gr.State([])
    avg_diff_x = gr.State()
    recalc_directions = gr.State(False)

    with gr.Row(elem_classes="container"):
        # Left Column - Controls
        with gr.Column(scale=4):
            with gr.Group(elem_classes="main-panel"):
                gr.Markdown("### Image Generation Controls")
                with gr.Group(elem_classes="controls-panel"):
                    prompt = gr.Textbox(
                        label=english_labels["Prompt"],
                        info="Enter the description",
                        placeholder="A dog in the park",
                        lines=2
                    )
                    with gr.Row():
                        with gr.Column(scale=1):
                            concept_1 = gr.Textbox(
                                label=english_labels["1st direction to steer"],
                                info="Initial state",
                                placeholder="Fresh"
                            )
                        with gr.Column(scale=1):
                            concept_2 = gr.Textbox(
                                label=english_labels["2nd direction to steer"],
                                info="Final state",
                                placeholder="Rotten"
                            )
                    with gr.Row(elem_classes="slider-container"):
                        x = gr.Slider(
                            minimum=0,
                            value=1.75,
                            step=0.1,
                            maximum=4.0,
                            label=english_labels["Strength"],
                            info="Maximum strength for each direction (above 2.5 may be unstable)"
                        )
                    submit = gr.Button(english_labels["Generate directions"], size="lg", variant="primary")
            with gr.Accordion(label=english_labels["Advanced options"], open=False, elem_classes="advanced-panel"):
                with gr.Row():
                    with gr.Column(scale=1):
                        interm_steps = gr.Slider(
                            label=english_labels["Num of intermediate images"],
                            minimum=3,
                            value=7,
                            maximum=65,
                            step=2
                        )
                    with gr.Column(scale=1):
                        guidance_scale = gr.Slider(
                            label=english_labels["Guidance scale"],
                            minimum=0.1,
                            maximum=10.0,
                            step=0.1,
                            value=3.5
                        )
                with gr.Row():
                    with gr.Column(scale=1):
                        iterations = gr.Slider(
                            label=english_labels["Num iterations for clip directions"],
                            minimum=0,
                            value=200,
                            maximum=400,
                            step=1
                        )
                    with gr.Column(scale=1):
                        steps = gr.Slider(
                            label=english_labels["Num inference steps"],
                            minimum=1,
                            value=3,
                            maximum=4,
                            step=1
                        )
                with gr.Row():
                    with gr.Column(scale=1):
                        randomize_seed = gr.Checkbox(
                            True,
                            label=english_labels["Randomize seed"]
                        )
                    with gr.Column(scale=1):
                        seed = gr.Slider(
                            minimum=0,
                            maximum=MAX_SEED,
                            step=1,
                            label=english_labels["Seed"],
                            interactive=True,
                            randomize=True
                        )
        # Right Column - Output
        with gr.Column(scale=8):
            with gr.Group(elem_classes="main-panel"):
                gr.Markdown("### Generated Results")
                # Swapped order: Image strip on top, video below (video is larger)
                image_strip = gr.Image(label="Image Strip", type="filepath", elem_id="strip", height=200)
                output_video = gr.Video(label=english_labels["Looping video"], elem_id="video", loop=True, autoplay=True, height=600)
                with gr.Row():
                    post_generation_image = gr.Image(
                        label=english_labels["Generated Images"],
                        type="filepath",
                        elem_id="interactive",
                        elem_classes="image-display"
                    )
                    post_generation_slider = gr.Slider(
                        minimum=-10,
                        maximum=10,
                        value=0,
                        step=1,
                        label=english_labels["From 1st to 2nd direction"]
                    )
    
    # Examples: 예제 클릭 시 입력란에 값이 바로 삽입됨
    gr.Examples(
        examples=examples,
        inputs=[prompt, concept_1, concept_2, x]
    )
    
    # Event Handlers
    submit.click(
        fn=generate,
        inputs=[
            prompt, concept_1, concept_2, x, randomize_seed, seed,
            recalc_directions, iterations, steps, interm_steps,
            guidance_scale, x_concept_1, x_concept_2, avg_diff_x, total_images
        ],
        outputs=[
            x_concept_1, x_concept_2, avg_diff_x,
            output_video,      # video output
            image_strip,       # canvas (image strip)
            total_images,
            post_generation_image,
            post_generation_slider,
            seed
        ]
    )
    
    iterations.change(fn=reset_recalc_directions, outputs=[recalc_directions])
    seed.change(fn=reset_recalc_directions, outputs=[recalc_directions])
    post_generation_slider.change(
        fn=update_pre_generated_images,
        inputs=[post_generation_slider, total_images],
        outputs=[post_generation_image],
        queue=False,
        show_progress="hidden",
        concurrency_limit=None
    )

if __name__ == "__main__":
    demo.launch()