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):
    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[1:], e[1:], color=red, linewidth=2, label=r"$E_{i,D}$")
    # 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"$E_{i,D}$", family='sans-serif', color=red, size=22)
    ax2.tick_params(axis='y', labelcolor=red)

    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=(4, 3))
    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'$E_{i,D}$', 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(int(min(e)*0.9), int(max(e)*1.1), 4))
    ax.tick_params(labelsize=14)
    ax.set_ylim(0.6, 1)
    ax.set_yticks(sp.linspace(0.6, 1, 5))
    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, isp1):
    coopr = np.loadtxt("outputs/CR_%s.csv" % mode, delimiter=',')
    x = np.arange(1, 16)
    e = np.loadtxt("outputs/EID_%s.csv" % mode, delimiter=',')

    if isp1:
        p1(x, coopr, e, mode, show)
    else:
        if mode == 'SURVIVE':
            p2(e[1:12], coopr[1:12], mode, show)
        else:
            p2c(e[1:], coopr[1:], mode, show)

if __name__ == '__main__':
    # eid_plot c/s 1/2 t/f
    mode = 'CLASSIC' if argv[1] == 'c' else 'SURVIVE'
    isp1 = argv[2] == '1'
    show = argv[3] == 't'
    plot(mode, show, isp1)

    """
SURVIVE
残差： [ 0.10769174]
Model parameter: [  1.33390242e-04   7.29761895e-01]
Other parameters: rank=2, sv=[ 1.37254846  0.34075025], rcond=3.33066907388e-15
error= 0.107692
pearson: 0.709599, p-value: 0.003045


CLASSIC
残差： [ 0.00992721]
Model parameter: [  3.14661037e-05   7.35284315e-01]
Other parameters: rank=2, sv=[ 1.41231835  0.07319068], rcond=3.33066907388e-15
error= 0.009927
pearson: 0.309326, p-value: 0.261915
    """