"""
    TRAINER SCRIPT - contains custom code for training your model
"""
import os
import argparse
import pickle

import pandas as pd
from sklearn.model_selection import train_test_split
from sklearn.preprocessing import StandardScaler
from sklearn.svm import SVC

from src.container.lifecycle import training_container

@training_container.wrap
def run_training(environment, logger):
    """Run training worflow"""
    parser = argparse.ArgumentParser()
    # get parameters
    parser.add_argument(
        '--iris', type=str,
        default=os.path.join(environment.input_data_dir, 'iris'),
        help='Expects a .csv file, representing the lookup.'
    )
    parser.add_argument(
        '--test_size', type=float,
        default=environment.hyperparameters.get('test_size', 0.33),
        help='size of test sample to evaluate on.'
    )
    # Get args
    args, _ = parser.parse_known_args()

    data = pd.read_csv(os.path.join(args.iris, 'train.csv'))
    X, y = data.drop(['target'], axis=1), data['target']
    X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=args.test_size, random_state=123, shuffle=True)

    transformer = StandardScaler()
    transformer.fit(X_train)
    X_train_std = transformer.transform(X_train)
    X_test_std = transformer.transform(X_test)

    logger.info('Running Training')

    model = SVC(kernel='rbf', random_state=0, gamma=.10, C=1.0)
    model.fit(X_train_std, y_train)

    logger.info('Running Evaluation')
    logger.info(f'train_accuracy={model.score(X_train_std, y_train)};')
    logger.info(f'test_accuracy={model.score(X_test_std, y_test)};')

    logger.info('Save Model artifacts')
    transformer_path = os.path.join(environment.model_dir, 'iris/transformer.pkl')

    with open(transformer_path, 'wb') as file_:
        pickle.dump(transformer, file_)

    model_path = os.path.join(environment.model_dir, 'iris/model.pkl')

    with open(model_path, 'wb') as file_:
        pickle.dump(model, file_)
