import pandas as pd
from sentence_transformers import SentenceTransformer
import faiss
import numpy as np
import streamlit as st
import requests
from PIL import Image
from io import BytesIO


st.title('Рекомендации сериалов по описанию пользователя')
st.divider()
st.write('Особенности: семантический поиск производится по косинусному сходству')
df = pd.read_csv('clean_series_data.csv')
embeddings = np.load('embeddings.npy')


def load_image_from_url(url):
    try:
        response = requests.get(url)
        response.raise_for_status()
        return Image.open(BytesIO(response.content))
    except Exception as e:
        st.error(f"Не удалось загрузить изображение: {e}")
        return None


model = SentenceTransformer('cointegrated/rubert-tiny2')
model.cpu()
# embeddings_desc = df['Описание'].apply(lambda x: model.encode(x))
# embeddings_gan = df['Жанры'].apply(lambda x: model.encode(x))

# embeddings = embeddings_desc + embeddings_gan

embeddings = np.array(embeddings).astype(np.float32)
faiss.normalize_L2(embeddings)

dimension = embeddings.shape[1]

index = faiss.IndexFlatIP(dimension)

index.add(embeddings)

query = [st.text_area('Введите описание сериала')]


button = st.button('Вывести результаты')
k = st.slider('Сколько сериалов рекомендовать?',
                      min_value=1, max_value=10, value=3, step=1)
if button:
    if query:
        query_embedding = model.encode(query).astype(np.float32)
        # Две строки ниже можно будет убрать
        # query_embedding = np.array(
        #     query_embedding, dtype=np.float32).reshape(1, -1)
        # faiss.normalize_L2(query_embedding)

        distances, indices = index.search(query_embedding, k)

        st.subheader('Похожие сериалы:')
        for i in range(k):
            url = df.loc[indices[0][i]]["Изображение"]
            image = load_image_from_url(url)
            st.image(image)
            st.write(f'Название: {df.loc[indices[0][i]]["Название"]}')
            st.write(f'Рейтинг: {df.loc[indices[0][i]]["Рейтинг"]}')
            st.write(f'Страна: {df.loc[indices[0][i]]["Страна"]}')
            st.write(
                f'Длительность одной серии: {df.loc[indices[0][i]]["Длительность"]}')
            st.write(
                f'Количество серий: {df.loc[indices[0][i]]["Количество серий"]}')
            st.write(f'Описание: {df.loc[indices[0][i]]["Описание"]}')
            st.write(f'Косинусное сходство: {distances[0][i]:.4f}')
            st.divider()