# -*- coding: utf-8 -*-
"""
Created on Fri Jul 28 12:12:47 2023

@author: Andrei Sontag
"""
# importing required libraries
import numpy as np
import pandas as pd
import os, fnmatch
import matplotlib.pyplot as plt
import seaborn as sns

plt.style.use('ggplot')
plt.close('all')

# make sure the script is in the same folder as the data for the empirical distributions
os.chdir(r'./')

#%%
# first analysing the data for the sample of 9 votes
file1 = "sample_9.csv"
# load data
df = pd.read_csv(file1)

# define the headings of the table we want to read: the election votes in each round
test=[]
for k in np.arange(1,122):
    test = np.append(test,r"empirical_dist.{0:.0f}.player.election_vote".format(k))
    
# get the votes in each round
elec_vote = df[test].to_numpy()

# this is the sample order used in the experiment and shown to the participants
sample = np.array([[2, 0, 7], [0, 3, 6], [0, 1, 8], [2, 1, 6], [5, 1, 3], [4, 1, 4], [3, 3, 3],
        [6, 1, 2], [2, 2, 5], [5, 4, 0], [3, 2, 4], [5, 4, 0], [4, 4, 1], [1, 3, 5], [0, 3, 6],
        [3, 5, 1], [3, 5, 1], [2, 6, 1], [0, 5, 4], [1, 1, 7], [1, 0, 8], [5, 3, 1], [0, 2, 7],
        [6, 2, 1], [7, 0, 2], [4, 0, 5], [8, 0, 1], [3, 1, 5], [1, 8, 0], [0, 6, 3], [7, 1, 1], 
        [7, 0, 2], [0, 7, 2], [1, 4, 4], [3, 6, 0], [5, 1, 3], [6, 3, 0], [0, 0, 9], [9, 0, 0],
        [2, 4, 3], [6, 0, 3], [1, 1, 7], [1, 2, 6], [2, 7, 0], [1, 8, 0], [3, 3, 3], [0, 0, 9],
        [6, 0, 3], [2, 3, 4], [1, 7, 1], [4, 2, 3], [2, 5, 2], [0, 0, 9], [2, 2, 5], [1, 0, 8],
        [3, 0, 6], [4, 5, 0], [4, 1, 4], [7, 1, 1], [1, 1, 7], [1, 1, 7], [1, 7, 1], [5, 2, 2], 
        [3, 4, 2], [2, 6, 1], [8, 1, 0], [3, 4, 2], [4, 3, 2], [0, 4, 5], [8, 1, 0], [0, 4, 5],
        [3, 6, 0], [8, 0, 1], [4, 0, 5], [2, 5, 2], [7, 2, 0], [6, 3, 0], [7, 2, 0], [0, 8, 1],
        [3, 0, 6], [0, 1, 8], [1, 6, 2], [1, 2, 6], [2, 2, 5], [4, 5, 0], [2, 2, 5], [2, 0, 7],
        [4, 4, 1], [4, 4, 1], [9, 0, 0], [1, 5, 3], [1, 4, 4], [4, 2, 3], [2, 7, 0], [6, 1, 2],
        [0, 5, 4], [3, 2, 4], [2, 4, 3], [0, 7, 2], [6, 2, 1], [4, 3, 2], [2, 3, 4], [0, 0, 9],
        [5, 0, 4], [5, 3, 1], [0, 8, 1], [2, 1, 6], [0, 9, 0], [3, 3, 3], [3, 3, 3], [4, 4, 1],
        [1, 5, 3], [1, 6, 2], [0, 2, 7], [3, 1, 5], [0, 9, 0], [5, 2, 2], [1, 3, 5], [0, 6, 3],
        [5, 0, 4]])

# array with the number of A votes in each sample
asample = sample[:,0]

# array with the number of B votes in each sample
bsample = sample[:,1]

# array with the number of Abstentions in each sample
abssample = sample[:,2]

# define the difference between A and B votes in each sample
vote_diff = asample-bsample

