from transformers import (AutoModelForTokenClassification,

                          AutoModelForSequenceClassification,

                          TrainingArguments,

                          AutoTokenizer,

                          AutoConfig,

                          Trainer)

from biobert_ner.utils_ner import (convert_examples_to_features, get_labels, NerTestDataset)

from biobert_ner.utils_ner import InputExample as NerExample

from biobert_re.utils_re import RETestDataset

from bilstm_crf_ner.model.config import Config as BiLSTMConfig

from bilstm_crf_ner.model.ner_model import NERModel as BiLSTMModel

from bilstm_crf_ner.model.ner_learner import NERLearner as BiLSTMLearner

import en_ner_bc5cdr_md

import numpy as np

import os

from torch import nn

from ehr import HealthRecord

from generate_data import scispacy_plus_tokenizer

from annotations import Entity

import logging

from typing import List, Tuple

logger = logging.getLogger(__name__)

BIOBERT_NER_SEQ_LEN = 128

BILSTM_NER_SEQ_LEN = 512

BIOBERT_RE_SEQ_LEN = 128

logging.getLogger('matplotlib.font_manager').disabled = True

BIOBERT_NER_MODEL_DIR = "biobert_ner/output_full"

BIOBERT_RE_MODEL_DIR = "biobert_re/output_full"

# =====BioBERT Model for NER======

biobert_ner_labels = get_labels('biobert_ner/dataset_full/labels.txt')

biobert_ner_label_map = {i: label for i, label in enumerate(biobert_ner_labels)}

num_labels_ner = len(biobert_ner_labels)

biobert_ner_config = AutoConfig.from_pretrained(

    os.path.join(BIOBERT_NER_MODEL_DIR, "config.json"),

    num_labels=num_labels_ner,

    id2label=biobert_ner_label_map,

    label2id={label: i for i, label in enumerate(biobert_ner_labels)})

biobert_ner_tokenizer = AutoTokenizer.from_pretrained(

    "dmis-lab/biobert-base-cased-v1.1")

biobert_ner_model = AutoModelForTokenClassification.from_pretrained(

    os.path.join(BIOBERT_NER_MODEL_DIR, "pytorch_model.bin"),

    config=biobert_ner_config)

biobert_ner_training_args = TrainingArguments(output_dir="/tmp", do_predict=True)

biobert_ner_trainer = Trainer(model=biobert_ner_model, args=biobert_ner_training_args)

label_ent_map = {'DRUG': 'Drug', 'STR': 'Strength',

                 'DUR': 'Duration', 'ROU': 'Route',

                 'FOR': 'Form', 'ADE': 'ADE',

                 'DOS': 'Dosage', 'REA': 'Reason',

                 'FRE': 'Frequency'}

# =====BiLSTM + CRF model for NER=========

bilstm_config = BiLSTMConfig()

bilstm_model = BiLSTMModel(bilstm_config)

bilstm_learn = BiLSTMLearner(bilstm_config, bilstm_model)

bilstm_learn.load("ner_15e_bilstm_crf_elmo")

scispacy_tok = en_ner_bc5cdr_md.load().tokenizer

scispacy_plus_tokenizer.__defaults__ = (scispacy_tok,)

# =====BioBERT Model for RE======

re_label_list = ["0", "1"]

re_task_name = "ehr-re"

biobert_re_config = AutoConfig.from_pretrained(

    os.path.join(BIOBERT_RE_MODEL_DIR, "config.json"),

    num_labels=len(re_label_list),

    finetuning_task=re_task_name)

biobert_re_model = AutoModelForSequenceClassification.from_pretrained(

    os.path.join(BIOBERT_RE_MODEL_DIR, "pytorch_model.bin"),

    config=biobert_re_config,)

biobert_re_training_args = TrainingArguments(output_dir="/tmp", do_predict=True)

biobert_re_trainer = Trainer(model=biobert_re_model, args=biobert_re_training_args)

