import os
import time

import pandas as pd
import torch

from datasets import Dataset, load_dataset
from sentence_transformers import util
from transformers import AutoTokenizer, AutoModel

import gradio as gr

device = torch.device('cpu')

# Helpers
def get_model_size(model):
    param_size = 0
    for param in model.parameters():
        param_size += param.nelement() * param.element_size()
    buffer_size = 0
    for buffer in model.buffers():
        buffer_size += buffer.nelement() * buffer.element_size()
    return (param_size + buffer_size) / 1024**2

# Load model
checkpoint = 'sberbank-ai/sbert_large_mt_nlu_ru'
tokenizer = AutoTokenizer.from_pretrained(checkpoint)
model = AutoModel.from_pretrained(checkpoint)
model = model.to(device)

def mean_pooling(token_embeddings, attention_mask):
    input_mask_expanded = attention_mask.unsqueeze(-1).expand(token_embeddings.size()).float()
    sum_embeddings = torch.sum(token_embeddings * input_mask_expanded, 1)
    sum_mask = torch.clamp(input_mask_expanded.sum(1), min=1e-9)
    return sum_embeddings / sum_mask

def get_embeddings(input):
    encoded_input = tokenizer(input, padding=True, truncation=True, max_length=50, return_tensors='pt').to(device)
    with torch.no_grad():
        model_output = model(**encoded_input)
    return mean_pooling(model_output[0], encoded_input['attention_mask']).cpu().numpy()

# Load data
ds_name = 'AresEkb/prof_standards_sbert_large_mt_nlu_ru'
domains_ds = load_dataset(ds_name, 'domains', split='train').add_faiss_index(column='embeddings')
generalized_functions_ds = load_dataset(ds_name, 'generalized_functions', split='train').add_faiss_index(column='embeddings')
jobs_ds = load_dataset(ds_name, 'jobs', split='train').add_faiss_index(column='embeddings')
particular_functions_ds = load_dataset(ds_name, 'particular_functions', split='train').add_faiss_index(column='embeddings')
actions_ds = load_dataset(ds_name, 'actions', split='train').add_faiss_index(column='embeddings')
skills_ds = load_dataset(ds_name, 'skills', split='train').add_faiss_index(column='embeddings')
knowledges_ds = load_dataset(ds_name, 'knowledges', split='train').add_faiss_index(column='embeddings')

indices = {'reg_number', 'generalized_function_id', 'particular_function_id'}

entity_kinds = {
    'Предметная область': domains_ds,
    'Процесс': generalized_functions_ds,
    'Подпроцесс': particular_functions_ds,
    'Функция': actions_ds,
    'Должность': jobs_ds,
    'Навык': skills_ds,
    'Знание': knowledges_ds,
}

# Main search logic
def search(context_entity_kind, context_entity_name, context_entity_count,
           target_entity_kind, target_entity_name, target_entity_count):
    # Find similar context entities
    start_time = time.perf_counter_ns()
    context_ds = entity_kinds[context_entity_kind]
    context_embedding = get_embeddings(context_entity_name)
    scores, samples = context_ds.get_nearest_examples(
        'embeddings', context_embedding, k=context_entity_count
    )
    cos_scores = util.cos_sim(context_embedding, samples['embeddings'])[0]
    cos_scores = [round(x, 4) for x in cos_scores.tolist()]
    results = pd.DataFrame({'name': samples['name'], 'score': cos_scores}).sort_values('score', ascending=False)
    search_time = round((time.perf_counter_ns() - start_time) / 10**6)

    # Get related entities
    start_time = time.perf_counter_ns()
    context_df = pd.DataFrame(samples).drop(columns=['embeddings']).rename(columns={'name': 'context_name'})
    context_df['context_score'] = cos_scores
    target_df = entity_kinds[target_entity_kind].to_pandas()
    common_indices = list(indices.intersection(context_df.columns).intersection(target_df.columns))
    target_ds = Dataset.from_pandas(context_df.merge(target_df, on=common_indices))
    target_ds.add_faiss_index(column='embeddings')

    # Find similar target entities
    target_embedding = get_embeddings(target_entity_name)
    scores, samples = target_ds.get_nearest_examples(
        'embeddings', target_embedding, k=target_entity_count
    )
    cos_scores = util.cos_sim(target_embedding, samples['embeddings'])[0]
    cos_scores = (cos_scores + torch.tensor(samples['context_score'])) / 2
    cos_scores = [round(x, 4) for x in cos_scores.tolist()]
    results2 = pd.DataFrame({'name': samples['name'], 'context_name': samples['context_name'], 'score': cos_scores}).sort_values('score', ascending=False)
    search_time2 = round((time.perf_counter_ns() - start_time) / 10**6)

    return [results.to_numpy(), search_time, results2.to_numpy(), search_time2]

