import json

import faiss
import flax
import gradio as gr
import jax
import numpy as np
import pandas as pd
import requests

from Models.CLIP import CLIP


def danbooru_id_to_url(image_id, selected_ratings, api_username="", api_key=""):
    headers = {"User-Agent": "image_similarity_tool"}
    ratings_to_letters = {
        "General": "g",
        "Sensitive": "s",
        "Questionable": "q",
        "Explicit": "e",
    }

    acceptable_ratings = [ratings_to_letters[x] for x in selected_ratings]

    image_url = f"https://danbooru.donmai.us/posts/{image_id}.json"
    if api_username != "" and api_key != "":
        image_url = f"{image_url}?api_key={api_key}&login={api_username}"

    r = requests.get(image_url, headers=headers)
    if r.status_code != 200:
        return None

    content = json.loads(r.text)
    image_url = content["large_file_url"] if "large_file_url" in content else None
    image_url = image_url if content["rating"] in acceptable_ratings else None
    return image_url


class Predictor:
    def __init__(self):
        self.base_model = "wd-v1-4-convnext-tagger-v2"

        with open(f"data/{self.base_model}/clip.msgpack", "rb") as f:
            data = f.read()

        self.params = flax.serialization.msgpack_restore(data)["model"]
        self.model = CLIP()

        self.tags_df = pd.read_csv("data/selected_tags.csv")

        self.images_ids = np.load("index/cosine_ids.npy")

        self.knn_index = faiss.read_index("index/cosine_knn.index")

        config = json.loads(open("index/cosine_infos.json").read())["index_param"]
        faiss.ParameterSpace().set_index_parameters(self.knn_index, config)

    def predict(self, positive_tags, negative_tags, n_neighbours=5):
        tags_df = self.tags_df
        model = self.model

        num_classes = len(tags_df)

        positive_tags = positive_tags.split(",")
        negative_tags = negative_tags.split(",")

        positive_tags_idxs = tags_df.index[tags_df["name"].isin(positive_tags)].tolist()
        negative_tags_idxs = tags_df.index[tags_df["name"].isin(negative_tags)].tolist()

        tags = np.zeros((1, num_classes), dtype=np.float32)
        tags[0][positive_tags_idxs] = 1
        emb_from_logits = model.apply(
            {"params": self.params},
            tags,
            method=model.encode_text,
        )
        emb_from_logits = jax.device_get(emb_from_logits)

        if len(negative_tags_idxs) > 0:
            tags = np.zeros((1, num_classes), dtype=np.float32)
            tags[0][negative_tags_idxs] = 1

            neg_emb_from_logits = model.apply(
                {"params": self.params},
                tags,
                method=model.encode_text,
            )
            neg_emb_from_logits = jax.device_get(neg_emb_from_logits)
            emb_from_logits = emb_from_logits - neg_emb_from_logits

        faiss.normalize_L2(emb_from_logits)

        dists, indexes = self.knn_index.search(emb_from_logits, k=n_neighbours)
        neighbours_ids = self.images_ids[indexes][0]
        neighbours_ids = [int(x) for x in neighbours_ids]

        captions = []
        image_urls = []
        for image_id, dist in zip(neighbours_ids, dists[0]):
            current_url = danbooru_id_to_url(
                image_id,
                [
                    "General",
                    "Sensitive",
                    "Questionable",
                    "Explicit",
                ],
            )
            if current_url is not None:
                image_urls.append(current_url)
                captions.append(f"{image_id}/{dist:.2f}")
        return list(zip(image_urls, captions))


def main():
    predictor = Predictor()

    with gr.Blocks() as demo:
        with gr.Row():
            positive_tags = gr.Textbox(label="Positive tags")
            negative_tags = gr.Textbox(label="Negative tags")

        find_btn = gr.Button("Find similar images")

        similar_images = gr.Gallery(label="Similar images", columns=[5])

        find_btn.click(
            fn=predictor.predict,
            inputs=[positive_tags, negative_tags],
            outputs=[similar_images],
        )

    demo.queue()
    demo.launch()


if __name__ == "__main__":
    main()