#!/usr/bin/env python

import os
import random
import uuid
import base64
import gradio as gr
import numpy as np
from PIL import Image
import spaces
import torch
import glob
from datetime import datetime
import json

from diffusers import StableDiffusionXLPipeline, EulerAncestralDiscreteScheduler

DESCRIPTION = """# DALL•E 3 XL v2 High Fi"""

VOTE_FILE = "vote_counts.json"

def load_vote_counts():
    if os.path.exists(VOTE_FILE):
        with open(VOTE_FILE, "r") as f:
            return json.load(f)
    return {}

def save_vote_counts(vote_counts):
    with open(VOTE_FILE, "w") as f:
        json.dump(vote_counts, f)

vote_counts = load_vote_counts()

def create_download_link(filename):
    with open(filename, "rb") as file:
        encoded_string = base64.b64encode(file.read()).decode('utf-8')
        download_link = f'<a href="data:image/png;base64,{encoded_string}" download="{filename}">Download Image</a>'
        return download_link
        
def save_image(img, prompt):
    timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
    filename = f"{timestamp}_{prompt[:50]}.png"  # Limit filename length
    img.save(filename)
    return filename

def randomize_seed_fn(seed: int, randomize_seed: bool) -> int:
    if randomize_seed:
        seed = random.randint(0, MAX_SEED)
    return seed

def get_image_gallery():
    image_files = glob.glob("*.png")
    image_files.sort(key=lambda x: calculate_score(x), reverse=True)
    return [(file, f"{file}\n👍 {vote_counts.get(file, {}).get('likes', 0)} 👎 {vote_counts.get(file, {}).get('dislikes', 0)} ❤️ {vote_counts.get(file, {}).get('hearts', 0)}") for file in image_files]

def calculate_score(filename):
    counts = vote_counts.get(filename, {})
    return (counts.get('hearts', 0) * 5) + counts.get('likes', 0) - counts.get('dislikes', 0)

def delete_all_images():
    for file in glob.glob("*.png"):
        os.remove(file)
    vote_counts.clear()
    save_vote_counts(vote_counts)
    return get_image_gallery()

def vote(filename, vote_type):
    if filename not in vote_counts:
        vote_counts[filename] = {'likes': 0, 'dislikes': 0, 'hearts': 0}
    vote_counts[filename][vote_type] += 1
    save_vote_counts(vote_counts)
    return get_image_gallery()

def get_random_style():
    styles = [
        "Impressionist", "Cubist", "Surrealist", "Abstract Expressionist",
        "Pop Art", "Minimalist", "Baroque", "Art Nouveau", "Pointillist", "Fauvism"
    ]
    return random.choice(styles)

MAX_SEED = np.iinfo(np.int32).max

if not torch.cuda.is_available():
    DESCRIPTION += "\n<p>Running on CPU 🥶 This demo may not work on CPU.</p>"

USE_TORCH_COMPILE = 0
ENABLE_CPU_OFFLOAD = 0

if torch.cuda.is_available():
    pipe = StableDiffusionXLPipeline.from_pretrained(
        "fluently/Fluently-XL-v4",
        torch_dtype=torch.float16,
        use_safetensors=True,
    )
    pipe.scheduler = EulerAncestralDiscreteScheduler.from_config(pipe.scheduler.config)
    
    pipe.load_lora_weights("ehristoforu/dalle-3-xl-v2", weight_name="dalle-3-xl-lora-v2.safetensors", adapter_name="dalle")
    pipe.set_adapters("dalle")

    pipe.to("cuda")

@spaces.GPU(enable_queue=True)
def generate(
    prompt: str,
    negative_prompt: str = "",
    use_negative_prompt: bool = False,
    seed: int = 0,
    width: int = 1024,
    height: int = 1024,
    guidance_scale: float = 3,
    randomize_seed: bool = False,
    progress=gr.Progress(track_tqdm=True),
):
    seed = int(randomize_seed_fn(seed, randomize_seed))

    if not use_negative_prompt:
        negative_prompt = ""

    images = pipe(
        prompt=prompt,
        negative_prompt=negative_prompt,
        width=width,
        height=height,
        guidance_scale=guidance_scale,
        num_inference_steps=20,
        num_images_per_prompt=1,
        cross_attention_kwargs={"scale": 0.65},
        output_type="pil",
    ).images
    image_paths = [save_image(img, prompt) for img in images]
    download_links = [create_download_link(path) for path in image_paths]

    return image_paths, seed, download_links, get_image_gallery()

