import pickle
import numpy as np
import os
import sys
import re
from levels.utils import convert_to_prompt
from overcooked import World
import pandas as pd
from scipy.stats import f_oneway
from scipy.stats import ttest_ind

def load_data(path):
    with open(path, 'rb') as f:
        data = pickle.load(f)
    return data

directory = '/home/nikepupu/Desktop/human_data'
files = os.listdir(directory)
files = [os.path.join(directory ,file) for file in files if file.endswith('.pkl')]
df = pd.read_csv( '/home/nikepupu/Desktop/human_data/data_v2.csv')
print(df.columns.tolist())
cnt = 0

group1 = {}
group2 = {}

single_agent = {}
gpt_1_agents = {}
gpt_2_agents = {}
gpt_3_agents = {}


for file_path in files:
    match = re.search(r'([a-f0-9]{8}-[a-f0-9]{4}-[a-f0-9]{4}-[a-f0-9]{4}-[a-f0-9]{12})', file_path)
    uuid_str = match.group(1)
    data = load_data(file_path)
    try:
        data['world_play_mode']
        game_mode = data['world_play_mode']
    except:
        game_mode = 'play_with_human'
     
    mask = df["user_id: "] == uuid_str
    resulting_rows = df[mask]
    if resulting_rows.empty:
        print('no user id')
        continue

    
    num_agents = resulting_rows.iloc[:, -4].tolist()[0]

    cnt += 1
    assisting = resulting_rows.iloc[:, 2].tolist()[0]
    trust = resulting_rows.iloc[:, 3].tolist()[0]
    productivity = resulting_rows.iloc[:, 4].tolist()[0]
    predictability = resulting_rows.iloc[:, 5].tolist()[0]
    dependability = resulting_rows.iloc[:, 6].tolist()[0]
    enjoyment = resulting_rows.iloc[:, 7].tolist()[0]
    more_fun = resulting_rows.iloc[:, 12].tolist()[0]



    env = World(num_agents=data['num_agents'], level=data['level_name'], 
                task_filename='./assets/tasks_level_iclr.json', 
                recipe_filename='./assets/recipe.json', use_task_lifetime_interval_oracle=True, 
                alpha=data['world_alpha'], beta=data['world_beta'],
                override_agent=True, play_mode=game_mode, seed=data['world_seed'])
    env.reset()
    for action in data['action_history']:
        state, success, info = env.step(actions=action)

    if num_agents > 0:
        if game_mode == 'play_with_human':
            group1['assisting'] = group1.get('assisting', []) + [assisting]
            group1['trust'] = group1.get('trust', []) + [trust]
            group1['productivity'] = group1.get('productivity', []) + [productivity]
            group1['predictability'] = group1.get('predictability', []) + [predictability]
            group1['dependability'] = group1.get('dependability', []) + [dependability]
            group1['enjoyment'] = group1.get('enjoyment', []) + [enjoyment]
            group1['more_fun'] = group1.get('more_fun', []) + [more_fun]
            group1['success'] = group1.get('success', []) + [env.success_count/(env.success_count+env.failed_count )]
        
        else:
            group2['assisting'] = group2.get('assisting', []) + [assisting]
            group2['trust'] = group2.get('trust', []) + [trust]
            group2['productivity'] = group2.get('productivity', []) + [productivity]
            group2['predictability'] = group2.get('predictability', []) + [predictability]
            group2['dependability'] = group2.get('dependability', []) + [dependability]
            group2['enjoyment'] = group2.get('enjoyment', []) + [enjoyment]
            group2['more_fun'] = group2.get('more_fun', []) + [more_fun]
            group2['success'] = group2.get('success', []) + [env.success_count/(env.success_count+env.failed_count)]
    
    else:
        single_agent['assisting'] = single_agent.get('assisting', []) + [assisting]
        single_agent['trust'] = single_agent.get('trust', []) + [trust]
        single_agent['productivity'] = single_agent.get('productivity', []) + [productivity]
        single_agent['predictability'] = single_agent.get('predictability', []) + [predictability]
        single_agent['dependability'] = single_agent.get('dependability', []) + [dependability]
        single_agent['enjoyment'] = single_agent.get('enjoyment', []) + [enjoyment]
        single_agent['more_fun'] = single_agent.get('more_fun', []) + [more_fun]
        single_agent['success'] = single_agent.get('success', []) + [env.success_count/(env.success_count+env.failed_count)]



    if game_mode == 'play_with_human':
        if num_agents == 1:
            gpt_1_agents['assisting'] = gpt_1_agents.get('assisting', []) + [assisting]
            gpt_1_agents['trust'] = gpt_1_agents.get('trust', []) + [trust]
            gpt_1_agents['productivity'] = gpt_1_agents.get('productivity', []) + [productivity]
            gpt_1_agents['predictability'] = gpt_1_agents.get('predictability', []) + [predictability]
            gpt_1_agents['dependability'] = gpt_1_agents.get('dependability', []) + [dependability]
            gpt_1_agents['enjoyment'] = gpt_1_agents.get('enjoyment', []) + [enjoyment]
            gpt_1_agents['more_fun'] = gpt_1_agents.get('more_fun', []) + [more_fun]
            gpt_1_agents['success'] = gpt_1_agents.get('success', []) + [env.success_count/(env.success_count+env.failed_count)]

        elif num_agents == 2:
            gpt_2_agents['assisting'] = gpt_2_agents.get('assisting', []) + [assisting]
            gpt_2_agents['trust'] = gpt_2_agents.get('trust', []) + [trust]
            gpt_2_agents['productivity'] = gpt_2_agents.get('productivity', []) + [productivity]
            gpt_2_agents['predictability'] = gpt_2_agents.get('predictability', []) + [predictability]
            gpt_2_agents['dependability'] = gpt_2_agents.get('dependability', []) + [dependability]
            gpt_2_agents['enjoyment'] = gpt_2_agents.get('enjoyment', []) + [enjoyment]
            gpt_2_agents['more_fun'] = gpt_2_agents.get('more_fun', []) + [more_fun]
            gpt_2_agents['success'] = gpt_2_agents.get('success', []) + [env.success_count/(env.success_count+env.failed_count)]
        elif num_agents == 3:
            gpt_3_agents['assisting'] = gpt_3_agents.get('assisting', []) + [assisting]
            gpt_3_agents['trust'] = gpt_3_agents.get('trust', []) + [trust]
            gpt_3_agents['productivity'] = gpt_3_agents.get('productivity', []) + [productivity]
            gpt_3_agents['predictability'] = gpt_3_agents.get('predictability', []) + [predictability]
            gpt_3_agents['dependability'] = gpt_3_agents.get('dependability', []) + [dependability]
            gpt_3_agents['enjoyment'] = gpt_3_agents.get('enjoyment', []) + [enjoyment]
            gpt_3_agents['more_fun'] = gpt_3_agents.get('more_fun', []) + [more_fun]
            gpt_3_agents['success'] = gpt_3_agents.get('success', []) + [env.success_count/(env.success_count+env.failed_count)]


