import json
from matplotlib import pyplot as plt
from island.match import Match
from island.matches import Matches
from numpy import mean, std
import numpy as np


matches = Matches('wos-data-new-2')
max_round = 15

survivals = {}
with open('survivals-2.json', 'r') as f:
    survivals = json.load(f)

neighbors = {}

cmean = []
dmean = []
cstd = []
dstd = []


def is_cooperator(rows, pid):
    m = 0
    mall = 0
    for r in rows:
        if pid == r['a'] and r['act_a'] == 'C':
            m += 1
            mall += 1
        elif pid == r['a'] and r['act_a'] == 'D':
            mall += 1
        elif pid == r['b'] and r['act_b'] == 'C':
            m += 1
            mall += 1
        elif pid == r['b'] and r['act_b'] == 'D':
            mall += 1
    return m*2 >= mall


for i in range(len(matches.data)):
    m = matches.data[i]
    n = {}
    for r in m.query('neighbor', 'create').raw_data:
        if r['a'] in n:
            n[r['a']].append(r['b'])
        else:
            n[r['a']] = [r['b']]

        if r['b'] in n:
            n[r['b']].append(r['a'])
        else:
            n[r['b']] = [r['a']]    
    neighbors[matches.names[i]] = n


for i in range(max_round):
    cneigh = []
    dneigh = []
    for j in range(len(matches.data)):
        rows = matches.data[j].query('action', 'done').where(lambda x: x['rno']==i+1).raw_data
        calced = set()
        for row in rows:
            if row['a'] not in calced:
                nn = 0
                for k in neighbors[matches.names[j]][row['a']]:
                    if k in survivals[matches.names[j]][str(i+1)]:
                        nn += 1
                
                if is_cooperator(rows, row['a']):
                    cneigh.append(nn)
                else:
                    dneigh.append(nn)

                calced.add(row['a'])
            if row['b'] not in calced:
                nn = 0
                for k in neighbors[matches.names[j]][row['b']]:
                    if k in survivals[matches.names[j]][str(i+1)]:
                        nn += 1
                if is_cooperator(rows, row['b']):
                    cneigh.append(nn)
                else:
                    dneigh.append(nn)
                calced.add(row['b'])

    if cneigh:
        cm = mean(cneigh)
        cs = std(cneigh)
    else:
        cm = 0
        cs = 0
    cmean.append(cm)
    cstd.append(cs)
    if dneigh:
        dm = mean(dneigh)
        ds = std(dneigh)
    else:
        dm = 0
        ds = 0
    dmean.append(dm)
    dstd.append(ds)

fig = plt.figure(figsize=(6.4, 3.6))
ax = fig.gca()
index = np.arange(max_round)
bar_width = 0.45
opacity = 1
error_config = dict(ecolor='#2d3436', capsize=0, elinewidth=1)
rects1 = ax.bar(index, cmean, bar_width,
                alpha=opacity, color='#0984e3',
                yerr=cstd, error_kw=error_config,
                label='Cooperator')
rects2 = ax.bar(index + bar_width, dmean, bar_width,
                alpha=opacity, color='#d63031',
                yerr=dstd, error_kw=error_config,
                label='Defector')

ax.set_xlabel('Round')
ax.set_ylabel('Size of the Neighborhood')
# ax.set_title('Scores by group and gender')
ax.set_xticks(index + bar_width / 2)
ax.set_xticklabels(index+1)
ax.legend()
fig.tight_layout()
plt.show()
# plt.savefig('graph/neigh_per_round.eps')