# we are going to look at the probabilities of voting A or B based on the difference seen in the sample
counts = np.zeros(19)
prob = np.zeros((19,3))
for rou in np.arange(0,120):
    for part in np.arange(0,len(elec_vote[:,1])):
        if ~np.isnan(elec_vote[part,rou]):
            counts[vote_diff[rou]+9] = counts[vote_diff[rou]+9]+1
            if elec_vote[part,rou+1] == 0:
                prob[vote_diff[rou]+9,0] = prob[vote_diff[rou]+9,0]+1
            elif elec_vote[part,rou+1] == 1:
                prob[vote_diff[rou]+9,1] = prob[vote_diff[rou]+9,1]+1
            else:
                prob[vote_diff[rou]+9,2] = prob[vote_diff[rou]+9,2]+1

# estimated probability = counts for each choosen option / total number of counts for that combination of A-B
prob[:,0] = prob[:,0]/counts
prob[:,1] = prob[:,1]/counts
prob[:,2] = prob[:,2]/counts

# define the probability of voting A, B or abstaining for each presented sample
prob2d = np.zeros((10,10,3))
counts2d = np.zeros((10,10))
for rou in np.arange(0,120):
    for part in np.arange(0,len(elec_vote[:,1])):
        if ~np.isnan(elec_vote[part,rou]):
            counts2d[int(asample[rou]),int(bsample[rou])] += 1
            if elec_vote[part,rou+1] == 0:
                prob2d[int(asample[rou]),int(bsample[rou]),0] = prob2d[int(asample[rou]),int(bsample[rou]),0]+1
            elif elec_vote[part,rou+1] == 1:
                prob2d[int(asample[rou]),int(bsample[rou]),1] = prob2d[int(asample[rou]),int(bsample[rou]),1]+1
            else:
                prob2d[int(asample[rou]),int(bsample[rou]),2] = prob2d[int(asample[rou]),int(bsample[rou]),2]+1

# estimate the probability
prob2d[:,:,0] = prob2d[:,:,0]/counts2d
prob2d[:,:,1] = prob2d[:,:,1]/counts2d
prob2d[:,:,2] = prob2d[:,:,2]/counts2d

### Plotting data ####
xax = np.arange(-9,10)

plt.figure(figsize=(8,6))
plt.bar(xax,prob[:,0])
plt.bar(xax,prob[:,1],bottom=prob[:,0])
plt.bar(xax,prob[:,2],bottom=prob[:,0]+prob[:,1])
plt.legend(["A", "B", "Abstain"])
plt.xlabel("A-B votes")
plt.title("Probability of voting A/B given difference in votes")
plt.show()

plt.figure(figsize=(8,6))
plt.bar(xax,prob[:,0],alpha=0.75)
plt.bar(xax,prob[:,1],alpha=0.75)
plt.legend(["A", "B"])
plt.xlabel("A-B votes")
plt.title("Probability of voting A/B given difference in votes")
plt.show()

plt.figure(figsize=(8,6))
plt.bar(xax,prob[:,0],alpha=0.75)
plt.bar(xax,prob[::-1,1],alpha=0.75)
plt.legend(["A", "B"])
plt.xlabel("A-B votes")
plt.title("Probability of voting A/B given difference in votes")
plt.show()

### 3D plot ###
axis = np.arange(0,10)
xv,yx = np.meshgrid(axis,axis)
xv = xv.reshape(1,100)[0]
yx = yx.reshape(1,100)[0]
zv = np.zeros(100)
dx = np.ones(100)
dy = np.ones(100)
dz = prob2d[:,:,0].reshape(1,100)[0]

# creating figure
fig = plt.figure(figsize=(8,6))
ax = fig.add_subplot(projection='3d')

# creating the plot
plot_geeks = ax.bar3d(xv,yx,zv,dx,dy,dz,color='royalblue')

# setting title and labels
ax.set_title("3D plot")
ax.set_xlabel('x-axis')
ax.set_ylabel('y-axis')
ax.set_zlabel('z-axis')
ax.set_zlim(0, 1)

# rotate the axes and update
for angle in range(0, 360):
   ax.view_init(30, 40)

# displaying the plot
plt.show()

