#!/usr/bin/env python3
#
# Trains a baseline Gaussian process classifier with an RBF kernel on
# the features generated by the logistic regression.

from ismb2020_maldi.datasets import AntibioticResistanceDataset
from ismb2020_maldi.datasets import EcoliAntibioticResistanceDataset
from ismb2020_maldi.datasets import KpneuAntibioticResistanceDataset
from ismb2020_maldi.datasets import SaureusAntibioticResistanceDataset

from maldi_learn.preprocessing import TotalIonCurrentNormalizer
from maldi_learn.preprocessing import SubsetPeaksTransformer
from maldi_learn.vectorization import BinningVectorizer

from imblearn.over_sampling import RandomOverSampler

from sklearn.gaussian_process import GaussianProcessClassifier
from sklearn.gaussian_process.kernels import RBF

from sklearn.exceptions import ConvergenceWarning
from sklearn.linear_model import LogisticRegression
from sklearn.metrics import average_precision_score
from sklearn.metrics import accuracy_score
from sklearn.model_selection import GridSearchCV
from sklearn.model_selection import StratifiedKFold
from sklearn.preprocessing import StandardScaler
from sklearn.pipeline import Pipeline

from joblib import parallel_backend

import numpy as np
import json_tricks as jt

import argparse
import os
import warnings


if __name__ == '__main__':

    parser = argparse.ArgumentParser()
    parser.add_argument('-s', '--species', type=str, required=True)
    parser.add_argument('-a', '--antibiotic', type=str, required=True)
    parser.add_argument('-S', '--seed', type=int, required=False, default=2020)
    parser.add_argument('-o', '--output', type=str)
    parser.add_argument('-n', '--normalize', action='store_true')

    args = parser.parse_args()

    species_to_dataset = {
        'ecoli': EcoliAntibioticResistanceDataset,
        'kpneu': KpneuAntibioticResistanceDataset,
        'saureus': SaureusAntibioticResistanceDataset
    }

    dataset = species_to_dataset[args.species](
                test_size=0.20,
                antibiotic=args.antibiotic,
                random_seed=args.seed,
                suffix='_peaks_warped'
            )

    X_train, y_train = dataset.training_data
    X_test, y_test = dataset.testing_data

    # Perform random oversampling in order to ensure class balance. This
    # is strictly speaking not required but we do it for the GP as well,
    # so in the interest of comparability, we have to do it here.

    ros = RandomOverSampler(random_state=args.seed)

    X_indices = np.asarray(
        [i for i in range(0, len(X_train))]).reshape(-1, 1)

    X_indices, y_train = ros.fit_sample(X_indices, y_train)
    X_train = np.take(X_train, X_indices.ravel())

    # Static information about the data set; will be extended later on
    # with information about the training itself.
    data = {
        'seed': args.seed,
        'species': args.species,
        'antibiotic': args.antibiotic,
        'spectra_path': os.getenv('ANTIBIOTICS_SPECTRA_PATH'),
        'endpoint_path': os.getenv('ANTIBIOTICS_ENDPOINT_PATH'),
        'normalize': args.normalize,
    }

    # This is the mode of the experiments run for the logistic
    # regression pipeline.
    n_bins = 3600

    data['n_bins'] = n_bins

    # Define pipeline and cross-validation setup

    pipeline = Pipeline(
        [
            ('bv', BinningVectorizer(
                    n_bins=n_bins,
                    min_bin=2000,
                    max_bin=20000)),
            ('gp', GaussianProcessClassifier(
                    kernel=RBF(),
                   )
            )
        ],
        memory=os.getenv('TMPDIR', default=None),
    )

    with warnings.catch_warnings():
        warnings.filterwarnings('ignore', category=ConvergenceWarning)
        warnings.filterwarnings('ignore', category=UserWarning)

        # Let's do the fitting in parallel, but the prediction can be done
        # without additional threading.
        with parallel_backend('loky', n_jobs=-1):
            pipeline.fit(X_train, y_train)

        data['kernel'] = repr(pipeline['gp'].kernel_.theta)

    # AUPRC

    y_pred = pipeline.predict_proba(X_test)
    average_precision = average_precision_score(y_test, y_pred[:, 1])

    data['average_precision'] = 100 * average_precision

    # Accuracy

    y_pred = pipeline.predict(X_test)
    accuracy = accuracy_score(y_test, y_pred)

    data['accuracy'] = 100 * accuracy

    if args.output is not None:
        with open(args.output, 'w') as f:
            jt.dump(data, f, indent=4)
    else:
        print(jt.dumps(data, indent=4))
