import streamlit as st
import pandas as pd
import numpy as np
import torch
from sentence_transformers import SentenceTransformer
from resources.functions import recommend, find_rows_with_genres, get_mask_in_range


st.markdown(f"<h1 style='text-align: center;'>Семантический поиск фильмов", unsafe_allow_html=True)

df = pd.read_csv('resources/DF_FINAL.csv')
genre_lists = df['ganres'].apply(lambda x: x.split(', ') if isinstance(x, str) else [])
all_genres = sorted(list(set([genre for sublist in genre_lists for genre in sublist])))

st.write(f'<p style="text-align: center; font-family: Arial, sans-serif; font-size: 20px; color: white;">Количество фильмов \
         для поиска {len(df)}</p>', unsafe_allow_html=True)

st.header(':wrench: Панель инструментов')
col1, col2, col3 = st.columns([1, 2, 2]) 
with col1:
    top_k = st.selectbox("Сколько фильмов?", options=[5, 10, 15, 20])
with col2:
    model_type = st.selectbox("Какой моделью пользуемся?\n ", options=['rubert-tiny2', 'msmarco-MiniLM-L-12-v3'])
with col3:
    genres_list = st.multiselect("Какого жанра фильм?\n ", options=all_genres)

if model_type == 'rubert-tiny2':
    model = SentenceTransformer('cointegrated/rubert-tiny2')
    emb = torch.load('resources/corpus_embeddings_rub.pth', map_location=torch.device('cpu'))
else:
    model = SentenceTransformer('msmarco-MiniLM-L-12-v3')
    emb = torch.load('resources/corpus_embeddings_ms.pth', map_location=torch.device('cpu'))

range_years = st.slider("В каком году вышел фильм?", min_value=df['year'].unique().min(), 
                         max_value=df['year'].unique().max(), 
                         value=(df['year'].unique().min(), df['year'].unique().max()))

text = st.text_input('Что будем искать?')
button = st.button('Начать поиск', type="primary")
    
if text and button:
    if len(genres_list) == 0:
        mask = get_mask_in_range(df=df, range_values=range_years)
    else:
        mask1 = find_rows_with_genres(df=df, genres_list=genres_list)
        mask2 = get_mask_in_range(df=df, range_values=range_years)
        mask = mask1 & mask2
    try:
        # emb = emb[mask]
        # df = df[mask]
        hits = recommend(model, text, emb, len(df))
        st.write(f'<p style="font-family: Arial, sans-serif; font-size: 24px; color: pink; font-weight: bold;"><strong>\
            {top_k} лучших рекомендаций</strong></p>', unsafe_allow_html=True)
        st.write('\n')
        mask_ind = df[mask].index.tolist()
        fil_hits = [hits[0][i] for i in range(len(hits[0])) if hits[0][i]['corpus_id'] in mask_ind]
        for i in range(top_k):
            col4, col5 = st.columns([3, 4]) 
            with col4:
                try:
                    st.image(df['poster'][fil_hits[i]['corpus_id']], width=300)
                except:
                    st.image('https://cdnn11.img.sputnik.by/img/104126/36/1041263627_235:441:1472:1802_1920x0_80_0_0_fc2acc893b618b7c650d661fafe178b8.jpg', width=300)
            with col5:
                st.write(f"***Название:*** {df['title'][fil_hits[i]['corpus_id']]}")
                st.write(f"***Жанр:*** {(df['ganres'][fil_hits[i]['corpus_id']])}")
                st.write(f"***Описание:*** {df['description'][fil_hits[i]['corpus_id']]}")
                st.write(f"***Год:*** {df['year'][fil_hits[i]['corpus_id']]}")
                st.write(f"***Актерский состав:*** {df['cast'][fil_hits[i]['corpus_id']]}")
                st.write(f"***Косинусное сходство:*** {round(fil_hits[i]['score'], 2)}")
                st.write(f"***Ссылка на фильм : {df['url'][fil_hits[i]['corpus_id']]}***")

            st.markdown(
            "<hr style='border: 2px solid #000; margin-top: 10px; margin-bottom: 10px;'>",
            unsafe_allow_html=True
            )
    except:
        message = '<p style="font-family: Arial, sans-serif; font-size: 24px; color: pink; font-weight: bold;"><strong>\
            Подходящих вариантов нет. Измените критерии поиска или отключайте интернет и читайте\
            <a href="https://huggingface.co/spaces/sakoser/rec_sys_books">книги</a>.</strong></p>'
        st.write(message, unsafe_allow_html=True)