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

def error(f,x,y):
    return sp.sum((f(x)-y)**2)

if __name__ == '__main__':

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

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

    neighbors = {}
    coopr = []
    x = np.arange(1, max_round)
    mwCo = {} # Match-wise frequency of cooperation
    mwTau = {} # Match-wise Tau
    bx = []
    tau = []

    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):
        co = []
        for j in range(len(matches.data)):
            coop = 0
            rows = matches.data[j].query('action', 'done').where(lambda x: x['rno']==i+1).raw_data
            for row in rows:
                if row['act_a'] == 'C':
                    coop += 1
                if row['act_b'] == 'C':
                    coop += 1

            if rows:
                co.append(float(coop) / float(len(rows)*2))
                mwCo["%s-%d"%(j,i)] = co[-1]
        bx.append(co)
        if co:
            coopr.append(np.average(co))


    for i in range(max_round-1):
        tp = []
        for j in range(len(matches.data)):
            if i == 0:
                for r in matches.data[j].query('player', 'join').raw_data:
                    t = 0
                    k = r['pid']
                    if k not in neighbors[matches.names[j]]:
                        print("[%s(%d)] alone: %d" % (matches.names[j], i+1, k))
                    else:
                        t = 1440 * len(neighbors[matches.names[j]][k])
                    tp.append(t if t < 1440 else 1440)
                    mwTau["%s-%d"%(j,i)] = tp[-1]
            else:
                if str(i) not in survivals[matches.names[j]]:
                    continue
                for k in survivals[matches.names[j]][str(i)]:
                    t = 0
                    if k not in neighbors[matches.names[j]]:
                        print("[%s(%d)] alone: %d" % (matches.names[j], i+1, k))
                    else:
                        trs = matches.data[j].get_tr(i, k, neighbors[matches.names[j]][k], survivals[matches.names[j]][str(i)])
                        for l in neighbors[matches.names[j]][k]:
                            if l in trs:
                                t += trs[l]
                    tp.append(t if t < 1440 else 1440)
                    mwTau["%s-%d"%(j,i)] = tp[-1]

        if tp:
            tau.append(np.average(tp))
        else:
            tau.append(0)


    blue = '#0984e3'
    red = '#d63031'

    # p1折线图
    # fig = plt.figure(figsize=(6.4, 3.6))
    # ax = fig.gca()
    # ax.plot(x, coopr, color=blue, linewidth=3)
    # ax.set_ylim(0.5, 1)
    # ax2 = ax.twinx()
    # ax2.plot(x, tau, color=red, linewidth=3)
    # ax2.set_ylim(0,1440)
    # ax.set_xlim(1,14)
    # ax.set_xlabel("Rounds")
    # ax.set_ylabel("Frequency of Cooperation", color=blue)
    # ax.tick_params(axis='y', labelcolor=blue)
    # ax2.set_ylabel("$\\tau_{p}$", family='sans-serif', color=red)
    # ax2.tick_params(axis='y', labelcolor=red)

    # plt.tight_layout()
    # # plt.show()
    # plt.savefig('graph/tau_p_co_plot.eps')

    tau2 = []
    coopr2 = []
    tau_r = []
    coopr_r = []
    for i in range(len(tau)):
        if tau[i] <= 720:
            tau2.append(tau[i])
            coopr2.append(coopr[i])
        else:
            tau_r.append(tau[i])
            coopr_r.append(coopr[i])

    # p2散点图
    fig = plt.figure(figsize=(6.4, 3.6))
    ax = fig.gca()
    # ax.set_ylim(0.5, 1)
    fp1,residuals,rank,sv,rcond = sp.polyfit(tau2, coopr2, 1, full=True)
    print("残差：",residuals)
    print('Model parameter:',fp1)
    f1 = sp.poly1d(fp1)
    print("error= %f" % error(f1, tau2, coopr2))
    # fx = sp.linspace(0,max(tau2),1000)
    fx = sp.linspace(0,720,2)

    plt.plot(fx,f1(fx),linewidth=2,color=red, ls='--', zorder=0)
    plt.scatter(tau2, coopr2, color=blue, linewidths=2, zorder=100)
    plt.scatter(tau_r, coopr_r, color='white', edgecolors=blue, linewidths=2, zorder=101)
    ax.set_xlabel('$\\tau_{p}$', family='sans-serif')
    ax.set_ylabel('Frequency of Cooperation')
    ax.set_xlim(0, 1440)
    ax.set_xticks(sp.linspace(0, 1440, 13))
    ax.set_ylim(0.5, 1)
    plt.tight_layout()
    plt.show()
    # plt.savefig('graph/tau_p_co_sca.eps')



    # 皮尔逊相关系数
    print("pearson: %f, p-value: %f" % pearsonr(tau2, coopr2))