def align_predictions(predictions: np.ndarray, label_ids: np.ndarray) -> List[List[str]]:

    """

    Get the list of labelled predictions from model output

    Parameters

    ----------

    predictions : np.ndarray

        An array of shape (num_examples, seq_len, num_labels).

    label_ids : np.ndarray

        An array of shape (num_examples, seq_length).

        Has -100 at positions which need to be ignored.

    Returns

    -------

    preds_list : List[List[str]]

        Labelled output.

    """

    preds = np.argmax(predictions, axis=2)

    batch_size, seq_len = preds.shape

    preds_list = [[] for _ in range(batch_size)]

    for i in range(batch_size):

        for j in range(seq_len):

            if label_ids[i, j] != nn.CrossEntropyLoss().ignore_index:

                preds_list[i].append(biobert_ner_label_map[preds[i][j]])

    return preds_list

def get_chunk_type(tok: str) -> Tuple[str, str]:

    """

    Args:

        tok: Label in IOB format

    Returns:

        tuple: ("B", "DRUG")

    """

    tag_class = tok.split('-')[0]

    tag_type = tok.split('-')[-1]

    return tag_class, tag_type

def get_chunks(seq: List[str]) -> List[Tuple[str, int, int]]:

    """

    Given a sequence of tags, group entities and their position

    Args:

        seq: ["O", "O", "B-DRUG", "I-DRUG", ...] sequence of labels

    Returns:

        list of (chunk_type, chunk_start, chunk_end)

    Example:

        seq = ["B-DRUG", "I-DRUG", "O", "B-STR"]

        result = [("DRUG", 0, 1), ("STR", 3, 3)]

    """

    default = "O"

    chunks = []

    chunk_type, chunk_start = None, None

    for i, tok in enumerate(seq):

        # End of a chunk 1

        if tok == default and chunk_type is not None:

            # Add a chunk.

            chunk = (chunk_type, chunk_start, i - 1)

            chunks.append(chunk)

            chunk_type, chunk_start = None, None

        # End of a chunk + start of a chunk!

        elif tok != default:

            tok_chunk_class, tok_chunk_type = get_chunk_type(tok)

            if chunk_type is None:

                chunk_type, chunk_start = tok_chunk_type, i

            elif tok_chunk_type != chunk_type or tok_chunk_class == "B":

                chunk = (chunk_type, chunk_start, i - 1)

                chunks.append(chunk)

                chunk_type, chunk_start = tok_chunk_type, i

        else:

            continue

    # end condition

    if chunk_type is not None:

        chunk = (chunk_type, chunk_start, len(seq))

        chunks.append(chunk)

    return chunks

# noinspection PyTypeChecker

def get_biobert_ner_predictions(test_ehr: HealthRecord) -> List[Tuple[str, int, int]]:

    """

    Get predictions for a single EHR record using BioBERT

    Parameters

    ----------

    test_ehr : HealthRecord

        The EHR record, this object should have a tokenizer set.

    Returns

    -------

    pred_entities : List[Tuple[str, int, int]]

        List of predicted Entities each with the format

        ("entity", start_idx, end_idx).

    """

    split_points = test_ehr.get_split_points(max_len=BIOBERT_NER_SEQ_LEN - 2)

    examples = []

    for idx in range(len(split_points) - 1):

        words = test_ehr.tokens[split_points[idx]:split_points[idx + 1]]

        examples.append(NerExample(guid=str(split_points[idx]),

                                   words=words,

                                   labels=["O"] * len(words)))

    input_features = convert_examples_to_features(

        examples,

        biobert_ner_labels,

        max_seq_length=BIOBERT_NER_SEQ_LEN,

        tokenizer=biobert_ner_tokenizer,

        cls_token_at_end=False,

        cls_token=biobert_ner_tokenizer.cls_token,

        cls_token_segment_id=0,

        sep_token=biobert_ner_tokenizer.sep_token,

        sep_token_extra=False,

        pad_on_left=bool(biobert_ner_tokenizer.padding_side == "left"),

        pad_token=biobert_ner_tokenizer.pad_token_id,

        pad_token_segment_id=biobert_ner_tokenizer.pad_token_type_id,

        pad_token_label_id=nn.CrossEntropyLoss().ignore_index,

        verbose=0)

    test_dataset = NerTestDataset(input_features)

    predictions, label_ids, _ = biobert_ner_trainer.predict(test_dataset)

    predictions = align_predictions(predictions, label_ids)

    # Flatten the prediction list

    predictions = [p for ex in predictions for p in ex]

    input_tokens = test_ehr.get_tokens()

    prev_pred = ""

    final_predictions = []

    idx = 0

    for token in input_tokens:

        if token.startswith("##"):

            if prev_pred == "O":

                final_predictions.append(prev_pred)

            else:

                pred_typ = prev_pred.split("-")[-1]

                final_predictions.append("I-" + pred_typ)

        else:

            prev_pred = predictions[idx]

            final_predictions.append(prev_pred)

            idx += 1

    pred_entities = []

    chunk_pred = get_chunks(final_predictions)

    for ent in chunk_pred:

        pred_entities.append((ent[0],

                              test_ehr.get_char_idx(ent[1])[0],

                              test_ehr.get_char_idx(ent[2])[1]))

    return pred_entities