# print(cnt)
keys = ['assisting', 'trust', 'productivity', 'predictability', 'dependability', 'enjoyment', 'more_fun', 'success']
for key in keys:
    t_statistic, p_value = ttest_ind(group1[key], group2[key])
    print(key, p_value)
    # print('group1: ', np.mean(group1[key]), np.std(group1[key]))
    # print('group2: ', np.mean(group2[key]), np.std(group2[key]))
    # print()
print()
print('single vs grou1p')
for key in keys:
    f_statistic, p_value = ttest_ind(single_agent[key], group1[key])
    print(key, p_value)
print()
print('1 agent vs 2agent')
for key in keys:
    f_statistic, p_value = ttest_ind(gpt_1_agents[key], gpt_2_agents[key])
    print(key, p_value)
    if key == 'assisting':
        print('gpt_1_agents: ', np.mean(gpt_1_agents[key]), np.std(gpt_1_agents[key]))
        print('gpt_2_agents: ', np.mean(gpt_2_agents[key]), np.std(gpt_2_agents[key]))
        print()

print()

print('1 agent vs 3agent')
for key in keys:
    f_statistic, p_value = ttest_ind(gpt_1_agents[key], gpt_3_agents[key])
    print(key, p_value)
    if key == 'assisting':
        print('gpt_1_agents: ', np.mean(gpt_1_agents[key]), np.std(gpt_1_agents[key]))
        print('gpt_3_agents: ', np.mean(gpt_3_agents[key]), np.std(gpt_3_agents[key]))
        print()
key = 'success'
print('single_agent: ', np.mean(single_agent[key]), np.std(single_agent[key]))
print('gpt_1_agents: ', np.mean(gpt_1_agents[key]), np.std(gpt_1_agents[key]))
print('gpt_2_agents: ', np.mean(gpt_2_agents[key]), np.std(gpt_2_agents[key]))
print('gpt_3_agents: ', np.mean(gpt_3_agents[key]), np.std(gpt_3_agents[key]))