import ast
import pandas as pd
from tqdm import tqdm
from ruamel import yaml


class nlu_generator():
    def __init__(self, mode, med_dataset):
        super().__init__()
        df = pd.read_csv(med_dataset)
        if mode == 'drug':
            col = 'drug_name'
            aka_col = 'Brand names'
        else:
            col = 'Lab test'
            aka_col = 'Also Known As'
        self.list = df[col].str.lower()
        aka_list = df[aka_col].str.lower()
        self.aka_list = aka_list.str.replace('®', '')
        self.mode = mode

    # Synonym key generator
    def __syn_generator(self, item, aka_list):
        inp = f"""\
nlu:
- synonym: {item}
  examples: |
"""
        code = yaml.load(inp, Loader=yaml.RoundTripLoader)
        if self.mode == 'lab':
            aka_list = ast.literal_eval(aka_list)
        else:
            aka_list = [aka_list]
        for aka in aka_list:
            code['nlu'][0]['examples'] += f'- {aka}\n'
        return code

    # First keys generator
    def __block_generator(self):
        inp = f"""\
version: "2.0"

nlu:
"""
        code = yaml.load(inp, Loader=yaml.RoundTripLoader)
        code['nlu'] = self.__syn_generator(
            self.list[0], self.aka_list[0])['nlu']
        return code

    # NaN check
    def __isnan(self, num):
        isnan = (num != num)
        return isnan

    # YAML generator
    def generate(self):
        # Generate first keys
        code = self.__block_generator()
        for i, item in tqdm(enumerate(self.list[1:], start=1)):
            # Check if aka_list for that item exists or not
            if not self.__isnan(self.aka_list[i]):
                # Append aka_list to YAML
                code['nlu'].append(self.__syn_generator(
                    item, self.aka_list[i])['nlu'][0])
        return code

    # Write data in the disk
    def write_data(self, code, nlu_file):
        with open(nlu_file, 'w') as f:
            yaml.dump(code, f, Dumper=yaml.RoundTripDumper)


def main():
    mode = 'lab'
    #mode = 'drug'
    generator = nlu_generator(
        mode, '../labtestonline_datasets/dataset_files/labtest_dataset.csv')
    #generator = nlu_generator(
        #mode, '../medlineplus_drug_dataset/dataset_files/MedlinePlus_2.csv')
    code = generator.generate()
    generator.write_data(code, 'data/synonym_lab.yml')
    #generator.write_data(code, 'data/synonym_drug.yml')


if __name__ == "__main__":
    main()