# -*- coding: utf-8 -*-
"""
Command line tool to do classification and valuation using different methods

Authors: Pietro Zambelli

Date: October 2014
"""
from __future__ import print_function
from __future__ import absolute_import

# from python standard library
from builtins import str
from builtins import range
import imp
import sys
import os
import argparse
from datetime import datetime
import pickle as pkl
from pprint import pprint

# import scientific libraires
import numpy as np
from sklearn.preprocessing import StandardScaler
from sklearn.decomposition import PCA
from sklearn.feature_selection import RFECV
from sklearn.ensemble import ExtraTreesClassifier
from sklearn.svm import LinearSVC
from sklearn.svm import SVC

# import local libraries
from .feature_selection import SSF
from .machine_learning import MLToolBox, SEP, NODATA, get_cv


MODELS = {# UC_26 Support Vector Machine (SVM)
          'SVM': [],
          # UC_27 Classificazione minima distanza: KNeighborsClassifier
          # UC_28 Classificazione massima verosimiglianza: EmpiricalCovariance
          # UC_29 Spectral Angle Mapper (SAM) i.spec.sam (G6)
          # UC_30 Stimatore lineare: LinearRegression
          # UC_32 Attribuzione/modifica delle classi tematiche: SVR
          'SVR': []}

TRANS = {'sqrt': np.sqrt, 'log': np.log}
UNTRANS = {'sqrt': np.exp2, 'log': np.exp}


# =======================================================================
# Parser functions
def pca_components(string):
    try:
        return int(string)
    except ValueError:
        return None if string == '' else string


def importable(path):
    return imp.load_source("classifiers", path)


def index(string, sep=',', rangesep='-'):
    """
    >>> indx = '1-5,34-36,40'
    >>> [i for i in get_indexes(indx)]
    [1, 2, 3, 4, 5, 34, 35, 36, 40]
    """
    if string == '':
        return
    for ind in string.split(sep):
        if rangesep in ind:
            start, stop = ind.split(rangesep)
            for i in range(int(start), int(stop) + 1):
                yield i
        else:
            yield int(ind)


def indexstr(string):
    return [i for i in index(string)]


