import os
import pandas as pd
import numpy as np
import sys
import warnings
if not sys.warnoptions:
    warnings.simplefilter("ignore")

def change_label_names(labels_kr, labels_eng):
    convert_dict = {}
    for num in range(len(labels_kr)):
        convert_dict[labels_kr[num]] = labels_eng[num]
    return convert_dict

def show_label_ratio(stats, stats_keys):
    for i in range(len(stats)):
        print("Ratio {}: {}".format(stats_keys[i], stats[i]/ len(preprocessed_train_df)))

if __name__ == '__main__':
    
    TRAIN_IMAGE_DIR = "/media/daitran/Data_SSD/data/SmartinsideAI/2021AIChamp/11.한국전력공사_데이터/train"
    DIR = "/media/daitran/Data_SSD/data/SmartinsideAI/2021AIChamp/11.한국전력공사_데이터/"
    SAVE_DIR = "/media/daitran/Data_SSD/code/SmartInsideAI/github/2021AIChamp_Submission/dataset"
    train_df_path = os.path.join(DIR, "train_output.csv")
    train_df = pd.read_csv(train_df_path)
    # display(train_df.head())

    print("Train sample: {}".format(len(train_df)))
    
    preprocessed_train_df = train_df
    
    status_df = preprocessed_train_df["status"]
    type_df = preprocessed_train_df["type"]
    fault_df = preprocessed_train_df["fault"]
    
    # STATUS
    print("====Convert STATUS===")
    status_label_kr = preprocessed_train_df["status"].unique()
    status_label_eng = [0,
                    1]
    status_convert_dict = change_label_names(status_label_kr, status_label_eng)
    print(status_convert_dict)
    preprocessed_train_df["status"].replace(status_convert_dict, inplace=True)
    print('STATUS classes ratio')
    status_nums = preprocessed_train_df["status"].value_counts()
    show_label_ratio(status_nums, status_nums.keys())
    
    # TYPE
    print("===Convert TYPE===")
    type_label_kr = preprocessed_train_df["type"].unique()
    type_label_eng = []
    for type_num in range(len(type_label_kr)):
        type_label_eng.append(type_num)
    type_convert_dict = change_label_names(type_label_kr, type_label_eng)
    print(type_convert_dict)
    preprocessed_train_df["type"].replace(type_convert_dict, inplace=True)
    print('TYPE classes ratio')
    type_nums = preprocessed_train_df["type"].value_counts()
    show_label_ratio(type_nums, type_nums.keys())
    
    # FAULT
    print("===Convert FAULT===")
    fault_label_kr = preprocessed_train_df["fault"].unique()
    fault_label_eng = []
    for fault_num in range(len(fault_label_kr)):
        fault_label_eng.append(fault_num)
    fault_convert_dict = change_label_names(fault_label_kr, fault_label_eng)
    print(fault_convert_dict)
    preprocessed_train_df["fault"].replace(fault_convert_dict, inplace=True)
    print('FAULT classes ratio')
    fault_nums = preprocessed_train_df["fault"].value_counts()
    show_label_ratio(fault_nums, fault_nums.keys())
    
    # Save df
    df_path = os.path.join(SAVE_DIR, 'preprocessed_train_df.csv')
    print("===Save train df at: {}===".format(df_path))
    preprocessed_train_df.to_csv(df_path)