'''
CODE FOR PREPROCESSING AND DISCRETIZING DATA SETS. 
The pipeline uses a .json file in the format generated by 
the notebook "data_analysis.ipynb"

Processed dbs are saved as:   *_prep.csv
Discretized dbs are saved as: *_disc.csv
'''

import pandas as pd
import numpy as np
import json
import glob
import argparse
from sklearn.preprocessing import KBinsDiscretizer


# global variables
VAR_TIME_NAME = 'survival_time'
VAR_EVENT_NAME = 'survival_status'
N_BINS = 5
ENCODE = 'ordinal'
STRATEGY = 'kmeans'


def _save_db(_db, _db_name, _path=None):
    
    with open(_path+'{}_dtypes.json'.format(_db_name), 'w') as f:
        json.dump(_db.dtypes.apply(lambda x: x.name).to_dict(),f)
    
    _db.to_csv(_path+'{}.csv'.format(_db_name), index=False)
    print('>> saved '+_db_name)    
    return


def preprocessing(_data_path, json_obj, _db_name='', _discretization=True, _save_path=None):
    
    # read db with [adjusted missing values]
    if json_obj['_dictMV']:
        mv_values = set([item for l in json_obj['_dictMV'].values() for item in l])
        db = pd.read_csv(_data_path, header=0, sep=',', na_values=mv_values)
    else:
        db = pd.read_csv(_data_path, header=0, sep=',')
    
    # drop columns
    cols2drop = json_obj['_emptCols'] + json_obj['_colsNot2use']
    if cols2drop:
        db.drop(columns=cols2drop, inplace=True)
    
    # rename survival attributes
    surv2change = {json_obj['_survivalAttr']['survivalTime_name']: VAR_TIME_NAME,
                   json_obj['_survivalAttr']['survivalEvent_name']: VAR_EVENT_NAME}
    db.rename(columns=surv2change, inplace=True)
    
    # adjust event variable
    if db[VAR_EVENT_NAME].isna().sum():                                                  # drop missing values on event feature
        db.dropna(subset=[VAR_EVENT_NAME], inplace=True)
        db.reset_index(drop=True, inplace=True)
    if (json_obj['_statusRepr']['censValue'] and json_obj['_statusRepr']['eventValue']): # if event representation adjustment
        repr2change = {VAR_EVENT_NAME:{}}
        for value in json_obj['_statusRepr']['censValue']:
            repr2change[VAR_EVENT_NAME][value] = 'False'
        for value in json_obj['_statusRepr']['eventValue']:
            repr2change[VAR_EVENT_NAME][value] = 'True'
        db.replace(to_replace=repr2change, inplace=True)      
        mask = {'False': 0, 'True': 1}
        db[VAR_EVENT_NAME].replace(to_replace=mask, inplace=True)
    
    # input categories
    if json_obj['_nan2replace']:
        db.fillna(value=json_obj['_nan2replace'], inplace=True)
    
    # dropping missing values
    db.dropna(inplace=True, how='any')
    db.reset_index(drop=True, inplace=True)
    
    # adjust dtypes
    dtypes = {}.fromkeys(json_obj['_colsType'].keys(), 'category')
    if json_obj['_survivalAttr']['survivalEvent_name'] in dtypes:                       # remove old VAR_EVENT_NAME
        del dtypes[json_obj['_survivalAttr']['survivalEvent_name']]                     # remove old VAR_TIME_NAME
    if json_obj['_survivalAttr']['survivalTime_name'] in dtypes:
        del dtypes[json_obj['_survivalAttr']['survivalTime_name']]
    dtypes[VAR_EVENT_NAME] = 'bool'                                        # add VAR_EVENT_NAME as bool
    numerical = dict.fromkeys(list(set(dtypes.keys()) ^ set(db.columns)),'int64')
    db = db.astype({**dtypes,**numerical}, copy=True)
    
    # reorder columns with survival features first
    cols = list(db.columns)
    cols.remove(VAR_EVENT_NAME)
    cols.remove(VAR_TIME_NAME)
    cols = [VAR_TIME_NAME, VAR_EVENT_NAME] + cols
    db = db.reindex(columns=cols)
    
    # save preprocessed database
    _save_db(db, _db_name+'_prep', _save_path)
    
    # discretization
    if _discretization:
        discretization(db, json_obj, _db_name, _save_path)
    return


