# *********************************************************
# AWS XGBOOST SCRIPT
# (C) David Renaudie 2019
# (C) Remi Domingues 2019
# *********************************************************
import logging
import os
import abc
import time
from typing import Dict, List

import boto3
import numpy as np
import pandas as pd
import sagemaker
import sagemaker.transformer

from load_dataset import is_s3_file, download_file_from_s3, load_preproc_dataset
from models.model import Model
from save_datasets import upload_file_to_s3


def endpoint_name(model_name):
    return model_name + '-endpoint'


def create_endpoint_config(boto3_sm, model_name, instance_type):
    """
    Create an Amazon SageMaker endpoint configuration by specifying the ML compute instances
    that you want to deploy your model to
    """
    endpoint_config_name = model_name + '-config'
    logging.info('Creating end point config {}'.format(endpoint_config_name))
    try:
        boto3_sm.delete_endpoint_config(EndpointConfigName=endpoint_config_name)
    except Exception:
        pass
    create_endpoint_config_response = boto3_sm.create_endpoint_config(
        EndpointConfigName=endpoint_config_name,
        ProductionVariants=[{
            'InstanceType': instance_type,
            'InitialVariantWeight': 1,
            'InitialInstanceCount': 1,
            'ModelName': model_name,
            'VariantName': 'AllTraffic'}])
    logging.info('Endpoint Config Arn: ' + create_endpoint_config_response['EndpointConfigArn'])

    return endpoint_config_name


def create_endpoint(boto3_sm, model_name, config_name):
    _endpoint_name = endpoint_name(model_name)
    logging.info('Creating endpoint {}'.format(_endpoint_name))
    create_endpoint_response = boto3_sm.create_endpoint(
        EndpointName=_endpoint_name,
        EndpointConfigName=config_name)
    logging.info('Endpoint Arn: {}'.format(create_endpoint_response['EndpointArn']))

    resp = boto3_sm.describe_endpoint(EndpointName=_endpoint_name)
    status = resp['EndpointStatus']
    logging.info("Status: " + status + '...')

    while status=='Creating':
        time.sleep(60)
        resp = boto3_sm.describe_endpoint(EndpointName=_endpoint_name)
        status = resp['EndpointStatus']
        logging.info("Status: " + status + '...')

    logging.info("Arn: " + resp['EndpointArn'])


def load_preds_file(file: str, parsing_fct) -> pd.DataFrame:
    """ Load the local prediction file in memory (result of the BT job for the prediction step)
    Each algorithm has its own format for the output data """
    logging.info('Loading from csv: {}...'.format(file))
    data = []

    with open(file) as f:
        for line in f:
            data.append(parsing_fct(line))

    return np.array(data)


