# -*- coding: utf-8 -*-
"""
Created on Thu Nov 18 14:37:58 2021

@author: Mohsen.Foroughifar17
"""
import pandas as pd
import os
import yaml
import numpy as np
from scipy import stats
import statsmodels.api as sm
import matplotlib.pyplot as plt
import numpy.matlib
from sklearn.linear_model import LogisticRegression
from scipy.optimize import curve_fit
currentDirectory=os.getcwd()
direct = currentDirectory+'\\_static\\global\\data_w_info\\'
sessions = os.listdir(direct)



def txt_to_num(inp):
    for i in range(len(inp)):
        if (inp[i]=='Box T (Tail)'):
            inp[i] = 2
        elif (inp[i]=='Box H (Head)'):
            inp[i] = 1
        elif (inp[i]=='Picture B'):
            inp[i] = 2
        elif (inp[i]=='Picture A'):
            inp[i] = 1
    return inp


# loading the data from the main experiment
for session in sessions:
    directory = direct + session
    files = os.listdir(directory)
    try:
        participant_files_list = np.append(participant_files_list,files)
    except:
        participant_files_list = np.array(files)
    for file in files:
        data = pd.read_csv(directory+'/'+file)
        try:
            participant_vars_col = data['player.participant_vars_dump'][:]
            participant_vars_all = participant_vars_col[~participant_vars_col.isna()]
            
            for participant_vars in participant_vars_all:
                participant_data = yaml.load(participant_vars)
                try:
                    #treatments = np.append(treatments,[participant_data['treatment']],axis=0)
                    #Majors = np.append(Majors,[participant_data['Major']],axis=0)
                    Prob_courses = np.append(Prob_courses,[participant_data['Prob_course']],axis=0)
                    Educations = np.append(Educations,[participant_data['Education']],axis=0)
                    Genders = np.append(Genders,[participant_data['Gender']],axis=0)
                    Ages = np.append(Ages,[participant_data['Age']],axis=0)
                    Feedbacks = np.append(Feedbacks,[participant_data['Feedback']],axis=0)
                    Comment_balls = np.append(Comment_balls,[participant_data['Comment_ball']],axis=0)
                    Comment_neighbors = np.append(Comment_neighbors,[participant_data['Comment_neighbor']],axis=0)
                except:
                    #treatments = np.array([participant_data['treatment']])
                    #Majors = np.array([participant_data['Major']])
                    Prob_courses = np.array([participant_data['Prob_course']])
                    Educations= np.array([participant_data['Education']])
                    Genders = np.array([participant_data['Gender']])
                    Ages = np.array([participant_data['Age']])
                    Feedbacks = np.array([participant_data['Feedback']])
                    Comment_balls = np.array([participant_data['Comment_ball']])
                    Comment_neighbors = np.array([participant_data['Comment_neighbor']])
                    
                participant_data['Guesses_sorted'] = txt_to_num(participant_data['Guesses_sorted'])
                participant_data['indiv_irrational'] = np.absolute(np.array(participant_data['Guesses_sorted'])-np.array(participant_data['state_realized_sorted']))
                try:
                    indiv_irrational = np.append(indiv_irrational,[participant_data['indiv_irrational']],axis=0)
                    indiv_posteriors = np.append(indiv_posteriors,[participant_data['Guesses_conf_sorted']],axis=0)
                    indiv_draws = np.append(indiv_draws,[participant_data['draw_sorted']],axis=0)
                    indiv_guess = np.append(indiv_guess,[participant_data['Guesses_sorted']],axis=0)
                    
                except:
                    indiv_irrational = np.array([participant_data['indiv_irrational']])
                    indiv_posteriors = np.array([participant_data['Guesses_conf_sorted']])
                    indiv_draws = np.array([participant_data['draw_sorted']])
                    indiv_guess = np.array([participant_data['Guesses_sorted']])                    
                   
        except:
            for i in range(len(data)): 
                Prob_courses = np.append(Prob_courses,[data['Prob_course'][i]],axis=0)
                Educations = np.append(Educations,[data['Education'][i]],axis=0)
                Genders = np.append(Genders,[data['Gender'][i]],axis=0)
                Ages = np.append(Ages,[data['Age'][i]],axis=0)
                Feedbacks = np.append(Feedbacks,[data['Feedback'][i]],axis=0)
                Comment_balls = np.append(Comment_balls,[data['Comment_ball'][i]],axis=0)
                Comment_neighbors = np.append(Comment_neighbors,[data['Comment_neighbor'][i]],axis=0)
                
                data['Guesses_sorted'][i] = txt_to_num(yaml.load(data['Guesses_sorted'][i]))
                indiv_irrational_temp = np.absolute(np.array(data['Guesses_sorted'][i])-np.array(yaml.load(data['state_realized_sorted'][i])))
                try:
                    indiv_irrational = np.append(indiv_irrational,[indiv_irrational_temp],axis=0)
                    indiv_posteriors = np.append(indiv_posteriors,[yaml.load(data['Guesses_conf_sorted'][i])],axis=0)
                    indiv_draws = np.append(indiv_draws,[yaml.load(data['draw_sorted'][i])],axis=0)
                    indiv_guess = np.append(indiv_guess,[data['Guesses_sorted'][i]],axis=0)
                except:
                    indiv_irrational = np.array([indiv_irrational_temp])
                    indiv_posteriors = np.array([yaml.load(data['Guesses_conf_sorted'][i])])
                    indiv_draws = np.array([yaml.load(data['draw_sorted'][i])])
                    indiv_guess = np.array([data['Guesses_sorted'][i]])
       
         
