from allennlp.predictors.predictor import Predictor

from allennlp.predictors.text_classifier import TextClassifierPredictor

from allennlp.predictors.sentence_tagger import SentenceTaggerPredictor

from allennlp.models.archival import *

from allennlp.data import DatasetReader

from allennlp.common.params import Params

from allennlp.models import Model

import allennlp_models.classification

import allennlp_models.tagging

import torch

import os

import shutil

def rm_tmp(tmp_start):

    # Remove new directories in tmp

    tmp_now = os.listdir('/tmp')

    for i in tmp_now:

        if i not in tmp_start:

            print('removing directory /tmp/' + i)

            shutil.rmtree('/tmp/' + i)

def get_config(archive_path):

    archive = load_archive(archive_path)

    config = archive.config.duplicate()

    return config

def download_glove():

    # Saves TextClassifierPredictor object

    tmp_start = os.listdir('/tmp')

    glove_predictor = Predictor.from_path("https://storage.googleapis.com/allennlp-public-models/basic_stanford_sentiment_treebank-2020.06.09.tar.gz")

    glove_path = os.path.join(os.path.dirname(os.path.abspath(__file__)), 'glove_sentiment_predictor.txt')

    torch.save(glove_predictor, glove_path)

    try:

        rm_tmp(tmp_start)

    except:

        pass

    return glove_predictor

def download_roberta():

    tmp_start = os.listdir('/tmp')

    archive_path = "https://storage.googleapis.com/allennlp-public-models/sst-roberta-large-2020.06.08.tar.gz"

    config = get_config(archive_path)

    roberta_predictor = Predictor.from_path(archive_path)

    serialization_dir = os.path.join(os.path.dirname(os.path.abspath(__file__)), 'roberta', '')

    if os.path.exists(serialization_dir):

        shutil.rmtree(serialization_dir)

    os.makedirs(serialization_dir)

    # Create config and model files

    config.to_file(os.path.join(serialization_dir, 'config.json'))

    with open(os.path.join(serialization_dir, 'whole_model.pt'), 'wb') as file:

        torch.save(roberta_predictor._model, file)

    try:

        rm_tmp(tmp_start)

    except:

        pass

    return roberta_predictor

def download_ner():

    tmp_start = os.listdir('/tmp')

    archive_path = "https://storage.googleapis.com/allennlp-public-models/ner-model-2020.02.10.tar.gz"

    config = get_config(archive_path)

    ner_predictor = Predictor.from_path(archive_path)

    serialization_dir = os.path.join(os.path.dirname(os.path.abspath(__file__)), 'elmo-ner', '')

    if os.path.exists(serialization_dir):

        shutil.rmtree(serialization_dir)

    os.makedirs(serialization_dir)

    config.to_file(os.path.join(serialization_dir, 'config.json'))

    vocab = ner_predictor._model.vocab

    vocab.save_to_files(os.path.join(serialization_dir, 'vocabulary'))

    with open(os.path.join(serialization_dir, 'whole_model.pt'), 'wb') as file:

        torch.save(ner_predictor._model.state_dict(), file)

    try:

        rm_tmp(tmp_start)

    except:

        pass

    return ner_predictor

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 = 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)

    else:

        print('Downloading Roberta Sentiment Analysis Model')

        predictor = download_roberta()

    return predictor

def load_ner():

    serialization_dir = "../allennlp/elmo-ner"

    config_file = os.path.join(serialization_dir, 'config.json')

    weights_file = os.path.join(serialization_dir, 'whole_model.pt')

    loaded_params = Params.from_file(config_file)

    loaded_model = Model.load(loaded_params, serialization_dir, weights_file)

    dataset_reader = DatasetReader.from_params(loaded_params.get('dataset_reader'))

    predictor = SentenceTaggerPredictor(loaded_model, dataset_reader)

    return predictor

if __name__ == "__main__":

    download_glove()

    download_roberta()

    download_ner()