from base64 import b64encode
from io import BytesIO
from math import ceil

import clip
from multilingual_clip import legacy_multilingual_clip, pt_multilingual_clip
import numpy as np
import pandas as pd
from PIL import Image
import requests
import streamlit as st
import torch
from torchvision.transforms import ToPILImage
from transformers import AutoTokenizer, AutoModel, BertTokenizer

from CLIP_Explainability.clip_ import load, tokenize
from CLIP_Explainability.rn_cam import (
    # interpret_rn,
    interpret_rn_overlapped,
    rn_perword_relevance,
)
from CLIP_Explainability.vit_cam import (
    # interpret_vit,
    vit_perword_relevance,
    interpret_vit_overlapped,
)

from pytorch_grad_cam.grad_cam import GradCAM

RUN_LITE = False  # Load vision model for CAM viz explainability for M-CLIP only

MAX_IMG_WIDTH = 500
MAX_IMG_HEIGHT = 800

st.set_page_config(layout="wide")


# The `find_best_matches` function compares the text feature vector to the feature vectors of all images and finds the best matches. The function returns the IDs of the best matching images.
def find_best_matches(text_features, image_features, image_ids):
    # Compute the similarity between the search query and each image using the Cosine similarity
    similarities = (image_features @ text_features.T).squeeze(1)

    # Sort the images by their similarity score
    best_image_idx = (-similarities).argsort()

    # Return the image IDs of the best matches
    return [[image_ids[i], similarities[i].item()] for i in best_image_idx]


# The `encode_search_query` function takes a text description and encodes it into a feature vector using the CLIP model.
def encode_search_query(search_query, model_type):
    with torch.no_grad():
        # Encode and normalize the search query using the multilingual model
        if model_type == "M-CLIP (multilingual ViT)":
            text_encoded = st.session_state.ml_model.forward(
                search_query, st.session_state.ml_tokenizer
            )
            text_encoded /= text_encoded.norm(dim=-1, keepdim=True)
        elif model_type == "J-CLIP (日本語 ViT)":
            t_text = st.session_state.ja_tokenizer(
                search_query, padding=True, return_tensors="pt"
            )
            text_encoded = st.session_state.ja_model.get_text_features(**t_text)
            text_encoded /= text_encoded.norm(dim=-1, keepdim=True)
        else:  # model_type == legacy
            text_encoded = st.session_state.rn_model(search_query)
            text_encoded /= text_encoded.norm(dim=-1, keepdim=True)

    # Retrieve the feature vector
    return text_encoded


def clip_search(search_query):
    if st.session_state.search_field_value != search_query:
        st.session_state.search_field_value = search_query

    model_type = st.session_state.active_model

    if len(search_query) >= 1:
        text_features = encode_search_query(search_query, model_type)

        # Compute the similarity between the descrption and each photo using the Cosine similarity
        # similarities = list((text_features @ photo_features.T).squeeze(0))

        # Sort the photos by their similarity score
        if model_type == "M-CLIP (multilingual ViT)":
            matches = find_best_matches(
                text_features,
                st.session_state.ml_image_features,
                st.session_state.image_ids,
            )
        elif model_type == "J-CLIP (日本語 ViT)":
            matches = find_best_matches(
                text_features,
                st.session_state.ja_image_features,
                st.session_state.image_ids,
            )
        else:  # model_type == legacy
            matches = find_best_matches(
                text_features,
                st.session_state.rn_image_features,
                st.session_state.image_ids,
            )

        st.session_state.search_image_ids = [match[0] for match in matches]
        st.session_state.search_image_scores = {match[0]: match[1] for match in matches}


def string_search():
    if "search_field_value" in st.session_state:
        clip_search(st.session_state.search_field_value)