class AWSModel(Model):
    def __init__(self, dataset_name: str, hyperparameters: Dict, infra_s3: Dict, infra_sm: Dict,
                 features: List, target: str, data_dir: str, training_job_common_dir: str, training_job_dir: str,
                 aws_model_id: str=None, clean: bool=False):
        Model.__init__(self, dataset_name, hyperparameters, infra_s3, features, target,
                       data_dir, training_job_dir, clean)

        self.infra_sm = infra_sm
        self.aws_model_id = aws_model_id
        if not aws_model_id:
            # data-config-timestamp
            self.aws_model_id = '-'.join(training_job_dir.replace('_', '').split('/')[1:])
            s3_timestamp = training_job_dir.split('/')[-1]
        else:
            s3_timestamp = self.aws_model_id[-23:]

        logging.info('======== Model ID ========\n{}'.format(self.aws_model_id))

        self.training_job_common_dir = training_job_common_dir
        # data/ml/<dataset>/<config>
        self.s3_ml_data_common_dir = '/'.join([self.infra_s3['s3_folder_ml']] + training_job_dir.split('/')[1:-1])
        # A new local folder is created for the logs and plots to avoid erasing previous data
        # However, we need to use the S3 folder previously generated by the prediction step (e.g. if we run only eval)
        # data/ml/<dataset>/<config>/<timestamp>
        self.s3_ml_data_job_dir = '/'.join([self.s3_ml_data_common_dir, s3_timestamp])

        self.csv_training_filename = 'training.csv'
        self.csv_validation_filename = 'validation.csv'
        self.csv_testing_filename = 'testing.csv'

        self.s3_training_csv_path = self.s3_filepath(filetype='ml_data', filename=self.csv_training_filename, common=True, uri=True)
        self.s3_validation_csv_path = self.s3_filepath(filetype='ml_data', filename=self.csv_validation_filename, common=True, uri=True)
        self.s3_testing_csv_path = self.s3_filepath(filetype='ml_data', filename=self.csv_testing_filename, common=True, uri=True)

        self.local_training_preds_file = self.local_filepath('training_preds.csv')
        self.local_validation_preds_file = self.local_filepath('validation_preds.csv')
        self.local_testing_preds_file = self.local_filepath('testing_preds.csv')

        self.s3_preds_folder = self.s3_filepath(filetype='ml_data', uri=True)
        self.s3_training_preds_file = self.s3_filepath(filetype='ml_data', filename='training.csv.out')
        self.s3_validation_preds_file = self.s3_filepath(filetype='ml_data', filename='validation.csv.out')
        self.s3_testing_preds_file = self.s3_filepath(filetype='ml_data', filename='testing.csv.out')

        boto3_sess = boto3.Session()
        self.boto3_sm = boto3_sess.client('sagemaker')
        self.container = self._get_container(boto3_sess)

    @abc.abstractmethod
    def _get_container(self, boto3_session):
        """ Return the URI corresponding to the container of the algorithm """
        raise NotImplementedError('_get_container method not implemented')

    @abc.abstractmethod
    def _init_s3_train_files(self):
        """ Initialize the training and validation files (features + label) required for the training step """
        raise NotImplementedError('_init_s3_train_files method not implemented')

    @abc.abstractmethod
    def _parse_preds_line(self, preds_file):
        raise NotImplementedError('_parse_preds_line method not implemented')

    def local_filepath(self, filename, common=False):
        """
        If common to all AWS jobs on this dataset, the returned filepath will be <jobs>/<data>/<config_id>/file.ext
        otherwise, it will be <jobs>/<data>/<config_id>/<timestamp>/file.ext
        """
        return os.path.join(self.training_job_dir, filename)

    def s3_to_uri(self, s3_filepath):
        """
        Append a prefix to an existing S3 filepath to obtain a full S3 URI
        return: s3://<s3_bucket>/<filepath>
        """
        return 's3://{}/{}'.format(self.infra_s3['s3_bucket'], s3_filepath)

    def s3_filepath(self, filetype, filename='', common=False, uri=False):
        """
        S3 filepath creation
        For ML data files, note that training requires libsvm or csv files, while batch transform requires csv
        filetype in ['preproc_data', 'ml_data', 'model']
        common is only available for filetype='ml_data'
        Filepaths:
            filetype='preproc_data' => data/preprocessed/<filename>
            filetype='ml_data' => data/ml/<dataset>/<config>[/<timestamp> if not common]/<filename>
            filetype='model' => models
        uri=True will add the prefix s3://<bucket>/ to the previous filepaths
        """
        if filetype == 'preproc_data':
            # data/preprocessed/<filename>
            s3_filepath = '/'.join([self.infra_s3['s3_folder_data'], filename])

        elif filetype == 'ml_data':
            if common:
                # data/ml/<dataset>/<config>/<filename>
                s3_filepath = '/'.join([self.s3_ml_data_common_dir, filename])
            else:
                # data/ml/<dataset>/<config>/<timestamp>/<filename>
                s3_filepath = '/'.join([self.s3_ml_data_job_dir, filename])

        elif filetype == 'model':
            # models
            s3_filepath = self.infra_s3['s3_folder_models']

        else:
            raise ValueError('Invalid filetype received: {}'.format(filetype))

        # If filename was empty, the request is for a directory. We thus remove the '/' in the end
        if s3_filepath[-1] == '/':
            s3_filepath = s3_filepath[:-1]

        if uri:
            # s3://<s3_bucket>/<filepath>
            s3_filepath = self.s3_to_uri(s3_filepath)

        return s3_filepath

    def is_s3_file(self, s3_filepath, verbose=True):
        """
        Returns True if the corresponding filepath exist in the S3 bucket specified in the infra_s3 (config file)
        The filepath must not include the S3 bucket, nor the URI prefix (s3://<bucket_name>/)
        Example of valid filepath: data/ml/<dataset>/<config>/training.libsvm
        """
        return is_s3_file(self.infra_s3['s3_bucket'], s3_filepath, verbose=verbose)

    def train(self):
        logging.info('======== TRAINING ========')
        logging.info('Preparing training data...')
        if self.training is None or self.validation is None:
            self.load_training_data()

        sm_session = sagemaker.Session()
        xgb_estimator = sagemaker.estimator.Estimator(self.container,
                                                      self.infra_sm['sm_role'],
                                                      train_instance_count=self.infra_sm['train_instance_count'],
                                                      train_instance_type=self.infra_sm['train_instance_type'],
                                                      output_path=self.s3_filepath(filetype='model', filename='', uri=True),
                                                      sagemaker_session=sm_session, train_max_run=self.infra_sm['maxruntime'])

        s3_input_training, s3_input_validation = self._init_s3_train_files()

        xgb_estimator.set_hyperparameters(**self.hyperparameters)
        xgb_estimator.fit({'train': s3_input_training, 'validation': s3_input_validation})
        self.model = xgb_estimator.create_model(name=self.aws_model_id)
        # TODO: check if model was saved correctly (i.e. if training didn't fail)
        # TODO: Force model saving. Currently, the model may not saved after the training, it seems AWS
        #       may persist it only when we perform the batch transform job in predict

        logging.debug('Model object: {}'.format(self.model.__dict__))

    def predict(self):
        logging.info('======= PREDICTION =======')
        # logging.info('Checking if previous prediction files exist...')
        if (not self.clean and self.is_s3_file(self.s3_training_preds_file)
                and self.is_s3_file(self.s3_validation_preds_file)
                and self.is_s3_file(self.s3_testing_preds_file)):
            logging.warning('S3 csv predictions files found - prediction step skipped. If recomputating is needed, please delete prediction files on S3 or use --clean.')
            return

        self.prepare_csv_data()

        logging.info('Predicting...')

        # if there's a model: first attempt to save it in AWS, during first transform job
        if self.model is not None:
            logging.info('Model reference found. Launching Batch Transform job...')
            transformer = self.model.transformer(
                instance_count=1,
                instance_type=self.infra_sm['bt_instance_type'],
                output_path=self.s3_preds_folder
            )
        # if no model: try to launch directly from an existing model stored in AWS
        else:
            logging.info('No model reference in memory: launching Batch Transform job from model ID {}...'.format(self.aws_model_id))
            transformer = sagemaker.transformer.Transformer(
                model_name=self.aws_model_id,
                base_transform_job_name='Batch-Transform',
                instance_count=1,
                instance_type=self.infra_sm['bt_instance_type'],
                output_path=self.s3_preds_folder
            )

        transformer.transform(self.s3_training_csv_path, content_type='text/csv', split_type='Line')
        transformer.transform(self.s3_validation_csv_path, content_type='text/csv', split_type='Line')
        transformer.transform(self.s3_testing_csv_path, content_type='text/csv', split_type='Line')
        transformer.wait()

        # Wait a couple of seconds to have the csv.out files available in S3, for further processing
        self.wait_for_files([self.s3_testing_preds_file, self.s3_training_preds_file, self.s3_validation_preds_file, ])
        logging.info('Predicting done.')

    def prepare_csv_data(self):
        logging.info('Preparing prediction data...')

        if self.training is None or self.validation is None:
            self.load_training_data()

        self.prepare_csv_file(self.training, self.csv_training_filename)
        self.prepare_csv_file(self.validation, self.csv_validation_filename)

        if self.testing is None:
            self.load_testing_data()

        self.prepare_csv_file(self.testing, self.csv_testing_filename)

    def prepare_csv_file(self, dataset: pd.DataFrame, filename) -> str:
        logging.info('Preparing csv file for Batch Transform: {}'.format(filename))
        s3_csv_file = self.s3_filepath(filetype='ml_data', filename=filename, common=True)

        if self.clean or not self.is_s3_file(s3_csv_file):
            logging.info('S3 csv file does not exist')
            local_csv_file = self.local_filepath(filename, common=True)
            if self.clean or not os.path.isfile(local_csv_file):
                logging.info('Local csv file does not exist. Computing...')
                dataset[self.features].to_csv(local_csv_file, header=False, index=False)
            logging.info('Local csv file found. Uploading to S3...')
            upload_file_to_s3(local_csv_file, self.infra_s3['s3_bucket'], s3_csv_file)
        else:
            logging.info('S3 csv file already exists. Skipping step.')

        logging.info('S3 csv file path: {}'.format(s3_csv_file))

    def prepare_evaluation_data(self):
        logging.info('Preparing evaluation data...')
        # To compute performance metrics, the true labels are required for all datasets in order to compare with predictions
        if self.training is None or self.validation is None:
            self.load_training_data()
        if self.testing is None:
            self.load_testing_data()

        if self.clean or not (os.path.isfile(self.local_training_preds_file)
                and os.path.isfile(self.local_validation_preds_file)
                and os.path.isfile(self.local_testing_preds_file)):
            logging.debug('training predictions s3 file: {}'.format(self.s3_training_preds_file))
            logging.debug('validation predictions s3 file: {}'.format(self.s3_validation_preds_file))
            logging.debug('testing predictions s3 file: {}'.format(self.s3_testing_preds_file))
            if (self.is_s3_file(self.s3_training_preds_file)
                    and self.is_s3_file(self.s3_validation_preds_file)
                    and self.is_s3_file(self.s3_testing_preds_file)):
                logging.info('s3 csv predictions files found, downloading...')
                download_file_from_s3(self.local_training_preds_file, self.infra_s3['s3_bucket'], self.s3_training_preds_file)
                download_file_from_s3(self.local_validation_preds_file, self.infra_s3['s3_bucket'], self.s3_validation_preds_file)
                download_file_from_s3(self.local_testing_preds_file, self.infra_s3['s3_bucket'], self.s3_testing_preds_file)
            else:
                logging.error('Predictions file not found, neither in local nor s3. Exiting.')
                raise FileNotFoundError

        logging.info('Local csv predictions files found, loading...')
        self.training_y_pred = load_preds_file(self.local_training_preds_file, self._parse_preds_line)
        self.validation_y_pred = load_preds_file(self.local_validation_preds_file, self._parse_preds_line)
        self.testing_y_pred = load_preds_file(self.local_testing_preds_file, self._parse_preds_line)

        logging.info('training_y_pred shape: {}'.format(self.training_y_pred.shape))
        logging.info('validation_y_pred shape: {}'.format(self.validation_y_pred.shape))
        logging.info('testing_y_pred shape: {}'.format(self.testing_y_pred.shape))

    def load_training_data(self):
        logging.info('Loading training/validation datasets...')
        postfixed_dataset_name = self.dataset_name + '_preprocessed'
        training, validation, _, _ = load_preproc_dataset(postfixed_dataset_name, self.data_dir,
                                                       self.infra_s3['s3_bucket'], self.infra_s3['s3_folder_data'],
                                                       self.clean, True, self.target)
        self.training = training
        self.validation = validation
        self._init_features(training)

        self.training_x = self.training[self.features].values[:]
        self.validation_x = self.validation[self.features].values[:]

        self.training_y = np.ravel(self.training[self.target])
        self.validation_y = np.ravel(self.validation[self.target])
        self.n_classes = len(set(self.training_y))

        logging.info('training_x shape: {}'.format(self.training_x.shape))
        logging.info('training_y shape: {}'.format(self.training_y.shape))
        logging.info('validation_x shape: {}'.format(self.validation_x.shape))
        logging.info('validation_y shape: {}'.format(self.validation_y.shape))

    def load_testing_data(self):
        logging.info('Loading testing dataset...')
        postfixed_dataset_name = self.dataset_name + '_preprocessed'
        _, _, testing, _ = load_preproc_dataset(postfixed_dataset_name, self.data_dir,
                                             self.infra_s3['s3_bucket'], self.infra_s3['s3_folder_data'],
                                             self.clean, True, self.target)
        self.testing = testing
        self._init_features(testing)
        self.testing_x = self.testing[self.features].values[:]
        self.testing_y = np.ravel(self.testing[self.target])

        logging.info('testing_x shape: {}'.format(self.testing_x.shape))
        logging.info('testing_y shape: {}'.format(self.testing_y.shape))

    def wait_for_files(self, files, max_time=600):
        if files:
            logging.info('Waiting for S3 files to be accessible:\n* {}'.format('\n* '.join(files)))
        i = 0
        while files and i < max_time:
            i += 1
            time.sleep(2)
            if self.is_s3_file(files[0], verbose=False):
                del files[0]

    def deploy(self):
        logging.info('======== DEPLOYMENT ========')
        logging.info('Deploying endpoint for model {}...'.format(self.aws_model_id))

        if self.model is None:
            # AWS SDK for Python (Boto 3) version
            config_name = create_endpoint_config(self.boto3_sm, self.aws_model_id, self.infra_sm['deploy_instance_type'])
            create_endpoint(self.boto3_sm, self.aws_model_id, config_name)
        else:
            # Amazon SageMaker Python SDK version (not available to deploy previously trained models)
            self.model.deploy(initial_instance_count=1,
                              instance_type=self.infra_sm['deploy_instance_type'],
                              endpoint_name=endpoint_name(self.aws_model_id),
                              update_endpoint=False)

        logging.info('Deployment done')
        logging.info('Endpoint name: {}'.format(endpoint_name(self.aws_model_id)))

    def delete_endpoint(self):
        logging.info('======== ENDPOINT DELETION ========')
        logging.info('Deleting endpoint: {}...'.format(endpoint_name(self.aws_model_id)))

        self.boto3_sm.delete_endpoint(EndpointName=endpoint_name(self.aws_model_id))

        logging.info('Endpoind deletion done')