# User Interface
ui = gr.Interface(
    search,
    [
        gr.Radio(label='Тип объекта', choices=list(entity_kinds.keys()), value='Функция'),
        gr.Textbox(label='Название объекта'),
        gr.Slider(1, 20, 10, step=1, label='Кол-во объектов'),
        gr.Radio(label='Тип связанного объекта', choices=list(entity_kinds.keys()), value='Должность'),
        gr.Textbox(label='Название связанного объекта'),
        gr.Slider(1, 20, 10, step=1, label='Кол-во связанных объектов'),
    ],
    [
        gr.Dataframe(label='Похожие объекты', headers=['Название', 'Сходство'], datatype=['str', 'number']),
        gr.Textbox(label='Время поиска, миллисекунды'),
        gr.Dataframe(label='Похожие связанные объекты', headers=['Название', 'Контекст', 'Сходство']),
        gr.Textbox(label='Время поиска, миллисекунды'),
    ],
    allow_flagging='never',
    live=True,
    examples=[
        ['Функция', 'проектирование базы данных', 7, 'Должность', '', 7],
        ['Функция', 'написать руководство пользователя', 7, 'Должность', '', 7],
        ['Должность', 'программист', 12, 'Процесс', '', 7],
    ],
    title='Поиск по профстандартам',
    description='''Выберите тип объектов, который вы хотите найти, введите его название.
        Опционально укажите какие связанные объекты вы хотите найти.''',
    article=f'''<p>Поиск выполняется по
           <a href="https://profstandart.rosmintrud.ru/obshchiy-informatsionnyy-blok/natsionalnyy-reestr-professionalnykh-standartov/reestr-professionalnykh-standartov/">реестру</a>
           профессиональных стандартов минтруда.</p>
        <p>В базе есть следующие данные:</p>
        <table>
          <tr><th>Тип объектов</th><th>Кол-во</th></tr>
          <tr><td>Предметные области</td><td>{domains_ds.num_rows}</td></tr>
          <tr><td>Процессы</td><td>{generalized_functions_ds.num_rows}</td></tr>
          <tr><td>Подпроцессы</td><td>{particular_functions_ds.num_rows}</td></tr>
          <tr><td>Функции</td><td>{actions_ds.num_rows}</td></tr>
          <tr><td>Должности</td><td>{jobs_ds.num_rows}</td></tr>
          <tr><td>Навыки</td><td>{skills_ds.num_rows}</td></tr>
          <tr><td>Знания</td><td>{knowledges_ds.num_rows}</td></tr>
        </table>
        <p>Для вычисления векторных представлений используется следующая модель:</p>
        <table>
          <tr><th>Характеристика модели</th><th>Значение</th></tr>
          <tr><td>Модель</td><td><a href="https://huggingface.co/{checkpoint}">{checkpoint}</a></td></tr>
          <tr><td>Размер, Мб</td><td>{round(get_model_size(model))}</td></tr>
          <tr><td>Количество параметров, миллионы</td><td>{round(model.num_parameters()/10**6)}</td></tr>
          <tr><td>Размерность векторных представлений</td><td>{get_embeddings('').shape[1]}</td></tr>
        </table>
    ''',
    css='.w-full .col:nth-child(2) { flex-grow: 2 !important; }')

ui.launch()