import streamlit as st
import numpy as np
import base64
from io import BytesIO
from multilingual_clip import pt_multilingual_clip
from transformers import CLIPTokenizerFast, AutoTokenizer, CLIPModel
import torch
import logging
from os import environ
from parse import parse
from clickhouse_connect import get_client
environ['TOKENIZERS_PARALLELISM'] = 'true'


db_name_map = {
    "Unsplash Photos 25K": lambda feat: f"mqdb_demo.unsplash_25k_{feat}_indexer",
    "RSICD: Remote Sensing Images 11K": lambda feat: f"mqdb_demo.rsicd_{feat}_b_32",
}
feat_name_map = {
    'Vanilla CLIP': "clip",
    'CLIP finetuned on RSICD': "cliprsicd"
}


DB_NAME = "mqdb_demo.unsplash_25k_clip_indexer"
DIMS = 512
# Ignore some bad links (broken in the dataset already)
BAD_IDS = {'9_9hzZVjV8s', 'RDs0THr4lGs', 'vigsqYux_-8',
           'rsJtMXn3p_c', 'AcG-unN00gw', 'r1R_0ZNUcx0'}


@st.experimental_singleton(show_spinner=False)
def init_db():
    """ Initialize the Database Connection

    Returns:
        meta_field: Meta field that records if an image is viewed or not
        client:     Database connection object
    """
    r = parse("{http_pre}://{host}:{port}", st.secrets["DB_URL"])
    client = get_client(
        host=r['host'], port=r['port'], user=st.secrets["USER"], password=st.secrets["PASSWD"]
    )
    meta_field = {}
    return meta_field, client


@st.experimental_singleton(show_spinner=False)
def init_query_num():
    print("init query_num")
    return 0


def query(xq, top_k=10):
    """ Query TopK matched w.r.t a given vector

    Args:
        xq (numpy.ndarray or list of floats): Query vector
        top_k (int, optional): Number of matched vectors. Defaults to 10.

    Returns:
        matches: list of Records object. Keys referrring to selected columns
    """
    attempt = 0
    xq = xq / np.linalg.norm(xq)
    while attempt < 3:
        try:
            xq_s = f"[{', '.join([str(float(fnum)) for fnum in list(xq)])}]"

            print('Excluded pre:', st.session_state.meta)
            if len(st.session_state.meta) > 0:
                exclude_list = ','.join(
                    [f'\'{i}\'' for i, v in st.session_state.meta.items() if v >= 1])
                print("Excluded:", exclude_list)
                # Using PREWHERE allows you to do column filter before vector search
                xc = st.session_state.index.query(f"SELECT id, url, vector,\
                        distance(vector, {xq_s}) AS dist\
                        FROM {db_name_map[st.session_state.db_name_ref](feat_name_map[st.session_state.feat_name])} \
                        WHERE id NOT IN ({exclude_list}) ORDER BY dist LIMIT {top_k}").named_results()
            else:
                xc = st.session_state.index.query(f"SELECT id, url, vector,\
                        distance(vector, {xq_s}) AS dist\
                        FROM {db_name_map[st.session_state.db_name_ref](feat_name_map[st.session_state.feat_name])} \
                        ORDER BY dist LIMIT {top_k}").named_results()
            real_xc = st.session_state.index.query(f"SELECT id, url, vector,\
                        distance(vector, {xq_s}) AS dist \
                        FROM {db_name_map[st.session_state.db_name_ref](feat_name_map[st.session_state.feat_name])} \
                        ORDER BY dist LIMIT {top_k}").named_results()
            top_k = [{k: v for k, v in r.items()} for r in real_xc]
            xc = [xi for xi in xc if xi['id'] not in st.session_state.meta or
                  st.session_state.meta[xi['id']] < 1]
            logging.info(
                f'{len(xc)} records returned, {[_i["id"] for _i in xc]}')
            matches = xc
            break
        except Exception as e:
            # force reload if we have trouble on connections or something else
            logging.warning(str(e))
            _, st.session_state.index = init_db()
            attempt += 1
            matches = []
    if len(matches) == 0:
        logging.error(f"No matches found for '{DB_NAME}'")
    return matches, top_k