examples = [
    f"{get_random_style()} painting of a majestic lighthouse on a rocky coast. Use bold brushstrokes and a vibrant color palette to capture the interplay of light and shadow as the lighthouse beam cuts through a stormy night sky.",
    f"{get_random_style()} still life featuring a pair of vintage eyeglasses. Focus on the intricate details of the frames and lenses, using a warm color scheme to evoke a sense of nostalgia and wisdom.",
    f"{get_random_style()} depiction of a rustic wooden stool in a sunlit artist's studio. Emphasize the texture of the wood and the interplay of light and shadow, using a mix of earthy tones and highlights.",
    f"{get_random_style()} scene viewed through an ornate window frame. Contrast the intricate details of the window with a dreamy, soft-focus landscape beyond, using a palette that transitions from cool interior tones to warm exterior hues.",
    f"{get_random_style()} close-up study of interlaced fingers. Use a monochromatic color scheme to emphasize the form and texture of the hands, with dramatic lighting to create depth and emotion.",
    f"{get_random_style()} composition featuring a set of dice in motion. Capture the energy and randomness of the throw, using a dynamic color palette and blurred lines to convey movement.",
    f"{get_random_style()} interpretation of heaven. Create an ethereal atmosphere with soft, billowing clouds and radiant light, using a palette of celestial blues, golds, and whites.",
    f"{get_random_style()} portrayal of an ancient, mystical gate. Combine architectural details with elements of fantasy, using a rich, jewel-toned palette to create an air of mystery and magic.",
    f"{get_random_style()} portrait of a curious cat. Focus on capturing the feline's expressive eyes and sleek form, using a mix of bold and subtle colors to bring out the cat's personality.",
    f"{get_random_style()} abstract representation of toes in sand. Use textured brushstrokes to convey the feeling of warm sand, with a palette inspired by a sun-drenched beach."
]

css = '''
.gradio-container{max-width: 1024px !important}
h1{text-align:center}
footer {
visibility: hidden
}
'''

with gr.Blocks(css=css, theme="pseudolab/huggingface-korea-theme") as demo:
    gr.Markdown(DESCRIPTION)

    with gr.Group():
        with gr.Row():
            prompt = gr.Text(
                label="Prompt",
                show_label=False,
                max_lines=1,
                placeholder="Enter your prompt",
                container=False,
            )
            run_button = gr.Button("Run", scale=0)
        result = gr.Gallery(label="Result", columns=1, preview=True, show_label=False)
    with gr.Accordion("Advanced options", open=False):
        use_negative_prompt = gr.Checkbox(label="Use negative prompt", value=True)
        negative_prompt = gr.Text(
            label="Negative prompt",
            lines=4,
            max_lines=6,
            value="""(deformed, distorted, disfigured:1.3), poorly drawn, bad anatomy, wrong anatomy, extra limb, missing limb, floating limbs, (mutated hands and fingers:1.4), disconnected limbs, mutation, mutated, ugly, disgusting, blurry, amputation, (NSFW:1.25)""",
            placeholder="Enter a negative prompt",
            visible=True,
        )
        seed = gr.Slider(
            label="Seed",
            minimum=0,
            maximum=MAX_SEED,
            step=1,
            value=0,
            visible=True
        )
        randomize_seed = gr.Checkbox(label="Randomize seed", value=True)
        with gr.Row(visible=True):
            width = gr.Slider(
                label="Width",
                minimum=512,
                maximum=2048,
                step=8,
                value=1920,
            )
            height = gr.Slider(
                label="Height",
                minimum=512,
                maximum=2048,
                step=8,
                value=1080,
            )
        with gr.Row():
            guidance_scale = gr.Slider(
                label="Guidance Scale",
                minimum=0.1,
                maximum=20.0,
                step=0.1,
                value=20.0,
            )
    
    image_gallery = gr.Gallery(label="Generated Images", show_label=True, columns=4, height="auto")
    
    with gr.Row():
        delete_all_button = gr.Button("🗑️ Delete All Images")
        like_button = gr.Button("👍 Like")
        dislike_button = gr.Button("👎 Dislike")
        heart_button = gr.Button("❤️ Heart")
    
    selected_image = gr.State(None)

    gr.Examples(
        examples=examples,
        inputs=prompt,
        outputs=[result, seed],
        fn=generate,
        cache_examples=False,
    )

    use_negative_prompt.change(
        fn=lambda x: gr.update(visible=x),
        inputs=use_negative_prompt,
        outputs=negative_prompt,
        api_name=False,
    )

    delete_all_button.click(
        fn=delete_all_images,
        inputs=[],
        outputs=[image_gallery],
    )

    image_gallery.select(
        fn=lambda evt: evt,
        inputs=[gr.State("value")],
        outputs=[selected_image],
    )

    like_button.click(
        fn=lambda x: vote(x, 'likes') if x else None,
        inputs=[selected_image],
        outputs=[image_gallery],
    )

    dislike_button.click(
        fn=lambda x: vote(x, 'dislikes') if x else None,
        inputs=[selected_image],
        outputs=[image_gallery],
    )

    heart_button.click(
        fn=lambda x: vote(x, 'hearts') if x else None,
        inputs=[selected_image],
        outputs=[image_gallery],
    )

    def update_gallery():
        return gr.update(value=get_image_gallery())

    gr.on(
        triggers=[
            prompt.submit,
            negative_prompt.submit,
            run_button.click,
        ],
        fn=generate,
        inputs=[
            prompt,
            negative_prompt,
            use_negative_prompt,
            seed,
            width,
            height,
            guidance_scale,
            randomize_seed,
        ],
        outputs=[result, seed, gr.HTML(visible=False), image_gallery],
        api_name="run",
    )

    demo.load(fn=update_gallery, outputs=image_gallery)

if __name__ == "__main__":
    demo.queue(max_size=20).launch(show_api=False, debug=False)