# -*- coding: utf-8 -*-
"""
Created on Wed Jul  8 21:56:55 2015.

@author: alex, rc

This script contain code to generate lvl1 model prediction.
usage : python genPreds.py model_name mode
with mode = val for validation and val = test for test.

This script will read the model description from the yaml file, load
dependencies, create preprocessing and classification pipeline and apply them
on raw data independently on each subjects.

This script support caching of preprocessed data, in order to allow reuse of
preprocessing pipeline across model.
"""
import os
import sys
if __name__ == '__main__' and __package__ is None:
    filePath = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
    sys.path.append(filePath)

import numpy as np
import pandas as pd
from time import time
from copy import deepcopy
import yaml
from sklearn.pipeline import make_pipeline, Pipeline
from progressbar import Bar, ETA, Percentage, ProgressBar, RotatingMarker

from sklearn.metrics import roc_auc_score

from preprocessing.aux import getEventNames, load_raw_data

from multiprocessing import Pool
cols = getEventNames()


def _from_yaml_to_func(method, params):
    """go from yaml to method.

    Need to be here for accesing local variables.
    """
    prm = dict()
    if params is not None:
        for key, val in params.iteritems():
            prm[key] = eval(str(val))
    return eval(method)(**prm)


def doCols(col):
    """Train and Predict for one event."""
    p = []
    for clf in clfs:
        clf.fit(trainPreprocessed, labels_train[:, col])
        p.append(clf.predict_proba(testPreprocessed)[:, 1])
    return p


yml = yaml.load(open(sys.argv[1]))

# Import package
for pkg, functions in yml['imports'].iteritems():
    stri = 'from ' + pkg + ' import ' + ','.join(functions)
    exec(stri)

# meta settings
fileName = yml['Meta']['file']
cores = yml['Meta']['cores']
subsample = yml['Meta']['subsample']
cache_preprocessed = yml['Meta']['cachePreprocessed']

if 'subsample_test' in yml['Meta'].keys():
    subsample_test = yml['Meta']['subsample_test']
else:
    subsample_test = 1

if 'addPreprocessed' in yml['Meta']:
    addPreprocessed = yml['Meta']['addPreprocessed']
else:
    addPreprocessed = []

# preprocessing pipeline
pipe = []
for item in yml['Preprocessing']:
    for method, params in item.iteritems():
        pipe.append(_from_yaml_to_func(method, params))
preprocess_base = make_pipeline(*pipe)

# post preprocessing
postpreprocess_base = None
if 'PostPreprocessing' in yml.keys():
    pipe = []
    for item in yml['PostPreprocessing']:
        for method, params in item.iteritems():
            pipe.append(_from_yaml_to_func(method, params))
    postpreprocess_base = make_pipeline(*pipe)

# models
clfs = []
for mdl in yml['Models']:
    clfs.append('Pipeline([ %s ])' % mdl)

for i, clf in enumerate(clfs):
    clfs[i] = eval(clf)

# ## read arguments ###

mode = sys.argv[2]
if mode == 'val':
    test = False
elif mode == 'test':
    test = True
else:
    raise('Invalid mode. Please specify either val or test')

if test:
    folder = 'test/'
    prefix = 'test_'
else:
    folder = 'val/'
    prefix = 'val_'


print 'Running %s, to be saved in file %s' % (mode, fileName)

saveFolder = folder + fileName
if not os.path.exists(saveFolder):
    os.makedirs(saveFolder)

# #### define lists #####
subjects = range(1, 13)
widgets = ['Cross Val : ', Percentage(), ' ', Bar(marker=RotatingMarker()),
           ' ', ETA(), ' ']
pbar = ProgressBar(widgets=widgets, maxval=len(subjects))
pbar.start()

