import sys
import os
import torch
import re
import loader
# from allennlp.models.archival import *
from allennlp.data import DatasetReader
from allennlp.common.params import Params
from allennlp.predictors.text_classifier import TextClassifierPredictor
import time

sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
from ehrkit import ehrkit
# from config import USERNAME, PASSWORD


def load_glove():
    # Loads GLOVE model
    glove_path = os.path.join(os.path.dirname(os.path.abspath(__file__)), "glove_sentiment_predictor.txt")
    if os.path.exists(glove_path):  # same dir for github
        print('Loading Glove Sentiment Analysis Model')
        predictor = torch.load(glove_path)
    else:
        print('Downloading Glove Sentiment Analysis Model')
        predictor = loader.download_glove()
    return predictor


def load_roberta():
    # Loads Roberta model
    serialization_dir = os.path.join(os.path.dirname(os.path.abspath(__file__)), 'roberta', '')
    config_file = os.path.join(serialization_dir, 'config.json')
    if os.path.exists(config_file):
        print('Loading Roberta Sentiment Analysis Model')
        model_file = os.path.join(serialization_dir, 'whole_model.pt')
        model = torch.load(model_file)
        loaded_params = Params.from_file(config_file)
        dataset_reader = DatasetReader.from_params(loaded_params.get('dataset_reader'))

        # Gets predictor from model and dataset reader
        predictor = TextClassifierPredictor(model, dataset_reader)

        # weights_file = os.path.join(serialization_dir, 'weights.th')
        # loaded_model = Model.load(loaded_params, serialization_dir, weights_file) # Takes forever
        # archive = load_archive(os.path.join('roberta', 'model.tar.gz')) # takes forever
    else:
        print('Downloading Roberta Sentiment Analysis Model')
        predictor = loader.download_roberta()
    return predictor


def get_doc():
    doc_id = input("MIMIC Document ID [press Enter for random]: ")
    if doc_id == '':
        ehrdb.cur.execute("SELECT ROW_ID FROM mimic.NOTEEVENTS ORDER BY RAND() LIMIT 1")
        doc_id = ehrdb.cur.fetchall()[0][0]
        print('Document ID: %s' % doc_id)
    try:
        text = ehrdb.get_document(int(doc_id))
        clean_text = re.sub('[^A-Za-z0-9\.\,\-\/]+', ' ', text).lower()
        return doc_id, clean_text
    except:
        message = 'Error: There is no document with ID \'' + doc_id + '\' in mimic.NOTEEVENTS'
        sys.exit(message)


if __name__ == '__main__':
    # ehrdb = ehrkit.start_session(USERNAME, PASSWORD)
    ehrdb = ehrkit.start_session("jeremy.goldwasser@localhost", "mysql4710")
    doc_id, clean_text = get_doc()
    # print('LENGTH OF DOCUMENT: %d' % len(clean_text))

    x = input('GloVe or RoBERTa predictor [g=GloVe, r=RoBERTa]? ')
    if x == 'g':
        glove_predictor = load_glove()
        probs = glove_predictor.predict(clean_text)['probs']
    elif x == 'r':
        roberta_predictor = load_roberta()
        try:
            probs = roberta_predictor.predict(clean_text)['probs']
        except:
            print('Document too long for RoBERTa model. Using GLoVe instead.')
            glove_predictor = load_glove()
            probs = glove_predictor.predict(clean_text)['probs']
    else:
        sys.exit('Error: Must input \'g\' or  \'r\'')

    classification = 'Positive' if probs[0] >= 0.5 else 'Negative'
    print('Sentiment of document: %s' % classification)


    # # jeremy.goldwasser@localhost
    # # Save sentiment as json file
    # sentiment = {'text': clean_text, 'sentiment': classification, 'prob': probs[0]}
    # with open('predicted_sentiments/' + str(doc_id) + '.json', 'w', encoding='utf-8') as f:
    #     json.dump(sentiment, f, ensure_ascii=False, indent=4)