import streamlit as st
from sklearn.metrics.pairwise import cosine_similarity
from sentence_transformers import SentenceTransformer
from transformers import CLIPTokenizer, CLIPModel
from scipy.spatial import distance
import numpy as np
import torch
import json
import random

model_id = "openai/clip-vit-base-patch32"
clip_model = CLIPModel.from_pretrained(model_id)
clip_tokenizer = CLIPTokenizer.from_pretrained(model_id)
device = "cuda" if torch.cuda.is_available() else "cpu"
# move the model to the device
clip_model.to(device)

def get_single_text_embedding(text):
  inputs = clip_tokenizer(text, return_tensors = "pt")
  text_embeddings = clip_model.to(device).get_text_features(**inputs.to(device))
  embedding_as_np = text_embeddings.cpu().detach().numpy()
  return embedding_as_np

def get_top_N_images(query, image_vectors, top_K=4):
    query_vect = get_single_text_embedding(query)

    data = cosine_similarity(query_vect, image_vectors)[0]

    most_similar_images = sorted(range(len(data)), key=lambda i: data[i], reverse=True)[:top_K]

    return most_similar_images


image_data = []

with open('Property-images.jsonl', 'r', encoding="utf-8") as f:
    for line in f:
        image_data.append(json.loads(line))

with open('property-image-vectors.npy', 'rb') as f:
    image_vectors = np.load(f)



def combine(data):
    new_data = []

    window = 4  # number of sentences to combine
    stride = 1  # number of sentences to 'stride' over, used to create overlap

    for i in range(0, len(data), stride):
        i_end = min(len(data)-1, i+window)
        if data[i]['title'] != data[i_end]['title']:
            # in this case we skip this entry as we have start/end of two videos
            continue
        text = ' '
        
        for x in data[i:i_end]:
          text += x['text']
        new_data.append({
            'start': data[i]['start'],
            'end': data[i_end]['end'],
            'title': data[i]['title'],
            'text': text,
            'id': data[i]['id'],
            'url': data[i]['url'],
            'published': data[i]['published']
        })
    
    return new_data


model_id = "multi-qa-mpnet-base-dot-v1"
model = SentenceTransformer(model_id)


meta_data = []

with open('Property-transcription.jsonl', 'r', encoding="utf-8") as f:
    for line in f:
        meta_data.append(json.loads(line))

meta_data = combine(meta_data)
    
with open('property-vectors.npy', 'rb') as f:
    text_vector = np.load(f)


def card(thumbnail: str, title: str, urls: list, contexts: list, starts: list, ends: list):
    meta = [(e, s, u, c) for e, s, u, c in zip(ends, starts, urls, contexts)]
    meta.sort(reverse=False)
    text_content = []
    current_start = 0
    current_end = 0
    for end, start, url, context in meta:
        # reformat seconds to timestamp
        time = start / 60
        mins = f"0{int(time)}"[-2:]
        secs = f"0{int(round((time - int(mins))*60, 0))}"[-2:]
        timestamp = f"{mins}:{secs}"
        if start < current_end and start > current_start:
            # this means it is a continuation of the previous sentence
            text_content[-1][0] = text_content[-1][0].split(context[:10])[0]
            text_content.append([f"[{timestamp}] {context.capitalize()}", url])
        else:
            text_content.append(["xxLINEBREAKxx", ""])
            text_content.append([f"[{timestamp}] {context}", url])
        current_start = start
        current_end = end
    html_text = ""
    for text, url in text_content:
        if text == "xxLINEBREAKxx":
            html_text += "<br>"
        else:
            html_text += f"<small><a href={url}>{text.strip()}... </a></small><br>"
    print(html_text)
    html = f"""
    <div class="container-fluid">
        <div class="row align-items-start">
            <div class="col-md-4 col-sm-4">
                <div class="position-relative">
                    <a href={urls[0]}><img src={thumbnail} class="img-fluid" style="width: 192px; height: 106px"></a>
                </div>
            </div>
            <div  class="col-md-8 col-sm-8">
                <h2>{title}</h2>
            </div>
        <div>
            {html_text}
    <br><br>
    """
    return st.markdown(html, unsafe_allow_html=True)

    
