#!/usr/bin/env python
# coding: utf-8

# **Packages**

#Some import
import numpy as np
#from sklearn.tree import export_text
from tqdm import tqdm
import time
import sys

import my_tree as mt
import pandas as pd
from sklearn.ensemble import RandomForestClassifier
from sklearn.model_selection import LeaveOneGroupOut
import os
import functools
import operator
from random import shuffle
import copy
import my_forest as mf

import dask
from multiprocessing import freeze_support



import json

def fusion_dict(list_dict) :
    output = list_dict[0]
    for d in list_dict[1::] :
        for k in d.keys() :
            output[k] += d[k]
    for k in output.keys() :
        if type(output[k]) != list :
            output[k] /= len(list_dict)
    return output
 

def rf_cross_validation(X, Y, nb_tree, cv=10, max_depth = None, groups = None, nb_forest = None) :
    nb_instance = len(Y)
    quotient = nb_instance // cv
    reste = nb_instance % cv
    if nb_forest is None :
        nb_forest = cv
    if groups is None :
        groups = [quotient*[i] for i in range(1,cv+1)]
        groups = functools.reduce(operator.iconcat, groups, [])
        groups += [i for i in range(1,reste+1)]
        shuffle(groups)
    loo = LeaveOneGroupOut()
    score = 0
    forests = []
    i = 0
    for ind_train, ind_test in loo.split(X, Y, groups=groups):
        if i < nb_forest :
            i += 1
            x_train = [X[x] for x in ind_train]
            y_train = [Y[x] for x in ind_train]
            x_test = [X[x] for x in ind_test]
            y_test = [Y[x] for x in ind_test]
            rf = RandomForestClassifier(max_depth=max_depth, n_estimators=nb_tree)
            rf.fit(x_train, y_train)
            y_predict = rf.predict(x_test)
            accuracy = (np.sum(y_predict == y_test)/len(y_test))*100
            score += accuracy
            forests.append((copy.deepcopy(rf),ind_train,ind_test))
    score /= nb_forest
    return score, forests, groups    


#Set parameter for testing

def work_on_forest(forest, d , l, acc, data_name, name = None) :
    
    os.chdir("../scripts")
    
    total_record = {}
    indice = 0
    
    for f in forest :
    
        record = {}
            
        #Creation of a my_tree
        my_forest = []
        for t in range(len(f[0].estimators_)) :
            my_tree = mt.decision_tree()
            my_tree.from_DecisionTreeClassifier(f[0].estimators_[t])
            my_forest.append(my_tree)
        my_forest = mf.decision_forest(my_forest)
        my_forest.labels = f[0].classes_
        
        X_train = np.array([d[i] for i in f[1]])
        
        x_test = np.array([d[i] for i in f[2]])
        y_test = np.array([l[i] for i in f[2]])
        
        nb_echantillon = np.min([25,len(y_test)])
        
        
        #Creation of list to store some results
        
        len_dir_reason = np.zeros((nb_echantillon, 2))
        len_min_maj_reason = np.zeros((nb_echantillon, 2))
        len_pref_shap_reason = np.zeros((nb_echantillon, 2))
        len_pref_f_importance_reason = np.zeros((nb_echantillon, 2))
        len_pref_wordf_reason = np.zeros((nb_echantillon, 2))
        len_pref_strate_wordf_reason = np.zeros((nb_echantillon, 2))
        
        len_enum_min_maj_reason = []
        len_enum_pref_shap_reason = []
        len_enum_pref_f_importance_reason = []
        len_enum_wordf_importance_reason = []
        len_enum_strate_wordf_importance_reason = []
        
        to = 60
        
        list_reason = []
        
        list_instance = []
        
        nb_tricky_instance = 0
        
        list_ok = []
        
        new_bina = {}
        for k_bin in my_forest.bina :
            new_bina[my_forest.bina[k_bin][0]] = [k_bin[0], k_bin[1]]
        
        with open(f'{name}.json', 'w') as file:
            file.write(json.dumps(new_bina, indent = 4))
        
        for j in tqdm(range(0,nb_echantillon)) : #x_test.shape[0]) :
            
            # assert forests[0][0].predict(x_test[j].reshape(1,-1)) == my_forest.predict(x_test[j])[0]
            try :
            
                if my_forest.predict(x_test[j])[0]==y_test[j] :
                    if type(y_test[j]) == np.int64 :
                        list_ok.append([True, int(y_test[j])])
                    else :
                        list_ok.append([True, y_test[j]])
                else :
                    if type(y_test[j]) == np.int64 :
                        list_ok.append([False, int(y_test[j])])
                    else :
                        list_ok.append([False, y_test[j]])
                
                
                
                    
                list_instance.append(x_test[j].tolist())
                
                #Just one reason
                tps = -time.time()
                direct_reason = my_forest.find_direct_reason(x_test[j])
                tps += time.time()
                len_dir_reason[j] = [len(direct_reason), tps]
                list_reason.append([direct_reason])
                
                
                #enumeration
                min_maj_reasons, tps = my_forest.find_pseudo_min_reason(x_test[j], nb=np.inf, time_out=to)
                len_vec = []
                for r in min_maj_reasons :
                    len_vec.append(len(r))
                if min_maj_reasons != [] :
                    len_enum_min_maj_reason.append([len_vec, tps, to > tps])
                    list_reason[-1].append(min_maj_reasons[0])
                else :
                    len_enum_pref_shap_reason.append([[-1], tps, False])
                    list_reason[-1].append([])
                
                prefered_reasons, tps = my_forest.find_prefered_majoritary_reason(x_test[j], rf=f[0], method = 'shapely', nb=np.inf, time_out = to, solver = "openwbo")
                len_vec = []
                for r in prefered_reasons :
                    len_vec.append(len(r))
                if prefered_reasons != [] :
                    len_enum_pref_shap_reason.append([len_vec, tps, prefered_reasons[-1][0] != 0])
                    list_reason[-1].append(prefered_reasons[0])
                else :
                    len_enum_pref_shap_reason.append([[-1], tps, False])
                    list_reason[-1].append([])
                    
                
                prefered_reasons, tps = my_forest.find_prefered_majoritary_reason(x_test[j], rf=f[0], method = 'f_importance', nb=np.inf, time_out=to, solver = "openwbo")
                len_vec = []
                for r in prefered_reasons :
                    len_vec.append(len(r))
                if prefered_reasons != [] :
                    len_enum_pref_f_importance_reason.append([len_vec, tps, prefered_reasons[-1][0] != 0])
                    list_reason[-1].append(prefered_reasons[0])
                else :
                    len_enum_pref_shap_reason.append([[-1], tps, False])
                    list_reason[-1].append([])
                
                prefered_reasons, tps = my_forest.find_prefered_majoritary_reason(x_test[j], rf=f[0], col=data_name, method = 'wordfreq', nb=np.inf, time_out=to, solver = "openwbo")
                len_vec = []
                for r in prefered_reasons :
                    len_vec.append(len(r))
                if prefered_reasons != [] :
                    len_enum_wordf_importance_reason.append([len_vec, tps, prefered_reasons[-1][0] != 0])
                    list_reason[-1].append(prefered_reasons[0])
                else :
                    len_enum_pref_shap_reason.append([[-1], tps, False])
                    list_reason[-1].append([])
                
                
            except TimeoutError :
                nb_tricky_instance += 1
                print(nb_tricky_instance)
            
        
        record["acc"] = acc
        record["nb_instance_passed"] = nb_tricky_instance
        record["instance"] = list_instance
        record["classified"] = list_ok
        record["len_bin"] = [len(my_forest.bina.keys()), x_test.shape[1]]
        record["ordre"]=[name]
        
        record["dir_r"] = len_dir_reason.tolist()
        record["min_maj_r"] = len_min_maj_reason.tolist()
        record["pref_shap_r"] = len_pref_shap_reason.tolist()
        record["pref_f_import_r"] = len_pref_f_importance_reason.tolist()
        record["pref_wordf_r"] = len_pref_wordf_reason.tolist()
        record["pref_strate_wordf_r"] = len_pref_strate_wordf_reason.tolist()
        
        record["enum_min_maj_r"] = len_enum_min_maj_reason
        record["enum_pref_shap_r"] = len_enum_pref_shap_reason
        record["enum_pref_f_import_r"] = len_enum_pref_f_importance_reason
        record["enum_wordf_import_r"] = len_enum_wordf_importance_reason
        
        record["reason"] = list_reason
        
        total_record[f"forest_{indice}"] = record
        indice += 1
    
    #record["dir_proto_r"] = len_dir_proto_reason.tolist()
    
    
    return record 