def get_bilstm_ner_predictions(test_ehr: HealthRecord) -> List[Tuple[str, int, int]]:

    """

    Get predictions for a single EHR record using BiLSTM

    Parameters

    ----------

    test_ehr : HealthRecord

        The EHR record, this object should have a tokenizer set.

    Returns

    -------

    pred_entities : List[Tuple[str, int, int]]

        List of predicted Entities each with the format

        ("entity", start_idx, end_idx).

    """

    split_points = test_ehr.get_split_points(max_len=BILSTM_NER_SEQ_LEN)

    examples = []

    for idx in range(len(split_points) - 1):

        words = test_ehr.tokens[split_points[idx]:split_points[idx + 1]]

        examples.append(words)

    predictions = bilstm_learn.predict(examples)

    pred_entities = []

    for idx in range(len(split_points) - 1):

        chunk_pred = get_chunks(predictions[idx])

        for ent in chunk_pred:

            pred_entities.append((ent[0],

                                  test_ehr.get_char_idx(split_points[idx] + ent[1])[0],

                                  test_ehr.get_char_idx(split_points[idx] + ent[2])[1]))

    return pred_entities

# noinspection PyTypeChecker

def get_ner_predictions(ehr_record: str, model_name: str = "biobert", record_id: str = "1") -> HealthRecord:

    """

    Get predictions for NER using either BioBERT or BiLSTM

    Parameters

    --------------

    ehr_record : str

        An EHR record in text format.

    model_name : str

        The model to use for prediction. Default is biobert.

    record_id : str

        The record id of the returned object. Default is 1.

    Returns

    -----------

    A HealthRecord object with entities set.

    """

    if model_name.lower() == "biobert":

        test_ehr = HealthRecord(record_id=record_id,

                                text=ehr_record,

                                tokenizer=biobert_ner_tokenizer.tokenize,

                                is_bert_tokenizer=True,

                                is_training=False)

        predictions = get_biobert_ner_predictions(test_ehr)

    elif model_name.lower() == "bilstm":

        test_ehr = HealthRecord(text=ehr_record,

                                tokenizer=scispacy_plus_tokenizer,

                                is_bert_tokenizer=False,

                                is_training=False)

        predictions = get_bilstm_ner_predictions(test_ehr)

    else:

        raise AttributeError("Accepted model names include 'biobert' "

                             "and 'bilstm'.")

    ent_preds = []

    for i, pred in enumerate(predictions):

        ent = Entity("T%d" % i, label_ent_map[pred[0]], [pred[1], pred[2]])

        ent_text = test_ehr.text[ent[0]:ent[1]]

        if not any(letter.isalnum() for letter in ent_text):

            continue

        ent.set_text(ent_text)

        ent_preds.append(ent)

    test_ehr.entities = ent_preds

    return test_ehr

def get_re_predictions(test_ehr: HealthRecord) -> HealthRecord:

    """

    Get predictions for Relation Extraction.

    Parameters

    -----------

    test_ehr : HealthRecord

        A HealthRecord object with entities set.

    Returns

    --------

    HealthRecord

        The original object with relations set.

    """

    test_dataset = RETestDataset(test_ehr, biobert_ner_tokenizer,

                                 BIOBERT_RE_SEQ_LEN, re_label_list)

    if len(test_dataset) == 0:

        test_ehr.relations = []

        return test_ehr

    re_predictions = biobert_re_trainer.predict(test_dataset=test_dataset).predictions

    re_predictions = np.argmax(re_predictions, axis=1)

    idx = 1

    rel_preds = []

    for relation, pred in zip(test_dataset.relation_list, re_predictions):

        if pred == 1:

            relation.ann_id = "R%d" % idx

            idx += 1

            rel_preds.append(relation)

    test_ehr.relations = rel_preds

    return test_ehr