import gradio as gr
from transformers import pipeline, BertTokenizer, BertForSequenceClassification
import os
import pickle
import pandas as pd
from service_dops_api.dops_config import ServiceDopsConfig
from service_dops_api.dops_classifier import DopsClassifier
import json
HF_TOKEN = os.getenv('HF_TOKEN')
tokenizer_cat = BertTokenizer.from_pretrained("etadevosyan/service_categorizer_v2",
                                          token=HF_TOKEN)
model_cat = BertForSequenceClassification.from_pretrained('etadevosyan/service_categorizer_v2',token=HF_TOKEN)

clf_cat = pipeline("text-classification", model=model_cat, tokenizer=tokenizer_cat)


tokenizer_spec = BertTokenizer.from_pretrained("etadevosyan/specialists_categorizer_model",
                                          token=HF_TOKEN)
model_spec = BertForSequenceClassification.from_pretrained('etadevosyan/specialists_categorizer_model',token=HF_TOKEN)

clf_spec = pipeline("text-classification", model=model_spec, tokenizer=tokenizer_spec)

id2label_spec = pd.read_pickle('id2label_spec_categoriser.pickle')
def categoriser_predict(input_text):
    predictions = clf_cat(input_text)
    text_label = predictions[0]['label']
    return text_label
def doctor_spec_predict(input_text):
    predictions = clf_spec(input_text)
    numeric_label = int(predictions[0]['label'].split("_")[1])
    text_label = id2label_spec[numeric_label]
    return text_label
def dops_predict(input_text):
    cfg = ServiceDopsConfig()
    model = DopsClassifier(config=cfg)
    result = model.run_all_dops(input_text)
    return result
def convert_dops_to_dif_fields(dops_result):
    dops_values = []
    for dop in dops_result:
        temp_values = ','.join(dop['values'])
        dops_values.append(temp_values)
    return dops_values
def service_pipeline(input_text):
    categoriser_result = categoriser_predict(input_text)
    if categoriser_result!='Консультация специалиста':
        return 'Эта услуга не относится к приему специалиста','-','-','-','-','-','-'
    else:
        doctor_spec_result = doctor_spec_predict(input_text)
        dops_result = dops_predict(input_text)
        dops_values = convert_dops_to_dif_fields(dops_result)
        return categoriser_result,doctor_spec_result,dops_values[0],dops_values[1],dops_values[2],dops_values[3],dops_values[4]
demo = gr.Interface(fn=service_pipeline,inputs=gr.components.Textbox(label='Название услуги'),
                    outputs=[gr.components.Textbox(label='Относится ли данная услуга к приёму специалиста'),
                             gr.components.Textbox(label='Специальность врача'),
                             gr.components.Textbox(label='Место оказания услуги'),
                             gr.components.Textbox(label='Учёная степень'),
                             gr.components.Textbox(label='Возрастная категория'),
                             gr.components.Textbox(label='Вид приёма'),
                             gr.components.Textbox(label='Расстояние в км от М(КАД)')],
                    examples=[
                        ['Врач-офтальмолог (высшая категория/кандидат медицинских наук), первичный приём'],
                        ['Прием (осмотр, консультация) - врача -оториноларинголога Первичный, рекомендации'],
                        ['Прием врача специалиста ЛОР'],
                        ['Прием (осмотр, консультация) врача-терапевта на дому, повторный (в пределах 5 км от М(КАД)/административной границы города)']])

if __name__ == "__main__":
    demo.launch()