# 3 модуль - Валидация и восстановление сущностей

from modules.NER import NER
from modules.paraGenerator import ParaphraseGenerator
from pymorphy3 import MorphAnalyzer

ner = NER()
pg = ParaphraseGenerator()
morph = MorphAnalyzer()

def compare_entities(original, generated):
    """
    Сравнивает два списка сущностей. Новые сущности допускаются.

    Использование: compare_entities(<сущности исходного текста>, <сущности сгенерированного текста>)

    Возвращает:
    - True - всё на месте
    - False - что то потерялось
    """
    def normalize(text):
        """
        Приводит слово к нормальной форме. Пример: Ивану --> Иван

        Возвращает:
        - Совпадают или нет сущности (bool)
        - Какие сущности потерялись (set)
        """
        return morph.parse(text)[0].normal_form

    if original.issubset(generated): # если original является подмножеством множества generated
        return True, set()
    
    # если нет, проверяем дополнительно, вдруг кейс по типу "Иван" - "Ивану"
    orig_norm = {normalize(e) for e in original} # нормализуем списки
    gen_norm = {normalize(e) for e in generated}

    if orig_norm.issubset(gen_norm):
        return True, set()
    
    # если по прежнему false, то ищем потерянные сущности
    lost = set()
    for o in original:
        if (normalize(o) not in gen_norm):
            lost.add(o)
    return False, lost

def validator(srcText, paraphrase, srcEntities, paraEntities):
    """
    Использование: validator(<исходный текст>, <парафразированный текст> <сущности исходного текста>, <сущности перефразированного текста>)

    Возвращает:
    - Исходный текст если сущности сохранены
    - Изменённый текст, если сущности не сохранены и были восстановлены
    - None, если сущности не удалось восстановить с трёх раз
    """
    ce = compare_entities(srcEntities, paraEntities)
    if ce[0]:
        return paraphrase # если всё нормально, возвращаем текст в неизменном виде
    
    print(f'Произошла потеря сущностей!')
    print(f'Исходные сущности: {srcEntities}')
    print(f'Сгенерируемые сущности: {paraEntities}')
    print(f'Потеря сущностей: {ce[1]}')
    # даём 5 попыток на восстановление
    for i in range(5):
        regen_prompt = (
        f'При перефразировании текста "{srcText}" из списка элементов "{', '.join(entity for entity in srcEntities)}"'
        f'были утеряны или изменены следующие важные элементы: "{', '.join(e for e in ce[1])}". '
         'Перефразируй исходный текст заново, обратив особое внимание на сохранение этих элементов. Выведи только текст.'
        )
        newParaphrase = pg.generateByPrompt(regen_prompt)
        paraEntities = ner.extract_entities(newParaphrase)
        ce = compare_entities(srcEntities, paraEntities)
        print(f'Попытка восстановления: {i+1}')
        print(f'Исходные сущности: {srcEntities}')
        print(f'Сгенерируемые сущности: {paraEntities}')
        print(f'Потеря сущностей: {ce[1]}')
        if (ce[0]): # если сравнение дало True, выходим из цикла
            return newParaphrase
    return None

# a = set(['a','b', 'c', 'd'])
# b = set(['a','b'])
# validator('123', '123', a, b)