import csv
from math import floor
import numpy as np
from matplotlib import pyplot as plt
from matplotlib import colors
from matplotlib import cm

data = np.loadtxt('outputs/decision.csv', delimiter=',')
# print(data)
print(data.shape)
ndata = np.zeros(shape=(6, 45))
for i in range(data.shape[0]):
    ans = floor(data[i, 0]) - 1
    if ans < 0:
        ans = 0
    elif ans >= 45:
        ans = 44
    ndata[int(data[i, 1]),  ans] += 1

print(ndata[4,11])


for i in range(ndata.shape[0]):
    for j in range(ndata.shape[1]):
        if ndata[i,j] == 0:
            ndata[i,j] = np.nan

# ndata = np.log(ndata+1)
# ndata += 1
fig = plt.figure()
ax = fig.gca()
# ax.set_xlabel()
ax.set_facecolor('white')
ax.set_xticks(np.arange(0,45,11))
colors = cm.get_cmap('OrRd')
colors.set_bad('white')
ax.set_xticklabels(np.arange(1, 46, 11))
dct = ax.imshow(ndata, cmap=colors, origin='lower')#, norm=colors.LogNorm())

ax.plot(np.arange(7, 45), np.floor(45 / np.arange(8, 46)),
        marker='x', color='black', markersize=3, linestyle='none', markeredgewidth=0.5)
ax.set_xlabel("Decision time")
ax.set_ylabel("Number of opponents")
fig.colorbar(dct, ax=ax, orientation='horizontal', aspect=40)

# plt.show()
plt.savefig('graph/decision_time.eps')