def load_image_features():
    # Load the image feature vectors
    if st.session_state.vision_mode == "tiled":
        ml_image_features = np.load("./image_features/tiled_ml_features.npy")
        ja_image_features = np.load("./image_features/tiled_ja_features.npy")
        rn_image_features = np.load("./image_features/tiled_rn_features.npy")
    elif st.session_state.vision_mode == "stretched":
        ml_image_features = np.load("./image_features/resized_ml_features.npy")
        ja_image_features = np.load("./image_features/resized_ja_features.npy")
        rn_image_features = np.load("./image_features/resized_rn_features.npy")
    else:  # st.session_state.vision_mode == "cropped":
        ml_image_features = np.load("./image_features/cropped_ml_features.npy")
        ja_image_features = np.load("./image_features/cropped_ja_features.npy")
        rn_image_features = np.load("./image_features/cropped_rn_features.npy")

    # Convert features to Tensors: Float32 on CPU and Float16 on GPU
    device = st.session_state.device
    if device == "cpu":
        ml_image_features = torch.from_numpy(ml_image_features).float().to(device)
        ja_image_features = torch.from_numpy(ja_image_features).float().to(device)
        rn_image_features = torch.from_numpy(rn_image_features).float().to(device)
    else:
        ml_image_features = torch.from_numpy(ml_image_features).to(device)
        ja_image_features = torch.from_numpy(ja_image_features).to(device)
        rn_image_features = torch.from_numpy(rn_image_features).to(device)

    st.session_state.ml_image_features = ml_image_features / ml_image_features.norm(
        dim=-1, keepdim=True
    )
    st.session_state.ja_image_features = ja_image_features / ja_image_features.norm(
        dim=-1, keepdim=True
    )
    st.session_state.rn_image_features = rn_image_features / rn_image_features.norm(
        dim=-1, keepdim=True
    )

    string_search()


def init():
    st.session_state.current_page = 1

    device = "cuda" if torch.cuda.is_available() else "cpu"
    st.session_state.device = device

    # Load the open CLIP models

    with st.spinner("Loading models and data, please wait..."):
        ml_model_name = "M-CLIP/XLM-Roberta-Large-Vit-B-16Plus"
        ml_model_path = "./models/vit_b_16_plus_240-laion400m_e32-699c4b84.pt"

        st.session_state.ml_image_model, st.session_state.ml_image_preprocess = load(
            ml_model_path, device=device, jit=False
        )

        st.session_state.ml_model = (
            pt_multilingual_clip.MultilingualCLIP.from_pretrained(ml_model_name)
        )
        st.session_state.ml_tokenizer = AutoTokenizer.from_pretrained(ml_model_name)

        ja_model_name = "hakuhodo-tech/japanese-clip-vit-h-14-bert-wider"
        ja_model_path = "./models/ViT-H-14-laion2B-s32B-b79K.bin"

        if not RUN_LITE:
            st.session_state.ja_image_model, st.session_state.ja_image_preprocess = (
                load(ja_model_path, device=device, jit=False)
            )

        st.session_state.ja_model = AutoModel.from_pretrained(
            ja_model_name, trust_remote_code=True
        ).to(device)
        st.session_state.ja_tokenizer = AutoTokenizer.from_pretrained(
            ja_model_name, trust_remote_code=True
        )

        if not RUN_LITE:
            st.session_state.rn_image_model, st.session_state.rn_image_preprocess = (
                clip.load("RN50x4", device=device)
            )

        st.session_state.rn_model = legacy_multilingual_clip.load_model(
            "M-BERT-Base-69"
        )
        st.session_state.rn_tokenizer = BertTokenizer.from_pretrained(
            "bert-base-multilingual-cased"
        )

        # Load the image IDs
        st.session_state.images_info = pd.read_csv("./metadata.csv")
        st.session_state.images_info.set_index("filename", inplace=True)

        with open("./images_list.txt", "r", encoding="utf-8") as images_list:
            st.session_state.image_ids = list(images_list.read().strip().split("\n"))

        st.session_state.active_model = "M-CLIP (multilingual ViT)"

    st.session_state.vision_mode = "tiled"
    st.session_state.search_image_ids = []
    st.session_state.search_image_scores = {}
    st.session_state.activations_image = None
    st.session_state.text_table_df = None

    with st.spinner("Loading models and data, please wait..."):
        load_image_features()


if "images_info" not in st.session_state:
    init()


