# 🚀 Import all necessary libraries
import os
import argparse
from functools import partial
from pathlib import Path
import sys
import random
from omegaconf import OmegaConf
from PIL import Image
import torch
from torch import nn
from torch.nn import functional as F
from torchvision import transforms
from torchvision.transforms import functional as TF
from tqdm import trange
from cloob_training import model_pt, pretrained
import ldm.models.autoencoder
from diffusion import sampling, utils
import train_latent_diffusion as train
from huggingface_hub import hf_hub_url, cached_download
import gradio as gr  # 🎨 The magic canvas for AI-powered image generation!

# 🖼️ Download the necessary model files
# These files are loaded from HuggingFace's repository
checkpoint = cached_download(hf_hub_url("huggan/distill-ccld-wa", filename="model_student.ckpt"))
ae_model_path = cached_download(hf_hub_url("huggan/ccld_wa", filename="ae_model.ckpt"))
ae_config_path = cached_download(hf_hub_url("huggan/ccld_wa", filename="ae_model.yaml"))

# 📐 Utility Functions: Math and images, what could go wrong?
# These functions help parse prompts and resize/crop images to fit nicely

def parse_prompt(prompt, default_weight=3.):
    """
    🎯 Parses a prompt into text and weight.
    """
    if prompt.startswith('http://') or prompt.startswith('https://'):
        vals = prompt.rsplit(':', 2)
        vals = [vals[0] + ':' + vals[1], *vals[2:]]
    else:
        vals = prompt.rsplit(':', 1)
    vals = vals + ['', default_weight][len(vals):]
    return vals[0], float(vals[1])

def resize_and_center_crop(image, size):
    """
    ✂️ Resize and crop image to center it beautifully.
    """
    fac = max(size[0] / image.size[0], size[1] / image.size[1])
    image = image.resize((int(fac * image.size[0]), int(fac * image.size[1])), Image.LANCZOS)
    return TF.center_crop(image, size[::-1])


# 🧠 Model loading: the brain of our operation! 🔥
# Load all the models: autoencoder, diffusion, and CLOOB

device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
print('Using device:', device)
print('loading models... 🛠️')

# 🔧 Autoencoder Setup: Let’s decode the madness into images
ae_config = OmegaConf.load(ae_config_path)
ae_model = ldm.models.autoencoder.AutoencoderKL(**ae_config.model.params)
ae_model.eval().requires_grad_(False).to(device)
ae_model.load_state_dict(torch.load(ae_model_path))
n_ch, side_y, side_x = 4, 32, 32

# 🌀 Diffusion Model Setup: The artist behind the scenes
model = train.DiffusionModel(192, [1,1,2,2], autoencoder_scale=torch.tensor(4.3084))
model.load_state_dict(torch.load(checkpoint, map_location='cpu'))
model = model.to(device).eval().requires_grad_(False)

# 👁️ CLOOB Setup: Our vision model to understand art in human style
cloob_config = pretrained.get_config('cloob_laion_400m_vit_b_16_16_epochs')
cloob = model_pt.get_pt_model(cloob_config)
checkpoint = pretrained.download_checkpoint(cloob_config)
cloob.load_state_dict(model_pt.get_pt_params(cloob_config, checkpoint))
cloob.eval().requires_grad_(False).to(device)


# 🎨 The key function: Where the magic happens!
# This is where we generate images based on text and image prompts

def generate(n=1, prompts=['a red circle'], images=[], seed=42, steps=15, method='plms', eta=None):
    """
    🖼️ Generates a list of PIL images based on given text and image prompts.
    """
    zero_embed = torch.zeros([1, cloob.config['d_embed']], device=device)
    target_embeds, weights = [zero_embed], []

    # Parse text prompts
    for prompt in prompts:
        txt, weight = parse_prompt(prompt)
        target_embeds.append(cloob.text_encoder(cloob.tokenize(txt).to(device)).float())
        weights.append(weight)

    # Parse image prompts
    for prompt in images:
        path, weight = parse_prompt(prompt)
        img = Image.open(utils.fetch(path)).convert('RGB')
        clip_size = cloob.config['image_encoder']['image_size']
        img = resize_and_center_crop(img, (clip_size, clip_size))
        batch = TF.to_tensor(img)[None].to(device)
        embed = F.normalize(cloob.image_encoder(cloob.normalize(batch)).float(), dim=-1)
        target_embeds.append(embed)
        weights.append(weight)

    # Adjust weights and set seed
    weights = torch.tensor([1 - sum(weights), *weights], device=device)
    torch.manual_seed(seed)

    # 💡 Model function with classifier-free guidance
    def cfg_model_fn(x, t):
        n = x.shape[0]
        n_conds = len(target_embeds)
        x_in = x.repeat([n_conds, 1, 1, 1])
        t_in = t.repeat([n_conds])
        embed_in = torch.cat([*target_embeds]).repeat_interleave(n, 0)
        vs = model(x_in, t_in, embed_in).view([n_conds, n, *x.shape[1:]])
        v = vs.mul(weights[:, None, None, None, None]).sum(0)
        return v

    # 🎞️ Run the sampler to generate images
    def run(x, steps):
        if method == 'ddpm':
            return sampling.sample(cfg_model_fn, x, steps, 1., {})
        if method == 'ddim':
            return sampling.sample(cfg_model_fn, x, steps, eta, {})
        if method == 'plms':
            return sampling.plms_sample(cfg_model_fn, x, steps, {})
        assert False

    # 🏃‍♂️ Generate the output images
    batch_size = n
    x = torch.randn([n, n_ch, side_y, side_x], device=device)
    t = torch.linspace(1, 0, steps + 1, device=device)[:-1]
    pil_ims = []
    for i in trange(0, n, batch_size):
        cur_batch_size = min(n - i, batch_size)
        out_latents = run(x[i:i + cur_batch_size], steps)
        outs = ae_model.decode(out_latents * torch.tensor(2.55).to(device))
        for j, out in enumerate(outs):
            pil_ims.append(utils.to_pil_image(out))

    return pil_ims


