#!/usr/bin/env python3

# Copyright (c) 2016 Aleksey Cheusov <vle@gmx.net>
#
# Permission is hereby granted, free of charge, to any person obtaining
# a copy of this software and associated documentation files (the
# "Software"), to deal in the Software without restriction, including
# without limitation the rights to use, copy, modify, merge, publish,
# distribute, sublicense, and/or sell copies of the Software, and to
# permit persons to whom the Software is furnished to do so, subject to
# the following conditions:
#
# The above copyright notice and this permission notice shall be
# included in all copies or substantial portions of the Software.
#
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND,
# EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF
# MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND
# NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE
# LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION
# OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION
# WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.

import pickle
import sys
import getopt

from svmlight_loader import load_svmlight_file
from sklearn import ensemble

def usage():
    print('''scikit_rf-train builds the model using sklearn.ensemble.RandomForestClassifier
usage: scikit_rf-train [OPTIONS] <dataset> <model>
OPTIONS:

   Most options below directly correspond to the arguments of
   sklearn.ensemble.RandomForestClassifier constructor. So, for
   detailed information see scikit learn documentation.

   -h   --help                     Display this screen.
        --estimators <num>         The number of estimator (default: 10).
        --criterion <criterion>    Function to measure the quality of a split
                                   (default: gini).
        --max_features <features>  The number of features to consider
                                   when looking for the best split.
        --max_depth <depth>        The maximum depth of the tree.
        --min_samples_split <num>  The minimum number of samples required
                                   to split an internal node.
        --min_samples_leaf <num>   The minimum number of samples required
                                   to split an internal node.
        --max_leaf_nodes <num>     Grow trees with max_leaf_nodes in
                                   best-first fashion. Best nodes are defined
                                   as relative reduction in impurity.
                                   If None then unlimited number of leaf nodes
                                   (default: None).
        --min_impurity_split <t>   Threshold for early stopping in tree growth.
                                   A node will split if its impurity is above
                                   the threshold, otherwise it is a leaf
                                   (default: 1e-7).
        --bootstrap <flag>         Whether bootstrap samples are used
                                   when building trees (default: True).
        --oob_score <flag>         Whether to use out-of-bag samples to
                                   estimate the generalization accuracy
                                   (default: False).
        --jobs <num>               The number of jobs to run in parallel
                                   for both fit and predict. If -1, then the
                                   number of jobs is set to the number
                                   of cores (default: 1).
        --random_state <seed>      Random seed (default: None).
        --verbose <num>            Controls the verbosity of the tree building
                                   process (default: 0).
        --warm_start <flag>        When set to True, reuse the solution of
                                   the previous call to fit and add more
                                   estimators to the ensemble, otherwise, just
                                   fit a whole new forest (default: False).
        --class_weight <weights>   Ignored (default: None).
        --extra <name=value>       Extra parameter to classifier constructor.
''')

def string_to_value(s):
    if s == "None":
        return None
    elif s == "":
        return ""
    elif s[0].isdigit():
        try:
            return int(s)
        except:
            return float(s)
    else:
        return s

def coerce(x):
    try:
        return int(x)
    except:
        return float(x)

def str2bool(v):
  return v.lower() in ("yes", "true", "y", "t", "1")

kwargs_constructor = {}
kwargs_fit = {}

#
opts,args = getopt.getopt(sys.argv[1:], 'h',
                          ['help', 'estimators=', "criterion=",
                           'max_features=', 'max_depth=', 'min_samples_split=',
                           'min_samples_leaf=', 'min_weight_fraction_leaf=',
                           'max_leaf_nodes=', 'min_impurity_split=',
                           'bootstrap=', 'oob_score=', 'jobs=', 'random_state=',
                           'verbose=', 'warm_start=', 'class_weight=', "extra="]
)

for o,a in opts:
    if o in ("-h", "--help", ):
        usage()
        sys.exit()
    if o in ("--estimators", ):
        kwargs_constructor['n_estimators'] = int(a)
    if o in ("--criterion", ):
        kwargs_constructor['criterion'] = a
    if o in ("--max_features", ):
        kwargs_constructor['max_features'] = a
        if a in ('auto', 'sqrt', 'log2'):
            pass
        elif a == 'None':
            kwargs_constructor['max_features'] = None
        else:
            kwargs_constructor['max_features'] = coerce(a)
    if o in ("--max_depth", ):
        kwargs_constructor['max_depth'] = a
        if a == 'None':
            kwargs_constructor['max_depth'] = None
        else:
            kwargs_constructor['max_depth'] = int(a)
    if o in ("--min_samples_split", ):
        kwargs_constructor['min_samples_split'] = coerce(a)
    if o in ("--min_samples_leaf", ):
        kwargs_constructor['min_samples_leaf'] = coerce(a)
    if o in ("--min_weight_fraction_leaf", ):
        kwargs_constructor['min_weight_fraction_leaf'] = float(a)
    if o in ("--max_leaf_nodes", ):
        if a == 'None':
            kwargs_constructor['max_leaf_nodes'] = None
        else:
            kwargs_constructor['max_leaf_nodes'] = int(a)
    if o in ("--min_impurity_split", ):
        kwargs_constructor['min_impurity_split'] = float(a)
    if o in ("--bootstrap", ):
        kwargs_constructor['bootstrap'] = str2bool(a)
    if o in ("--oob_score", ):
        kwargs_constructor['oob_score'] = str2bool(a)
    if o in ("--jobs", ):
        kwargs_constructor['jobs'] = int(a)
    if o in ("--random_state", ):
        if a == 'None':
            kwargs_constructor['random_state'] = None
        else:
            kwargs_constructor['random_state'] = int(a)
    if o in ("--verbose", ):
        kwargs_constructor['verbose'] = int(a)
    if o in ("--warm_start", ):
        kwargs_constructor['warm_start'] = str2bool(a)
    if o in ("--extra"):
        sep = a.find("=")
        kwargs_constructor[a[:sep]] = string_to_value(a[sep+1:])

X_train, y_train = load_svmlight_file(args[0])

clf = ensemble.RandomForestClassifier(**kwargs_constructor)

clf.fit(X_train, y_train, **kwargs_fit)
with open(args[1], 'wb') as f:
    pickle.dump(clf, f)
    pickle.dump(X_train.shape[1], f)
