import re
from ruamel import yaml
import numpy as np
import sys


def load_entities(entity_names):
    entities_dict = dict()
    for entity in entity_names:
        entity_file_adr = f'./{entity}.yml'
        entities_dict[entity] = load_entities_list(entity_file_adr)
    return entities_dict
 

def load_entities_list(file_adr):
    with open(file_adr, 'r') as f:
        entities = yaml.load(f, Loader=yaml.RoundTripLoader)
    entities = entities['nlu'][0]
    entities = entities['examples'].split('\n')
    entities = [row[2:] for row in entities]
    return entities


def randomize_entity_names(nlu_dict, entities_dict):
    for intent in nlu_dict['nlu']:
        examples = ''
        for example in intent['examples'].split('\n'):
            entity_match = dict()
            for entity in entities_dict:
                entity_match = re.search(f'\[[^\]]*\]\({entity}\)', example)
                if entity_match:
                    start, end = entity_match.span()
                    random_entity_name = f'[{np.random.choice(entities_dict[entity]).strip()}]({entity})'
                    example = example.replace(example[start:end] , random_entity_name)
                    
            examples += example + '\n'
        intent['examples'] = examples[:-1]  #removing last \n to avoid \n\n after the last example
    return nlu_dict



if __name__=='__main__':
    
    NLU_FILE = './nlu_cleaned.yml'
    ENTITY_NAMES = ['drug', 'lab'] #lookup files should be in same directoty as this file
    OUTPUT_FILE = 'nlu_random.yml'
    #load files
    with open(NLU_FILE, 'r') as f:
        nlu = yaml.load(f, Loader=yaml.RoundTripLoader)
    entities = load_entities(ENTITY_NAMES)

    randomized_nlu = randomize_entity_names(nlu, entities)
    #save
    with open(OUTPUT_FILE, 'w') as f:
        yaml.dump(randomized_nlu, f, Dumper=yaml.RoundTripDumper, default_flow_style=None)