@st.experimental_singleton(show_spinner=False)
def init_random_query():
    xq = np.random.rand(DIMS).tolist()
    return xq, xq.copy()


class Classifier:
    """ Zero-shot Classifier
    This Classifier provides proxy regarding to the user's reaction to the probed images.
    The proxy will replace the original query vector generated by prompted vector and finally
    give the user a satisfying retrieval result.

    This can be commonly seen in a recommendation system. The classifier will recommend more 
    precise result as it accumulating user's activity.
    """

    def __init__(self, xq: list):
        # initialize model with DIMS input size and 1 output
        # note that the bias is ignored, as we only focus on the inner product result
        self.model = torch.nn.Linear(DIMS, 1, bias=False)
        # convert initial query `xq` to tensor parameter to init weights
        init_weight = torch.Tensor(xq).reshape(1, -1)
        self.model.weight = torch.nn.Parameter(init_weight)
        # init loss and optimizer
        self.loss = torch.nn.BCEWithLogitsLoss()
        self.optimizer = torch.optim.SGD(self.model.parameters(), lr=0.1)

    def fit(self, X: list, y: list, iters: int = 5):
        # convert X and y to tensor
        X = torch.Tensor(X)
        y = torch.Tensor(y).reshape(-1, 1)
        for i in range(iters):
            # zero gradients
            self.optimizer.zero_grad()
            # Normalize the weight before inference
            # This will constrain the gradient or you will have an explosion on query vector
            self.model.weight.data = self.model.weight.data / \
                torch.norm(self.model.weight.data, p=2, dim=-1)
            # forward pass
            out = self.model(X)
            # compute loss
            loss = self.loss(out, y)
            # backward pass
            loss.backward()
            # update weights
            self.optimizer.step()

    def get_weights(self):
        xq = self.model.weight.detach().numpy()[0].tolist()
        return xq


class NormalizingLayer(torch.nn.Module):
    def forward(self, x):
        return x / torch.norm(x, dim=-1, keepdim=True)


def card(i, url):
    return f'<img id="img{i}" src="{url}" width="200px;">'


def card_with_conf(i, conf, url):
    conf = "%.4f" % (conf)
    return f'<img id="img{i}" src="{url}" width="200px;" style="margin:50px 50px"><div><p><b>Relevance: {conf}</b></p></div>'


def get_top_k(xq, top_k=9):
    """ wrapper function for query

    Args:
        xq (numpy.ndarray or list of floats): Query vector
        top_k (int, optional): Number of returned vectors. Defaults to 9.

    Returns:
        matches: See `query()`
    """
    matches = query(
        xq, top_k=top_k
    )
    return matches


def tune(X, y, iters=2):
    """ Train the Zero-shot Classifier

    Args:
        X (numpy.ndarray): Input vectors (retreived vectors)
        y (list of floats or numpy.ndarray): Scores given by user
        iters (int, optional): iterations of updates to be run
    """
    assert len(X) == len(y)
    # train the classifier
    st.session_state.clf.fit(X, y, iters=iters)
    # extract new vector
    st.session_state.xq = st.session_state.clf.get_weights()


def refresh_index():
    """ Clean the session
    """
    del st.session_state["meta"]
    st.session_state.meta = {}
    st.session_state.query_num = 0
    logging.info(f"Refresh for '{st.session_state.meta}'")
    init_db.clear()
    # refresh session states
    st.session_state.meta, st.session_state.index = init_db()
    del st.session_state.clf, st.session_state.xq


def calc_dist():
    xq = np.array(st.session_state.xq)
    orig_xq = np.array(st.session_state.orig_xq)
    return np.linalg.norm(xq - orig_xq)


