import numpy as np
from matplotlib import pyplot as plt
import scipy as sp
from scipy.stats import pearsonr
from matplotlib import markers
from sys import argv
blue = '#0984e3'
red = '#d63031'

def error(f,x,y):
    return sp.sum((f(x)-y)**2)
def p1(x, coopr, e, postfix, show=True, max_round=15):
    fig = plt.figure(figsize=(3, 2))
    ax = fig.gca()
    ax.plot(x, coopr, color=blue, linewidth=2, label=r"f$_{\rm c}$")
    ax.set_ylim(0, 1)
    ax.set_yticks(sp.linspace(0, 1, 5))
    ax2 = ax.twinx()
    ax2.plot(x[:-1], e[:-1], color=red, linewidth=2, label=r"${\rm \theta}$")
    ax2.set_ylim(-1500, 200)
    ax2.set_yticks(sp.linspace(-1500, 200, 5))
    ax2.tick_params(labelsize=10)
    ax.tick_params(labelsize=10)
    ax.set_xlim(1, max_round)
    ax.set_xlabel("Rounds", size=11)
    ax.set_ylabel(r"f$_{\rm c}$", color=blue, size=11)
    ax.tick_params(axis='y', labelcolor=blue)
    ax2.set_ylabel(r"${\rm \theta}$", color=red, size=11)
    ax2.tick_params(axis='y', labelcolor=red)
    ax.tick_params(direction='in')
    ax2.tick_params(direction='in')

    plt.tight_layout()
    if show:
        plt.show()
    else:
        plt.savefig("graph/eid_co_plot_%s.eps" % postfix)

def p2(e, coopr, postfix, show=True, showline=True):
    # p2散点图
    fig = plt.figure(figsize=(3, 2))
    ax = fig.gca()
    if showline:
        fp1,residuals,rank,sv,rcond = sp.polyfit(e, coopr, 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, e, coopr))
        fx = sp.linspace(np.min(e), np.max(e), 2)
        plt.plot(fx,f1(fx),linewidth=2,color=red, ls='--', zorder=0)

    plt.scatter(e, coopr, color='white', edgecolors=blue, linewidths=2, zorder=101)
    ax.set_xlabel(r"${\rm \theta}$", size=12)
    ax.set_ylabel(r"f$_{\rm c}$", size=12)
    # ax.set_xlim(0, 1440)
    # ax.set_xticks(sp.linspace(int(min(e)*0.9), int(max(e)*1.1), 4))
    ax.tick_params(labelsize=10)
    ax.set_ylim(0.6, 1)
    ax.set_yticks(sp.linspace(0.6, 1, 5))
    ax.tick_params(direction='in')
    plt.tight_layout()
    if show:
        plt.show()
    else:
        plt.savefig("graph/eid_co_sca_%s.eps" % postfix)

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

def p2c(e,c,m,s):
    p2(e,c,m,s,False)

def plot(mode, show, network_type, isp1):
    coopr = np.loadtxt("outputs/CR_%s_%s.csv" % (mode, network_type), delimiter=',')
    e = np.loadtxt("outputs/EID_%s_%s.csv" % (mode, network_type), delimiter=',')

    if isp1:
        if mode == 'NEW':
            x = np.arange(1, 29)
        else:
            x = np.arange(1, 16)
        p1(x, coopr[:28], e, mode, show, max_round=28)
    else:
        if mode == 'SURVIVE':
            p2(e[:-1], coopr[:-1], mode, show)
        elif mode == 'CLASSIC':
            p2c(e[:-1], coopr[:-1], mode, show)
        else:
            p2(e[:-1], coopr[:27], f'{mode}_{network_type}', show)

if __name__ == '__main__':
    # eid_plot c/s/nw/nb 1/2 t/f
    network_type = None
    if argv[1] == 'c':
        mode = 'CLASSIC'
    elif argv[1] == 's':
        mode = 'SURVIVE'
    elif argv[1] == 'nb':
        mode = 'NEW'
        network_type = 'BA'
    else:
        mode = 'NEW'
        network_type = 'WS'
    isp1 = argv[2] == '1'
    show = argv[3] == 't'
    plot(mode, show, network_type, isp1)

    """
SURVIVE
残差： [ 0.06081364]
Model parameter: [ -3.54739130e-04   4.41956696e-01]
Other parameters: rank=2, sv=[ 1.40523958  0.15906514], rcond=3.10862446895e-15
error= 0.060814
pearson: -0.837958, p-value: 0.000183


CLASSIC
pearson: 0.384519, p-value: 0.174626
    """