import gradio as gr
from transformers import pipeline

NER_MODEL_NAME = "Babelscape/wikineural-multilingual-ner"


def create_ner_pipeline():
    return pipeline(
        "ner",
        model=NER_MODEL_NAME,
        aggregation_strategy="simple",
    )


ner = create_ner_pipeline()


def normalize_entity_label(label: str) -> str:
    label = label.replace("B-", "").replace("I-", "").upper()

    label_mapping = {
        "FIRST_NAME": "PER",
        "LAST_NAME": "PER",
        "MIDDLE_NAME": "PER",
        "PERSON": "PER",
        "PER": "PER",

        "CITY": "LOC",
        "COUNTRY": "LOC",
        "LOCATION": "LOC",
        "LOC": "LOC",

        "ORGANIZATION": "ORG",
        "COMPANY": "ORG",
        "ORG": "ORG",
    }

    return label_mapping.get(label, label)


def analyze(text: str) -> str:
    if not text or not text.strip():
        return "Введите текст для анализа"

    results = ner(text)

    if not results:
        return "Сущности не найдены"

    answer = []

    for result in results:
        word = result["word"]
        entity_group = normalize_entity_label(result["entity_group"])
        score = round(float(result["score"]), 4)

        answer.append(f"{word} {entity_group} — уверенность {score}")

    return "\n".join(answer)


def run_gradio_ner_test():
    demo = gr.Interface(
        fn=analyze,
        inputs=gr.Textbox(
            label="Текст для анализа",
            placeholder="Напишите текст, в котором нужно определить персону или локацию",
            lines=3,
        ),
        outputs=gr.Textbox(label="Результат"),
        title="Определение персон и локаций",
        description="Введите текст на русском языке — модель определит персону и/или локацию",
        examples=[
            ["Яндекс основан Аркадием Воложем в Москве в 1997 году."],
            ["Иван Павлов родился в Рязани и учился в Санкт-Петербургском университете."],
            ["Т-Банк запустил сервис для клиентов в Екатеринбурге."],
        ],
    )

    demo.launch()


if __name__ == "__main__":
    run_gradio_ner_test()