indiv_irr_avg = np.round(np.mean(indiv_irrational,axis=0)*100)
indiv_irr_avg_W = np.empty(21)*np.nan
indiv_irr_avg_B = np.empty(21)*np.nan
indiv_irr_n_W = np.empty(21)*np.nan
indiv_irr_n_B = np.empty(21)*np.nan
for i in range(len(indiv_irrational[0,])):
    indiv_irr_avg_W[i] = np.round(np.mean(indiv_irrational[indiv_guess[:,i]==1,i]),2)
    indiv_irr_n_W[i] = np.sum(indiv_draws[:,i]==1) # 
    indiv_irr_avg_B[i] = np.round(np.mean(indiv_irrational[indiv_guess[:,i]==2,i]),2)
    indiv_irr_n_B[i] = np.sum(indiv_draws[:,i]==2) #

plt.plot(np.linspace(1,21,21),np.mean(indiv_irrational,axis=0),'g-')
plt.plot(np.linspace(1,21,21),indiv_irr_avg_W,'r--')
plt.plot(np.linspace(1,21,21),indiv_irr_avg_B,'b-.')
plt.grid(True)



import pandas as pd
import os
import yaml
import numpy as np
from scipy import stats
import statsmodels.api as sm
import matplotlib.pyplot as plt
import numpy.matlib
from sklearn.linear_model import LogisticRegression
from scipy.optimize import curve_fit

def txt_to_num(inp):
    for i in range(len(inp)):
        if (inp[i]=='Box T (Tail)'):
            inp[i] = 2
        elif (inp[i]=='Box H (Head)'):
            inp[i] = 1
        elif (inp[i]=='Picture B'):
            inp[i] = 2
        elif (inp[i]=='Picture A'):
            inp[i] = 1
    return inp


currentDirectory=os.getcwd()
data = pd.read_csv(currentDirectory+'/Interaction_II_info_2021-11-22.csv')
participant_vars_col = data['player.participant_vars_dump'][:]
participant_vars_all = participant_vars_col[~participant_vars_col.isna()]