def submit():
    """ Tune the model w.r.t given score from user.
    """
    st.session_state.query_num += 1
    matches = st.session_state.matches
    velocity = 1  # st.session_state.velocity
    scores = {}
    states = [
        st.session_state[f"input{i}"] for i in range(len(matches))
    ]
    for i, match in enumerate(matches):
        scores[match['id']] = float(states[i])
    # reset states to 1.0
    for i in range(len(matches)):
        st.session_state[f"input{i}"] = 1.0
    # get training data and labels
    X = list([match['vector'] for match in matches])
    y = [v for v in list(scores.values())]
    tune(X, y, iters=int(st.session_state.iters))
    # update record metadata after training
    for match in matches:
        st.session_state.meta[match['id']] = 1
    logging.info(f"Exclude List: {st.session_state.meta}")


def delete_element(element):
    del element


@st.experimental_singleton(show_spinner=False)
def init_clip_mlang():
    """ Initialize CLIP Model

    Returns:
        Tokenizer: CLIPTokenizerFast (which convert words into embeddings)
    """
    MODEL_ID = 'M-CLIP/XLM-Roberta-Large-Vit-B-32'
    clip = pt_multilingual_clip.MultilingualCLIP.from_pretrained(MODEL_ID)
    tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
    return tokenizer, clip


@st.experimental_singleton(show_spinner=False)
def init_clip_vanilla():
    """ Initialize CLIP Model

    Returns:
        Tokenizer: CLIPTokenizerFast (which convert words into embeddings)
    """
    MODEL_ID = "openai/clip-vit-base-patch32"
    tokenizer = CLIPTokenizerFast.from_pretrained(MODEL_ID)
    clip = CLIPModel.from_pretrained(MODEL_ID)
    return tokenizer, clip


@st.experimental_singleton(show_spinner=False)
def init_clip_rsicd():
    """ Initialize CLIP Model

    Returns:
        Tokenizer: CLIPTokenizerFast (which convert words into embeddings)
    """
    MODEL_ID = "flax-community/clip-rsicd"
    tokenizer = CLIPTokenizerFast.from_pretrained(MODEL_ID)
    clip = CLIPModel.from_pretrained(MODEL_ID)
    return tokenizer, clip


def prompt2vec_mlang(prompt: str, tokenizer, clip):
    """ Convert prompt into a computational vector

    Args:
        prompt (str): Text to be tokenized

    Returns:
        xq: vector from the tokenizer, representing the original prompt
    """
    out = clip.forward(prompt, tokenizer)
    xq = out.squeeze(0).cpu().detach().numpy().tolist()
    return xq


def prompt2vec_vanilla(prompt: str, tokenizer, clip):
    inputs = tokenizer(prompt, return_tensors='pt')
    out = clip.get_text_features(**inputs)
    xq = out.squeeze(0).cpu().detach().numpy().tolist()
    return xq


st.markdown("""
<link
  rel="stylesheet"
  href="https://fonts.googleapis.com/css?family=Roboto:300,400,500,700&display=swap"
/>
""", unsafe_allow_html=True)

messages = [
    f"""
    Find most relevant examples from a large visual dataset by combining text query and few-shot learning.
    """,
    f"""
    Then then you can adjust the weight on each image. Those weights should **represent how much it 
    can meet your preference**. You can either choose the images that match your prompt or change 
    your mind.

    You might notice that there is a iteration slide bar on the top of all retrieved images. This will
    control the speed of changes on vectors. More **iterations** will change the vector faster while 
    lower values on **iterations** will make the retrieval smoother.
    """,
    f"""
    This example will manage to train a classifier to distinguish between samples you want and samples 
    you don't want. By initializing the weight from prompt, you can get a good enough classifier to cluster
    images you want to search. If you think the result is not as perfect as you expected, you can also 
    supervise the classifer with **Relevance** annotation. If you cannot see any difference in Top-K 
    Retrieved results, try to enlarge **Number of Iteration**
    """,
    # TODO @ fangruil: fill the link with our tech blog
    f"""
    The app uses the [MyScale](http://mqdb.page.moqi.ai/mqdb-docs/) to store and query images 
    using vector search. All images are sourced from the 
    [Unsplash Lite dataset](https://unsplash-datasets.s3.amazonaws.com/lite/latest/unsplash-research-dataset-lite-latest.zip) 
    and encoded using [OpenAI's CLIP](https://huggingface.co/openai/clip-vit-base-patch32). We explain how
    it all works [here]().
    """
]

