import argparse
import os
import re
from typing import Dict, List

from deidentify.base import Annotation, Document
from deidentify.dataset import brat


def readlines(filename):
    with open(filename) as file:
        lines = file.readlines()
    return lines


def documents_iter(notes):
    lines = readlines(notes)

    record_lines = []
    for line in lines:
        if line.startswith('START_OF_RECORD'):
            record_lines = []
            patient_id, record_id = re.findall(r'\d+', line)
        elif line.startswith('||||END_OF_RECORD'):
            yield Document(
                name='note-{}-{}'.format(patient_id, record_id),
                text=''.join(record_lines).rstrip(),
                annotations=[]
            )
        else:
            record_lines.append(line)


def annotations_iter(annotations):
    lines = readlines(annotations)

    current_pid, current_rid = lines[0].split(maxsplit=5)[0:2]

    annotations = []
    i = 1
    for line in lines:
        pid, rid, start, end, tag, text = line.strip().split(maxsplit=5)
        if pid != current_pid or rid != current_rid:
            yield annotations
            annotations = []
            i = 1
            current_pid = pid
            current_rid = rid

        annotations.append(Annotation(
            text=text,
            start=int(start),
            end=int(end),
            tag=tag,
            ann_id='T{}'.format(i),
            doc_id='note-{}-{}'.format(current_pid, current_rid)
        ))
        i += 1

    yield annotations


def _map_annotations(annotations):
    # Mapping: doc_id -> List[Annotation]
    mapping: Dict[str, List[Annotation]] = {}

    for doc_anns in annotations:
        mapping[doc_anns[0].doc_id] = doc_anns

    return mapping


def main(args):
    documents = documents_iter(args.notes_file)
    annotations = annotations_iter(args.phi_file)

    doc_annotations_mapping = _map_annotations(annotations)

    for doc in documents:
        anns = doc_annotations_mapping.get(doc.name, [])
        brat.write_brat_document(args.output_dir, doc_name=doc.name,
                                 text=doc.text, annotations=anns)


def arg_parser():
    parser = argparse.ArgumentParser()
    parser.add_argument("notes_file", help="Full path to raw notes file (notes-raw.txt)")
    parser.add_argument("phi_file", help="Full path to annotations file (id-phi.phrase)")
    parser.add_argument("output_dir", help="Path to output directory.")
    return parser.parse_args()


if __name__ == '__main__':
    ARGS = arg_parser()
    os.makedirs(ARGS.output_dir, exist_ok=True)
    main(ARGS)