for participant_vars in participant_vars_all:
    participant_data = yaml.load(participant_vars)
    try:
        treatments = np.append(treatments,[participant_data['treatment']],axis=0)
        #Majors = np.append(Majors,[participant_data['Major']],axis=0)
        Prob_courses = np.append(Prob_courses,[participant_data['Prob_course']],axis=0)
        Educations = np.append(Educations,[participant_data['Education']],axis=0)
        Genders = np.append(Genders,[participant_data['Gender']],axis=0)
        Ages = np.append(Ages,[participant_data['Age']],axis=0)
        Feedbacks = np.append(Feedbacks,[participant_data['Feedback']],axis=0)
        Comment_balls = np.append(Comment_balls,[participant_data['Comment_ball']],axis=0)
        Comment_neighbors = np.append(Comment_neighbors,[participant_data['Comment_neighbor']],axis=0)
    except:
        treatments = np.array([participant_data['treatment']])
        #Majors = np.array([participant_data['Major']])
        Prob_courses = np.array([participant_data['Prob_course']])
        Educations= np.array([participant_data['Education']])
        Genders = np.array([participant_data['Gender']])
        Ages = np.array([participant_data['Age']])
        Feedbacks = np.array([participant_data['Feedback']])
        Comment_balls = np.array([participant_data['Comment_ball']])
        Comment_neighbors = np.array([participant_data['Comment_neighbor']])
        
    participant_data['Guesses_sorted'] = txt_to_num(participant_data['Guesses_sorted'])
    participant_data['indiv_irrational'] = np.absolute(np.array(participant_data['Guesses_sorted'])-np.array(participant_data['state_realized_sorted']))
    try:
        indiv_irrational = np.append(indiv_irrational,[participant_data['indiv_irrational']],axis=0)
        indiv_posteriors = np.append(indiv_posteriors,[participant_data['Guesses_conf_sorted']],axis=0)
        indiv_draws = np.append(indiv_draws,[participant_data['draw_sorted']],axis=0)
        indiv_guess = np.append(indiv_guess,[participant_data['Guesses_sorted']],axis=0)
        
    except:
        indiv_irrational = np.array([participant_data['indiv_irrational']])
        indiv_posteriors = np.array([participant_data['Guesses_conf_sorted']])
        indiv_draws = np.array([participant_data['draw_sorted']])
        indiv_guess = np.array([participant_data['Guesses_sorted']]) 
    
    participant_data['Guesses_neighbor_sorted'] = txt_to_num(participant_data['Guesses_neighbor_sorted'])
    participant_data['neighbor_guess_sorted'] = txt_to_num(participant_data['neighbor_guess_sorted'])
    participant_data['neighbor_draw_sorted'] = txt_to_num(participant_data['neighbor_draw_sorted'])
    participant_data['neighbor_rational'] = np.logical_not(np.absolute(np.array(participant_data['neighbor_guess_sorted'])-np.array(participant_data['neighbor_draw_sorted'])))
    participant_data['social_irrational'] = np.absolute(np.array(participant_data['Guesses_neighbor_sorted'])-np.array(participant_data['neighbor_guess_sorted']))
    participant_data['social_irrational_ball'] = np.absolute(np.array(participant_data['Guesses_neighbor_sorted'])-np.array(participant_data['neighbor_draw_sorted']))         
    participant_data['social_irrational_corrected'] = np.logical_and(np.array(participant_data['neighbor_rational']),np.array(participant_data['social_irrational']))
   
    try:
        if np.logical_or(participant_data['treatment']=="D_ID_wo",participant_data['treatment']=="ID_D_wo"):
            neighbor_rational = np.append(neighbor_rational,[participant_data['neighbor_rational']],axis=0)
            social_irrational_corrected = np.append(social_irrational_corrected,[participant_data['social_irrational_corrected']],axis=0)
        else:
            neighbor_rational = np.append(neighbor_rational,[np.logical_or(participant_data['neighbor_rational'],1)],axis=0)
            social_irrational_corrected = np.append(social_irrational_corrected,[participant_data['social_irrational']],axis=0)       
        social_irrational = np.append(social_irrational,[participant_data['social_irrational']],axis=0)
        social_irrational_ball = np.append(social_irrational_ball,[participant_data['social_irrational_ball']],axis=0)                   
        social_posteriors = np.append(social_posteriors,[participant_data['Guesses_conf_neighbor_sorted']],axis=0)
        social_ndraws = np.append(social_ndraws,[participant_data['neighbor_draw_sorted']],axis=0)
        social_nguess = np.append(social_nguess,[participant_data['neighbor_guess_sorted']],axis=0)
        social_guess = np.append(social_guess,[participant_data['Guesses_neighbor_sorted']],axis=0)

    except:
        if np.logical_or(participant_data['treatment']=="D_ID_wo",participant_data['treatment']=="ID_D_wo"):
            neighbor_rational = np.array([participant_data['neighbor_rational']])
            social_irrational_corrected = np.array([participant_data['social_irrational_corrected']])
        else:
            neighbor_rational = np.array([np.logical_or(participant_data['neighbor_rational'],1)])
            social_irrational_corrected = np.array([participant_data['social_irrational']])
        social_irrational = np.array([participant_data['social_irrational']])
        social_irrational_ball = np.array([participant_data['social_irrational_ball']])                   
        social_posteriors = np.array([participant_data['Guesses_conf_neighbor_sorted']])
        social_ndraws = np.array([participant_data['neighbor_draw_sorted']])
        social_nguess = np.array([participant_data['neighbor_guess_sorted']])
        social_guess = np.array([participant_data['Guesses_neighbor_sorted']])