# =======================================================================
# Define the parser options
def get_parser():
    """Create the parser for running as script"""
    parser = argparse.ArgumentParser(description='Test several machine-'
                                                 'learning models.')
    parser.add_argument('vector', type=str, metavar='VECTOR',
                        help='Training vector format supported by OGR')
    parser.add_argument('column', type=str, metavar='COLUMN',
                        help='Column name with the values/classes'
                             ' used as training')
    parser.add_argument('-u', '--use-columns', type=str, metavar='COLUMNS',
                        default=None, nargs='+', dest='use_columns',
                        help='Name of the columns containing the data that '
                             'will be used instad of the raster input map.')
    parser.add_argument('-r', '--raster', type=str, default=None, dest='raster',
                        help='Multi/Hyper spectral raster file supported by GDAL')
    parser.add_argument('-t', '--csv-training', type=str, dest='csvtraining',
                        default='training.csv',
                        help='Name of the CSV file with the training values'
                             ' extracted by the raster and vector map.')
    parser.add_argument('-cr', '--csv-cross', type=str, dest='csvcross',
                        default='cross_validation.csv',
                        help='Name of the CSV file with the models results.')
    parser.add_argument('-cd', '--csv-delimiter', type=str, default=SEP,
                        dest='csvdelimiter', help='CSV delimiter.')
    parser.add_argument('-m', '--models', type=importable,
                        dest='models', default='classifiers.py',
                        help='A python file containing a list of dictionary'
                             ' with the models that will be tested.')
    parser.add_argument('-n', '--list-name', type=str,
                        dest='listname', default='ALL',
                        help='The python name variable with the model list')
    parser.add_argument('-k', '--scoring', default='accuracy', dest='scoring',
                        choices=['accuracy', 'f1', 'precision',
                                 'recall', 'roc_auc', 'adjusted_rand_score',
                                 'mean_absolute_error', 'mean_squared_error',
                                 'r2'],
                        help='Set the method that will be used for the model'
                             ' evaluation. See:\nhttp://scikit-learn.org/'
                             'stable/modules/model_evaluation.html')
    parser.add_argument('-i', '--mod-indexes', type=indexstr,
                        dest='imod', default='',
                        help='Define which model will be tested'
                             ' by index, for example with the following string'
                             ' "1-5,34-36,40", only model with index:'
                             '[1, 2, 3, 4, 5, 34, 35, 36, 40] will be used.')
    parser.add_argument('-x', '--cols-indexes', type=str,
                        dest='icols', default='',
                        help='Npy input file with the columns index to use'
                              'could be use to select a subset of the dataset')
    parser.add_argument('-o', '--output-dir', type=str,
                        dest='odir', default='',
                        help='Specify the directory that will contain all the'
                             'outputs')
    parser.add_argument('-of', '--output-file', type=str,
                        dest='ofile', default=None,
                        help='Specify the output filename')
    parser.add_argument('-w', '--output-raster-name', type=str,
                        dest='rname', default='%s',
                        help='Specify the name of the output raster map that'
                             ' will be written.')
    parser.add_argument('-c', '--cross-validation', action='store_true',
                        dest='cross', default=False,
                        help='Apply the models to the input data and '
                             'write the output raster maps.')
    parser.add_argument('-e', '--execute', action='store_true',
                        dest='execute', default=False,
                        help='Apply the models to the input data and '
                             'write the output raster maps.')
    parser.add_argument('--o', '--overwrite', action='store_true',
                        dest='overwrite', default=False,
                        help='overwrite')
    parser.add_argument('-b', '--n-best', type=int,
                        dest='n_best', default=1,
                        help='Number of best model to apply on the raster.')
    parser.add_argument('-bs', '--best-strategy', default='mean',
                        dest='best_strategy', choices=['min', 'mean', 'median'],
                        help='Strategy to use to select the best classifier.')
    parser.add_argument('-bp', '--best-pickle',
                        default='best_model.pickle',
                        dest='best_pickle', type=str,
                        help='All the best models are saved and serialized '
                             'using python pickle.')
    parser.add_argument('-nf', '--n-folds', type=int, dest='n_folds',
                        default=5, help='Number of folds that we want to use'
                                        ' to cross validate the dataset.')
    parser.add_argument('-nj', '--n-jobs', type=int, dest='n_jobs',
                        default=1, help='Number of jobs that we want to use'
                                        ' during the cross validations.')
    parser.add_argument('-s', '--scale', type=bool, nargs=2, dest='scale',
                        metavar='S', help='StandardScaler options: '
                                          'with_mean, with_std')
    parser.add_argument('-p', '--pca', type=pca_components, dest='pca',
                        metavar='n_components', default=NODATA, nargs='?',
                        help='See PCA n_componets documentation')
    parser.add_argument('-f', '--feature-selection', default='None',
                        dest='fs', help='Feature selection',
                        choices=['None', 'SSF', 'RFECV', 'LinearSVC',
                                 'ExtraTreesClassifier'], )
    parser.add_argument('-ff', '--feature-selection-file', default=None,
                        dest='ff', help='File with the Feature selected.')
    parser.add_argument('-fs', '--SSF-strategy', default='min',
                        dest='SSF_strategy', choices=['min', 'mean', 'median'],
                        help='Sequential Forward Floating Feature Selection'
                             ' (SSF) strategy to select the best.')
    parser.add_argument('-tr', '--test-raster', type=str, default=None,
                        dest='traster',
                        help='Multi/Hyper spectral raster file supported by'
                             ' GDAL')
    parser.add_argument('-tv', '--test-vector', type=str, default=None,
                        dest='tvector',
                        help='Training vector format supported by OGR')
    parser.add_argument('-tc', '--test-column', type=str, default=None,
                        dest='tcolumn',
                        help='Column name with the values/classes'
                             ' used as training')
    parser.add_argument('-tp', '--test-pickle',
                        default='test_model.pickle',
                        dest='test_pickle', type=str,
                        help='All the test models results saved and serialized '
                             'using python pickle.')
    parser.add_argument('-tf', '--csv-test', type=str, dest='csvtest',
                        default='test.csv',
                        help='Name of the CSV file with the test values'
                             ' extracted by the raster and vector map.')
    parser.add_argument('--transform', choices=['sqrt', 'log', ],
                        dest='transform', default=None,
                        help='Transform regression target before to apply the'
                             'regression.')
    parser.add_argument('-l', '--labels', type=pca_components, dest='labels',
                        metavar='LABEL', default=None, nargs='+',
                        help='Label tag')
    parser.add_argument('-v', '--verbosity', type=int, choices=(0, 1, 2),
                        default=0, dest='verbosity', help='Output verbosity.')
    parser.add_argument('-nd', '--no-data', type=float, default=NODATA,
                        dest='nodata',
                        help='Name of the columns containing the data that '
                             'will be used instad of the raster input map.')
    parser.add_argument('--memory-factor', type=float, dest='memory',
                        default=1.,
                        help='Memory factor to reduce the memory consumption')
    parser.add_argument('--log', type=str, dest='log',
                        help='Log file with features data selection log')
    return parser


