#!/usr/local/bin/python

# A sample training component that trains a simple scikit-learn model with optuna.
# This implementation works in File mode and makes no assumptions about the input file names.
# Input is specified as CSV with a data point in each row and the labels in the first column.

from __future__ import print_function

import os
import json
import pickle
import sys
import traceback
import optuna

import pandas as pd

from sklearn.ensemble import GradientBoostingRegressor
from sklearn.model_selection import train_test_split
from sklearn.model_selection import cross_validate
from sklearn.model_selection import KFold
from sklearn.metrics import mean_squared_error
from sklearn.metrics import r2_score

# These are the paths to where SageMaker mounts interesting things in your container.

prefix = '/opt/ml'

input_path = os.path.join(prefix, 'input/data')
output_path = os.path.join(prefix, 'output')
model_path = os.path.join(prefix, 'model')
param_path = os.path.join(prefix, 'input/config/hyperparameters.json')

# This algorithm has a single channel of input data called 'training'. Since we run in
# File mode, the input files are copied to the directory specified here.
channel_name = 'training'
training_path = os.path.join(input_path, channel_name) # /opt/ml/input/data/training

class Objective:
    def __init__(self, X, y):
        self.X = X
        self.y = y

    def __call__(self, trial):
        params = {
            'n_estimators': trial.suggest_int('n_estimators', 1, 100),
            'learning_rate': trial.suggest_uniform('learning_rate', 0.01, 1.0)
        }
        # 5-Fold Stratified CV with GradientBoostingRegressor, R2 score validation
        clf = GradientBoostingRegressor(**params)
        kf = KFold(n_splits=5, shuffle=True, random_state=42)
        scores = cross_validate(clf, X=self.X, y=self.y, cv=kf, scoring='r2', n_jobs=-1)
        return scores['test_score'].mean()

# The function to execute the training.
def train():
    print('Starting the training.')
    try:
        # Read in any parameters that the user passed for Optuna
        with open(param_path, 'r') as tc:
            params = json.load(tc)

        # Take the set of files and read all into a single pandas dataframe
        input_files = [os.path.join(training_path, file) for file in os.listdir(training_path)]
        if len(input_files) == 0:
            raise ValueError(('There are no files in {}.\n' +
                              'This usually indicates that the channel ({}) was incorrectly specified,\n' +
                              'the data specification in S3 was incorrectly specified or the role specified\n' +
                              'does not have permission to access the data.').format(training_path, channel_name))
        raw_data = [pd.read_csv(file, header=None) for file in input_files]
        train_data = pd.concat(raw_data)

        # Note that parameters are always passed in as strings, so we need to do any necessary conversions.
        seconds = params.get('seconds', None)
        if seconds is not None:
            seconds = int(seconds)
        else:
            seconds = 300

        # labels are at the last column
        X = train_data.iloc[:, :-1]
        y = train_data.iloc[:, [-1]]
        print('X shape: (%i,%i)' % X.shape)
        print('y shape: (%i,%i)' % y.shape)

        # train and test split with validate data
        X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, random_state=39)

        # Use Optuna to find the hyperparameters with the highest score
        print('Hyperparameter optimization starting.')
        objective = Objective(X_train, y_train.values.ravel())
        study = optuna.create_study(direction='maximize')
        study.optimize(objective, timeout=seconds)
        hist_df = study.trials_dataframe()
        hist_df.to_csv(os.path.join(output_path, 'optimization.csv'))
        print("Optimized params: {}".format(study.best_params))

        clf = GradientBoostingRegressor(n_estimators=study.best_params['n_estimators'], learning_rate=study.best_params['learning_rate'])
        clf = clf.fit(X_train, y_train.values.ravel())
        # Use test dataset to evaluate the model
        y_pred_test = clf.predict(X_test)
        print("Evaluation for test data: {}".format(r2_score(y_test, y_pred_test)))

        # save the model
        with open(os.path.join(model_path, 'optimized.pkl'), 'wb') as out:
            pickle.dump(clf, out)
        print('Training completed.')
    
    except Exception as e:
        # Write out an error file. This will be returned as the failureReason in the
        # DescribeTrainingJob result.
        trc = traceback.format_exc()
        with open(os.path.join(output_path, 'failure'), 'w') as s:
            s.write('Exception during training: ' + str(e) + '\n' + trc)
        # Printing this causes the exception to be in the training job logs, as well.
        print('Exception during training: ' + str(e) + '\n' + trc, file=sys.stderr)
        # A non-zero exit code causes the training job to be marked as Failed.
        sys.exit(255)

if __name__ == '__main__':
    # train execution
    train()

    # A zero exit code causes the job to be marked as Succeeded.
    sys.exit(0)

