from collections import defaultdict
import random
from typing import List
import tarfile
import shutil
import typer
from pathlib import Path
import spacy
from spacy.language import Language
from spacy.tokens import Doc, DocBin, Span
from spacy.util import filter_spans
from wasabi import msg
from spacy.tokenizer import Tokenizer
from spacy.util import compile_prefix_regex, compile_infix_regex, compile_suffix_regex
import functools 

random.seed(42)


def main(
    input_dir: Path = typer.Argument(..., exists=True),
    output_dir: Path = typer.Argument(...),
    beth_train_tar_name: str = "i2b2_Beth_Train_Release.tar.gz",
    partners_train_tar_name: str = "i2b2_Partners_Train_Release.tar.gz",
    test_zip_name: str = "Task_1C.zip",
    merge_docs: bool = True,
):
    """Extract and preprocess raw n2c2 2011 Challenge data into spaCy DocBin format.
    input_dir (Path): Input directory with raw downloads from Harvard DBMI Portal.
    output_dir (Path): Output directory to save spaCy .docbin files to.
    beth_train_tar_name (str): Filename of downloaded tarfile for Beth Training Data.
    partners_train_tar_name (str): Filename of downloaded tarfile for Partners Training Data.
    test_zip_name (str): Filename of downloaded tarfile for n2c2 Test Data.
    merge_docs (bool): If False, create spaCy docs for each line of each medical record
    """
    # Unpack compressed data files
    msg.info("Extracting raw data.")
    beth_train_tar_path = input_dir / beth_train_tar_name
    partners_train_tar_path = input_dir / partners_train_tar_name
    test_zip_path = input_dir / test_zip_name

    #for path in [beth_train_tar_path, partners_train_tar_path]:
    #    if path.name.endswith("tar.gz"):
    #        msg.text(f"Extracting {path}")
    #        tar = tarfile.open(path, "r:gz")
    #        tar.extractall(path.parent)
    #        tar.close()

    #shutil.unpack_archive(test_zip_path, input_dir / test_zip_name.replace(".zip", ""))

    # preprocess data
    msg.info("Converting to spaCy Doc objects.")
    
    with open(r'/notebooks/Clinical_NER/difference.txt', 'a') as fp:
            fp.write(str((input_dir / "Beth_Train").stem)+'\n')
    
    beth_train_docs = docs_from_many_clinical_records(
        input_dir / "Beth_Train", merge_docs=merge_docs
    )
    
    with open(r'/notebooks/Clinical_NER/difference.txt', 'a') as fp:
            fp.write(str((input_dir / "Partners_Train").stem)+'\n')
    
    partners_train_docs = docs_from_many_clinical_records(
        input_dir / "Partners_Train", merge_docs=merge_docs
    )
    train_docs = beth_train_docs + partners_train_docs

    with open(r'/notebooks/Clinical_NER/difference.txt', 'a') as fp:
            fp.write(str((input_dir / "Task_1C/i2b2_Test/i2b2_Beth_Test").stem)+'\n')
    
    beth_test_docs = docs_from_many_clinical_records(
        input_dir / "Task_1C/i2b2_Test/i2b2_Beth_Test", merge_docs=merge_docs
    )
    
    with open(r'/notebooks/Clinical_NER/difference.txt', 'a') as fp:
            fp.write(str((input_dir / "Task_1C/i2b2_Test/i2b2_Partners_Test").stem)+'\n')
    
    partners_test_docs = docs_from_many_clinical_records(
        input_dir / "Task_1C/i2b2_Test/i2b2_Partners_Test", merge_docs=merge_docs
    )
    test_docs = beth_test_docs + partners_test_docs

    random.shuffle(train_docs)
    split_idx = int(len(train_docs) * 0.8)
    train_docs, dev_docs = train_docs[:split_idx], train_docs[split_idx:]

    msg.good(f"Num Train Docs: {len(train_docs)}")
    msg.good(f"Num Dev Docs: {len(dev_docs)}")
    msg.good(f"Num Test Docs: {len(test_docs)}")

    with msg.loading(f"Saving docs to: {output_dir}..."):
        DocBin(docs=train_docs).to_disk(output_dir / "train.spacy")
        DocBin(docs=dev_docs).to_disk(output_dir / "dev.spacy")
        DocBin(docs=test_docs).to_disk(output_dir / "test.spacy")
        msg.good("Done.")