def visualize_gradcam(viz_image_id):
    if "search_field_value" not in st.session_state:
        return

    header_cols = st.columns([80, 20], vertical_alignment="bottom")
    with header_cols[0]:
        st.title("Image + query details")
    with header_cols[1]:
        if st.button("Close"):
            st.rerun()

    st.markdown(
        f"**Query text:** {st.session_state.search_field_value} | **Image relevance:** {round(st.session_state.search_image_scores[viz_image_id], 3)}"
    )

    with st.spinner("Calculating..."):
        # info_text = st.text("Calculating activation regions...")

        image_url = st.session_state.images_info.loc[viz_image_id]["image_url"]
        image_response = requests.get(image_url)
        image = Image.open(BytesIO(image_response.content), formats=["JPEG", "GIF"])
        image = image.convert("RGB")

        img_dim = 224
        if st.session_state.active_model == "M-CLIP (multilingual ViT)":
            img_dim = 240
        elif st.session_state.active_model == "Legacy (multilingual ResNet)":
            img_dim = 288

        orig_img_dims = image.size

        ##### If the features are based on tiled image slices
        tile_behavior = None

        if st.session_state.vision_mode == "tiled":
            scaled_dims = [img_dim, img_dim]

            if orig_img_dims[0] > orig_img_dims[1]:
                scale_ratio = round(orig_img_dims[0] / orig_img_dims[1])
                if scale_ratio > 1:
                    scaled_dims = [scale_ratio * img_dim, img_dim]
                    tile_behavior = "width"
            elif orig_img_dims[0] < orig_img_dims[1]:
                scale_ratio = round(orig_img_dims[1] / orig_img_dims[0])
                if scale_ratio > 1:
                    scaled_dims = [img_dim, scale_ratio * img_dim]
                    tile_behavior = "height"

            resized_image = image.resize(scaled_dims, Image.LANCZOS)

            if tile_behavior == "width":
                image_tiles = []
                for x in range(0, scale_ratio):
                    box = (x * img_dim, 0, (x + 1) * img_dim, img_dim)
                    image_tiles.append(resized_image.crop(box))

            elif tile_behavior == "height":
                image_tiles = []
                for y in range(0, scale_ratio):
                    box = (0, y * img_dim, img_dim, (y + 1) * img_dim)
                    image_tiles.append(resized_image.crop(box))

            else:
                image_tiles = [resized_image]

        elif st.session_state.vision_mode == "stretched":
            image_tiles = [image.resize((img_dim, img_dim), Image.LANCZOS)]

        else:  # vision_mode == "cropped"
            if orig_img_dims[0] > orig_img_dims[1]:
                scale_factor = orig_img_dims[0] / orig_img_dims[1]
                resized_img_dims = (round(scale_factor * img_dim), img_dim)
                resized_img = image.resize(resized_img_dims)
            elif orig_img_dims[0] < orig_img_dims[1]:
                scale_factor = orig_img_dims[1] / orig_img_dims[0]
                resized_img_dims = (img_dim, round(scale_factor * img_dim))
            else:
                resized_img_dims = (img_dim, img_dim)

            resized_img = image.resize(resized_img_dims)

            left = round((resized_img_dims[0] - img_dim) / 2)
            top = round((resized_img_dims[1] - img_dim) / 2)
            x_right = round(resized_img_dims[0] - img_dim) - left
            x_bottom = round(resized_img_dims[1] - img_dim) - top
            right = resized_img_dims[0] - x_right
            bottom = resized_img_dims[1] - x_bottom

            # Crop the center of the image
            image_tiles = [resized_img.crop((left, top, right, bottom))]

        image_visualizations = []

        if st.session_state.active_model == "M-CLIP (multilingual ViT)":
            # Sometimes used for token importance viz
            tokenized_text = st.session_state.ml_tokenizer.tokenize(
                st.session_state.search_field_value
            )

            text_features = st.session_state.ml_model.forward(
                st.session_state.search_field_value, st.session_state.ml_tokenizer
            )

            image_model = st.session_state.ml_image_model

            for altered_image in image_tiles:
                p_image = (
                    st.session_state.ml_image_preprocess(altered_image)
                    .unsqueeze(0)
                    .to(st.session_state.device)
                )

                vis_t = interpret_vit_overlapped(
                    p_image.type(st.session_state.ml_image_model.dtype),
                    text_features,
                    image_model.visual,
                    st.session_state.device,
                    img_dim=img_dim,
                )

                image_visualizations.append(vis_t)

        elif st.session_state.active_model == "J-CLIP (日本語 ViT)":
            # Sometimes used for token importance viz
            tokenized_text = st.session_state.ja_tokenizer.tokenize(
                st.session_state.search_field_value
            )

            t_text = st.session_state.ja_tokenizer(
                st.session_state.search_field_value, return_tensors="pt"
            )
            text_features = st.session_state.ja_model.get_text_features(**t_text)

            image_model = st.session_state.ja_image_model

            for altered_image in image_tiles:
                p_image = (
                    st.session_state.ja_image_preprocess(altered_image)
                    .unsqueeze(0)
                    .to(st.session_state.device)
                )

                vis_t = interpret_vit_overlapped(
                    p_image.type(st.session_state.ja_image_model.dtype),
                    text_features,
                    image_model.visual,
                    st.session_state.device,
                    img_dim=img_dim,
                )

                image_visualizations.append(vis_t)

        else:  # st.session_state.active_model == Legacy
            # Sometimes used for token importance viz
            tokenized_text = st.session_state.rn_tokenizer.tokenize(
                st.session_state.search_field_value
            )

            text_features = st.session_state.rn_model(
                st.session_state.search_field_value
            )

            image_model = st.session_state.rn_image_model

            for altered_image in image_tiles:
                p_image = (
                    st.session_state.rn_image_preprocess(altered_image)
                    .unsqueeze(0)
                    .to(st.session_state.device)
                )

                vis_t = interpret_rn_overlapped(
                    p_image.type(st.session_state.rn_image_model.dtype),
                    text_features,
                    image_model.visual,
                    GradCAM,
                    st.session_state.device,
                    img_dim=img_dim,
                )

                image_visualizations.append(vis_t)

        transform = ToPILImage()

        vis_images = [transform(vis_t) for vis_t in image_visualizations]

        if st.session_state.vision_mode == "cropped":
            resized_img.paste(vis_images[0], (left, top))
            vis_images = [resized_img]

        if orig_img_dims[0] > orig_img_dims[1]:
            scale_factor = MAX_IMG_WIDTH / orig_img_dims[0]
            scaled_dims = [MAX_IMG_WIDTH, int(orig_img_dims[1] * scale_factor)]
        else:
            scale_factor = MAX_IMG_HEIGHT / orig_img_dims[1]
            scaled_dims = [int(orig_img_dims[0] * scale_factor), MAX_IMG_HEIGHT]

        if tile_behavior == "width":
            vis_image = Image.new("RGB", (len(vis_images) * img_dim, img_dim))
            for x, v_img in enumerate(vis_images):
                vis_image.paste(v_img, (x * img_dim, 0))
            st.session_state.activations_image = vis_image.resize(scaled_dims)

        elif tile_behavior == "height":
            vis_image = Image.new("RGB", (img_dim, len(vis_images) * img_dim))
            for y, v_img in enumerate(vis_images):
                vis_image.paste(v_img, (0, y * img_dim))
            st.session_state.activations_image = vis_image.resize(scaled_dims)

        else:
            st.session_state.activations_image = vis_images[0].resize(scaled_dims)

        image_io = BytesIO()
        st.session_state.activations_image.save(image_io, "PNG")
        dataurl = "data:image/png;base64," + b64encode(image_io.getvalue()).decode(
            "ascii"
        )

        st.html(
            f"""<div style="display: flex; flex-direction: column; align-items: center;">
                    <img src="{dataurl}" />
                </div>"""
        )

    tokenized_text = [tok.replace("▁", "") for tok in tokenized_text if tok != "▁"]
    tokenized_text = [
        tok for tok in tokenized_text if tok not in ["s", "ed", "a", "the", "an", "ing"]
    ]

    if (
        len(tokenized_text) > 1
        and len(tokenized_text) < 25
        and st.button(
            "Calculate text importance (may take some time)",
        )
    ):
        search_tokens = []
        token_scores = []

        progress_text = f"Processing {len(tokenized_text)} text tokens"
        progress_bar = st.progress(0.0, text=progress_text)

        for t, tok in enumerate(tokenized_text):
            token = tok

            if st.session_state.active_model == "Legacy (multilingual ResNet)":
                word_rel = rn_perword_relevance(
                    p_image,
                    st.session_state.search_field_value,
                    image_model,
                    tokenize,
                    GradCAM,
                    st.session_state.device,
                    token,
                    data_only=True,
                    img_dim=img_dim,
                )
            else:
                word_rel = vit_perword_relevance(
                    p_image,
                    st.session_state.search_field_value,
                    image_model,
                    tokenize,
                    st.session_state.device,
                    token,
                    data_only=True,
                    img_dim=img_dim,
                )
            avg_score = np.mean(word_rel)
            if avg_score == 0 or np.isnan(avg_score):
                continue
            search_tokens.append(token)
            token_scores.append(1 / avg_score)

            progress_bar.progress(
                (t + 1) / len(tokenized_text),
                text=f"Processing token {t+1} of {len(tokenized_text)}",
            )
        progress_bar.empty()

        normed_scores = torch.softmax(torch.tensor(token_scores), dim=0)

        token_scores = [f"{round(score.item() * 100, 3)}%" for score in normed_scores]
        st.session_state.text_table_df = pd.DataFrame(
            {"token": search_tokens, "importance": token_scores}
        )

        st.markdown("**Importance of each text token to relevance score**")
        st.table(st.session_state.text_table_df)