def main() :
    freeze_support()

    dir = "../"
    os.chdir(dir+"datasets")
    fichiers = os.listdir()
    
    database = {}
    
    fichiers = [arg +".csv" for arg in sys.argv[1:]]#, "dorothea.csv", "dexter.csv", "mnist38.csv", "adult.csv"]
    
    for f in fichiers :
        database[f.split(".")[0]] = pd.read_csv(f)
    
    #fichiers = os.listdir()
    #fichiers = fichiers[76::]
    
    os.chdir(dir+"scripts")
    
    info = "info_data_RF.json"
    with open(info) as json_data:
        info_nb_tree = json.load(json_data)
        
    for dataset in database.keys() :
        
        print(f"Work on {dataset}")
        
        nb_tree = info_nb_tree[dataset]
        if nb_tree % 2 == 0:
            nb_tree += 1
        cv = 10
        print("nb forest = ", nb_tree)
        data = database[dataset].copy()
        size = data.shape
        labels = data[data.columns[-1]]
        labels = np.array(labels)
        data = data.drop(columns=[data.columns[-1]])
        data_name = list(data.columns)
        data = data.values
        
        score, forests, groups = rf_cross_validation(data, labels, nb_tree, cv=cv, nb_forest = cv)
        
        
        list_record = []
        num_rf = 0
        print(score)
        #Generate a place to save data
        for f in range(0, cv) : 
            record = dask.delayed(work_on_forest)([forests[f]], data, labels, score, data_name, name = dataset + "_" + str(num_rf))
            num_rf += 1
            list_record.append(record)
            
        final_record = dask.delayed(fusion_dict)(list_record)
        #final_record.visualize()
        final_results = final_record.compute(scheduler="processes")
        
        
        
        os.chdir(dir+"result_RF_IJCAI-local")
        
                
        with open(f'{dataset}_RF_stat.json', 'w') as file:
            file.write(json.dumps(final_results, indent = 4))
        
        #done.append(dask.delayed(work_on_forest, pure=True)(deepcopy(forests[0]), deepcopy(data), name = deepcopy(dataset)))
    #result = dask.delayed(print)(done)
    #with ProgressBar():
        #result.compute(scheduler="processes", optimize_graph=False, num_workers=5)
    #print(result)
    
if __name__ == '__main__':
    main()
            
        