def docs_from_clinical_record(
    lines: List[str], annotations: List[str], nlp: Language, merge_docs: bool = False
) -> List[Doc]:
    """Create spaCy docs from a single annotated medical record in the n2c2 2011 format
    lines (List[str]): Text of the clinical record as a list separated by newlines
    annotations (List[str]): Raw entity annotations in the n2c2 2011 format
    nlp (Language): spaCy Language object. Defaults to spacy.blank("en").
    merge_docs (bool): If True: merge all lines into a single spaCy doc so
        there is only 1 element in the output array.
        If False: create a spaCy doc for each line in the original record
    RETURNS (List[Doc]): List of spaCy Doc objects with entity spans set
    """
    difference = []
    docs = []
    spans_by_line = defaultdict(list)
    nlp.Defaults.prefixes = [signs for signs in nlp.Defaults.prefixes if ':' not in signs and '#' not in signs and '+' not in signs and '(' not in signs and ')' not in signs and '*' not in signs and "'" not in signs and "%" not in signs and "_" not in signs and ";" not in signs and ">" not in signs and "," not in signs and "&" not in signs and '"' not in signs and "<" not in signs ]
    infixes = nlp.Defaults.prefixes + [r"[-]~"]

    infix_re = spacy.util.compile_infix_regex(infixes)

    def custom_tokenizer(nlp):
        return Tokenizer(nlp.vocab, infix_finditer=infix_re.finditer)

    nlp.tokenizer = custom_tokenizer(nlp)

    entities = {}
    
    for row in annotations:
        row = row.split("||")
        text_info = row[0]
        type_info = row[1]

        offset_start = text_info.split(" ")[-2]
        offset_end = text_info.split(" ")[-1]

        start_line, word_start = offset_start.split(":")
        end_line, word_end = offset_end.split(":")

        label = type_info.split('"')[-2]

        if start_line != end_line:
            # This happens very infrequently (only about 10 times in total)
            # so we just skip these annotations
            continue
        else:
            spans_by_line[int(start_line)].append(
                (int(word_start), int(word_end), label)
            )
            
            if start_line in entities :
                entities[start_line].append(text_info.split('"')[1])
            else :
                entities[start_line] = [text_info.split('"')[1]]
    
    extracted_entities = {}
    
    for i, line in enumerate(lines):
        n = i + 1
        line = line.replace("  "," ")
        doc = nlp.make_doc(line)
        if n in spans_by_line:
            ents = [
                Span(doc, start, end + 1, label=label)
                for (start, end, label) in spans_by_line[n]
            ]
            ents = [
                e for e in ents if bool(e.text.strip()) and e.text.strip() == e.text
            ]
            doc.ents = filter_spans(ents)
            extracted_entities[str(n)] =  [ e.text for e in ents if bool(e.text.strip()) and e.text.strip() == e.text ]

        docs.append(doc)
    
    for key , value in entities.items() :
        if key in extracted_entities :
            if functools.reduce(lambda x, y : x and y, map(lambda p, q: p.lower() != q.lower(),entities[key],extracted_entities[key]), True): 
                difference = difference+[key]+entities[key]+extracted_entities[key]
        else :
            difference = difference+[key+" Key not present"]+entities[key]
    
    with open(r'/notebooks/Clinical_NER/difference.txt', 'a') as fp:
        fp.write('\n'.join(difference))
    return [Doc.from_docs(docs)] if merge_docs else docs


def docs_from_many_clinical_records(
    base_path: Path, nlp: Language = spacy.blank("en"), merge_docs: bool = True
) -> List[Doc]:
    """Convert raw n2c2 annotated clinical records into a list of
        spaCy Doc objects to be ready to be used in training
    base_path (Path): Root path to the raw data
    nlp (Language): spaCy Language object. Defaults to spacy.blank("en").
    merge_docs (bool): If True: merge all lines into a single spaCy doc so
        there is only 1 element in the output array.
        If False: create a spaCy doc for each line in the original record

    RETURNS (List[Doc]): List of spaCy Doc objects with entity spans set
    """
    all_docs = []
    concept_paths = sorted((base_path / "concepts").glob("*.txt.con"))
    document_paths = sorted((base_path / "docs").glob("*.txt"))

    for con_path, doc_path in zip(concept_paths, document_paths):
        with open(r'/notebooks/Clinical_NER/difference.txt', 'a') as fp:
            fp.write('\n'+str(con_path.stem))
        annotations = con_path.open().read().splitlines()
        lines = doc_path.open().read().splitlines()

        docs = docs_from_clinical_record(lines, annotations, nlp, merge_docs=merge_docs)
        all_docs += docs

    return all_docs


if __name__ == "__main__":
    typer.run(main)