indiv_irr_avg = np.round(np.mean(indiv_irrational,axis=0)*100)
indiv_irr_avg_W = np.empty(21)*np.nan
indiv_irr_avg_B = np.empty(21)*np.nan
indiv_irr_n_W = np.empty(21)*np.nan
indiv_irr_n_B = np.empty(21)*np.nan
for i in range(len(indiv_irrational[0,])):
    indiv_irr_avg_W[i] = np.round(np.mean(indiv_irrational[indiv_guess[:,i]==1,i]),2)
    indiv_irr_n_W[i] = np.sum(indiv_draws[:,i]==1) # 
    indiv_irr_avg_B[i] = np.round(np.mean(indiv_irrational[indiv_guess[:,i]==2,i]),2)
    indiv_irr_n_B[i] = np.sum(indiv_draws[:,i]==2) #

plt.plot(np.linspace(1,21,21),np.mean(indiv_irrational,axis=0),'g-')
plt.plot(np.linspace(1,21,21),indiv_irr_avg_W,'r--')
plt.plot(np.linspace(1,21,21),indiv_irr_avg_B,'b-.')
plt.grid(True)

social_irr_avg = np.round(np.mean(social_irrational,axis=0)*100)
social_irr_avg_W = np.empty(21)*np.nan
social_irr_avg_B = np.empty(21)*np.nan
social_irr_n_W = np.empty(21)*np.nan
social_irr_n_B = np.empty(21)*np.nan
for i in range(len(indiv_irrational[0,])):
    social_irr_avg_W[i] = np.round(np.mean(social_irrational[social_nguess[:,i]==1,i]),2)
    social_irr_n_W[i] = np.sum(social_nguess[:,i]==1) # 
    social_irr_avg_B[i] = np.round(np.mean(social_irrational[social_nguess[:,i]==2,i]),2)
    social_irr_n_B[i] = np.sum(social_nguess[:,i]==2) #

plt.plot(np.linspace(1,21,21),np.mean(social_irrational,axis=0),'g-')
plt.plot(np.linspace(1,21,21),social_irr_avg_W,'r--')
plt.plot(np.linspace(1,21,21),social_irr_avg_B,'b-.')
plt.grid(True)


plt.plot(np.linspace(1,21,21),np.mean(indiv_irrational,axis=0),'g-')
plt.plot(np.linspace(1,21,21),np.mean(social_irrational,axis=0),'k-')
plt.grid(True)
######################################
                    
# keeping the whole data in temporary variables before removing the unuseful rounds      
indiv_irrational_temp = indiv_irrational
social_irrational_temp = social_irrational
social_irrational_corrected_temp = social_irrational_corrected
social_irrational_ball_temp = social_irrational_ball
neighbor_rational_temp = neighbor_rational

