import json
import logging
import os
from functools import partial

import gradio as gr
from datasets import Dataset, load_dataset
from dotenv import load_dotenv

logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)
logger.setLevel(logging.INFO)

load_dotenv()

# dataset = load_dataset("detection-datasets/coco")
it_dataset = load_dataset(
    "imagenet-1k", split="train", streaming=True, trust_remote_code=True
).shuffle(42)


def gen_from_iterable_dataset(iterable_ds):
    """
    Convert an iterable dataset to a generator
    """
    yield from iterable_ds


# imagenet_categories_data.json is a JSON file containing a hierarchy of ImageNet categories.
# We want to take all categories under "artifact, artefact".
# Each node has this structure:
# {
#     "id": 1,
#     "name": "entity",
#     "children": ...
# }
with open("imagenet_categories_data.json") as f:
    data = json.load(f)

    # Recursively find all categories under "artifact, artefact".
    # We want to get all the "index" values of the leaf nodes. Nodes that are not leaf nodes have a "children" key.
    def find_categories(node):
        if "children" in node:
            for child in node["children"]:
                yield from find_categories(child)
        elif "index" in node:
            yield node["index"]

    broad_categories = data["children"]
    artifact_category = next(
        filter(lambda x: x["name"] == "artifact, artefact", broad_categories)
    )
    artifact_categories = list(find_categories(artifact_category))
    logger.info(f"Artifact categories: {artifact_categories}")


def filter_imgs_by_label(x):
    """
    Filter out the images that have label -1
    """
    print(f'label: {x["label"]}')
    return x["label"] in artifact_categories


dataset = it_dataset.take(1000).filter(filter_imgs_by_label)
dataset = Dataset.from_generator(
    partial(gen_from_iterable_dataset, it_dataset), features=it_dataset.features
)
dataset_iterable = iter(dataset)


def get_user_prompt():
    # Pick the first 3 images and labels
    images = []
    machine_labels = []
    human_labels = []
    for i in range(3):
        data = next(dataset_iterable)
        logger.info(f"Data: {data}")
        images.append(data["image"])
        # Get the label as a human readable string
        machine_labels.append(data["label"])
        logger.info(dataset)
        human_label = dataset.features["label"].int2str(data["label"]) + str(
            data["label"]
        )
        human_labels.append(human_label)
    return {
        "images": images,
        "machine_labels": machine_labels,
        "human_labels": human_labels,
    }


hf_writer = gr.HuggingFaceDatasetSaver(
    hf_token=os.environ["HF_TOKEN"], dataset_name="acmc/maker-faire-bot", private=True
)
csv_writer = gr.CSVLogger(simplify_file_data=True)

theme = gr.themes.Default(primary_hue="cyan", secondary_hue="fuchsia")

with gr.Blocks(theme=theme) as demo:
    with gr.Row() as header:
        gr.Image(
            "maker-faire-logo.webp",
            show_download_button=False,
            show_label=False,
            show_share_button=False,
            container=False,
            # height=100,
            scale=0.2,
        )
        gr.Markdown(
            """
            # Maker Faire Bot
            """,
            visible=False,
        )

    user_prompt = gr.State(get_user_prompt())

    gr.Markdown("""# Think about these objects...""")
    gr.Markdown(
        """We want to teach the Maker Faire Bot some creativity. Help us get ideas on what you'd build!"""
    )
    image_components = []
    with gr.Row(variant="panel") as row:
        for i in range(len(user_prompt.value["images"])):
            with gr.Column(variant="default") as col:
                img = gr.Image(
                    user_prompt.value["images"][i],
                    label=user_prompt.value["human_labels"][i],
                    interactive=False,
                    show_download_button=False,
                    show_share_button=False,
                )
                image_components.append(img)
                btn = gr.Button("Change", variant="secondary")

                def change_image(user_prompt):
                    data = next(dataset_iterable)
                    logger.info(user_prompt)
                    user_prompt = user_prompt.copy()
                    user_prompt["images"][i] = data["image"]
                    user_prompt["machine_labels"][i] = data["label"]
                    user_prompt["human_labels"][i] = dataset.features["label"].int2str(
                        data["label"]
                    )
                    logger.info(user_prompt)
                    return (
                        user_prompt,
                        user_prompt["images"][i],
                        gr.update(
                            label=user_prompt["human_labels"][i],
                        ),
                    )

                btn.click(
                    change_image,
                    inputs=[user_prompt],
                    outputs=[user_prompt, img, img],
                    preprocess=True,
                    postprocess=True,
                )

    user_answer_object = gr.Textbox(
        autofocus=True,
        placeholder="(example): An digital electronic guitar",
        label="What would you build?",
    )
    user_answer_explanation = gr.TextArea(
        autofocus=True,
        label="How would you build it?",
        # The example uses a roll of string, a camera, and a loudspeaker to build an electronic guitar.
        placeholder="""To build an electronic guitar, I would:
1. Use the roll of string to create the strings of the guitar.
2. Use the camera to capture a live video of the hand movements. That way, I can use an AI model to predict the chords.
3. Using a computer vision model, identify where the fingers are placed on the strings.
4. Calculate the sounds that the loudspeaker should produce based on the finger placements.
5. Play the sound through the loudspeaker.
""",
    )

    csv_writer.setup(
        components=[user_prompt, user_answer_object, user_answer_explanation],
        flagging_dir="user_data_csv",
    )
    hf_writer.setup(
        components=[user_prompt, user_answer_object, user_answer_explanation],
        flagging_dir="user_data_hf",
    )

    submit_btn = gr.Button("Submit", variant="primary")

    def log_results(prompt, object, explanation):
        csv_writer.flag([prompt, object, explanation])
        hf_writer.flag([prompt, object, explanation])

    submit_btn.click(
        log_results,
        inputs=[user_prompt, user_answer_object, user_answer_explanation],
        preprocess=False,
    )

    new_prompt_btn = gr.Button("New Prompt", variant="secondary")
    new_prompt_btn.click(
        get_user_prompt,
        outputs=[user_prompt],
        preprocess=False,
    )

    gr.Markdown(
        """
        This is an experimental project. Your data is anonymous and will be used to train an AI model. By using this tool, you agree to our [policy](https://makerfaire.com/privacy).
        """
    )

if __name__ == "__main__":
    demo.launch()