# coding: utf-8

# Base Dependencies

# ------------------

import json

import xml.etree.ElementTree as ET

from os.path import join as pjoin

from pathlib import Path

from tqdm import tqdm

from typing import List, Tuple, Set, Dict

# Local Dependencies

# --------------------

from models import Document, Entity, RelationN2C2, RelationDDI, RelationCollection

from utils import files_ddi, files_n2c2, doc_id_n2c2, make_dir

# 3rd-Party Dependencies

# ----------------------

from PyRuSH import PyRuSHSentencizer

# Constants

# ---------

from constants import N2C2_PATH, DDI_PATH

# Auxiliar Functions

# ------------------

def read_txt(file: Path) -> str:

    """Reads a .txt file

    Args:

        file (Path): path to the .txt file

    """

    # read text file

    with open(file, "r", encoding="utf-8") as fin:

        text = fin.read()

    return text

def read_json(file: Path) -> str:

    """Reads a .json file

    Args:

        file (Path): path to the .json file

    """

    return json.loads(read_txt(file))

def read_annotations_n2c2(file: Path) -> Tuple[List[Entity], Set[str]]:

    """Reads a n2c2 .ann file and extracts the entities and the relations

    Args:

        file (Path): path to the n2c2 annotation file

    """

    # read file

    with open(file, "r", encoding="utf-8") as fin:

        annotations: List[str] = fin.readlines()

    # process file

    doc_id: str = doc_id_n2c2(file)

    entities: List[Entity] = list()

    gt_relations: Set[str] = set()  # ground-truth relations

    for line in annotations:

        if line.startswith("T"):  # process entity

            entities.append(Entity.from_n2c2_annotation(doc_id, line))

        elif line.startswith("R"):  # process relation

            id, definition = line.strip().split("\t")

            type, entity1_id, entity2_id = definition.split()

            entity1_id = entity1_id.split(":")[1]

            entity2_id = entity2_id.split(":")[1]

            gt_relations.add("{}-{}".format(entity1_id, entity2_id))

        else:  # ignore annotator's note

            continue

    # sort entities by their end character

    entities.sort(key=lambda ent: ent.end)

    return entities, gt_relations

# Main Functions

# ---------------

def generate_relations(

    dataset: str, save_to_disk: bool = True

) -> Dict[str, RelationCollection]:

    """Generates relations of a given dataset and saves them to disk

    Args:

        dataset (str): dataset's name

        save_to_disk (bool, optional): the relation collections are saved to disk in a datading or not. Defaults to True.

    Raises:

        ValueError: unsupported dataset

    Returns:

        Dict[str, RelationCollection]: train and test relation collections

    """

    if dataset == "n2c2":

        return generate_relations_n2c2(save_to_disk=save_to_disk)

    elif dataset == "ddi":

        return generate_relations_ddi(save_to_disk=save_to_disk)

    else:

        raise ValueError("unsupported dataset '{}'".format(dataset))

def generate_relations_n2c2(save_to_disk: bool = True) -> Dict[str, RelationCollection]:

    """Generates relations of the n2c2 dataset

        1. Per document

        2. Read all entities, all true relations

        3. Separate in to drugs and per attribute

        4. For each relation type, combine each drug with each attribute within the same sentence

    Args:

        save_to_disk (bool): the relation collections are saved to disk in a datading or not. Default to True.

    Returns:

        Dict[str, RelationCollection]: train and test relation collections

    """

    print("Generating relations for the n2c2 dataset...\n")

    dataset = files_n2c2()

    collections = {}

    for split, files in dataset.items():

        print(split, ": ")

        split_entities = []

        split_relations = []

        for basepath in tqdm(files):

            # process clinical text, split in sentences

            document: Document = Document.from_json(read_txt(basepath + ".json"))

            # read annotation file

            entities, gt_relations = read_annotations_n2c2(basepath + ".ann")

            # generate relations

            relations = RelationN2C2.generate_relations_n2c2(

                document, entities, gt_relations, (split == "test")

            )

            split_entities.extend(entities)

            split_relations.extend(relations)

        # create collection

        collection = RelationCollection(split_relations)

        # remove invalid relations

        collection = collection[collection.valid_indexes()]

        # write to databing

        if save_to_disk:

            make_dir(pjoin(N2C2_PATH, "{}_datading".format(split)))

            collection.to_datading(

                pjoin(N2C2_PATH, "{}_datading".format(split), "relations.msgpack")

            )

        collections[split] = collection

    return collections

def generate_relations_ddi(save_to_disk: bool = True) -> Dict[str, RelationCollection]:

    """Generates relations of the ddi dataset

    Args:

        save_to_disk (bool): the relation collections are saved to disk in a datading or not. Default to True.

    Returns:

        Dict[str, RelationCollection]: train and test relation collections

    """

    print("Generating relations for the DDI Extraction corpus...")

    dataset = files_ddi()

    collections = {}

    for split, files in dataset.items():

        print(split, ": ")

        split_relations = []

        for file in tqdm(files):

            xml_tree = ET.parse(file)

            relations = RelationDDI.generate_relations_ddi(xml_tree)

            split_relations.extend(relations)

        # create collection

        collection = RelationCollection(split_relations)

        # remove invalid relations

        collection = collection[collection.valid_indexes()]

        # write to databing

        if save_to_disk:

            make_dir(pjoin(DDI_PATH, "{}_datading".format(split)))

            collection.to_datading(

                pjoin(DDI_PATH, "{}_datading".format(split), "relations.msgpack")

            )

        collections[split] = collection

    return collections