import numpy as np

def weighted_linear_fit(w,x,y):
    A = ((w*x*y).sum() - (w*x).sum()*(w*y).sum()/w.sum())/((w*x*x).sum() - (w*x).sum()**2/w.sum())
    B = ((w*x*y).sum()-(w*x*x).sum()*A)/(w*x).sum()
    return A,B

def find_peaks(x):
    index = []
    for i in range(1,x.size-1):
        pre = x[i] - x[i-1]
        post = x[i] - x[i+1]
        if pre>0 and post>0:
            index.append(i)
    return index

def find_valleys(x):
    index = []
    for i in range(1,x.size-1):
        pre = x[i] - x[i-1]
        post = x[i] - x[i+1]
        if pre<0 and post<0:
            index.append(i)
    return index 

def weighted_cov(x,y,w):
    return np.average((x - np.average(x, weights=w)) * (y - np.average(y, weights=w)), weights=w)

def weighted_cor(x,y,w):
    return weighted_cov(x, y, w) / np.sqrt(weighted_cov(x, x, w) * weighted_cov(y, y, w))


def my_round(x):
    if type(x) is np.ndarray:
        return (2*x).round()/2
    else:
        return round(x*2)/2

def my_round_n(x,n):
    if type(x) is np.ndarray:
        return (x*int(1/n)).round()/int(1/n)
    else:
        return round(x*int(1/n))/int(1/n)

#print(my_round_n(np.linspace(0,10,333),1/np.pi))

def complex_arg(x):
    if type(x) is np.ndarray:
        phi = np.zeros_like(x)
        phi = np.arccos(x.real/np.sqrt(x.real**2+x.imag**2))
        xxx = x.imag < 0 
        phi[xxx] =  - phi[xxx]
        
    else:
        phi = np.arccos(x.real/np.sqrt(x.real**2+x.imag**2))
        if x.real < 0:
            phi = - phi
    return phi