text_model_map = {
    'Multi Lingual': {'Vanilla CLIP': [prompt2vec_mlang, ]},
    'English': {'Vanilla CLIP': [prompt2vec_vanilla, ],
                'CLIP finetuned on RSICD': [prompt2vec_vanilla, ],
                }
}


with st.spinner("Connecting DB..."):
    st.session_state.meta, st.session_state.index = init_db()

with st.spinner("Loading Models..."):
    # Initialize CLIP model
    if 'xq' not in st.session_state:
        text_model_map['Multi Lingual']['Vanilla CLIP'].append(
            init_clip_mlang())
        text_model_map['English']['Vanilla CLIP'].append(init_clip_vanilla())
        text_model_map['English']['CLIP finetuned on RSICD'].append(
            init_clip_rsicd())
        st.session_state.query_num = 0

if 'xq' not in st.session_state:
    # If it's a fresh start
    if st.session_state.query_num < len(messages):
        msg = messages[st.session_state.query_num]
    else:
        msg = messages[-1]
    prompt = ''
    # Basic Layout
    with st.container():
        if 'prompt' in st.session_state:
            del st.session_state.prompt
        st.title("Visual Dataset Explorer")
        start = [st.empty(), st.empty(), st.empty(), st.empty(),
                 st.empty(), st.empty(), st.empty(), st.empty()]
        start[0].info(msg)
        start_col = start[1].columns(3)
        st.session_state.db_name_ref = start_col[0].selectbox(
            "Select Database:", list(db_name_map.keys()))
        st.session_state.lang = start_col[1].selectbox(
            "Select Language:", list(text_model_map.keys()))
        st.session_state.feat_name = start_col[2].selectbox("Select Image Feature:",
                                                            list(text_model_map[st.session_state.lang].keys()))
        if st.session_state.db_name_ref == "RSICD: Remote Sensing Images 11K":
            start[2].warning('If you are searching for Remote Sensing Images, \
                        try to use prompt "An aerial photograph of <your-real-query>" \
                        to obtain best search experience!')
        if len(prompt) > 0:
            st.session_state.prompt = prompt.replace(' ', '_')
        start[4].markdown(
            '<p style="color:gray;"> Don\'t know what to search? Try <b>Random</b>!</p>\
            <p>🌟 We also support multi-language search. Type any language you know to search! ⌨️ </p>',
            unsafe_allow_html=True)
        upld_model = start[6].file_uploader(
            "Or you can upload your previous run!", type='onnx')
        upld_btn = start[7].button(
            "Use Loaded Weights", disabled=upld_model is None)
        prompt = start[3].text_input(
            "Prompt:",
            value="An aerial photograph of "if st.session_state.db_name_ref == "RSICD: Remote Sensing Images 11K" else "",
            placeholder="Examples: playing corgi, 女人举着雨伞, mouette volant au-dessus de la mer, ガラスの花瓶の花 ...",)
        with start[5]:
            col = st.columns(8)
            has_no_prompt = (len(prompt) == 0 and upld_model is None)
            prompt_xq = col[6].button("Prompt", disabled=len(prompt) == 0)
            random_xq = col[7].button("Random", disabled=not (
                len(prompt) == 0 and upld_model is None))

    if random_xq:
        # Randomly pick a vector to query
        xq, orig_xq = init_random_query()
        st.session_state.xq = xq
        st.session_state.orig_xq = orig_xq
        _ = [elem.empty() for elem in start]
    elif prompt_xq or upld_btn:
        if upld_model is not None:
            # Import vector from a file
            import onnx
            from onnx import numpy_helper
            _model = onnx.load(upld_model)
            weights = _model.graph.initializer
            assert len(weights) == 1
            xq = numpy_helper.to_array(weights[0]).tolist()
            assert len(xq) == DIMS
            st.session_state.prompt = upld_model.name.split(".onnx")[
                0].replace(' ', '_')
        else:
            print(f"Input prompt is {prompt}")
            # Tokenize the vectors
            p2v_func, args = text_model_map[st.session_state.lang][st.session_state.feat_name]
            xq = p2v_func(prompt, *args)
        st.session_state.xq = xq
        st.session_state.orig_xq = xq
        _ = [elem.empty() for elem in start]

