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_p和合作频率之间的关系
'''

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

def p1(x, coopr, tau, postfix, show=True):
    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()
    if show:
        plt.show()
    else:
        plt.savefig("graph/tau_p_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=(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)
    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('$\\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()
    if show:
        plt.show()
    else:
        plt.savefig("graph/tau_p_co_sca_%s.eps" % postfix)

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

if __name__ == '__main__':
  
    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)
    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-1):
        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-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折线图
    # p1(x, coopr, tau, 'survive', False)

    p2(tau[:12], coopr[:12], 720, 'survive', False)


    '''
    classic
    残差： [ 0.00317365]
    Model parameter: [ -1.24291981e-04   9.70766132e-01]
    Other parameters: rank=2, sv=[ 1.41276801  0.06392615], rcond=2.6645352591e-15
    error= 0.003174
    pearson: -0.607866, p-value: 0.036010

    survive
    残差： [ 0.00548837]
    Model parameter: [ -2.71223484e-04   1.00851422e+00]
    Other parameters: rank=2, sv=[ 1.36321571  0.37635479], rcond=1.99840144433e-15
    error= 0.005488
    pearson: -0.829102, p-value: 0.005723
    '''