indiv_posteriors_temp = indiv_posteriors
social_posteriors_temp = social_posteriors

# removing the rounds in which both boxes have 5 black balls and 5 white balls
indiv_irrational = np.delete(indiv_irrational, 0, axis=1)
social_irrational = np.delete(social_irrational, 0, axis=1)
social_irrational_corrected = np.delete(social_irrational_corrected, 0, axis=1)
social_irrational_ball = np.delete(social_irrational_ball, 0, axis=1)
neighbor_rational = np.delete(neighbor_rational, 0, axis=1)

indiv_posteriors = np.delete(indiv_posteriors, 0, axis=1)
social_posteriors = np.delete(social_posteriors, 0, axis=1)

##########################################
# compute the irrationalities observed in the data
# =============================================================================
# N_social_irr_corrected = np.sum(neighbor_noball_wo) + np.sum(neighbor_noball_w) +\
#                         np.sum(neighbor_machine) + np.sum(neighbor_rational[np.where(neighbor_ball_wo)[0],])/20
# N_social_irr_indiv_upfront_corrected = np.sum(treatments=="D_I_wo") + np.sum(treatments=="D_I_w") +\
#                                          np.sum(treatments=="D_ID_machine") + np.sum(neighbor_rational[np.where(treatments=="D_ID_wo")[0],])/20
# N_social_irr_social_upfront_corrected = np.sum(treatments=="I_D_wo") + np.sum(treatments=="I_D_w") +\
#                                          np.sum(treatments=="ID_D_machine") + np.sum(neighbor_rational[np.where(treatments=="ID_D_wo")[0],])/20
# 
# =============================================================================
N_social_irr_corrected = np.sum(treatments == "D_I_w_errorrate")

social_irr = np.mean(social_irrational)
social_irr_corrected = np.sum(social_irrational_corrected)/(N_social_irr_corrected*20) # special
indiv_irr = np.mean(indiv_irrational)



indiv_rational_given_posterior = indiv_posteriors>49
indiv_error_posterior = np.logical_and(indiv_irrational,indiv_rational_given_posterior)
indiv_error_rationality = np.logical_and(indiv_irrational,np.logical_not(indiv_rational_given_posterior))
err_indiv = np.mean(indiv_irrational)
err_post_indiv = np.mean(indiv_error_posterior)
err_rationality_indiv = np.mean(indiv_error_rationality)


social_rational_given_posterior = social_posteriors>49
social_error_posterior = np.logical_and(social_irrational_corrected,social_rational_given_posterior)
social_error_rationality = np.logical_and(social_irrational_corrected,np.logical_not(social_rational_given_posterior))
err_social = np.sum(social_irrational_corrected)/(N_social_irr_corrected*20)
err_post_social = np.sum(social_error_posterior)/(N_social_irr_corrected*20)
err_rationality_social = np.sum(social_error_rationality)/(N_social_irr_corrected*20)




def graph_bar_corrected(colors,x_axis,y_axis,title,means,raw_data,name_fig):
    materials = x_axis 
    x_pos = np.arange(len(materials))
    CTEs = means
    error = []
    for i in range(len(materials)):
        error.append(np.sqrt(means[i]*(1-means[i])/((raw_data[i])*20)))
    
    # Build the plot
    fig, ax = plt.subplots()
    ax.bar(x_pos, CTEs, yerr=error, align='center', alpha=0.5, ecolor='black', capsize=10, color = colors)
    ax.set_ylabel(y_axis)
    ax.set_xticks(x_pos)
    ax.set_xticklabels(materials)
    ax.set_title(title)
    ax.yaxis.grid(True)  
    # Save the figure and show
    plt.tight_layout()
    #plt.savefig(name_fig)
    plt.show()
    