def format_vision_mode(mode_stub):
    return mode_stub.capitalize()


@st.dialog(" ", width="large")
def image_modal(vis_image_id):
    visualize_gradcam(vis_image_id)


st.title("Explore Japanese visual aesthetics with CLIP models")

st.markdown(
    """
    <style>
    [data-testid=stImageCaption] {
        padding: 0 0 0 0;
    }
    [data-testid=stVerticalBlockBorderWrapper] {
        line-height: 1.2;
    }
    [data-testid=stVerticalBlock] {
        gap: .75rem;
    }
    [data-testid=baseButton-secondary] {
        min-height: 1rem;
        padding: 0 0.75rem;
        margin: 0 0 1rem 0;
    }
    div[aria-label="dialog"]>button[aria-label="Close"] {
        display: none;
    }
    [data-testid=stFullScreenFrame] {
        display: flex;
        flex-direction: column;
        align-items: center;
    }
    </style>
    """,
    unsafe_allow_html=True,
)

search_row = st.columns([45, 8, 8, 10, 1, 8, 20], vertical_alignment="center")
with search_row[0]:
    search_field = st.text_input(
        label="search",
        label_visibility="collapsed",
        placeholder="Type something, or click a suggested search below.",
        on_change=string_search,
        key="search_field_value",
    )