if 'xq' in st.session_state:
    # If it is not a fresh start
    if st.session_state.query_num+1 < len(messages):
        msg = messages[st.session_state.query_num+1]
    else:
        msg = messages[-1]
    # initialize classifier
    if 'clf' not in st.session_state:
        st.session_state.clf = Classifier(st.session_state.xq)

    # if we want to display images we end up here
    st.info(msg)
    # first retrieve images from pinecone
    st.session_state.matches, st.session_state.top_k = get_top_k(
        st.session_state.clf.get_weights(), top_k=9)

    # export the model into executable ONNX
    st.session_state.dnld_model = BytesIO()
    torch.onnx.export(torch.nn.Sequential(NormalizingLayer(), st.session_state.clf.model),
                      torch.as_tensor(st.session_state.xq).reshape(1, -1),
                      st.session_state.dnld_model,
                      input_names=['input'],
                      output_names=['output'])

    with st.container():
        with st.sidebar:
            with st.container():
                st.header("Top K Nearest in Database")
                for i, k in enumerate(st.session_state.top_k):
                    url = k["url"]
                    url += "?q=75&fm=jpg&w=200&fit=max"
                    if k["id"] not in BAD_IDS:
                        disabled = False
                    else:
                        disable = True
                    dist = np.matmul(st.session_state.clf.get_weights() / np.linalg.norm(st.session_state.clf.get_weights()),
                                     np.array(k["vector"]).T)
                    st.markdown(card_with_conf(i, dist, url),
                                unsafe_allow_html=True)
        dnld_nam = st.text_input('Download Name:',
                                 f'{(st.session_state.prompt if "prompt" in st.session_state else "model")}.onnx',
                                 max_chars=50)
        dnld_btn = st.download_button('Download your classifier!',
                                      st.session_state.dnld_model,
                                      dnld_nam,)
        # once retrieved, display them alongside checkboxes in a form
        with st.form("batch", clear_on_submit=False):
            st.session_state.iters = st.slider(
                "Number of Iterations to Update", min_value=0, max_value=10, step=1, value=2)
            col = st.columns([1, 9])
            col[0].form_submit_button("Train!", on_click=submit)
            col[1].form_submit_button(
                "Choose a new prompt", on_click=refresh_index)
            # we have three columns in the form
            cols = st.columns(3)
            for i, match in enumerate(st.session_state.matches):
                # find good url
                url = match["url"]
                url += "?q=75&fm=jpg&w=200&fit=max"
                if match["id"] not in BAD_IDS:
                    disabled = False
                else:
                    disable = True
                # the card shows an image and a checkbox
                cols[i % 3].markdown(card(i, url), unsafe_allow_html=True)
                # we access the values of the checkbox via st.session_state[f"input{i}"]
                cols[i % 3].slider(
                    "Relevance",
                    min_value=0.0,
                    max_value=1.0,
                    value=1.0,
                    step=0.05,
                    key=f"input{i}",
                    disabled=disabled
                )