import tensorflow as tf
from tensorflow import keras
from tensorflow.keras import backend as K


def running_mean(x,n=100,K=K,tf=tf):
    """
    Calculate the running mean of an array
    """
    cumsum = tf.cumsum(x) 
    return (cumsum[n:] - cumsum[:-n]) / float(n)
def running_variance(x,n=100,K=K,tf=tf):
    """
    Calculate the running variance of an array
    """
    return running_mean(x**2,n=n,K=K,tf=tf)-running_mean(x,n=n,K=K,tf=tf)**2
def running_std(x,n=100,K=K,tf=tf):
    """
    Calculate the running standard deviation of an array
    """
    return K.sqrt(running_variance(x,n=n,K=K,tf=tf))

def loss2d(a,b,n=25,K=K,tf=tf):
    q=b
    x,y=q[:,0],q[:,1]
    #sort by x, evaluate stds on y
    dex=tf.argsort(x)
    yy=tf.gather(y,dex)

    ss=running_std(yy,n=n)
    s=K.std(yy)

    return K.mean(ss)/s

def numpyloss2d(a,b,n=25):
    import numpy as np
    q=np.concatenate((np.expand_dims(a,1),np.expand_dims(b,1)),axis=1)
    np.gather=np.take
    return loss2d(q,q,n=n,K=np,tf=np)



if __name__=='__main__':
    import numpy as np
    x=np.random.uniform(-1,1,size=(1000,2))
    print(numpyloss2d(x[:,0],x[:,1],n=25))