with search_row[1]:
    st.button(
        "Search", on_click=string_search, use_container_width=True, type="primary"
    )
with search_row[2]:
    st.markdown("**Vision mode:**")
with search_row[3]:
    st.selectbox(
        "Vision mode",
        options=["tiled", "stretched", "cropped"],
        key="vision_mode",
        help="How to consider images that aren't square",
        on_change=load_image_features,
        format_func=format_vision_mode,
        label_visibility="collapsed",
    )
with search_row[4]:
    st.empty()
with search_row[5]:
    st.markdown("**CLIP model:**")
with search_row[6]:
    st.selectbox(
        "CLIP Model:",
        options=[
            "M-CLIP (multilingual ViT)",
            "J-CLIP (日本語 ViT)",
            "Legacy (multilingual ResNet)",
        ],
        key="active_model",
        on_change=string_search,
        label_visibility="collapsed",
    )

canned_searches = st.columns([12, 22, 22, 22, 22], vertical_alignment="top")
with canned_searches[0]:
    st.markdown("**Suggested searches:**")
if st.session_state.active_model == "J-CLIP (日本語 ViT)":
    with canned_searches[1]:
        st.button(
            "間",
            on_click=clip_search,
            args=["間"],
            use_container_width=True,
        )
    with canned_searches[2]:
        st.button("奥", on_click=clip_search, args=["奥"], use_container_width=True)
    with canned_searches[3]:
        st.button("山", on_click=clip_search, args=["山"], use_container_width=True)
    with canned_searches[4]:
        st.button(
            "花に酔えり 羽織着て刀 さす女",
            on_click=clip_search,
            args=["花に酔えり 羽織着て刀 さす女"],
            use_container_width=True,
        )