report = pd.DataFrame(index=[fileName])
start_time = time()
# #### generate predictions #####
for subject in subjects:
    print 'Loading data for subject %d...' % subject
    # ############### READ DATA ###############################################
    data_train, labels_train, data_test, labels_test = load_raw_data(subject,
                                                                     test)

    trainPreprocessed = None
    testPreprocessed = None
    cacheFile = '%s/train_sub%d.npy' % (saveFolder, subject)
    # copy processing pipeline to start fresh
    preprocess = deepcopy(preprocess_base)
    if postpreprocess_base is not None:
        postpreprocess = deepcopy(postpreprocess_base)
    else:
        postpreprocess = None

    # ### preprocessing ####
    print 'Preprocessing Training data...'

    if cache_preprocessed and os.path.isfile(cacheFile):
        # if cache activated + file exist, load file
        trainPreprocessed = np.load(cacheFile)
    else:
        # if not, do preprocessing
        trainPreprocessed = preprocess.fit_transform(data_train, labels_train)
        # if cache activated but no file, save
        if cache_preprocessed:
            np.save(cacheFile, trainPreprocessed)
            trainPreprocessed = None
    data_train = None

    print 'Preprocessing Test data...'
    cacheFile = '%s/test_sub%d.npy' % (saveFolder, subject)

    # update subsampling for test Preprocessing
    for name, step in preprocess.steps:
        if hasattr(step, 'update_subsample'):
            step.update_subsample(subsample, subsample_test)

    if cache_preprocessed and os.path.isfile(cacheFile):
        # if cache activated + file exist, load file
        testPreprocessed = np.load(cacheFile)
    else:
        # if not, do preprocessing
        testPreprocessed = preprocess.transform(data_test)
        # if cache activated but no file, save
        if cache_preprocessed:
            np.save(cacheFile, testPreprocessed)
    data_test = None

    print 'Post Preprocessing data...'
    if cache_preprocessed and (trainPreprocessed is None):
        # if cache activated load file
        cacheFile = '%s/train_sub%d.npy' % (saveFolder, subject)
        trainPreprocessed = np.load(cacheFile)

    # Add preprocessed feature if they have been set in the config file
    for feat_name in addPreprocessed:
        featFile = '%s/%s/train_sub%d.npy' % (folder, feat_name, subject)
        if os.path.isfile(featFile):
            feat = np.load(featFile)
            if trainPreprocessed is None:
                trainPreprocessed = feat
            else:
                trainPreprocessed = np.c_[trainPreprocessed, feat]
            feat = None
        else:
            raise ValueError("File %s does not exist" % featFile)

    # Add preprocessed feature if they have been set in the config file
    for feat_name in addPreprocessed:
        featFile = '%s/%s/test_sub%d.npy' % (folder, feat_name, subject)
        if os.path.isfile(featFile):
            feat = np.load(featFile)
            if testPreprocessed is None:
                testPreprocessed = feat
            else:
                testPreprocessed = np.c_[testPreprocessed, feat]
            feat = None
        else:
            raise ValueError('File %s does not exist' % featFile)

    trainPreprocessed[np.isnan(trainPreprocessed)] = 0
    testPreprocessed[np.isnan(testPreprocessed)] = 0

    if postpreprocess is not None:
        trainPreprocessed = postpreprocess.fit_transform(trainPreprocessed,
                                                         labels_train)
        for name, step in postpreprocess.steps:
            if hasattr(step, 'update_subsample'):
                step.update_subsample(subsample, subsample_test)

        testPreprocessed = postpreprocess.transform(testPreprocessed)

    print 'Training models...'
    labels_train = labels_train[::subsample]
    if cores == 1:
        preds = [doCols(i) for i in range(len(cols))]
    else:
        pool = Pool(processes=cores)
        preds = pool.map(doCols, range(len(cols)))
        pool.close()
    # ### results #####
    print 'Aggregating results...'
    for i in range(len(clfs)):
        pred_i = [j[i] for j in preds]
        pred_i = np.array(np.vstack(pred_i)).transpose()
        np.save('%s/sub%d_clf%d.npy' % (saveFolder, subject, i), pred_i)
        if not test:
            auc = np.mean([roc_auc_score(trueVals, p) for trueVals, p in
                          zip(labels_test[::subsample_test].T, pred_i.T)])
            print '%d, clf %d: %.5f' % (subject, i, auc)

    # clear memory
    preds = None
    trainPreprocessed = None
    testPreprocessed = None

    # update progress Bar
    pbar.update(subject)


if not test:
    labels = np.load('../infos_val.npy')[:, :-2]

# ## AGGREGATE HERE
preds_tot = []

for i in range(len(clfs)):
    preds_tot.append([])
    for subject in subjects:
        preds_tot[i].append(np.load('%s/sub%d_clf%d.npy' % (saveFolder,
                                                            subject, i)))
    preds_tot[i] = np.concatenate(preds_tot[i])
    if not test:
        auc = [roc_auc_score(trueVals, p) for trueVals, p in
               zip(labels[::subsample_test].T, preds_tot[i].T)]
        report['AUC'] = np.mean(auc)
        print np.mean(auc)

# ## save the model ###
np.save(folder + prefix + fileName + '.npy', preds_tot)

# ## save report
end_time = time()
report['Time'] = end_time - start_time
report.to_csv("report/%s_%s.csv" % (prefix, fileName))
print report