import argparse

parser = argparse.ArgumentParser(description='The backend of the specified frontend. Service obtains sentences and predicts entities.')

parser.add_argument('-l', '--length', type=int, default=128,

                    help='Choose the maximum length of the model\'s input layer.')

parser.add_argument('-m', '--model', type=str, default='../models/medcondbert.pth',

                    help='Choose the directory of the model to be used for prediction.')

parser.add_argument('-tr', '--transfer_learning', type=bool, default=False,

                    help='Choose whether the given model has been trained on BioBERT or not. \

                    Careful: It will not work if wrongly specified!')

parser.add_argument('-p', '--port', type=int, default=5000,

                    help='The port on which the model is going to run.')

parser.add_argument('-t', '--type', type=str, required=True,

                    help='Specify the type of annotation to process. Type of annotation needs to be one of the following: Medical Condition, Symptom, Medication, Vital Statistic, Measurement Value, Negation Cue, Medical Procedure')

args = parser.parse_args()

max_length = args.length

model_path = args.model

transfer_learning = args.transfer_learning

port = args.port

print("Preparing model...")

from gevent.pywsgi import WSGIServer # Imports the WSGIServer

from gevent import monkey; monkey.patch_all()

from flask import Flask, request, jsonify

from flask_cors import CORS

from utils.dataloader import Dataloader

from utils.BertArchitecture import BertNER, BioBertNER

from utils.metric_tracking import MetricsTracking

import torch

from torch.optim import SGD

from torch.utils.data import DataLoader

import numpy as np

import pandas as pd

from tqdm import tqdm

from transformers import BertTokenizer,BertForTokenClassification

import spacy

# initializing backend

if not args.transfer_learning:

    print("Training base BERT model...")

    model = BertNER(3) #O, B-, I- -> 3 entities

    if args.type == 'Medical Condition':

        type = 'MEDCOND'

    elif args.type == 'Symptom':

        type = 'SYMPTOM'

    elif args.type == 'Medication':

        type = 'MEDICATION'

    elif args.type == 'Vital Statistic':

        type = 'VITALSTAT'

    elif args.type == 'Measurement Value':

        type = 'MEASVAL'

    elif args.type == 'Negation Cue':

        type = 'NEGATION'

    elif args.type == 'Medical Procedure':

        type = 'PROCEDURE'

    else:    

        raise ValueError('Type of annotation needs to be one of the following: Medical Condition, Symptom, Medication, Vital Statistic, Measurement Value, Negation Cue, Medical Procedure')

    tokenizer = BertTokenizer.from_pretrained('bert-base-uncased')

    tokenizer.add_tokens(['B-' + args.type, 'I-' + args.type])

else:

    print("Training BERT model based on BioBERT diseases...")

    if not args.type == 'Medical Condition':

        raise ValueError('Type of annotation needs to be Medical Condition when using BioBERT as baseline.')

    model = BioBertNER(3) #O, B-, I- -> 3 entities

    tokenizer = BertTokenizer.from_pretrained('alvaroalon2/biobert_diseases_ner')

    type = 'DISEASE'

label_to_ids = {

    'B-' + type: 0,

    'I-' + type: 1,

    'O': 2

    }

ids_to_label = {

    0:'B-' + type,

    1:'I-' + type,

    2:'O'

    }

model.load_state_dict(torch.load(model_path))

model.eval()

app = Flask(__name__)

CORS(app)  # Initialize CORS

sentence_detector = spacy.load("en_core_web_sm")

print("Serving API now...")

def predict_sentence(sentence):

    t_sen = tokenizer.tokenize(sentence)

    sen_code = tokenizer.encode_plus(sentence,

        return_tensors='pt',

        add_special_tokens=True,

        max_length = max_length,

        padding='max_length',

        return_attention_mask=True,

        truncation = True

        )

    inputs = {key: torch.as_tensor(val) for key, val in sen_code.items()}

    attention_mask = inputs['attention_mask'].squeeze(1)

    input_ids = inputs['input_ids'].squeeze(1)

    outputs = model(input_ids, attention_mask)

    predictions = outputs.logits.argmax(dim=-1)

    predictions = [ids_to_label.get(x) for x in predictions.numpy()[0]]

    #beware special tokens

    cutoff = min(len(predictions)-1, len(t_sen))

    predictions = predictions[1:cutoff+1]

    t_sen = t_sen[:cutoff]

    return t_sen, predictions

def clean(tokens, labels):

    cleaned_tokens = []

    cleaned_labels = []

    cnt = 1

    for i in range(len(tokens)): #same length

        if tokens[i].startswith("##") and len(cleaned_tokens) > 0:

            cleaned_tokens[i-cnt] = cleaned_tokens[i-cnt] + tokens[i][2:]

            cnt = cnt + 1

        else:

            cleaned_tokens.append(tokens[i])

            cleaned_labels.append(labels[i])

    return cleaned_tokens, cleaned_labels

def handle_request(data):

    sentences = sentence_detector(data).sents

    tokens = []

    labels = []

    for sentence in sentences:

        new_tokens, new_labels = predict_sentence(sentence.text)

        tokens = tokens + new_tokens

        labels = labels + new_labels

    cleaned_tokens, cleaned_labels = clean(tokens, labels)

    return cleaned_tokens, cleaned_labels

@app.route('/extract_entities', methods=['POST'])

def main():

    text = request.get_data(as_text=True)

    result = handle_request(text)

    return jsonify({'tokens': result[0], 'entities': result[1]})

if __name__ == '__main__':

    LISTEN = ('0.0.0.0',port)

    http_server = WSGIServer( LISTEN, app )

    http_server.serve_forever()