# Licensed under an MIT open source license - see LICENSE
from __future__ import print_function, absolute_import, division

import numpy as np
import astropy.units as u
from astropy.utils import NumpyRNGContext

'''
Routines for fitting a line with errors in both variables.
'''


def leastsq_linear(x, y, x_err, y_err, verbose=False):
    '''
    Fit a line with errors in both variables using least squares fitting.
    Specifically, this uses orthogonal distance regression (in scipy) since
    there is no clear "independent" variable here and the covariance is
    unknown (or at least difficult to estimate).

    Parameters
    ----------
    x : `~numpy.ndarray`
        x data.
    y : `~numpy.ndarray`
        y data.
    x_err : `~numpy.ndarray`
        x errors.
    y_err : `~numpy.ndarray`
        y errors.
    verbose : bool, optional
        Plot the resulting fit.

    Returns
    -------
    params : `~numpy.ndarray`
        Fit parameters (slope, intercept)
    errors : `~numpy.ndarray`
        1-sigma errors from the covariance matrix (slope, intercept).
    '''

    import scipy.odr as odr
    from scipy.stats import linregress

    def fit_func(B, x):
        return B[0] * x + B[1]

    linear = odr.Model(fit_func)

    mydata = odr.Data(x, y, wd=np.power(x_err, -2), we=np.power(y_err, -2))

    # Set the initial guess for ODR from a normal linear regression
    beta0 = linregress(x, y)[:2]

    # beta sets the initial parameters
    myodr = odr.ODR(mydata, linear, beta0=beta0)

    output = myodr.run()

    params = output.beta
    errors = output.sd_beta

    # found a source saying this equivalent to reduced chi-square. Not sure if
    # this is true... Bootstrapping is likely a better way to go.
    # gof = output.res_var

    if verbose:
        output.pprint()

        import matplotlib.pyplot as plt