import numpy as np
from scipy.optimize import minimize

class GeneralizedLeastSquares:
    """
    This class is used to compute the parameters (m,b) that produce
    the best-fitting line through a set of (x,y) points. Here, (m,b)
    refer to the slope and y-intercept of y = mx + b.
    (Eventually, this will be a generalized routine.)
    """

    def __init__(self, x=None, y=None, parameters=[], weights=None, method='minimize residuals', minimization='Nelder-Mead'):
        self.x = x
        self.y = y
        self.parameters = parameters
        self.weights = weights
        self.method = method
        self.minimization = minimization

    def __str__(self, parameters):
        print("nslope = %.2f   ny-intercept = %.2f    n" %(self.parameters[0], self.parameters[1]))

    @classmethod
    def get_fitting_func(self):
        """ This function returns y_new = mx + b """
        return [self.parameters[0] * self.x[idx] + self.parameters[1] for idx in range(len(self.x))]

    @classmethod
    def get_residuals(self):
        """ This function returns the residuals """
        y_trial = get_fitting_func(self)
        return [(self.y[idx] - y_trial[idx])**2 for idx in range(len(self.y)) if len(self.y) == len(y_trial)]

    @classmethod
    def update_weights(self):
        """ This function returns the weights to be used """
        if self.weights is None:
            res = np.ones(len(self.x))
        else:
            resid = get_residuals(self)
            inv_resid = [abs(1 / resid[idx]) for idx in range(len(resid))]
            tot_resid = sum(inv_resid)
            res = [inv_resid[idx] / tot_resid for idx in range(len(inv_resid))]
        return res

    @classmethod
    def get_error_func(self):
        """ This function returns the sum of the weighted inverse residuals """
        resid = get_residuals(self)
        return sum([weights[idx] * resid[idx] for idx in range(len(resid)) if len(resid) == len(weights)])

    @classmethod
    def perform_gls_fit(self):
        """ This function returns the fitting parameters of the linear fit """
        if self.parameters == 'estimate':
            slope_est = (self.y[-1] - self.y[0]) / (self.x[-1] - self.x[0])
            y_int_est = np.mean([self.y[-1] - slope_est * self.x[-1], self.y[0] - slop_est * self.x[0]])
            self.parameters = [slope_est, y_int_est]
        elif not isinstance(self.parameters, (list, np.ndarray)):
            raise ValueError("parameters = 'estimate' or [slope, y-intercept]")
        if self.method == 'minimize residuals':
            res = minimize(get_error_func, x0=self.parameters, args=(self,), method=self.minimization)
            self.parameters = res.x
            return self.parameters

# y = 2x+3
x_data = np.linspace(1, 10, 10)
y_data = np.array([5, 7.1, 9.2, 10.8, 12.9, 15.1, 17, 19.2, 20.8, 23])
gg = GeneralizedLeastSquares(x_data, y_data, 'estimate').perform_gls_fit
print(gg)

## produces this output
<bound method GeneralizedLeastSquares.perform_gls_fit of <class '__main__.GeneralizedLeastSquares'>>