import numpy as np

from choosenext import choosenext

from plt import plt


def draw(p,dat):
    mp=np.mean(p)
    d=np.square(p-mp)
    sx=[(xx,dd) for xx,dd in zip(dat,d)]
    sx.sort(key=lambda x:x[1])
    sx=[zw[0] for zw in sx]
    sx=np.array(sx)

    col1=[1.0,0.0,0.0]
    col2=[0.0,1.0,0.0]
    col1,col2=np.array(col1),np.array(col2)
    ln=len(sx)
    cols=[col1*(i/ln)+col2*(1-i/ln) for i in range(ln)]

    plt.scatter(sx[:,0],sx[:,1],c=cols)


def combine(a,b):
    a=1+(a-np.mean(a))/np.std(a)
    b=1+(b-np.mean(b))/np.std(b)
    toc=np.concatenate((np.expand_dims(a,axis=1),np.expand_dims(b,axis=1)),axis=1)
    toc=np.sqrt(np.mean(toc**2,axis=1))
    return toc

if __name__=="__main__":
    f=np.load("merged.npz")
    dat=f["x"]
    x=f["ps"]
    np.random.shuffle(x)
    given=x[0]
    possble=x[1:5]

    ac=choosenext(given,possble)
    nextbest=ac[:,0]
    remainder=ac[:,1:]

    for row in np.transpose(ac):
        print(np.corrcoef(given,row)[0,1])
    #as you see: the correlation is the lowest for the first row.
    #so lets combine it
    updated=combine(given,nextbest)


    draw(given,dat)
    plt.savefig("before.png")
    plt.show()
    draw(nextbest,dat)
    plt.savefig("suggestion.png")
    plt.show()
    draw(updated,dat)
    plt.savefig("updated.png")
    plt.show()