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

'''
计算tau_f和合作频率之间的关系
'''

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

def p1(x, coopr, tau, postfix, show=True):
    fig = plt.figure(figsize=(4.5, 3))
    ax = fig.gca()
    ax.plot(x, coopr, color=blue, linewidth=2, label="$f_c$")
    ax.set_ylim(0, 1)
    ax.set_yticks(sp.linspace(0, 1, 5))
    ax2 = ax.twinx()
    ax2.plot(x, tau, color=red, linewidth=2, label=r"$\tau_f$")
    ax2.set_ylim(0, 1440)
    ax2.set_yticks(sp.linspace(0, 1440, 5))
    ax2.tick_params(labelsize=18)
    ax.tick_params(labelsize=18)
    ax.set_xlim(1, 15)
    ax.set_xlabel("Rounds", size=22)
    ax.set_ylabel(r"$f_c$", family='sans-serif', color=blue, size=22)
    ax.tick_params(axis='y', labelcolor=blue)
    ax2.set_ylabel(r"$\tau_{f}$", family='sans-serif', color=red, size=22)
    ax2.tick_params(axis='y', labelcolor=red)

    # fig.legend(fontsize=18)
    plt.tight_layout()
    if show:
        plt.show()
    else:
        plt.savefig("graph/tau_f_co_plot_%s.eps" % postfix)

def p2(tau, coopr, limited, postfix, show=True):
    tau2 = []
    coopr2 = []
    tau_r = []
    coopr_r = []
    for i in range(len(tau)):
        if tau[i] <= limited:
            tau2.append(tau[i])
            coopr2.append(coopr[i])
        else:
            tau_r.append(tau[i])
            coopr_r.append(coopr[i])

    # p2散点图
    fig = plt.figure(figsize=(4, 3))
    ax = fig.gca()
    if tau2:
        # ax.set_ylim(0.5, 1)
        fp1,residuals,rank,sv,rcond = sp.polyfit(tau2, coopr2, 1, full=True)
        print("残差：",residuals)
        print('Model parameter:',fp1)
        print("Other parameters: rank=%s, sv=%s, rcond=%s"%(str(rank), str(sv), str(rcond)))
        f1 = sp.poly1d(fp1)
        print("error= %f" % error(f1, tau2, coopr2))
        fx = sp.linspace(0, limited, 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(r'$\tau_{f}$', family='sans-serif', size=20)
    ax.set_ylabel(r'$f_{c}$', family='sans-serif', size=20)
    ax.set_xlim(0, 1440)
    ax.set_xticks(sp.linspace(0, 1440, 5))
    ax.tick_params(labelsize=14)
    ax.set_ylim(0.5, 1)
    ax.set_yticks(sp.linspace(0.5, 1, 5))
    plt.tight_layout()
    if show:
        plt.show()
    else:
        plt.savefig("graph/tau_f_co_sca_%s.eps" % postfix)

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

if __name__ == '__main__':
  
    # matches = Matches.from_profile_expr(lambda r: 'CLASSIC' in r)
    matches = Matches.from_profile_expr(lambda r: 'SURVIVE' in r)
    max_round = 15

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

    neighbors = {}
    coopr = []
    x = np.arange(1, max_round+1)
    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))
        else:
            coopr.append(0)


    for i in range(max_round):
        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)
                    tp.append(t)
                    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 and trs[l] > 0:
                                t += trs[l] if trs[l] < 1440 else 1440
                    tp.append(t if t < 1440 else 1440)
                    # tp.append(t)
                    mwTau["%s-%d"%(j,i)] = tp[-1]

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


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

    # p1折线图
    # p1(x, coopr, tau, 'classic', False)
    # p1(x, coopr, tau, 'survive', True)

    p2(tau[:12], coopr[:12], 1322, 'survive', True)
    # p2(tau, coopr, 1399, 'classic', True)


    '''

    survive
    残差： [ 0.0054399]
    Model parameter: [ -2.69062302e-04   1.00822767e+00]
    Other parameters: rank=2, sv=[ 1.36265876  0.37836635], rcond=1.99840144433e-15
    error= 0.005440
    pearson: -0.830766, p-value: 0.005540

    classic
    pearson: -0.622448, p-value: 0.013207
    '''