import csv
import numpy as np
from os import makedirs
from os.path import realpath, exists, isabs, dirname, join
from optparse import OptionParser
from sklearn.ensemble import RandomForestClassifier

from context import *
from classfication.log import log

train_dataset = dict()
target_dataset = dict()


def conver_list_to_float(source):
    result = []
    if type(source) == str:
        token = source.split(' ')
    else:
        token = source
    for t in token:
        result.append(float(t))
    return result


def train_set_dataset(row):
    features = conver_list_to_float(row[2:-2][0])
    return {
        'ClusterNumber': int(row[0]),
        'FileName': str(row[1]),
        'Features': np.array(features),
        'IsMalware': float(row[-1])
    }


def target_set_dataset(row):
    features = conver_list_to_float(row[2:])
    return {
        'name': row[0],
        'feature': np.array(features)
    }


def load_train_data(train_file_path):
    with open(train_file_path) as file_object:
        reader = csv.reader(file_object, delimiter=',', quoting=csv.QUOTE_MINIMAL)
        for index, row in enumerate(reader):
            if index > 0:
                train_dataset[index] = train_set_dataset(row)
    train_cluster = train_label = (np.array([row[1]['ClusterNumber'] for row in train_dataset.items()]))
    train_morb_name = (np.array([row[1]['FileName'] for row in train_dataset.items()]))
    train_feature = (np.array([row[1]['Features'] for row in train_dataset.items()]))
    train_label = (np.array([row[1]['IsMalware'] for row in train_dataset.items()]))

    return train_cluster, train_morb_name, train_feature, train_label


def load_test_data(test_file_path):
    with open(test_file_path, 'r') as csvfile:
        reader = csv.reader(csvfile, delimiter=',', quoting=csv.QUOTE_MINIMAL)
        for index, row in enumerate(reader):
            if index > 0:
                target_dataset[index] = target_set_dataset(row)

    target_morb_name = (np.array([row[1]['name'] for row in target_dataset.items()]))
    target_feature = (np.array([row[1]['feature'] for row in target_dataset.items()]))
    return target_morb_name, target_feature


def main(train_file_path, test_file_path):
    current_module_path = dirname(realpath(__file__))
    train_cluster, train_morb_name, train_feature, train_label = load_train_data(train_file_path)
    answer_table = dict()

    # print("Train_Feature: %s" %train_feature)
    train_cluster_list = train_cluster.tolist()
    train_label_list = train_label.tolist()

    # print(train_cluster_list)
    for k in range(0, len(train_cluster_list)):
        if train_cluster_list[k] in answer_table:
            continue
        else:
            answer_table[train_cluster_list[k]] = train_label_list[k]

    target_morb_name, target_feature = load_test_data(test_file_path)

    clf = RandomForestClassifier(n_estimators=100, max_depth=2, random_state=0)
    log().info(train_feature)
    log().info(train_label)
    clf = clf.fit(train_feature, train_cluster)

    predict_result = clf.predict(target_feature)

    # make output dir
    output_base_dir = join(current_module_path, 'exported/randomforest')
    if exists(output_base_dir) is False:
        makedirs(output_base_dir)

    # with cluster label
    result_file_with_cn_path = join(output_base_dir, 'andro_result_with_cluster_number.csv')
    cn_file_object = open(result_file_with_cn_path, 'w')  # output file write
    cn_file_object.write("file,  cluster_number, class\n")

    for j in range(0, len(target_feature)):
        answer_qry = answer_table[predict_result[j]]
        cn_file_object.write(target_morb_name[j] + ", " + str(predict_result[j]) + ", " + str(answer_qry) + "\n")
    cn_file_object.close()

    # not cluster label
    result_file_path = join(output_base_dir, 'androd_result.csv')
    file_object = open(result_file_path, 'w')  # output file write
    file_object.write("file,  class\n")
    for j in range(0, len(target_feature)):
        answer_qry = answer_table[predict_result[j]]
        file_object.write(target_morb_name[j] + ", " + str(answer_qry) + "\n")


if __name__ == '__main__':
    opt_parser = OptionParser()
    opt_parser.add_option(
        '-t', '--train_file_path', dest='train_file',
        help='csv input feature table.')
    opt_parser.add_option(
        '-f', '--test_file_path', dest='test_file',
        help='csv input feature table.')
    options, _ = opt_parser.parse_args()
    if options.test_file is None or exists(options.test_file) is False:
        opt_parser.print_help()
        exit(-1)
    if options.train_file is None or exists(options.train_file) is False:
        opt_parser.print_help()
        exit(-1)
    if isabs(options.test_file) is False:
        test_file_path = realpath(options.test_file)
    else:
        test_file_path = options.test_file
    if isabs(options.train_file) is False:
        train_file_path = realpath(options.train_file)
    else:
        train_file_path = options.train_file
    main(train_file_path, test_file_path)