'''
samples from all raw data;
by default samples in a non-iid manner; namely, randomly selects users from 
raw data until their cumulative amount of data exceeds the given number of 
datapoints to sample (specified by --fraction argument);
ordering of original data points is not preserved in sampled data
'''

import argparse
import json
import os
import random
import time

from collections import OrderedDict

from constants import DATASETS, SEED_FILES
from util import iid_divide

parser = argparse.ArgumentParser()

parser.add_argument('--name',
                help='name of dataset to parse; default: sent140;',
                type=str,
                choices=DATASETS,
                default='sent140')
parser.add_argument('--iid',
                help='sample iid;',
                action="store_true")
parser.add_argument('--niid',
                help="sample niid;",
                dest='iid', action='store_false')
parser.add_argument('--fraction',
                help='fraction of all data to sample; default: 0.1;',
                type=float,
                default=0.1)
parser.add_argument('--u',
                help=('number of users in iid data set; ignored in niid case;'
                      'represented as fraction of original total number of users; '
                      'default: 0.01;'),
                type=float,
                default=0.01)
parser.add_argument('--seed',
                help='seed for random sampling of data',
                type=int,
                default=None)
parser.set_defaults(iid=False)

args = parser.parse_args()

print('------------------------------')
print('sampling data')

parent_path = os.path.dirname(os.path.dirname(os.path.realpath(__file__)))
data_dir = os.path.join(parent_path, args.name, 'data')
subdir = os.path.join(data_dir, 'all_data')
files = os.listdir(subdir)
files = [f for f in files if f.endswith('.json')]

rng_seed = (args.seed if (args.seed is not None and args.seed >= 0) else int(time.time()))
print ("Using seed {}".format(rng_seed))
rng = random.Random(rng_seed)
print (os.environ.get('LEAF_DATA_META_DIR'))
if os.environ.get('LEAF_DATA_META_DIR') is not None:
    seed_fname = os.path.join(os.environ.get('LEAF_DATA_META_DIR'), SEED_FILES['sampling'])
    with open(seed_fname, 'w+') as f:
        f.write("# sampling_seed used by sampling script - supply as "
                "--smplseed to preprocess.sh or --seed to utils/sample.py\n")
        f.write(str(rng_seed))
    print ("- random seed written out to {file}".format(file=seed_fname))
else:
    print ("- using random seed '{seed}' for sampling".format(seed=rng_seed))

new_user_count = 0 # for iid case
for f in files:
    file_dir = os.path.join(subdir, f)
    with open(file_dir, 'r') as inf:
        # Load data into an OrderedDict, to prevent ordering changes
        # and enable reproducibility
        data = json.load(inf, object_pairs_hook=OrderedDict)

    num_users = len(data['users'])

    tot_num_samples = sum(data['num_samples'])
    num_new_samples = int(args.fraction * tot_num_samples)

    hierarchies = None
    cluster_ids = None

    if(args.iid):
        raw_list = list(data['user_data'].values())
        raw_x = [elem['x'] for elem in raw_list]
        raw_y = [elem['y'] for elem in raw_list]
        x_list = [item for sublist in raw_x for item in sublist] # flatten raw_x
        y_list = [item for sublist in raw_y for item in sublist] # flatten raw_y

        num_new_users = int(round(args.u * num_users))
        if num_new_users == 0:
            num_new_users += 1

        indices = [i for i in range(tot_num_samples)]
        new_indices = rng.sample(indices, num_new_samples)
        users = [str(i+new_user_count) for i in range(num_new_users)]

        user_data = {}
        for user in users:
            user_data[user] = {'x': [], 'y': []}
        all_x_samples = [x_list[i] for i in new_indices]
        all_y_samples = [y_list[i] for i in new_indices]
        x_groups = iid_divide(all_x_samples, num_new_users)
        y_groups = iid_divide(all_y_samples, num_new_users)
        for i in range(num_new_users):
            user_data[users[i]]['x'] = x_groups[i]
            user_data[users[i]]['y'] = y_groups[i]
        
        num_samples = [len(user_data[u]['y']) for u in users]
        if 'cluster_ids' in data:
            cluster_ids = [data['cluster_ids'][int(u)] for u in users]

        new_user_count += num_new_users

    else:

        ctot_num_samples = 0

        users = data['users']
        users_hiers_cids = None
        users_hiers = None
        users_cids = None
        if 'hierarchies' in data and 'cluster_ids' in data:
            users_hiers_cids = list(zip(users, data['hierarchies'], data['cluster_ids']))
            rng.shuffle(users_hiers_cids)
        elif 'hierarchies' in data:
            users_hiers = list(zip(users, data['hierarchies']))
            rng.shuffle(users_hiers)
        elif 'cluster_ids' in data:
            users_cids = list(zip(users, data['cluster_ids']))
            rng.shuffle(users_cids)
        else:
            rng.shuffle(users)
        user_i = 0
        num_samples = []
        user_data = {}

        if 'hierarchies' in data:
            hierarchies = []
        
        if 'cluster_ids' in data:
            cluster_ids = []

        while(ctot_num_samples < num_new_samples):
            hierarchy = None
            if users_hiers_cids is not None:
                user, hier, cid = users_hiers_cids[user_i]
            elif users_hiers is not None:
                user, hier = users_hiers[user_i]
            elif users_cids is not None:
                user, cid = users_cids[user_i]
            else:
                user = users[user_i]

            cdata = data['user_data'][user]

            cnum_samples = len(data['user_data'][user]['y'])

            if (ctot_num_samples + cnum_samples > num_new_samples):
                cnum_samples = num_new_samples - ctot_num_samples
                indices = [i for i in range(cnum_samples)]
                new_indices = rng.sample(indices, cnum_samples)
                x = []
                y = []
                for i in new_indices:
                    x.append(data['user_data'][user]['x'][i])
                    y.append(data['user_data'][user]['y'][i])
                cdata = {'x': x, 'y': y}
            
            if 'hierarchies' in data:
                hierarchies.append(hier)

            if 'cluster_ids' in data:
                cluster_ids.append(cid)

            num_samples.append(cnum_samples)
            user_data[user] = cdata

            ctot_num_samples += cnum_samples
            user_i += 1

        if 'hierarchies' in data and 'cluster_ids' in data:
            users = [u for u, h, cid in users_hiers_cids][:user_i]
        elif 'hierarchies' in data:
            users = [u for u, h in users_hiers][:user_i]
        elif 'cluster_ids' in data:
            users = [u for u, h in users_cids][:user_i]
        else:
            users = users[:user_i]

    # ------------
    # create .json file

    all_data = {}
    all_data['users'] = users
    if hierarchies is not None:
        all_data['hierarchies'] = hierarchies
    all_data['num_samples'] = num_samples
    if cluster_ids is not None:
        all_data['cluster_ids'] = cluster_ids
    all_data['user_data'] = user_data

    slabel = ''
    if(args.iid):
        slabel = 'iid'
    else:
        slabel = 'niid'

    arg_frac = str(args.fraction)
    arg_frac = arg_frac[2:]
    arg_nu = str(args.u)
    arg_nu = arg_nu[2:]
    arg_label = arg_frac
    if(args.iid):
        arg_label = '%s_%s' % (arg_nu, arg_label)
    file_name = '%s_%s_%s.json' % ((f[:-5]), slabel, arg_label)
    ouf_dir = os.path.join(data_dir, 'sampled_data', file_name)

    print('writing %s' % file_name)
    with open(ouf_dir, 'w') as outfile:
        json.dump(all_data, outfile)
