import pandas as pd
import os
from glob import glob


def read_vocab(path):
    data = []
    with open(path, encoding='utf-8') as input_stream:
        for line in input_stream:
            data.append({'label': line.split('||')[0], 'text':line.strip().split('||')[1]})
    return pd.DataFrame(data)


def read_annotation_file(ann_file_path):
    data = []
    with open(ann_file_path, encoding='utf-8') as input_stream:
        for row_id, line in enumerate(input_stream):
            splitted_line = line.strip().split('||')
            mention = splitted_line[-2]
            concept_id = splitted_line[-1]
            data.append({'entity_text': mention, 'label':concept_id})
    return data


def read_dataset(dataset_folder):
    ann_file_pattern = os.path.join(dataset_folder, '*.concept')
    dataset = []
    for ann_file_path in glob(ann_file_pattern):
        dataset += read_annotation_file(ann_file_path)
    return dataset


def process_dataset(test_dataset, train_dataset):
    train_df = pd.DataFrame(train_dataset)
    test_df = pd.DataFrame(test_dataset)
    refined_set = test_df[~test_df.entity_text.isin(train_df.entity_text)]
    return refined_set.drop_duplicates().to_dict('records')


def save_dataset(dataset, path):
    if not os.path.exists(path): os.mkdir(path)
    fpath = os.path.join(path, '0.concept')
    with open(fpath, 'w', encoding='utf-8') as output_stream:
        for entity in dataset:
            output_stream.write(f"-1||0|0||Disease||{entity['entity_text']}||{entity['label']}\n")
