from FIDDLE_config import *
import pickle
import pandas as pd
import numpy as np
import time
import os

import argparse
from FIDDLE_helpers import str2bool

if __name__ == '__main__':

    parser = argparse.ArgumentParser(description='')
    parser.add_argument('--T',               type=float,   required=True)
    parser.add_argument('--dt',              type=float,   required=True)
    parser.add_argument('--theta_1',         type=float,   default=0.001)
    parser.add_argument('--theta_2',         type=float,   default=0.001)
    parser.add_argument('--theta_freq',      type=float,   default=1.0)
    parser.add_argument('--stats_functions', nargs='+',    default=['min', 'max', 'mean'])
    parser.add_argument('--binarize',        type=str2bool, default=True, nargs='?', const=True)

    parser.add_argument('--data_path',       type=str,     required=True)
    parser.add_argument('--input_fname',     type=str,     required=False)
    parser.add_argument('--population',      type=str,     required=True)
    parser.add_argument('--N',               type=int,     required=False)
    parser.add_argument('--Ds',              nargs='+',    type=int)

    parser.add_argument('--no_prefilter',    dest='prefilter',  action='store_false')
    parser.add_argument('--no_postfilter',   dest='postfilter', action='store_false')
    parser.set_defaults(prefilter=True, postfilter=True)

    args = parser.parse_args([
        "--data_path=./sample_output/out_labs/", 
        "--input_fname=./sample_input/labs.csv", 
        "--population=./sample_input/windows.csv", 
        "--T=240", "--dt=240", 
        "--no_prefilter", "--no_postfilter", "--theta_freq=1",
        "--stats_functions", 'min', 'max', 'mean',
    ])
    args.variables = sorted(pd.read_csv('metadata/labs/value_types.csv')['variable_name'])
    args.variables_num_freq = []

    #########
    data_path = args.data_path
    if not data_path.endswith('/'):
        data_path += '/'

    population = args.population
    T = int(args.T)
    dt = args.dt
    theta_1 = args.theta_1
    theta_2 = args.theta_2
    theta_freq = args.theta_freq
    stats_functions = args.stats_functions
    binarize = args.binarize

    df_population = pd.read_csv(population).set_index('ID')
    N = args.N or len(df_population)
    df_population = df_population.iloc[:args.N]
    L = int(np.floor(T/dt))

    args.df_population = df_population
    args.N = N
    args.L = L
    args.parallel = parallel

    if args.input_fname and os.path.isfile(args.input_fname):
        input_fname = args.input_fname
        if input_fname.endswith('.p' or '.pickle'):
            df_data = pd.read_pickle(input_fname)
        elif input_fname.endswith('.csv'):
            df_data = pd.read_csv(input_fname)
        else:
            assert False
    else:
        raise NotImplementedError

    ##########
    from FIDDLE_steps_2 import *

    print('Input data file:', input_fname)
    print()
    print('Input arguments:')
    print('    {:<6} = {}'.format('T', T))
    print('    {:<6} = {}'.format('dt', dt))
    print('    {:<6} = {}'.format('\u03B8\u2081', theta_1))
    print('    {:<6} = {}'.format('\u03B8\u2082', theta_2))
    print('    {:<6} = {}'.format('\u03B8_freq', theta_freq))
    print('    {:<6} = {} {}'.format('k', len(stats_functions), stats_functions))
    print('{} = {}'.format('binarize', {False: 'no', True: 'yes'}[binarize]))
    print()
    print('N = {}'.format(N))
    print('L = {}'.format(L))
    print('', flush=True)

    df_data = df_data[df_data['ID'].isin(df_population.index)]
    df_data = df_data.sort_values(by=['ID', 't'])
    df_data_time_invariant = df_data.drop_duplicates(subset=['ID', 'variable_name'], keep='last') # Keep most recent recording
    df_data_time_invariant = df_data_time_invariant.drop(columns={'t'})

    dir_path = data_path + '/'
    start_time = time.time()

    ## Create NxLxD^ table
    df_time_invariant = df_data_time_invariant.set_index(['ID', 'variable_name']).unstack()
    print('Done unstacking', flush=True)
    df_time_invariant.columns = df_time_invariant.columns.droplevel(0)
    df_time_invariant = df_time_invariant.reindex(columns=args.variables, fill_value=np.nan)
    df_time_invariant = df_time_invariant.add_suffix('_value')
    print('Done reindexing', flush=True)
    print('Time elapsed: %f seconds' % (time.time() - start_time), flush=True)

    sdf = df_time_invariant.astype('Sparse[object]')
    joblib.dump(sdf, args.data_path + 'sdf.joblib')

    #####
    import joblib
    import pandas as pd
    import numpy as np
    import sparse
    import json

    df_time_series = joblib.load('sample_output/out_labs/sdf.joblib')
    discretization_bins = json.load(open('metadata/labs/discretization.json', 'r'))
    feature_names = json.load(open('metadata/labs/X_all.feature_names.json', 'r'))

    cols_useful = list(discretization_bins.keys())
    df_time_series = df_time_series[cols_useful]
    ddf = df_time_series.sparse.to_dense()

    from FIDDLE_steps_2 import *
    from FIDDLE_helpers import *

    df = ddf

    pool = multiprocessing.Pool(multiprocessing.cpu_count())
    out_0 = list(tqdm(pool.imap_unordered(
        smart_qcut_bins,
        [(df[col], discretization_bins[col]) for col in df.columns]), total=len(df.columns)
    ))
    cols_data, dis_bins = zip(*out_0)
    out = list(tqdm(pool.imap_unordered(dummify, cols_data), total=len(df.columns)
    ))
    pool.close()
    pool.join()

    df_features = pd.concat(out, axis=1).sort_index(axis=1)

    # Drop any values='missing' for now
    df_features = df_features.loc[:, [col for col in df_features.columns if 'missing' not in col]]

    df_features = df_features.reindex(columns=feature_names, fill_value=0)
    sdf = df_features.astype(pd.SparseDtype(int, fill_value=0))
    X_all = sparse.COO(sdf.sparse.to_coo())

    sparse.save_npz('sample_output/out_labs/X_all.npz', X_all)
    df_time_series[[]].to_csv('sample_output/out_labs/X_all.IDs.csv')