#!/usr/bin/env python

# Trains a classifier and classifies new instances
import argparse
import csv
import logging

from common.configuration import ConfigurationHandler
from common.workspace import WorkspaceHandler
import pandas as pd
from sklearn.tree import DecisionTreeClassifier
from sklearn.ensemble import RandomForestClassifier
from sklearn.neighbors import KNeighborsClassifier
from sklearn.naive_bayes import MultinomialNB
from sklearn.tree import _tree
import numpy as np
from subprocess import call
import gc

logger = logging.getLogger('pyccflex')
logger.setLevel(logging.DEBUG)
ch = logging.StreamHandler()
ch.setLevel(logging.INFO)
logger.addHandler(ch)


def tree_to_code(tree, feature_names):
    result = ""
    tree_ = tree.tree_
    feature_name = [
        feature_names[i] if i != _tree.TREE_UNDEFINED else "undefined!"
        for i in tree_.feature
    ]
    result += "def tree({}):\n".format(", ".join(feature_names))

    def recurse(node, depth, result):
        indent = "  " * depth
        if tree_.feature[node] != _tree.TREE_UNDEFINED:
            name = feature_name[node]
            threshold = tree_.threshold[node]
            result += "{}if '{}' <= {}:\n".format(indent, name, threshold)
            result = recurse(tree_.children_left[node], depth + 1, result)
            result += "{}else:  # if '{}' > {}\n".format(indent, name, threshold)
            result = recurse(tree_.children_right[node], depth + 1, result)
        else:
            result += "{}return {}\n".format(indent, np.argmax(tree_.value[node][0]), result)
        return result

    result = recurse(0, 1, result)
    return result


sklearn_classifiers = ("CART", "RandomForest", "KNN", "MultinomialNB")

