#!/usr/bin/env python
'''
Gets experimental results out of the database and table described in config.txt
and writes the data to a json or csv file.

Usage: speriment-output [-j] filename [-e excluded]'''

from sqlalchemy import create_engine, MetaData, Table
import json
import pandas as pd
import sys
from psiturk.psiturk_config import PsiturkConfig
import argparse

def parse():
    parser = argparse.ArgumentParser(description='''Retrieve and format the
            data gathered in an experiment and write it to a csv file.''')
    parser.add_argument('filename', type=str, help = '''File
            to write experimental results to in csv format.''')
    parser.add_argument('-j', '--json', action = 'store_true', help = '''Write the output in JSON
    format.''')
    parser.add_argument('-e', '--exclude', nargs='*',
            default=[], help = '''Worker IDs of any participants whose data you don't
            want to write to the output file.''')
    return parser.parse_args()

def get_credentials():
    config = PsiturkConfig()
    config.load_config()
    DBURL = config.get('Database Parameters', 'database_url')
    TABLENAME = config.get('Database Parameters', 'table_name')
    return DBURL, TABLENAME

def retrieve(db_url, table_name, exclude = []):
    # status codes PsiTurk gives subjects who completed experiment
    statuses = [3,4,5,7]

    # boilerplace sqlalchemy setup
    engine = create_engine(db_url)
    metadata = MetaData()
    metadata.bind = engine
    table = Table(table_name, metadata, autoload=True)
    s = table.select()
    rows = s.execute()

    # filter participants
    complete_participants = [participant for participant in rows
            if participant['status'] in statuses
            and participant['workerid'] not in exclude]
    return complete_participants

def format_data(complete_participants):
    # column PsiTurk saves your data to
    data_column_name = 'datastring'
    # JSON property Speriment tells PsiTurk to log trial data to
    # PsiTurk also keeps questiondata and eventdata, which Speriment doesn't use
    data_property_name = 'data'

    # 'data' is a list of objects containing among other things 'uniqueid' and
    # 'trialdata'. push uniqueid into trialdata and then use just trialdata.
    # also push information outside of 'data' into 'trialdata'.
    # this way, each row contains all the study-level and participant-level
    # information.
    trials = []
    for participant in complete_participants:
        json_data = json.loads(participant[data_column_name])
        for trial in json_data[data_property_name]:
            trial['trialdata'].update({
                'UniqueID': trial['uniqueid'],
                'TrialNumber': trial['current_trial'],
                'Version': participant['cond'],
                'Permutation': participant['counterbalance'],
                'HIT': participant['hitid'],
                'WorkerID': participant['workerid'],
                'ExperimentVersion': participant['codeversion']
                })
            trials.append(trial['trialdata'])
    return trials

def python_dataframe(trials, filename):
    data_frame = pd.DataFrame(trials)
    data_frame['ReactionTime'] = data_frame['EndTime'] - data_frame['StartTime']
    data_frame.to_csv(filename)

def write_json(trials, filename):
    json_content = json.dumps(trials)
    with open(filename, 'w') as f:
        f.write(json_content)


if __name__ == '__main__':
    # usage: speriment-output filename exclude
    args = parse()
    filename = args.filename
    exclude = args.exclude
    json_output = args.json
    (db_url, table_name) = get_credentials()
    data = retrieve(db_url, table_name, exclude)
    formatted = format_data(data)
    if json_output:
        write_json(formatted, filename)
    else:
        python_dataframe(formatted, filename)