if __name__ == "__main__":
    parser = get_parser()
    args = parser.parse_args()

    best = None

    if args.log:
        import logging

        logging.basicConfig(filename=args.log, filemode='w',
                            level=logging.DEBUG)
    else:
        logging = None

    if args.odir and not os.path.isdir(args.odir):
        os.mkdir(args.odir)

    with open(os.path.join(args.odir, 'history.txt'), 'a') as cmd:
        cmd.write('# %s\n' % str(datetime.now()))
        cmd.write(' '.join(sys.argv) + '\n')

    # -----------------------------------------------------------------------
    # Get Models
    models = getattr(args.models, args.listname)
    # write indexes
    for i, mod in enumerate(models):
        mod['index'] = i

    # filter models by index
    if args.imod:
        models = [models[i] for i in args.imod]

    # =======================================================================
    # Instantiate the class
    mltb = MLToolBox(vector=args.vector, column=args.column,
                     use_columns=args.use_columns, raster=args.raster,
                     models=models, scoring=args.scoring,
                     n_folds=args.n_folds, n_jobs=args.n_jobs,
                     n_best=args.n_best, nodata=args.nodata,
                     tvector=args.tvector, tcolumn=args.tcolumn,
                     traster=args.traster,
                     best_strategy=getattr(np, args.best_strategy),
                     scaler=None, fselector=None, decomposer=None,
                     transform=None, untransform=None,
                     memory_factor=args.memory, logging=logging)

    # -----------------------------------------------------------------------
    # Extract training samples
    print('\nExtract training samples')
    trnpath = os.path.join(args.odir, args.csvtraining)
    if (not os.path.exists(trnpath) or args.overwrite):
        print('    From:')
        print('      - vector: %s' % mltb.vector)
        print('      - training column: %s' % mltb.column)
        if mltb.use_columns:
            print('      - use columns: %s' % mltb.use_columns)
        if mltb.raster:
            print('      - raster: %s' % mltb.raster)
        X, y = mltb.extract_training(csv_file=trnpath, delimiter=SEP,
                                     nodata=args.nodata, dtype=np.uint32)
    else:
        print('    Load from:')
        print('      - %s' % trnpath)
        dt = np.loadtxt(trnpath, delimiter=SEP, skiprows=1)
        X, y = dt[:, :-1], dt[:, -1]
    X = X.astype(float)
    print('\nTraining sample shape:', X.shape)

    fselect = {
           # Sequential Forward Floating Feature Selection
           # with Jeffries-Matusita Distance
           'SSF': SSF(strategy=getattr(np, args.SSF_strategy)),
           # Recursive Feature Elimination and Cross-Validated (RFECV)
           'RFECV': RFECV(estimator=SVC(), cv=get_cv(y, args.n_folds)),
           # Linear Support Vector Classification
           'LinearSVC': LinearSVC(C=0.01, penalty="l1", dual=False),
           # Tree-based feature selection
           'ExtraTreesClassifier': ExtraTreesClassifier(n_estimators=250,
                                                        random_state=0)
           }

    # -----------------------------------------------------------------------
    # Scale
    scaler = None
    if args.scale:
        wm, ws = args.scale
        scaler = StandardScaler(with_mean=wm, with_std=ws)

    # -----------------------------------------------------------------------
    # Feature selector
    fselector = None
    fscolumns = None
    fspath = None
    if args.ff:
        fspath = os.path.join(args.odir, args.ff)
        if (os.path.exists(fspath) and not args.overwrite):
            fscolumns = np.loadtxt(fspath)
    if args.fs != 'None':
        fselector = fselect[args.fs]

    # -----------------------------------------------------------------------
    # Decomposer
    decomposer = PCA(n_components=args.pca) if args.pca != NODATA else None

    # -----------------------------------------------------------------------
    # Transform the input data
    X = mltb.data_transform(X=X, y=y, scaler=scaler, fselector=fselector,
                            decomposer=decomposer, fscolumns=fscolumns,
                            fsfile=fspath, fsfit=True)

    # -----------------------------------------------------------------------
    # Transform the training/target
    if args.transform:
        transform = TRANS[args.transform]
        untransform = UNTRANS[args.transform]
    else:
        transform, untransform = None, None

    # -----------------------------------------------------------------------
    # Extract test samples
    Xtest, ytest = None, None
    if args.tvector and args.tcolumn:
        # extract_training(vector_file, column, csv_file, raster_file=None,
        #                  use_columns=None, delimiter=SEP, nodata=None,
        #                  dtype=np.uint32)
        print('\nExtract test samples')
        testpath = os.path.join(args.odir, args.csvtest)
        if (not os.path.exists(testpath) or args.overwrite):
            print('    From:')
            print('      - vector: %s' % mltb.tvector)
            print('      - training column: %s' % mltb.tcolumn)
            if mltb.use_columns:
                print('      - use columns: %s' % mltb.use_columns)
            if mltb.raster:
                print('      - raster: %s' % mltb.traster)
            Xtest, ytest = mltb.extract_test(csv_file=testpath,
                                             nodata=args.nodata)
            dt = np.concatenate((Xtest.T, ytest[None, :]), axis=0).T
            np.savetxt(testpath, dt, delimiter=SEP,
                       header="# last column is the training.")
        else:
            print('    Load from:')
            print('      - %s' % trnpath)
            dt = np.loadtxt(testpath, delimiter=SEP, skiprows=1)
            Xtest, ytest = dt[:, :-1], dt[:, -1]
        Xtest = Xtest.astype(float)
        print('Training sample shape:', Xtest.shape)

    # -----------------------------------------------------------------------
    # Cross Models
    if args.cross:
        print('\nCross-validation of the models')
        crosspath = os.path.join(args.odir, args.csvcross)
        bpkpath = os.path.join(args.odir, args.best_pickle)
        if (not os.path.exists(crosspath) or args.overwrite):
            cross = mltb.cross_validation(X=X, y=y, transform=transform)
            np.savetxt(crosspath, cross, delimiter=args.csvdelimiter, fmt='%s',
                       header=args.csvdelimiter.join(['id', 'name', 'mean', 'max',
                                                      'min', 'std', 'time']))
            mltb.find_best(models)
            best = mltb.select_best()
            with open(bpkpath, 'w') as bpkl:
                pkl.dump(best, bpkl)
        else:
            print('    Read cross-validation results from file:')
            print('      -  %s' % crosspath)
            with open(bpkpath, 'r') as bpkl:
                best = pkl.load(bpkl)
            order, models = mltb.find_best(models=best)
            best = mltb.select_best(best=models)
        print('\nBest models:')
        pprint(best)

    # -----------------------------------------------------------------------
    # test Models
    if Xtest is not None and ytest is not None:
        print('\nTest models with an indipendent dataset')
        testpath = os.path.join(args.odir, args.csvtest)
        bpkpath = os.path.join(args.odir, args.test_pickle)
        if (not os.path.exists(testpath) or args.overwrite):
            Xtest = mltb.data_transform(X=Xtest, y=ytest)
            test = mltb.test(Xtest=Xtest, ytest=ytest, X=X, y=y,
                             transform=transform)
            np.savetxt(testpath, test, delimiter=args.csvdelimiter, fmt='%s',
                       header=args.csvdelimiter.join(list(test[0].__dict__.keys())))
            mltb.find_best(models, strategy=lambda x: x, key='score_test')
            best = mltb.select_best()
            with open(bpkpath, 'w') as bpkl:
                pkl.dump(best, bpkl)
        else:
            with open(bpkpath, 'r') as bpkl:
                best = pkl.load(bpkl)
            order, models = mltb.find_best(models=best, strategy=lambda x: x, key='score_test')
            best = mltb.select_best(best=models)
        print('Best models:')
        pprint(best)

    # -----------------------------------------------------------------------
    # execute Models and save the output raster/vector map
    if args.execute:
        if best is None:
            order, models = mltb.find_best(models, strategy=lambda x: x, key='score')
            best = mltb.select_best(best=models)

        print('\Execute the model to the whole raster map.')
        mltb.execute(X=X, y=y, best=best, output_file=args.ofile,
                     transform=transform, untransform=untransform)

    print('Finished!')