if __name__ == '__main__':

    logger.info("\n#### Running: {}".format(__file__))

    # Parse input parameters
    parser = argparse.ArgumentParser()
    parser.add_argument("train_input_csv",
                        help="Path to input train csv file", type=str)
    parser.add_argument("classify_input_csv",
                        help="Path to input classify csv file", type=str)
    parser.add_argument('--classifier', type=str, required=True)
    parser.add_argument("--locations_config", help="Path to locations configuration file",
                        type=str, required=False, default="./locations.json")
    parser.add_argument("--files_format_config", help="Path to files format configuration file",
                        type=str, required=False, default="./files_format.json")
    parser.add_argument("--classes_config", help="Path to classes configuration file",
                        type=str, required=False, default="./classes.json")
    parser.add_argument("--classifiers_options", help="Path to classifiers options file",
                        type=str, required=False, default="./classifiers_options.json")
    parser.add_argument("--chunk_size", help="Number of lines to process in a batch",
                        type=int, required=False, default=10 ** 5)
    parser.add_argument("--output_prefix", help="Prefix used for the output file",
                        type=str, required=False, default="")

    args = vars(parser.parse_args())
    logger.info("Run parameters: {}".format(str(args)))

    locations_file_path = args['locations_config']
    files_format_file_path = args['files_format_config']
    classes_file_path = args['classes_config']
    classifiers_options_file_path = args['classifiers_options']
    chunk_size = args['chunk_size']
    output_prefix = args['output_prefix']

    classifier = args['classifier']

    try:
        locations_config = ConfigurationHandler(locations_file_path)
    except Exception as e:
        logger.error("Couldn't load configuration file {}".format(locations_file_path))
        exit(1)

    try:
        files_format_config = ConfigurationHandler(files_format_file_path)
    except Exception as e:
        logger.error("Couldn't load configuration file {}".format(files_format_file_path))
        exit(1)

    try:
        classifiers_options_config = ConfigurationHandler(classifiers_options_file_path)
    except Exception as e:
        logger.error("Couldn't load configuration file {}".format(classifiers_options_file_path))
        exit(1)

    try:
        classes_config = ConfigurationHandler(classes_file_path)
    except Exception as e:
        logger.error("Couldn't load configuration file {}".format(classes_file_path))
        exit(1)
    decision_classes = classes_config.get("classes", {})

    csv_separator = files_format_config.get("csv_sep", ",")

    workspace_dir_conf = locations_config.get('workspace_dir', None)
    workspace_dir_path = workspace_dir_conf.get("path", "")
    workspace_dir = WorkspaceHandler(workspace_dir_path)

    train_input_file = workspace_dir.get_processing_file_path(args['train_input_csv'])
    classify_input_file = workspace_dir.get_processing_file_path(args['classify_input_csv'])

    logger.info(">>>> Setting up paths complete!")

    # read and transform inputs
    logger.info(">>>> Loading and transforming inputs")
    input_raw = pd.read_csv(train_input_file, sep=csv_separator, encoding="utf-8")

    collumns_to_drop = ['id', 'class_name']
    if 'contents' in list(input_raw.columns):
        collumns_to_drop.append('contents')
    input_data = input_raw.drop(collumns_to_drop, inplace=False, axis=1)

    logger.info(">>> Preparing training data")
    Y = input_data['class_value']
    X = input_data.drop(['class_value'], inplace=False, axis=1)

    # train model
    model_options = classifiers_options_config.get(classifier, {})
    output_file_path = workspace_dir.get_results_file_path("{}classify-output-{}.csv".format(output_prefix, classifier))

    if classifier == "CART":
        logger.info(">>> {}: training model".format(classifier))
        model = DecisionTreeClassifier(**model_options)
        model.fit(X, Y)

        tree_text = tree_to_code(model, X.columns)
        output_tree_file_path = workspace_dir.get_results_file_path(
            "{}classify-output-{}-model.txt".format(output_prefix, classifier))
        with open(output_tree_file_path, "wt") as tree_file:
            tree_file.write(tree_text)

        logger.info(">>> {} model structure saved to the file {} ".format(classifier, output_tree_file_path))

    elif classifier == "RandomForest":
        logger.info(">>> {}: training model".format(classifier))
        model = RandomForestClassifier(**model_options)
        model.fit(X, Y)

    elif classifier == "KNN":
        model = KNeighborsClassifier(**model_options)
        model.fit(X, Y)

    elif classifier == "MultinomialNB":
        model = MultinomialNB(**model_options)
        model.fit(X, Y)

    # classify instances
    logger.info(">>> {}: classifying instances".format(classifier))

    # Running C50 algorithm training and classification in R
    if classifier == "C50":
        output_file_prefix = workspace_dir.get_results_file_path(
            "{}classify-output-{}".format(output_prefix, classifier))
        r_script_exec = locations_config.get("rscript_executable_path", "RScript.exe")

        call([r_script_exec, "./bin/classify_c5.R", train_input_file, classify_input_file, output_file_prefix,
              classifiers_options_file_path, classes_file_path, csv_separator])

    # Classifying instances for sklearn algorithms
    if classifier in sklearn_classifiers:
        classify_raw_chunks = pd.read_csv(classify_input_file, sep=csv_separator, encoding="utf-8",
                                          chunksize=chunk_size)

        for i, classify_raw in enumerate(classify_raw_chunks):

            collumns_to_drop = ['id']
            columns = list(classify_raw.columns)
            if 'class_value' in columns:
                collumns_to_drop.append('class_value')
                collumns_to_drop.append('class_name')
            if 'contents' in columns:
                collumns_to_drop.append('contents')
            classify_data = classify_raw.drop(collumns_to_drop, inplace=False, axis=1)

            pred = model.predict(classify_data)
            del classify_data
            gc.collect()

            # report
            output_full = classify_raw[['id', 'contents']]
            output_full = output_full.assign(pred_class=pred)
            if i == 0:
                with open(output_file_path, "w", newline='', encoding="utf-8") as f:
                    output_full.to_csv(f, sep=csv_separator, index=False, encoding="utf-8", header=True,
                                       quoting=csv.QUOTE_NONNUMERIC)
            else:
                with open(output_file_path, "a", newline='', encoding="utf-8") as f:
                    output_full.to_csv(f, sep=csv_separator, index=False, encoding="utf-8", header=False,
                                       quoting=csv.QUOTE_NONNUMERIC)
            logger.info(
                ">>> {}: {} lines of predictions saved to {}".format(classifier, output_full.shape[0],
                                                                     output_file_path))
            del output_full
            del pred
            del classify_raw
            gc.collect()

    # Reporting results for classes
    output_file_chunks = pd.read_csv(output_file_path, sep=csv_separator, encoding="utf-8",
                                     chunksize=chunk_size)

    for i, output_full in enumerate(output_file_chunks):

        for decision_class in decision_classes["labeled"]:
            output_file_path_class = workspace_dir.get_results_file_path(
                "{}classify-output-{}-{}.csv".format(output_prefix, classifier, decision_class['name']))
            output_class = output_full[(output_full['pred_class'] == decision_class['value'])]

            if i == 0:
                with open(output_file_path_class, "w", newline='', encoding="utf-8") as f:
                    output_class.to_csv(f, sep=csv_separator, index=False, encoding="utf-8", header=True,
                                        quoting=csv.QUOTE_NONNUMERIC)
            else:
                with open(output_file_path_class, "a", newline='', encoding="utf-8") as f:
                    output_class.to_csv(f, sep=csv_separator, index=False, encoding="utf-8", header=False,
                                        quoting=csv.QUOTE_NONNUMERIC)
            logger.info(
                ">>> {}: {} lines of predictions for the class \"{}\" saved to {}".format(classifier,
                                                                                          output_class.shape[0],
                                                                                          decision_class['name'],
                                                                                          output_file_path_class))
            del output_class
            gc.collect()

        output_file_path_class = workspace_dir.get_results_file_path(
            "{}classify-output-{}-{}.csv".format(output_prefix, classifier, decision_classes["default"]['name']))
        output_class = output_full[(output_full['pred_class'] == decision_classes["default"]['value'])]
        if i == 0:
            with open(output_file_path_class, "w", newline='', encoding="utf-8") as f:
                output_class.to_csv(f, sep=csv_separator, index=False, encoding="utf-8", header=True,
                                    quoting=csv.QUOTE_NONNUMERIC)
        else:
            with open(output_file_path_class, "a", newline='', encoding="utf-8") as f:
                output_class.to_csv(f, sep=csv_separator, index=False, encoding="utf-8", header=False,
                                    quoting=csv.QUOTE_NONNUMERIC)
        logger.info(
            ">>> {}: {} lines of predictions for the class \"{}\" saved to {}".format(classifier, output_class.shape[0],
                                                                                      decision_classes["default"][
                                                                                          'name'],
                                                                                      output_file_path_class))
        del output_class
        del output_full
        gc.collect()
