import csv
import os
from tqdm import tqdm
from importlib.resources import open_text
import json
from transformer_deid.label import Label
from transformer_deid.data import DeidDataset
from transformer_deid.tokenization import split_sequences, assign_tags, encode_tags


def load_label(filename):
    """
    Loads annotations from a CSV file.
    CSV file should have entity_type/start/stop columns.
    """
    # can change this, make it a variable? options: "base", "hipaa", "binary"
    label_map = get_label_map('base')

    with open(filename, 'r') as fp:
        csvreader = csv.reader(fp, delimiter=',', quotechar='"')
        header = next(csvreader)
        # identify which columns we want
        idx = [
            header.index('entity_type'),
            header.index('start'),
            header.index('stop'),
            header.index('entity')
        ]

        # iterate through the CSV and load in the labels
        if label_map is not None:
            labels = [
                Label(entity_type=label_map[row[idx[0]].upper()],
                      start=int(row[idx[1]]),
                      length=int(row[idx[2]]) - int(row[idx[1]]),
                      entity=row[idx[3]]) for row in csvreader
            ]
        else:
            labels = [
                Label(entity_type=row[idx[0].upper()],
                      start=int(row[idx[1]]),
                      length=int(row[idx[2]]) - int(row[idx[1]]),
                      entity=row[idx[3]]) for row in csvreader
            ]

    return labels


def load_data(path) -> dict:
    """Creates a dict the dataset."""
    examples = {'guid': [], 'text': [], 'ann': []}

    # for deid datasets, "path" is a folder containing txt/ann subfolders
    # "txt" subfolder has text files with the text of the examples
    # "ann" subfolder has annotation files with the labels
    txt_path = path / 'txt'
    ann_path = path / 'ann'
    for f in os.listdir(txt_path):
        if not f.endswith('.txt'):
            continue

        # guid is the file name
        guid = f[:-4]
        with open(txt_path / f, 'r') as fp:
            text = ''.join(fp.readlines())

        # load the annotations from disk
        # these datasets have consistent folder structures:
        #   root_path/txt/RECORD_NAME.txt - has text
        #   root_path/ann/RECORD_NAME.gs - has annotations
        if os.path.isdir(ann_path):
            labels = load_label(ann_path / f'{f[:-4]}.gs')
            examples['ann'].append(labels)

        examples['guid'].append(guid)
        examples['text'].append(text)

    return examples


def create_deid_dataset(data_dict, tokenizer, label2id=None) -> DeidDataset:
    """Creates a dataset from set of texts and labels.

       Args:
            data_dict: dict of text data with 'text', 'ann', and 'guid' keys; e.g., from load_data()
            tokenizer: HuggingFace tokenizer, e.g., loaded from AutoTokenizer.from_pretrained()
            label2id: dict to convert label (e.g., 'O,' 'DATE') to id (e.g., 0, 1)

       Returns: DeidDataset; see class definition in data.py
    """

    # specify dataset arguments
    split_long_sequences = True

    # split text/labels into multiple examples
    # (1) tokenize text
    # (2) identify split points
    # (3) output text as it was originally
    texts = data_dict['text']
    labels = data_dict['ann']
    ids = data_dict['guid']

    if split_long_sequences:
        split_dict = split_sequences(tokenizer, texts, labels, ids=ids)

    texts = split_dict['texts']
    labels = split_dict['labels']
    guids = split_dict['guids']

    encodings = tokenizer(texts,
                          is_split_into_words=False,
                          return_offsets_mapping=True,
                          padding=True,
                          truncation=True)

    if labels != []:
        # use the offset mappings in train_encodings to assign labels to tokens
        tags = assign_tags(encodings, labels)

        # encodings are dicts with three elements:
        #   'input_ids', 'attention_mask', 'offset_mapping'
        # these are used as kwargs to model training later
        tags = encode_tags(tags, encodings, label2id)

    else:
        tags = None

    # prepare a dataset compatible with Trainer module
    encodings.pop("offset_mapping")
    dataset = DeidDataset(encodings, tags, guids)

    return dataset


def get_label_map(transform):
    """ Gets dictionary of labels to convert from specific labels (may differ across datasets) to more general labels. """
    with open_text('transformer_deid', 'label.json') as fp:
        label_map = json.load(fp)

    # label_membership has different label transforms as keys
    if transform not in label_map:
        raise KeyError('Unable to find label transform %s in label.json' %
                       transform)
    label_map = label_map[transform]

    # label_map has items "harmonized_label": ["label 1", "label 2", ...]
    # invert this for the final label mapping
    return {
        label: harmonized_label
        for harmonized_label, original_labels in label_map.items()
        for label in original_labels
    }


def get_labels(labels_each_doc):
    """Gets the list of labels for this data set."""
    unique_labels = list(
        set(label.entity_type for labels in labels_each_doc
            for label in labels))

    unique_labels.sort()

    # add in the object tag - ensure it is the first element for label2id and id2label dict
    if 'O' in unique_labels:
        unique_labels.pop('O')
    unique_labels = ['O'] + unique_labels

    return unique_labels


def save_labels(labels, ids, out_path):
    if not os.path.exists(out_path):
        os.mkdir(out_path)
    header = ['start', 'stop', 'entity', 'entity_type']

    for doc, id in tqdm(zip(labels, ids), total=len(ids)):
        label_list = [[
            label.start, label.start + label.length, label.entity,
            label.entity_type
        ] for label in doc]
        with open(f'{out_path}/{id}.gs', 'w') as f:
            writer = csv.writer(f)
            writer.writerow(header)
            writer.writerows(label_list)