import json
from matplotlib import pyplot as plt
from island.match import Match
from island.matches import Matches
import numpy as np
from scipy.stats import pearsonr


matches = Matches('wos-data-new')

k = np.arange(2, 11)
succ = np.zeros(9)
total = np.zeros(9)

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

neighbors = {}
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 m_i in range(len(matches.data)):
    m = matches.data[m_i]
    info = m.query('game', 'created').select('info').first()['info']
    conf = json.loads(info['config'])
    game_end_at = int(info['game_end_at'])

    for p in m.query('player', 'join').raw_data:
        pid = p['pid']
        for i in range(2, game_end_at):
            neighborhood = []
            if pid not in neighbors[matches.names[m_i]]:
                break
            for j in neighbors[matches.names[m_i]][pid]:
                if j in survivals[matches.names[m_i]][str(i-1)]:
                    neighborhood.append(j)
            if len(neighborhood) < 2:
                break
            previous_round_partner = []
            for r in m.query('action', 'done').where(lambda x: x['rno']==i-1 and (x['a']==pid or x['b']==pid)).raw_data:
                if r['a'] == pid:
                    previous_round_partner.append(r['b'])
                else:
                    previous_round_partner.append(r['a'])
            new_partner_request = 0
            for r in m.query('action', 'request').where(lambda x: x['rno']==i and (x['from']==pid or x['to']==pid)).raw_data:
                if r['from'] == pid:
                    if r['to'] not in previous_round_partner:
                        new_partner_request += 1
                else:
                    if r['from'] not in previous_round_partner:
                        new_partner_request += 1
            if new_partner_request == 0:
                continue
            new_partner_succ = 0
            for r in m.query('action', 'done').where(lambda x: x['rno']==i and (x['a']==pid or x['b']==pid)).raw_data:
                if r['a'] == pid:
                    if r['b'] not in previous_round_partner:
                        new_partner_succ += 1
                else:
                    if r['a'] not in previous_round_partner:
                        new_partner_succ += 1
            succ[len(neighborhood)-2] += new_partner_succ
            total[len(neighborhood)-2] += new_partner_request


red = '#d63031'
fig = plt.figure(figsize=(6.4, 4))
ax = fig.gca()
bar_width = 0.35
opacity = 1
error_config = {'ecolor': '0.3', 'capsize': 4}
rects1 = ax.bar(k, succ, bar_width,
                alpha=opacity, color='#00b894',
                # yerr=c_req_suc_std, error_kw=error_config,
                label='Success')
rects3 = ax.bar(k + bar_width, total, bar_width,
                alpha=opacity, color='#fdcb6e',
                # yerr=d_req_suc_mean, error_kw=error_config,
                label='Requests')

ax.set_xlabel('k')
ax.set_ylabel('Count')
# ax.set_title('Scores by group and gender')
ax.set_xticks(k + bar_width / 2)
ax.set_xticklabels(k)
ax.legend()
ax2 = ax.twinx()
ax2.plot(k, succ/total,linewidth=2,color=red, ls='--')
ax2.set_ylabel("Frequency of new partners", family='sans-serif', color=red)
ax2.tick_params(axis='y', labelcolor=red)
ax2.set_ylim(0,1)
fig.tight_layout()
# plt.show()
# plt.savefig('graph/k_and_new_partner.eps')

print("[succ vs k]pearson: %f, p-value: %f" % pearsonr(succ, k))
print("[total vs k]pearson: %f, p-value: %f" % pearsonr(total, k))
print("[rate vs k]pearson: %f, p-value: %f" % pearsonr(succ/total, k))
print(np.average(succ/total))