import os
import time

import pandas as pd
import torch

from datasets import 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')

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

# Main logic
def calc(entity_kind, entity1_name, entity2_name, entity3_name, entity_count):
    start_time = time.perf_counter_ns()

    if not entity1_name or not entity2_name and not entity3_name:
        return [None, 0]

    embedding1 = get_embeddings(entity1_name)
    embedding2 = get_embeddings(entity2_name)
    embedding3 = get_embeddings(entity3_name)

    if entity2_name and entity3_name:
        embedding = embedding1 - embedding2 + embedding3
    elif entity2_name:
        embedding = (embedding1 + embedding2) / 2
    else:
        embedding = (embedding1 + embedding3) / 2

    context_ds = entity_kinds[entity_kind]
    scores, samples = context_ds.get_nearest_examples(
        'embeddings', embedding, k=entity_count
    )
    cos_scores = util.cos_sim(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)

    return [results.to_numpy(), search_time]

# User interface
ui = gr.Interface(
    calc,
    [
        gr.Radio(label='Тип искомого объекта', choices=list(entity_kinds.keys()), value='Функция'),
        gr.Textbox(label='Исходный объект'),
        gr.Textbox(label='Вычитаемый объект'),
        gr.Textbox(label='Добавляемый объект'),
        gr.Slider(1, 20, 10, step=1, label='Кол-во результатов'),
    ],
    [
        gr.Dataframe(label='Результат вычисления', headers=['Название', 'Сходство'], datatype=['str', 'number'], wrap=True),
        gr.Textbox(label='Время поиска, миллисекунды'),
    ],
    allow_flagging='never',
    examples=[
        ['Функция', 'проектирование баз данных', 'база данных', 'интерфейс', 10],
        ['Знание', 'управление проектами', 'проекты', 'производство', 10],
        ['Функция', 'проектирование баз данных', 'проектирование баз данных', 'диагностика оборудования', 10],
        ['Функция', 'проектирование автомобиля', 'управление автомобилем', 'управление данными', 10],
        ['Функция', 'оценка данных', 'управление данными', 'управление персоналом', 10],
        ['Функция', 'проектирование баз данных', 'диагностика оборудования', '', 10],
    ],
    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()