st.write("""
# Property Video Search 🏘️
Utilize AI to quickly locate amenities and features within your residence.🤗
""")

st.markdown("""
<link rel="stylesheet" href="https://cdn.jsdelivr.net/npm/bootstrap@4.0.0/dist/css/bootstrap.min.css" integrity="sha384-Gn5384xqQ1aoWXA+058RXPxPg6fy4IWvTNh0E263XmFcJlSAwiGgFAW/dAiS6JXm" crossorigin="anonymous">
""", unsafe_allow_html=True)

query = st.text_input("Search!", "")

if query != "":

    vector1 = model.encode(query)
    text_cosines = dict()
    for i,vec in enumerate(text_vector):
        text_cosines[str(i)+'-text'] = 1 - distance.cosine(vector1, vec)
    

    if random.randint(0, 1):
        vector1 = get_single_text_embedding(query)
        image_cosines = dict()
        for i,vec in enumerate(image_vectors):
            image_cosines[str(i)+'-image'] = 1 - distance.cosine(vector1.reshape((512,)), vec)
        
        text_cosines.update(image_cosines)

    sorted_cosines = sorted(text_cosines.items(), key=lambda x:x[1], reverse=True)
    converted_dict = dict(sorted_cosines)
    
    
    results = {}
    order = []

    for vec_index in list(converted_dict.keys())[:7]:
        idx = int(vec_index.split('-')[0])
        video_id = image_data[idx]['url'].split('/')[-1]

        if vec_index.split('-')[-1] == 'image':
            if video_id not in results:
                results[video_id] = {
                    'title': image_data[idx]['title'],
                    'urls': [f"{image_data[idx]['url']}?t={int(image_data[idx]['sec'])}"],
                    'contexts': ['Image-query'],
                    'starts': [int(image_data[idx]['sec'])],
                    'ends': [int(image_data[idx]['sec']+6)]
                }
                order.append(video_id)
            else:
                results[video_id]['urls'].append(
                    f"{image_data[idx]['url']}?t={int(image_data[idx]['sec'])}"
                )
                results[video_id]['contexts'].append('Image-query')
                results[video_id]['starts'].append(int(image_data[idx]['sec']))
                results[video_id]['ends'].append(int(image_data[idx]['sec']+6))

        elif vec_index.split('-')[-1] == 'text':
            if video_id not in results:
                results[video_id] = {
                    'title': meta_data[idx]['title'],
                    'urls': [f"{meta_data[idx]['url']}?t={int(meta_data[idx]['start'])}"],
                    'contexts': [meta_data[idx]['text']],
                    'starts': [int(meta_data[idx]['start'])],
                    'ends': [int(meta_data[idx]['end'])]
                }
                order.append(video_id)
            else:
                results[video_id]['urls'].append(
                    f"{meta_data[idx]['url']}?t={int(meta_data[idx]['start'])}"
                )
                results[video_id]['contexts'].append(
                    meta_data[idx]['text']
                )
                results[video_id]['starts'].append(int(meta_data[idx]['start']))
                results[video_id]['ends'].append(int(meta_data[idx]['end']))

    # now display cards
    for video_id in order:
        card(
            thumbnail=f"https://img.youtube.com/vi/{video_id}/maxresdefault.jpg",
            title=results[video_id]['title'],
            urls=results[video_id]['urls'],
            contexts=results[video_id]['contexts'],
            starts=results[video_id]['starts'],
            ends=results[video_id]['ends']
        )
else:
    st.warning('💡Try searching: huge balcony, swimming pool, spiral staircase, panoramic view ... etc')