import numpy as np

def LinearLeastSquaresFit(x,y):
    """Take in arrays representing (x,y) values for a set of linearly varying
    data and perform a linear least squares regression. Return the resulting
    slope and intercept parameters of the best fit line with their
    uncertainties (m,b,m_err,b_err).
    """
    n = len(x)
    x_avg = np.sum(x)/n
    y_avg = np.sum(y)/n
    xsq_avg = np.sum(x*x)/n
    xy_avg = np.sum(x*y)/n
    m = (xy_avg - (x_avg*y_avg))/(xsq_avg - (x_avg**2))
    b = ((xsq_avg*y_avg)-(x_avg*xy_avg))/(xsq_avg - (x_avg**2))
    
    u = np.array([])
    for i in range(n):
        w=y[i]-((m*x[i])+b)
        u = np.append(u,w)
    
    usq_avg = np.sum((u*u))/len(u)

    m_err = np.sqrt((1/n-2.)*(usq_avg/(xsq_avg -(x_avg**2))))
    b_err = np.sqrt((1/n-2.)*((usq_avg*xsq_avg)/(xsq_avg - (x_avg**2))))

    return m,b,m_err,b_err

def WeightedLinearLeastSquaresFit(x,y,w):
    """Take in arrays representing (x,y) values for a set of linearly varying
    data and an array of weights w. Perform a weighted linear least squares
    regression. Return the resulting slope and intercept parameters of the
    best fit line with their uncertainties. (m,m_err,b,b_err)
    """
    f = len(x)
    w_avg = np.sum(w)/f
    wx_avg = np.sum(w*x)/f
    wy_avg = np.sum(w*y)/f
    wxsq_avg = np.sum(w*(x*x))/f
    wxy_avg = np.sum(w*(x*y))/f
    
    m = ((w_avg*wxy_avg)-(wx_avg*wy_avg))/((w_avg*wxsq_avg)-(wx_avg**2))
    b = ((wxsq_avg*wy_avg)-(wx_avg*wxy_avg))/((w_avg*wxsq_avg)-(wx_avg**2))
    
    if np.sum(w)/len(w) == w[0] and w[1]:
        m_err = LinearLeastSquaresFit(x,y)[2]
        b_err = LinearLeastSquaresFit(x,y)[3]
    else:
        m_err = np.sqrt(w_avg/((w_avg*wxsq_avg)-(wx_avg**2)))
        b_err = np.sqrt(wxsq_avg/((w_avg*wxsq_avg)-(wx_avg**2)))
    
    return m,m_err,b,b_err
