import numpy as np
from ss_setPlotParams import setPlotParams
from scipy.optimize import minimize
from orthofit import *

def get_excess_data():
    data = np.loadtxt('datasets/Srinivasanetal2009_Crich_8micron_excess.csv', skiprows = 1, \
                      dtype = {'names': ('L', 'sigma_L', 'X8', 'sigma_X8'), \
                               'formats': ('f4', 'f4', 'f4', 'f4')}, \
                      delimiter = ',')
    return data

def func(pars, x, y, sigx, sigy):
    #Note: x, y are the log10 versions of the luminosity and excess
    Delta = -x * np.sin(pars[0]) + y * np.cos(pars[0]) - pars[1]
    Sigma2 = sigx**2 * np.sin(pars[0])**2 + sigy**2 * np.cos(pars[0])**2
    logL = (np.log(Sigma2 + pars[2]) + Delta**2/(Sigma2 + pars[2])).sum()
    return 0.5*logL #Note positive sign, because we want to MINIMISE this

def linear_fit(x, y, sigy):
    #Invert the data matrices to get the best-fit parameter values and uncertainties
    ### Define A and Y matrices and the inverse of the covariance matrix
    A = np.array([np.ones(len(x)), x]).transpose()
    Y = np.array(y)[:, np.newaxis]
    Sigma_1 = np.linalg.inv(np.diag(sigy**2))
    ### Compute best-fit parameters and parameter covariance matrix
    Temp = A.transpose().dot(Sigma_1)
    Sigma_Pars = np.linalg.inv(Temp.dot(A))
    X = Sigma_Pars.dot(Temp.dot(Y))
    return A, Y, X, Sigma_Pars

def dobootstrap(x, y, sigy, b, m):
    B = 100 #Number of bootstrap resamples
    bb = np.zeros(B)
    mm = np.zeros(B)
    Ndata = len(x)
    index = np.arange(Ndata)
    for i in range(B):
        j = np.random.choice(index, Ndata, replace = True)
        x = x[j]
        y = y[j]
        sigy = sigy[j]
        A, Y, X, Sigma_Pars = linear_fit(x, y, sigy)
        bb[i] = X[0][0]
        mm[i] = X[1][0]
    sigma_b = ((bb - b)**2).mean()
    sigma_m = ((mm - m)**2).mean()
    print("-----------------")
    print("Best-fit intercept: {}".format(np.round(b, decimals = 3)))
    print("Standard deviation in intercept: {}".format(np.round(sigma_b, decimals = 3)))
    print("Best-fit slope: {}".format(np.round(m, decimals = 3)))
    print("Standard deviation in slope: {}".format(np.round(sigma_m, decimals = 4)))
    print("-----------------")
    return sigma_b, sigma_m

def mains():
    data = get_excess_data()
    x = np.log10(data['L'])
    sigx = data['sigma_L']/data['L']/np.log(10.)
    y = np.log10(data['X8'])
    sigy = data['sigma_X8']/data['X8']/np.log(10.)

    A, Y, X, Sigma_Pars = linear_fit(x, y, sigy)
    print("Intercept and slope from linear fitting with y uncertainties: {} and {} respectively."\
          .format(np.round(X[0][0], decimals = 2), np.round(X[1][0], decimals = 2)))
    Sigma_Pars_1 = np.sqrt(Sigma_Pars[0][0])
    Sigma_Pars_2 = np.sqrt(Sigma_Pars[1][1])
    rho_12 = Sigma_Pars[0][1]/Sigma_Pars_1/Sigma_Pars_2
    print("Uncertainties in the intercept and slope are {} and {} respectively." \
          .format(np.round(Sigma_Pars_1, decimals = 5), np.round(Sigma_Pars_2, decimals = 5)))
    print("The correlation coefficient between the uncertainties of the intercept and the slope is {}." \
          .format(np.round(rho_12, decimals = 4)))

    #Compute a reduced chi^2
    Sigma_1 = np.linalg.inv(np.diag(sigy**2))
    chi2red_standard = (Y - A.dot(X)).transpose().dot(Sigma_1.dot(Y - A.dot(X)))/(len(data) - len(X))
    print("The reduced chi-square for the standard fit is {}." \
          .format(np.round(chi2red_standard[0][0], decimals = 2)))

    #sigb, sigm = dobootstrap(x, y, sigy, X[0][0], X[1][0])


    pars_init = np.array([np.arctan(X[1]), X[0]*np.cos(np.arctan(X[1])), 0.25])
    b, m, V = orthofit(x, y, sigx, sigy, initial_guess = pars_init)
    b_hi = b + np.sqrt(V)/(1/np.sqrt(1 + m**2))
    b_lo = b - np.sqrt(V)/(1/np.sqrt(1 + m**2))

    print("The improved intercept and slope are {} and {} respectively.".format(np.round(b, decimals = 2), np.round(m, decimals = 2)))

    xx = np.linspace(np.min(x)*0.8, np.max(x)*1.2, 1000)
    yy1 = X[0] + X[1] * xx
    yy2 = b + m * xx
    yy2_hi = b_hi + m * xx
    yy2_lo = b_lo + m * xx
    plt = setPlotParams()
    plt.figure(figsize = (8, 8))
    plt.plot(x, y, 'ro', label = "Data")
    plt.plot(xx, yy1, label = "Model 1")
    plt.plot(xx, yy2, label = "Model 2")
    plt.plot(xx, yy2_hi, label = "Model 2 + Sqrt(V)", ls = 'dashed')
    plt.plot(xx, yy2_lo, label = "Model 2 - Sqrt(V)", ls = 'dashed')
    plt.title(r'$8 \mu\textrm{m excess vs luminosity for LMC C-rich AGB stars}$')
    plt.xlabel(r"$\log{[\textrm{Luminosity} (L_\odot)]}$")
    plt.ylabel(r"$\log{[8 \mu\textrm{m Excess Flux (Jy)}]}$")
    plt.legend(loc = 'best')
    plt.savefig('orthofit_excess.png', bbox_inches = 'tight')