# %%
import sqlite3
import os
import pandas as pd
import numpy as np
sqlite3.register_adapter(np.int64, lambda val: int(val))

DATA_PATH = '/home/marek/Documents/marta/data/'

if __name__ == "__main__":
    
    try: 
        os.remove('/home/marek/Documents/marta/oTree/_static/iq_results/iq_data.db')
    except Exception:
        pass


    conn = sqlite3.connect('/home/marek/Documents/marta/oTree/_static/iq_results/iq_data.db')
    c = conn.cursor()

    c.execute('''CREATE TABLE IF NOT EXISTS results
                (ID text NOT NULL PRIMARY KEY,
                TOT_EARN real,
                IQ INTEGER NOT NULL,
                TASK INTEGER NOT NULL,
                BOX INTEGER NOT NULL,
                RANK INTEGER NOT NULL,
                S real,
                NUM real,
                MONEY_TASK real, 
                TEST_DATE text
                )''')

    conn.commit()

    # sort by date so that we have to order of participants right
    all_files  = os.listdir(DATA_PATH)
    all_files.sort() 

    all_files  = [filename for filename in all_files if (('.csv' in filename) &  ('_clean' in filename) & ('#' not in filename) )]
    # load participation labels 
    plabels = open("/home/marek/Documents/marta/oTree/_rooms/plabs.txt", "r").read().split('\n')
    
    # set participant counter
    pcounter = 0

    # columns to extract
    COLS2GET = [
    'final_payoff',    # TOT_EARN
    'payoff_IQ',    # IQ
    'task_paid',    # TASK
    'Guessing_Rank',    # BOX
    'Rank',    # RANK
    'S',    # S
    'random02',    # NUM
    'payment_task'    # MONEY_TASK
    ]

    # go over all data files 
    for filename in all_files:

        # load raw file
        raw_data = pd.read_csv(DATA_PATH+filename)

        # check if main or control
        is_main  = 'Belief2urn2' in raw_data.columns
        
        if is_main:
            raw_data['Guessing_Rank'] = raw_data['Guessing_Rank']+1
        else:
            raw_data['Guessing_Rank'] = (raw_data['box_drawn']+1)%2+1
        
        raw_data['final_payoff']  = round((raw_data['final_payoff'] +1e-8)* 2.0) / 2.0
        raw_data['S']  = round(raw_data['S'],2)
        raw_data['random02']  = round(raw_data['random02'],2)

        temp = filename.split('_')[0][2:]
        test_date = temp[2:]+'/'+temp[0:2]+'/2020'
        # core SQL command to insert values
        core_string1 = ''' INSERT INTO results VALUES (?,?,?,?,?,?,?,?,?,?) '''

        # go row by row, get participant code, get data, commint to sql
        for row in range(0,len(raw_data)):
            P_id = plabels[pcounter]

            values_list = raw_data.iloc[row,:][COLS2GET].to_list()
            values_list.append(test_date)
            values_list.insert(0,P_id)
            
            # Insert a row of data
            c.execute(core_string1,values_list)

            pcounter+=1

        #raw_data.iloc[row,1]
        # Save (commit) the changes
        conn.commit()



    # We can also close the connection if we are done with it.
    # Just be sure any changes have been committed or they will be lost.
    conn.close()