fig, axes = plt.subplots(1, 3, figsize=(24, 6), sharey=True)
sns.heatmap(prob2d[:,:,0],ax=axes[0])
axes[0].set_title("Probability of voting A")
axes[0].set(xlabel='B votes', ylabel='A votes')
sns.heatmap(prob2d[:,:,1],ax=axes[1])
axes[1].set_title("Probability of voting B")
axes[1].set(xlabel='B votes')
sns.heatmap(prob2d[:,:,2],ax=axes[2])
axes[2].set_title("Probability of abstaining")
axes[2].set(xlabel='B votes')
plt.show()

#%%
# same analysis but now for the data of a sample of 5 votes
file2 = "sample_5.csv"
df = pd.read_csv(file2)

test=[]
for k in np.arange(1,122):
    test = np.append(test,r"empirical_dist.{0:.0f}.player.election_vote".format(k))
elec_vote = df[test].to_numpy()

sample = np.array([[2, 1, 2], [1, 0, 4], [0, 3, 2], [0, 2, 3], [0, 4, 1], [4, 0, 1], [0, 0, 5], [0, 3, 2],
       [0, 3, 2], [0, 4, 1], [0, 0, 5], [2, 0, 3], [0, 2, 3], [0, 0, 5], [2, 3, 0], [1, 3, 1],
       [0, 3, 2], [3, 1, 1], [3, 0, 2], [3, 1, 1], [2, 2, 1], [1, 1, 3], [3, 0, 2], [1, 1, 3],
       [1, 1, 3], [3, 0, 2], [3, 2, 0], [0, 0, 5], [1, 1, 3], [4, 1, 0], [1, 1, 3], [0, 1, 4],
       [1, 3, 1], [3, 0, 2], [0, 5, 0], [1, 0, 4], [0, 0, 5], [5, 0, 0], [1, 1, 3], [0, 0, 5],
       [2, 3, 0], [3, 0, 2], [0, 3, 2], [3, 1, 1], [2, 0, 3], [4, 1, 0], [0, 5, 0], [2, 0, 3],
       [0, 1, 4], [2, 2, 1], [5, 0, 0], [1, 2, 2], [1, 4, 0], [4, 1, 0], [5, 0, 0], [2, 2, 1],
       [1, 4, 0], [1, 3, 1], [3, 2, 0], [0, 2, 3], [1, 2, 2], [1, 0, 4], [2, 0, 3], [2, 2, 1],
       [0, 0, 5], [0, 1, 4], [0, 4, 1], [0, 0, 5], [2, 1, 2], [3, 2, 0], [2, 2, 1], [2, 1, 2],
       [4, 1, 0], [0, 0, 5], [0, 0, 5], [1, 2, 2], [1, 1, 3], [1, 3, 1], [0, 4, 1], [3, 2, 0],
       [1, 2, 2], [2, 3, 0], [0, 5, 0], [5, 0, 0], [2, 2, 1], [0, 2, 3], [4, 0, 1], [0, 4, 1],
       [3, 1, 1], [1, 0, 4], [4, 0, 1], [1, 4, 0], [0, 2, 3], [4, 0, 1], [2, 2, 1], [1, 1, 3],
       [0, 1, 4], [2, 2, 1], [2, 2, 1], [1, 4, 0], [0, 5, 0], [2, 1, 2], [5, 0, 0], [4, 0, 1],
       [2, 1, 2], [2, 3, 0], [4, 1, 0], [1, 2, 2], [1, 0, 4], [3, 1, 1], [0, 1, 4], [3, 2, 0],
       [1, 1, 3], [2, 0, 3], [2, 3, 0], [1, 3, 1], [2, 2, 1], [1, 4, 0], [1, 1, 3], [0, 5, 0]])

asample = sample[:,0]
bsample = sample[:,1]
abssample = sample[:,2]
vote_diff = asample-bsample


counts = np.zeros(11)
prob = np.zeros((11,3))
for rou in np.arange(0,120):
    for part in np.arange(0,len(elec_vote[:,1])):
        if ~np.isnan(elec_vote[part,rou]):
            counts[vote_diff[rou]+5] = counts[vote_diff[rou]+5]+1
            if elec_vote[part,rou+1] == 0:
                prob[vote_diff[rou]+5,0] = prob[vote_diff[rou]+5,0]+1
            elif elec_vote[part,rou+1] == 1:
                prob[vote_diff[rou]+5,1] = prob[vote_diff[rou]+5,1]+1
            else:
                prob[vote_diff[rou]+5,2] = prob[vote_diff[rou]+5,2]+1

