import data
import torch
import gradio as gr
from models import imagebind_model
from models.imagebind_model import ModalityType


device = "cuda:0" if torch.cuda.is_available() else "cpu"
model = imagebind_model.imagebind_huge(pretrained=True)
model.eval()
model.to(device)


def image_text_zeroshot(image, text_list):
    image_paths = [image]
    labels = [label.strip(" ") for label in text_list.strip(" ").split("|")]
    inputs = {
        ModalityType.TEXT: data.load_and_transform_text(labels, device),
        ModalityType.VISION: data.load_and_transform_vision_data(image_paths, device),
    }

    with torch.no_grad():
        embeddings = model(inputs)

    scores = (
        torch.softmax(
            embeddings[ModalityType.VISION] @ embeddings[ModalityType.TEXT].T, dim=-1
        )
        .squeeze(0)
        .tolist()
    )

    score_dict = {label: score for label, score in zip(labels, scores)}

    return score_dict


def audio_text_zeroshot(audio, text_list):
    audio_paths = [audio]
    labels = [label.strip(" ") for label in text_list.strip(" ").split("|")]
    inputs = {
        ModalityType.TEXT: data.load_and_transform_text(labels, device),
        ModalityType.AUDIO: data.load_and_transform_audio_data(audio_paths, device),
    }

    with torch.no_grad():
        embeddings = model(inputs)

    scores = (
        torch.softmax(
            embeddings[ModalityType.AUDIO] @ embeddings[ModalityType.TEXT].T, dim=-1
        )
        .squeeze(0)
        .tolist()
    )

    score_dict = {label: score for label, score in zip(labels, scores)}

    return score_dict


def video_text_zeroshot(video, text_list):
    video_paths = [video]
    labels = [label.strip(" ") for label in text_list.strip(" ").split("|")]
    inputs = {
        ModalityType.TEXT: data.load_and_transform_text(labels, device),
        ModalityType.VISION: data.load_and_transform_video_data(video_paths, device),
    }

    with torch.no_grad():
        embeddings = model(inputs)

    scores = (
        torch.softmax(
            embeddings[ModalityType.VISION] @ embeddings[ModalityType.TEXT].T, dim=-1
        )
        .squeeze(0)
        .tolist()
    )

    score_dict = {label: score for label, score in zip(labels, scores)}

    return score_dict


def inference(
    task,
    text_list=None,
    image=None,
    audio=None,
    video=None,
):
    if task == "image-text":
        result = image_text_zeroshot(image, text_list)
    elif task == "audio-text":
        result = audio_text_zeroshot(audio, text_list)
    elif task == "video-text":
        result = video_text_zeroshot(video, text_list)
    else:
        raise NotImplementedError
    return result


def main():
    inputs = [
        gr.inputs.Radio(
            choices=[
                "image-text",
                "audio-text",
                "video-text",
            ],
            type="value",
            default="image-text",
            label="Task",
        ),
        gr.inputs.Textbox(lines=1, label="Candidate texts"),
        gr.inputs.Image(type="filepath", label="Input image"),
        gr.inputs.Audio(type="filepath", label="Input audio"),
        gr.inputs.Video(type=None, label="Input video"),
    ]

    iface = gr.Interface(
        inference,
        inputs,
        "label",
        examples=[
            ["image-text", "A dog|A car|A bird", "assets/dog_image.jpg", None, None],
            ["image-text", "A dog|A car|A bird", "assets/car_image.jpg", None, None],
            ["audio-text", "A dog|A car|A bird",  None, "assets/bird_audio.wav", None],
            ["video-text", "A dog|A car|A bird", None, None, "assets/dog_video.mp4"],
        ],
        description="""<p>This is a simple demo of ImageBind for zero-shot cross-modal understanding (now including image classification, audio classification, and video classification). Please refer to the original <a href='https://arxiv.org/abs/2305.05665' target='_blank'>paper</a> and <a href='https://github.com/facebookresearch/ImageBind' target='_blank'>repo</a> for more details.<br>
                    To test your own cases, you can upload an image, an audio or a video, and provide the candidate texts separated by "|".<br>
                    You can duplicate this space and run it privately: <a href='https://huggingface.co/spaces/OFA-Sys/chinese-clip-zero-shot-image-classification?duplicate=true'><img src='https://img.shields.io/badge/-Duplicate%20Space-blue?labelColor=white&style=flat&logo=data:image/png;base64,iVBORw0KGgoAAAANSUhEUgAAABAAAAAQCAYAAAAf8/9hAAAAAXNSR0IArs4c6QAAAP5JREFUOE+lk7FqAkEURY+ltunEgFXS2sZGIbXfEPdLlnxJyDdYB62sbbUKpLbVNhyYFzbrrA74YJlh9r079973psed0cvUD4A+4HoCjsA85X0Dfn/RBLBgBDxnQPfAEJgBY+A9gALA4tcbamSzS4xq4FOQAJgCDwV2CPKV8tZAJcAjMMkUe1vX+U+SMhfAJEHasQIWmXNN3abzDwHUrgcRGmYcgKe0bxrblHEB4E/pndMazNpSZGcsZdBlYJcEL9Afo75molJyM2FxmPgmgPqlWNLGfwZGG6UiyEvLzHYDmoPkDDiNm9JR9uboiONcBXrpY1qmgs21x1QwyZcpvxt9NS09PlsPAAAAAElFTkSuQmCC&logoWidth=14' alt='Duplicate Space'></a></p>""",
        title="ImageBind: Zero-shot Cross-modal Understanding",
    )

    iface.launch()


if __name__ == "__main__":
    main()