def discretization(_database, json_obj, _db_name='', _save_path=None):
    
    if isinstance(_database,str):
        _database = pd.read_csv(_database, header=0, sep=',')    
    
    db = _database.copy()
    
    # remove old VAR_TIME_NAME from features to discretize
    if json_obj['_survivalAttr']['survivalTime_name'] in json_obj['_cols2disc']:
        json_obj['_cols2disc'].remove(json_obj['_survivalAttr']['survivalTime_name'])
        
    # return if there is no col to discretize
    if not json_obj['_cols2disc']:
        # save db as db_disc
        _save_db(db, _db_name+'_disc', _save_path)
        # save log discretization
        with open(_save_path+'{}_log_discretization.json'.format(_db_name), 'w') as f:
            json.dump('! This data set has no features for discretization',f)
        return
    
    # discretization process: iterates over all coluns
    log_discretization = {}
    db_disc = db.copy()
    for col_name in json_obj['_cols2disc']:
        col_log = {}
        
        # data shape to discretize
        data = np.array(db[col_name]).reshape(-1, 1)
        
        try:
            discretizer = KBinsDiscretizer(n_bins=N_BINS, encode=ENCODE, strategy=STRATEGY)
            discretizer.fit(data)
            # data discretized/encoded
            data_encoded = pd.Series(discretizer.transform(data).reshape(1,-1)[0]).astype('int64')
        except:
            discretizer = KBinsDiscretizer(n_bins=N_BINS, encode=ENCODE, strategy='quantile')
            discretizer.fit(data)
            # data discretized/encoded
            data_encoded = pd.Series(discretizer.transform(data).reshape(1,-1)[0]).astype('int64')      
        
        
        # categories representation
        categories = sorted(data_encoded.unique().tolist())
        bin_edges = list(discretizer.bin_edges_[0])
        map_names = {}
        for idx,ctg in enumerate(categories):
            if idx+1 == len(categories):
                string = '[{:0.2f},{:0.2f}]'.format(bin_edges[idx],bin_edges[idx+1])
            else:
                string = '[{:0.2f},{:0.2f})'.format(bin_edges[idx],bin_edges[idx+1])
            map_names[ctg] = string
        
        # data decodification
        data_ctg = data_encoded.map(map_names)
        db_disc[col_name] = data_ctg.astype('category')

        # register log for discretized column
        col_log['discretizer_params'] = discretizer.get_params().copy()
        col_log['bin_edges'] = bin_edges
        col_log['dict_categories_names'] = map_names
        col_log['data_orig'] = data.tolist()
        col_log['data_encoded'] = data_encoded.tolist()
        col_log['data_categories'] = data_ctg.tolist()
        log_discretization[col_name] = col_log.copy()
    
    # save discretized database
    _save_db(db_disc, _db_name+'_disc', _save_path)
    
    # save log discretization
    with open(_save_path+'{}_log_discretization.json'.format(_db_name), 'w') as f:
        json.dump(log_discretization,f)
        
    return


def process_folder(folder_path, ext, json_folder, save_path, disc, prep, _sep):
    
    glob_path = folder_path+'*.{}'.format(ext)
    print('FOLDER TO PROCESS:{}'.format(glob_path))
    
    for file in glob.iglob(glob_path):
        
        print('>> process file: {}'.format(file))
        save_name = file.split(_sep)[-1].split('.')[0]
        
        # json-arg
        if not json_folder:
            json_file = folder_path.split('.')[0]+'{}_prep.json'.format(save_name)
        else:
            json_file = json_folder+'{}{}_prep.json'.format(_sep,save_name)
        with open(json_file, 'r') as f:
            json_obj = json.load(f)
        
        # save-arg
        if not save_path:
            save_path = folder_path
        
        # prep-arg
        if not prep:
            discretization(file, json_obj, _db_name=save_name, _save_path=save_path)
        else:
            preprocessing(file, json_obj, _db_name=save_name, _discretization=disc, _save_path=save_path)
    
    return


if __name__ == '__main__':
    
    # parse args setting
    parser = argparse.ArgumentParser(description='Script to perform data preprocessing and discretization [default].')
    
    group_f = parser.add_mutually_exclusive_group(required=True)
    group_f.add_argument("--fl", type=str,
                        help="File path")
    group_f.add_argument("--fd", type=str,
                        help="Folder path: executes the script to all files in the folder (or the ones defined by extension --ext)")
    
    group_j = parser.add_mutually_exclusive_group()
    group_j.add_argument("--jfl", type=str,
                        help="Json preprocessing file path: if not provided, uses the file folder")
    group_j.add_argument("--jfd", type=str,
                        help="Folder for json preprocessing files: if not provided, uses the file folder")
    
    parser.add_argument("--ext", type=str, default='csv',
                        help="Extension of the files to process")
    parser.add_argument("--s", type=str,
                        help="Save path: if not provided, saves processed data in the file folder")
    parser.add_argument("--d", type=bool, default=True,
                        help="Discretization [bool]: whether to perform or not")
    parser.add_argument("--p", type=bool, default=True,
                        help="Preprocessing [bool]: whether to perform or not")
    
    args = parser.parse_args()
    file_path = args.fl
    json_file = args.jfl
    folder_path = args.fd
    json_folder = args.jfd
    ext = args.ext
    save_path = args.s
    disc = args.d
    prep = args.p
    
    # string separator and folder-names adjustment
    if file_path:
        if '\\' in file_path: _sep = '\\'
        elif '//' in file_path: _sep = '//'
        else: _sep = ''
    else:
        if '\\' in folder_path: _sep = '\\'
        elif '//' in folder_path: _sep = '//'
        else: _sep = ''
        folder_path = folder_path + _sep
        if json_folder:
            json_folder = json_folder + _sep
    if save_path:
        save_path = save_path + _sep
    
       
    if folder_path:
        process_folder(folder_path, ext, json_folder, save_path, disc, prep, _sep)
    
    else:
        
        # json-arg
        if not json_file:
            json_file = file_path.split('.')[0]+'_prep.json'
        with open(json_file, 'r') as f:
            json_obj = json.load(f)
		
        # save-arg
        if not save_path:
            if not _sep:
                save_path = ''
                save_name = file_path.split('.')[0]
            else:
                save_path = '{}'.format(_sep).join(file_path.split(_sep)[:-1])+_sep
                save_name = file_path.split(_sep)[-1].split('.')[0]
        else:
            save_name = file_path.split(_sep)[-1].split('.')[0]
        
        # prep-arg
        if not prep:
            discretization(file_path, json_obj, _db_name=save_name, _save_path=save_path)
        else:
            preprocessing(file_path, json_obj, _db_name=save_name, _discretization=disc, _save_path=save_path)
