#!/usr/bin/env python
# coding: utf-8

# Uncomment to run on cpu
# import os
# os.environ["JAX_PLATFORM_NAME"] = "cpu"

import random

import gradio as gr
import jax
import numpy as np
from flax.jax_utils import replicate
from flax.training.common_utils import shard
from PIL import Image, ImageDraw, ImageFont

# ## CLIP Scoring
from transformers import BartTokenizer, CLIPProcessor, FlaxCLIPModel
from vqgan_jax.modeling_flax_vqgan import VQModel

from dalle_mini.model import CustomFlaxBartForConditionalGeneration

DALLE_REPO = "flax-community/dalle-mini"
DALLE_COMMIT_ID = "4d34126d0df8bc4a692ae933e3b902a1fa8b6114"

VQGAN_REPO = "flax-community/vqgan_f16_16384"
VQGAN_COMMIT_ID = "90cc46addd2dd8f5be21586a9a23e1b95aa506a9"

tokenizer = BartTokenizer.from_pretrained(DALLE_REPO, revision=DALLE_COMMIT_ID)
model = CustomFlaxBartForConditionalGeneration.from_pretrained(
    DALLE_REPO, revision=DALLE_COMMIT_ID
)
vqgan = VQModel.from_pretrained(VQGAN_REPO, revision=VQGAN_COMMIT_ID)


def captioned_strip(images, caption=None, rows=1):
    increased_h = 0 if caption is None else 48
    w, h = images[0].size[0], images[0].size[1]
    img = Image.new("RGB", (len(images) * w // rows, h * rows + increased_h))
    for i, img_ in enumerate(images):
        img.paste(img_, (i // rows * w, increased_h + (i % rows) * h))

    if caption is not None:
        draw = ImageDraw.Draw(img)
        font = ImageFont.truetype(
            "/usr/share/fonts/truetype/liberation2/LiberationMono-Bold.ttf", 40
        )
        draw.text((20, 3), caption, (255, 255, 255), font=font)
    return img


def custom_to_pil(x):
    x = np.clip(x, 0.0, 1.0)
    x = (255 * x).astype(np.uint8)
    x = Image.fromarray(x)
    if not x.mode == "RGB":
        x = x.convert("RGB")
    return x


def generate(input, rng, params):
    return model.generate(
        **input,
        max_length=257,
        num_beams=1,
        do_sample=True,
        prng_key=rng,
        eos_token_id=50000,
        pad_token_id=50000,
        params=params,
    )


def get_images(indices, params):
    return vqgan.decode_code(indices, params=params)


p_generate = jax.pmap(generate, "batch")
p_get_images = jax.pmap(get_images, "batch")

bart_params = replicate(model.params)
vqgan_params = replicate(vqgan.params)

clip = FlaxCLIPModel.from_pretrained("openai/clip-vit-base-patch32")
print("Initialize FlaxCLIPModel")
processor = CLIPProcessor.from_pretrained("openai/clip-vit-base-patch32")
print("Initialize CLIPProcessor")


def hallucinate(prompt, num_images=64):
    prompt = [prompt] * jax.device_count()
    inputs = tokenizer(
        prompt,
        return_tensors="jax",
        padding="max_length",
        truncation=True,
        max_length=128,
    ).data
    inputs = shard(inputs)

    all_images = []
    for i in range(num_images // jax.device_count()):
        key = random.randint(0, 1e7)
        rng = jax.random.PRNGKey(key)
        rngs = jax.random.split(rng, jax.local_device_count())
        indices = p_generate(inputs, rngs, bart_params).sequences
        indices = indices[:, :, 1:]

        images = p_get_images(indices, vqgan_params)
        images = np.squeeze(np.asarray(images), 1)
        for image in images:
            all_images.append(custom_to_pil(image))
    return all_images


def clip_top_k(prompt, images, k=8):
    inputs = processor(text=prompt, images=images, return_tensors="np", padding=True)
    outputs = clip(**inputs)
    logits = outputs.logits_per_text
    scores = np.array(logits[0]).argsort()[-k:][::-1]
    return [images[score] for score in scores]


def compose_predictions(images, caption=None):
    increased_h = 0 if caption is None else 48
    w, h = images[0].size[0], images[0].size[1]
    img = Image.new("RGB", (len(images) * w, h + increased_h))
    for i, img_ in enumerate(images):
        img.paste(img_, (i * w, increased_h))

    if caption is not None:
        draw = ImageDraw.Draw(img)
        font = ImageFont.truetype(
            "/usr/share/fonts/truetype/liberation2/LiberationMono-Bold.ttf", 40
        )
        draw.text((20, 3), caption, (255, 255, 255), font=font)
    return img


def top_k_predictions(prompt, num_candidates=32, k=8):
    images = hallucinate(prompt, num_images=num_candidates)
    images = clip_top_k(prompt, images, k=k)
    return images


def run_inference(prompt, num_images=32, num_preds=8):
    images = top_k_predictions(prompt, num_candidates=num_images, k=num_preds)
    predictions = captioned_strip(images)
    output_title = f"""
    <b>{prompt}</b>
    """
    return (output_title, predictions)


outputs = [
    gr.outputs.HTML(label=""),  # To be used as title
    gr.outputs.Image(label=""),
]

description = """
DALL·E-mini is an AI model that generates images from any prompt you give! Generate images from text:
"""
gr.Interface(
    run_inference,
    inputs=[gr.inputs.Textbox(label="What do you want to see?")],
    outputs=outputs,
    title="DALL·E mini",
    description=description,
    article="<p style='text-align: center'> Created by Boris Dayma et al. 2021 | <a href='https://github.com/borisdayma/dalle-mini'>GitHub</a> | <a href='https://wandb.ai/dalle-mini/dalle-mini/reports/DALL-E-mini--Vmlldzo4NjIxODA'>Report</a></p>",
    layout="vertical",
    theme="huggingface",
    examples=[
        ["an armchair in the shape of an avocado"],
        ["snowy mountains by the sea"],
    ],
    allow_flagging=False,
    live=False,
    # server_port=8999
).launch(share=True)