else:
    with canned_searches[1]:
        st.button(
            "negative space",
            on_click=clip_search,
            args=["negative space"],
            use_container_width=True,
        )
    with canned_searches[2]:
        st.button("間", on_click=clip_search, args=["間"], use_container_width=True)
    with canned_searches[3]:
        st.button("음각", on_click=clip_search, args=["음각"], use_container_width=True)
    with canned_searches[4]:
        st.button(
            "αρνητικός χώρος",
            on_click=clip_search,
            args=["αρνητικός χώρος"],
            use_container_width=True,
        )

controls = st.columns([35, 5, 35, 5, 20], gap="large", vertical_alignment="center")
with controls[0]:
    im_per_pg = st.columns([30, 70], vertical_alignment="center")
    with im_per_pg[0]:
        st.markdown("**Images/page:**")
    with im_per_pg[1]:
        batch_size = st.select_slider(
            "Images/page:", range(10, 50, 10), label_visibility="collapsed"
        )
with controls[1]:
    st.empty()
with controls[2]:
    im_per_row = st.columns([30, 70], vertical_alignment="center")
    with im_per_row[0]:
        st.markdown("**Images/row:**")
    with im_per_row[1]:
        row_size = st.select_slider(
            "Images/row:", range(1, 6), value=5, label_visibility="collapsed"
        )
num_batches = ceil(len(st.session_state.image_ids) / batch_size)
with controls[3]:
    st.empty()
with controls[4]:
    pager = st.columns([40, 60], vertical_alignment="center")
    with pager[0]:
        st.markdown(f"Page **{st.session_state.current_page}** of **{num_batches}** ")
    with pager[1]:
        st.number_input(
            "Page",
            min_value=1,
            max_value=num_batches,
            step=1,
            label_visibility="collapsed",
            key="current_page",
        )


if len(st.session_state.search_image_ids) == 0:
    batch = []
else:
    batch = st.session_state.search_image_ids[
        (st.session_state.current_page - 1) * batch_size : st.session_state.current_page
        * batch_size
    ]

grid = st.columns(row_size)
col = 0
for image_id in batch:
    with grid[col]:
        link_text = st.session_state.images_info.loc[image_id]["permalink"].split("/")[
            2
        ]
        # st.image(
        #     st.session_state.images_info.loc[image_id]["image_url"],
        #     caption=st.session_state.images_info.loc[image_id]["caption"],
        # )
        st.html(
            f"""<div style="display: flex; flex-direction: column; align-items: center">
                    <img src="{st.session_state.images_info.loc[image_id]['image_url']}" style="max-width: 100%; max-height: {MAX_IMG_HEIGHT}px" />
                    <div>{st.session_state.images_info.loc[image_id]['caption']} <b>[{round(st.session_state.search_image_scores[image_id], 3)}]</b></div>
                </div>"""
        )
        st.caption(
            f"""<div style="display: flex; flex-direction: column; align-items: center; position: relative; top: -12px">
                    <a href="{st.session_state.images_info.loc[image_id]['permalink']}">{link_text}</a>
                <div>""",
            unsafe_allow_html=True,
        )
        if not RUN_LITE or st.session_state.active_model == "M-CLIP (multilingual ViT)":
            st.button(
                "Explain this",
                on_click=image_modal,
                args=[image_id],
                use_container_width=True,
                key=image_id,
            )
    col = (col + 1) % row_size