prob[:,0] = prob[:,0]/counts
prob[:,1] = prob[:,1]/counts
prob[:,2] = prob[:,2]/counts
print(np.sum(prob,axis=1))

prob2d = np.zeros((6,6,3))
counts2d = np.zeros((6,6))
for rou in np.arange(0,120):
    for part in np.arange(0,len(elec_vote[:,1])):
        if ~np.isnan(elec_vote[part,rou]):
            counts2d[int(asample[rou]),int(bsample[rou])] += 1
            if elec_vote[part,rou+1] == 0:
                prob2d[int(asample[rou]),int(bsample[rou]),0] = prob2d[int(asample[rou]),int(bsample[rou]),0]+1
            elif elec_vote[part,rou+1] == 1:
                prob2d[int(asample[rou]),int(bsample[rou]),1] = prob2d[int(asample[rou]),int(bsample[rou]),1]+1
            else:
                prob2d[int(asample[rou]),int(bsample[rou]),2] = prob2d[int(asample[rou]),int(bsample[rou]),2]+1

prob2d[:,:,0] = prob2d[:,:,0]/counts2d
prob2d[:,:,1] = prob2d[:,:,1]/counts2d
prob2d[:,:,2] = prob2d[:,:,2]/counts2d

### Plotting figures ###
xax = np.arange(-5,6)

plt.figure(figsize=(8,6))
plt.bar(xax,prob[:,0])
plt.bar(xax,prob[:,1],bottom=prob[:,0])
plt.bar(xax,prob[:,2],bottom=prob[:,0]+prob[:,1])
plt.legend(["A", "B", "Abstain"])
plt.xlabel("A-B votes")
plt.title("Probability of voting A/B given difference in votes")
plt.show()

plt.figure(figsize=(8,6))
plt.bar(xax,prob[:,0],alpha=0.75)
plt.bar(xax,prob[:,1],alpha=0.75)
plt.legend(["A", "B"])
plt.xlabel("A-B votes")
plt.title("Probability of voting A/B given difference in votes")
plt.show()

plt.figure(figsize=(8,6))
plt.bar(xax,prob[:,0],alpha=0.75)
plt.bar(xax,prob[::-1,1],alpha=0.75)
plt.legend(["A", "B"])
plt.xlabel("A-B votes")
plt.title("Probability of voting A/B given difference in votes")
plt.show()

### 3D plot ###
axis = np.arange(0,6)
xv,yx = np.meshgrid(axis,axis)
xv = xv.reshape(1,36)[0]
yx = yx.reshape(1,36)[0]
zv = np.zeros(36)
dx = np.ones(36)
dy = np.ones(36)
dz = prob2d[:,:,0].reshape(1,36)[0]

# creating figure
fig = plt.figure(figsize=(8,6))
ax = fig.add_subplot(projection='3d')

# creating the plot
plot_geeks = ax.bar3d(xv,yx,zv,dx,dy,dz)

# setting title and labels
ax.set_title("3D plot")
ax.set_xlabel('x-axis')
ax.set_ylabel('y-axis')
ax.set_zlabel('z-axis')
ax.set_zlim(0, 1)

# rotate the axes and update
for angle in range(0, 360):
   ax.view_init(30, 40)

# displaying the plot
plt.show()

fig, axes = plt.subplots(1, 3, figsize=(24, 6), sharey=True)
sns.heatmap(prob2d[:,:,0],ax=axes[0])
axes[0].set_title("Probability of voting A")
axes[0].set(xlabel='B votes', ylabel='A votes')
sns.heatmap(prob2d[:,:,1],ax=axes[1])
axes[1].set_title("Probability of voting B")
axes[1].set(xlabel='B votes')
sns.heatmap(prob2d[:,:,2],ax=axes[2])
axes[2].set_title("Probability of abstaining")
axes[2].set(xlabel='B votes')
plt.show()