def grouped_graph_bar_corrected(x_axis,y_axis,title,means,raw_data,name_fig,labels):
    barWidth = 0.25
    materials = x_axis 
    x_pos = np.arange(len(materials))
    CTEs = means
    error = []
    for i in range(len(materials)*len(labels)):
        error.append(np.sqrt(means[i]*(1-means[i])/((raw_data[i])*20)))
    
    # Build the plot
    fig, ax = plt.subplots()
    for i in range(len(labels)):
        ax.bar(x_pos+barWidth*i, CTEs[i*len(materials):(i+1)*len(materials)], yerr=error[i*len(materials):(i+1)*len(materials)], width=barWidth, alpha=0.7, ecolor='black', capsize=5, label = labels[i])
    ax.set_ylabel(y_axis)
    ax.set_xticks(x_pos+barWidth*(len(labels)-1)/2)
    ax.set_xticklabels(materials)
    ax.set_title(title)
    ax.yaxis.grid(True) 
    ax.legend()

    # Save the figure and show
    plt.tight_layout()
    #plt.savefig(name_fig)
    plt.show()


# plot Figure 3    
graph_bar_corrected(['tab:blue','tab:orange'],['Individual', 'Social'],'Fraction of Irrational Trials','Irrationality',\
                    [indiv_irr, social_irr_corrected],[len(indiv_irrational),N_social_irr_corrected],'Figure3.eps')

# plot Figure 4
NNN = 151
plt.rcParams.update({'font.size': 12})
mean_ind = np.array([np.sum(np.sum(indiv_irrational[:,],axis=1)==0)/NNN,\
                 np.sum(np.sum(indiv_irrational[:,],axis=1)==1)/NNN,\
                 np.sum(np.sum(indiv_irrational[:,],axis=1)==2)/NNN,\
                 np.sum(np.sum(indiv_irrational[:,],axis=1)==3)/NNN,\
                 np.sum(np.sum(indiv_irrational[:,],axis=1)==4)/NNN,\
                 np.sum(np.sum(indiv_irrational[:,],axis=1)==5)/NNN,\
                 np.sum(np.sum(indiv_irrational[:,],axis=1)>5)/NNN])
mean_soc = np.array([np.sum(np.sum(social_irrational_corrected[:,],axis=1)==0)/NNN,\
                 np.sum(np.sum(social_irrational_corrected[:,],axis=1)==1)/NNN,\
                 np.sum(np.sum(social_irrational_corrected[:,],axis=1)==2)/NNN,\
                 np.sum(np.sum(social_irrational_corrected[:,],axis=1)==3)/NNN,\
                 np.sum(np.sum(social_irrational_corrected[:,],axis=1)==4)/NNN,\
                 np.sum(np.sum(social_irrational_corrected[:,],axis=1)==5)/NNN,\
                 np.sum(np.sum(social_irrational_corrected[:,],axis=1)>5)/NNN])  
barWidth = 0.35
materials = np.array(['0','1','2','3','4','5','>5'])
labels = ['individual','social']
x_pos = np.arange(len(materials))
CTEs = np.append(mean_ind,mean_soc)
fig, ax = plt.subplots()
for i in range(len(labels)):
    ax.bar(x_pos+barWidth*i, CTEs[i*len(materials):(i+1)*len(materials)], width=barWidth,\
                                  alpha=0.7, ecolor='black', capsize=5, label = labels[i])
ax.set_ylabel('Fraction of Subjects')
ax.set_xticks(x_pos+barWidth*(len(labels)-1)/2)
ax.set_xticklabels(materials)
ax.set_xlabel('Number of Irrational Choices')
ax.set_title('Distribution of Irrationality across Subjects')
ax.yaxis.grid(True)
plt.tight_layout()    
ax.legend()
plt.show()


# plot Figure 8
grouped_graph_bar_corrected(['Individual', 'Social'],'Fraction of Irrational Trials','Causes of Irrationality',\
                            [err_post_indiv, err_post_social, err_rationality_indiv, err_rationality_social],\
                            [len(indiv_irrational),N_social_irr_corrected,len(indiv_irrational),N_social_irr_corrected],\
                            'Figure8.eps',['posterior error', 'reasoning error'])