#!/usr/bin/env python2
import pandas as pd
import json
import ast
import re
import os



writevar = lambda *args: None

def to_snake_case(name):
    name = re.sub('(.)([A-Z][a-z]+)', r'\1_\2', name)
    name = re.sub(r'[.:\/]', '_', name)
    return re.sub('([a-z0-9])([A-Z])', r'\1_\2', name).lower()


class Labeler(object):
    """Assigns unique integer labels."""
    def __init__(self, init=()):
        self._labels = {}
        self._xs = []
        for x in init:
            self.label(x)

    def label(self, x):
        if x not in self._labels:
            self._labels[x] = len(self._labels)
            self._xs.append(x)
        return self._labels[x]

    def unlabel(self, label):
        return self._xs[label]

    __call__ = label


def fetch_data(experiment, code_version):
    data_path = 'data/{experiment}/human_raw/{code_version}/'.format(
        experiment=experiment,
        code_version=code_version,
    )


    pdf = pd.read_csv(data_path + 'questiondata.csv', header=None)
    # finished_turk_ids = pdf[0][pdf[1] == 'final_bonus'].apply(lambda x: x.split(':')[0])
    # Create participants dataframe (pdf)
    pdf = pdf[pdf[1] == 'params'].reset_index(drop=True)
    pdf['code_version'] = code_version
    expanded = pd.DataFrame.from_records(pdf[2].apply(ast.literal_eval))
    if len(expanded):
        pdf = expanded.join(pdf[0])

    # Create trials dataframe (tdf)
    tdf = pd.read_csv(data_path + 'trialdata.csv', header=None)
    tdf = pd.DataFrame.from_records(tdf[3].apply(json.loads)).join(tdf[0])
    
    n_trials = tdf.groupby(0).trialID.apply(lambda s: len(s.dropna()))
    complete = n_trials[n_trials >= n_trials.median()].index

    def clean_up(df):
        # remove debug trials and assgin participants ids
        df = df.set_index(0).ix[complete].reset_index()
        df['worker_id'] = df[0].apply(lambda x: x.split(':')[0])
        df['assignment_id'] = df[0].apply(lambda x: x.split(':')[1])
        df = df[df.worker_id.map(lambda x: not x.startswith('debug'))].copy()
        df['pid'] = df['worker_id'].apply(Labeler())
        return df.drop(0, axis=1)

    pdf, tdf = map(clean_up, (pdf, tdf))

    data = {'participants': pdf}
    for trial_type, df in tdf.groupby('trial_type'):
        df = df.dropna(axis=1)
        df = df.drop('internal_node_id', axis=1)
        df = df.drop('trial_index', axis=1)
        data[trial_type] = df

    return data


def write_data(experiment, code_version, data):
    path = 'data/{experiment}/human/{code_version}/'.format(
        experiment=experiment,
        code_version=code_version,
    )
    if not os.path.isdir(path):
        os.makedirs(path)
    for name, df in data.items():
        dest = path + name + '.csv'
        # Hack to make feedback the final column
        if 'feedback' in df:
            x = df.feedback
            del df['feedback']
            df['feedback'] = x
        df.to_csv(dest)
        print('wrote {} with {} rows.'.format(dest, len(df)))



if __name__ == '__main__':
    from argparse import ArgumentParser, ArgumentDefaultsHelpFormatter

    parser = ArgumentParser(
        formatter_class=ArgumentDefaultsHelpFormatter)
    parser.add_argument(
        '-e', '--experiment',
        default='1',
        help='Directory in data/ to put the results.')
    parser.add_argument(
        '-v', '--version',
        required=True,
        help=('Experiment version. This corresponds to the '
              'experiment_code_version parameter in the psiTurk config.txt file '
              'that was used when the data was collected.'))

    args = parser.parse_args()
    data = fetch_data(args.experiment, args.version)
    write_data(args.experiment, args.version, data)



# def misformatted(df):
#     def check(c):
#         return len(c.dropna()) !=138
    
#     idx =  df.groupby('pid').correct.apply(check).as_matrix()
#     return set(df.pid.unique()[idx])

# def failed_catch(df, max_fails=1):
#     fails = (df.query("kind == 'control'")
#              .groupby('pid')
#              .correct.agg(lambda c: len(c[c== False]))
#     )
#     failed_pids = fails[fails > max_fails].index
#     return set(failed_pids)

# misformat_pids = misformatted(data['ball'])
# caught_pids = failed_catch(data['ball'])
# bad_pids = misformat_pids | caught_pids

# writevar('N_CAUGHT', len(caught_pids))
# writevar('N_FORMAT', len(misformat_pids))



# caught = data['participants'].ix[caught_pids].turk_id
# with open('caught_wids.txt', 'w+') as f:
#     f.write('\n'.join(caught))



