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')
max_round = 14

c_req_suc_mean = []
d_req_suc_mean = []
c_req_suc_std = []
d_req_suc_std = []

c_req_fail_mean = []
d_req_fail_mean = []
c_req_fail_std = []
d_req_fail_std = []

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(max_round):
    cReSu = []
    cReFa = []
    dReSu = []
    dReFa = []
    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:
                r = matches.data[j].query('action', 'request').where(lambda x: x['rno'] == i+2 and x['from'] == row['a']).raw_data
                rs = matches.data[j].query('action', 'approve').where(lambda x: x['rno'] == i+2 and x['to'] == row['a']).raw_data
                if is_cooperator(rows, row['a']):
                    cReSu.append(len(rs))
                    cReFa.append(len(r) - len(rs))
                else:
                    dReSu.append(len(rs))
                    dReFa.append(len(r) - len(rs))
                calced.add(row['a'])

            if row['b'] not in calced:
                r = matches.data[j].query('action', 'request').where(lambda x: x['rno'] == i+2 and x['from'] == row['b']).raw_data
                rs = matches.data[j].query('action', 'approve').where(lambda x: x['rno'] == i+2 and x['to'] == row['b']).raw_data
                if is_cooperator(rows, row['a']):
                    cReSu.append(len(rs))
                    cReFa.append(len(r) - len(rs))
                else:
                    dReSu.append(len(rs))
                    dReFa.append(len(r) - len(rs))
                calced.add(row['b'])

    if cReSu:
        cm = mean(cReSu)
        cs = std(cReSu)
    else:
        cm = 0
        cs = 0
    c_req_suc_mean.append(cm)
    c_req_suc_std.append(cs)

    if cReFa:
        cm = mean(cReFa)
        cs = std(cReFa)
    else:
        cm = 0
        cs = 0
    c_req_fail_mean.append(cm)
    c_req_fail_std.append(cs)

    if dReSu:
        dm = mean(dReSu)
        ds = std(dReSu)
    else:
        dm = 0
        ds = 0
    d_req_suc_mean.append(dm)
    d_req_suc_std.append(ds)

    if dReFa:
        dm = mean(dReFa)
        ds = std(dReFa)
    else:
        dm = 0
        ds = 0
    d_req_fail_mean.append(dm)
    d_req_fail_std.append(ds)    

fig = plt.figure(figsize=(6.4, 4))
ax = fig.gca()
index = np.arange(max_round)
bar_width = 0.35
opacity = 1
error_config = {'ecolor': '0.3', 'capsize': 4}
rects1 = ax.bar(index, c_req_suc_mean, bar_width,
                alpha=opacity, color='#00b894',
                # yerr=c_req_suc_std, error_kw=error_config,
                label='Request from Cooperator Approved')
rects2 = ax.bar(index, c_req_fail_mean, bar_width,
                alpha=opacity, color='#0984e3',
                # yerr=c_req_fail_std, error_kw=error_config, 
                bottom=c_req_suc_mean,
                label='Request from Cooperator Rejected')

rects3 = ax.bar(index + bar_width, d_req_suc_mean, bar_width,
                alpha=opacity, color='#fdcb6e',
                # yerr=d_req_suc_mean, error_kw=error_config,
                label='Request from Defector Approved')

rects4 = ax.bar(index + bar_width, d_req_fail_mean, bar_width,
                alpha=opacity, color='#d63031',
                # yerr=d_req_fail_mean, error_kw=error_config, 
                bottom=d_req_suc_mean,
                label='Request from Defector Rejected')

ax.set_xlabel('Round')
ax.set_ylabel('Number of Requests')
# ax.set_title('Scores by group and gender')
ax.set_xticks(index + bar_width / 2)
ax.set_xticklabels(index+2)
ax.legend()
fig.tight_layout()
# plt.show()
plt.savefig('graph/request_success_per_round.eps')