from learn_multi_class import *
import numpy as np
import pandas as pd
import pickle
import time
import sys

'''
Main file for training random forest:

RUN CODE:
python main_learn.py <N number of networks> <bool (0/1) for balanced data> <maxL>
maxL: in the data set, maxL is the maximum number of leaves of an instance group.

EXAMPLE: 
python LearningCherries/main_learn.py 1000 1 100
'''

if __name__ == "__main__":
    N = int(sys.argv[1])
    balanced = int(sys.argv[2])
    max_l = int(sys.argv[3])

    # PREPARE METADATA OF ML MODEL
    metadata = pd.Series(index=["N", "maxL", "num_data", "data_non_cher", "data_cher",
                                "data_ret_cher", "data_no_ret_cher", "acc_all", "acc_non_cher", "acc_cher",
                                "acc_ret_cher", "acc_no_ret_cher", "train_dur", "balanced"], dtype=float)
    metadata["balanced"] = balanced
    metadata["maxL"] = max_l

    # rf model name
    data_set_string = f"N{N}_maxL{max_l}_random"
    if balanced:
        data_set_string += "_balanced"

    print(f"Start {data_set_string}")
    output_list = []
    data_set = f"../ALL_ML_data_N{N}_maxL{max_l}_random.pickle"
    with open(data_set, "rb") as handle:
        output_list.append(pickle.load(handle))

    # Clean data by deleting bad data points
    X_new = output_list[0]["X"]
    Y_new = output_list[0]["Y"]
    for output in output_list[1:]:
        X_new = pd.concat([X_new, output["X"]])
        Y_new = pd.concat([Y_new, output["Y"]])

    # make new index for X and Y first
    X_new.index = np.arange(len(X_new))
    Y_new.index = X_new.index
    X_new.replace([np.inf, -np.inf], np.nan, inplace=True)
    X_new.dropna(inplace=True)
    Y_new = Y_new.loc[X_new.index]

    # Balance data such that each class has equal size
    if balanced:
        X_new["class"] = Y_new.dot(np.array([0, 1, 2, 3]).transpose())
        g = X_new.groupby('class')
        g = pd.DataFrame(g.apply(lambda x: x.sample(int(g.size().min()*1.5), replace=True).reset_index(drop=False)))
        X_new = g
        X_new.index = g["index"]
        X_new.drop(["index", "class"], axis=1, inplace=True)
        Y_new = Y_new.loc[X_new.index]

    output_info = Y_new.describe()
    print(output_info)
    num_data = len(Y_new)
    metadata["N"] = N
    metadata["num_data"] = num_data
    metadata[["data_non_cher", "data_cher", "data_ret_cher", "data_no_ret_cher"]] = list(output_info.loc["mean"])
    print(f"Num data points = {num_data}")

    # TRAIN RF MODEL
    start_time = time.time()
    print("\nRF MODEL")
    metadata[["acc_all", "acc_non_cher", "acc_cher", "acc_ret_cher", "acc_no_ret_cher"]] = train_cherries_rf(X_new, Y_new, name=data_set_string)
    print(f"RF train duration (normal): {np.round(time.time() - start_time, 3)}s")
    metadata["train_dur"] = time.time() - start_time
    metadata.to_pickle(f"RFModels/model_metadata_{data_set_string}.pickle")