# 🖌️ Interface: Gradio's brush to paint the UI
# Gradio is used here to create a user-friendly interface for art generation.

def gen_ims(prompt, im_prompt=None, seed=None, n_steps=10, method='plms'):
    """
    💡 Gradio function to wrap image generation.
    """
    if seed is None:
        seed = random.randint(0, 10000)
    prompts = [prompt]
    im_prompts = []
    if im_prompt is not None:
        im_prompts = [im_prompt]
    pil_ims = generate(n=1, prompts=prompts, images=im_prompts, seed=seed, steps=n_steps, method=method)
    return pil_ims[0]

# 🖼️ Gradio UI: The interface where users can input text or image prompts
iface = gr.Interface(
    fn=gen_ims,
    inputs=[
        gr.Textbox(label="Text prompt"),
        gr.Image(optional=True, label="Image prompt", type='filepath')
    ],
    outputs=gr.Image(type="pil", label="Generated Image"),
    examples=[
  ["Virgin and Child, in the style of Jacopo Bellini"],
  ["Art Nouveau, in the style of John Singer Sargent"], 
  ["Neoclassicism, in the style of Gustav Klimt"], 
  ["Abstract Art, in the style of M.C. Escher"], 
  ['Surrealism, in the style of Salvador Dali'], 
   ["Romanesque Art, in the style of Leonardo da Vinci"], 
   ["landscape"],
    ["portrait"],
    ["sculpture"],
    ["photo"],
    ["figurative"],
    ["illustration"],
    ["still life"],
    ["cityscape"],
    ["marina"],
    ["animal painting"],
    ["graffiti"],
    ["mythological painting"],
    ["battle painting"],
    ["self-portrait"],
    ["Impressionism, oil on canvas"],
    ["Katsushika Hokusai, The Dragon of Smoke Escaping from Mount Fuji"],
    ["Moon Light Sonata by Basuki Abdullah"],
    ["Two Trees by M.C. Escher"],
    ["Futurism, in the style of Wassily Kandinsky"], 
    ["Surrealism, in the style of Edgar Degas"], 
    ["Expressionism, in the style of Wassily Kandinsky"], 
    ["Futurism, in the style of Egon Schiele"], 
    ["Cubism, in the style of Gustav Klimt"], 
    ["Op Art, in the style of Marc Chagall"], 
    ["Romanticism, in the style of M.C. Escher"], 
    ["Futurism, in the style of M.C. Escher"], 
    ["Mannerism, in the style of Paul Klee"], 
    ["High Renaissance, in the style of Rembrandt"], 
    ["Magic Realism, in the style of Gustave Dore"], 
    ["Realism, in the style of Jean-Michel Basquiat"], 
    ["Art Nouveau, in the style of Paul Gauguin"], 
    ["Avant-garde, in the style of Pierre-Auguste Renoir"], 
    ["Baroque, in the style of Edward Hopper"], 
    ["Post-Impressionism, in the style of Wassily Kandinsky"], 
    ["Naturalism, in the style of Rene Magritte"], 
    ["Constructivism, in the style of Paul Cezanne"], 
    ["Abstract Expressionism, in the style of Henri Matisse"], 
    ["Pop Art, in the style of Vincent van Gogh"], 
    ["Futurism, in the style of Zdzislaw Beksinski"], 
    ["Aaron Wacker, oil on canvas"]
    ],
    title='Art Generator and Style Mixer from 🧠 Cloob and 🎨 WikiArt - Visual Art Encyclopedia',
    description="Trained on images from the [WikiArt](https://www.wikiart.org/) dataset, comprised of visual arts",
    article='Model used is: [model card](https://huggingface.co/huggan/distill-ccld-wa).'
)

# 🚀 Launch the Gradio interface
iface.launch(enable_queue=True)