import json
import uuid

import comet_ml
import gradio as gr
import pandas as pd
from PIL import Image
from transformers import CLIPModel, CLIPProcessor

CLIP_MODEL_PATH = "openai/clip-vit-base-patch32"

clip_model = CLIPModel.from_pretrained(CLIP_MODEL_PATH)
clip_processor = CLIPProcessor.from_pretrained(CLIP_MODEL_PATH)

DESCRIPTION = """Glad to see you here 😄.
You can use this Space to log predictions to [Comet](https://www.comet.ml/site) from Spaces that use Text to Image Diffusion Models.

Keep track of all your prompts and generated images so that you remember the good ones!

Set your Comet credentials in the Comet Settings tab and create an Experiment for logging data.
If you want to continue logging to the same Experiment over multiple sessions, add in the

Then use the path to a Space to generate from in the Diffusion Model tab
"""


def create_experiment(
    comet_api_key,
    comet_workspace,
    comet_project_name,
    comet_experiment_name,
    experiment,
):
    if not comet_api_key:
        experiment = None
        return (
            experiment,
            """
            Please add your API key in order to log your predictions to a Comet Experiment.
            If you don't have a Comet account yet, you can sign up using the link below:

            https://www.comet.ml/signup
            """,
        )

    try:
        api_experiment = comet_ml.APIExperiment(
            api_key=comet_api_key,
            workspace=comet_workspace,
            project_name=comet_project_name,
            experiment_name=comet_experiment_name,
        )
        experiment = {
            "api_key": comet_api_key,
            "workspace": comet_workspace,
            "project_name": comet_project_name,
            "previous_experiment": api_experiment.id,
        }

        return experiment, f"Started {api_experiment.name}. Happy logging!😊"

    except Exception as e:
        return None, e


def get_experiment(kwargs) -> comet_ml.APIExperiment:
    try:
        return comet_ml.APIExperiment(**kwargs)
    except Exception as e:
        return None


def get_experiment_status(experiment_state):
    experiment = get_experiment(experiment_state)
    if experiment is not None:
        name = experiment.name
        return experiment_state, f"Currently logging to: {name}"

    return experiment_state, f"No Experiments found"


def predict(
    model,
    prompt,
    experiment_state,
):
    io = gr.Interface.load(model)
    image = io(prompt)
    pil_image = Image.open(image)

    inputs = clip_processor(
        text=[prompt],
        images=pil_image,
        return_tensors="pt",
        padding=True,
    )
    outputs = clip_model(**inputs)
    clip_score = outputs.logits_per_image.item() / 100.0

    experiment = get_experiment(experiment_state)
    if experiment is not None:
        image_id = uuid.uuid4().hex
        experiment.log_image(image, image_id)

        asset = pd.DataFrame.from_records(
            [
                {
                    "prompt": prompt,
                    "model": model,
                    "clip_model": CLIP_MODEL_PATH,
                    "clip_score": round(clip_score, 3),
                }
            ]
        )
        experiment.log_table(f"{image_id}.json", asset, orient="records")

    return image, experiment_state


def start_interface():
    demo = gr.Blocks()
    with demo:
        description = gr.Markdown(DESCRIPTION)
        with gr.Tabs():
            with gr.TabItem(label="Comet Settings"):
                # credentials
                comet_api_key = gr.Textbox(
                    label="Comet API Key",
                    placeholder="This is required if you'd like to create an Experiment",
                )
                comet_workspace = gr.Textbox(label="Comet Workspace")
                comet_project_name = gr.Textbox(label="Comet Project Name")
                comet_experiment_name = gr.Textbox(
                    label="Comet Experiment Name",
                    placeholder=(
                        "Set this if you'd like"
                        "to continue logging to an existing Experiment",
                    ),
                )

                with gr.Row():
                    start = gr.Button("Create Experiment", variant="primary")
                    status = gr.Button("Experiment Status")

                output = gr.Markdown(label="Status")
                experiment_state = gr.Variable(label="Experiment State")

                start.click(
                    create_experiment,
                    inputs=[
                        comet_api_key,
                        comet_workspace,
                        comet_project_name,
                        comet_experiment_name,
                        experiment_state,
                    ],
                    outputs=[experiment_state, output],
                )

                status.click(
                    get_experiment_status,
                    inputs=[experiment_state],
                    outputs=[experiment_state, output],
                )

            with gr.TabItem(label="Diffusion Model"):
                diff_description = gr.Markdown(
                    """The Model must be a path to any Space that accepts"
                    only text as input and produces an image as an output
                    """
                )
                model = gr.Textbox(label="Model", value="spaces/valhalla/glide-text2im")
                prompt = gr.Textbox(label="Prompt")

                outputs = gr.Image(label="Image")

                submit = gr.Button("Submit", variant="primary")
                submit.click(
                    predict,
                    inputs=[model, prompt, experiment_state],
                    outputs=[outputs, experiment_state],
                )

    demo.launch()


start_interface()