from scipy import optimize
import numpy as np
import copy
import collections

def maxlike_fitting(post, verbose=True):
    """Maximum Likelihood Fitting

    Perform a maximum likelihood fit.

    Args:
        post (radvel.Posterior): Posterior object with initial guesses
        verbose (bool): (optional) Print messages and fitted values?

    Returns: radvel.Posterior : Posterior object with parameters
        updated their maximum likelihood values

    """
    
    post0 = copy.deepcopy(post)
    if verbose:
        print "Initial loglikelihood = %f" % post0.logprob()
        print "Performing maximum likelihood fit..."

    res  = optimize.minimize(
        post.neglogprob_array, post.get_vary_params(), method='Powell',
        options=dict(maxiter=100,maxfev=100000,xtol=1e-8)
    )

    cpspost = copy.deepcopy(post)
    cpsparams = post.params.basis.to_cps(post.params)
    cpspost.params.update(cpsparams)

    if verbose:
        print "Final loglikelihood = %f" % post.logprob()
        print "Best-fit parameters:"
        print cpspost
        
    return post
    

def model_comp(post, verbose=False):
    """Model Comparison

    Fit for planets adding one at a time.  Save results as list of
    posterior objects.

    Args:
        post (radvel.Posterior): posterior object for final best-fit solution 
            with all planets
        verbose (bool): (optional) print out statistics
        
    Returns:

        list of dictionaries : List of dictionaries with fit
            statistics. Each value in the dictionary is a tuple with
            the statistic value as the first element and a description
            of that statistic in the second element.
    """

    ipost = copy.deepcopy(post)
        
    num_planets = post.likelihood.model.num_planets

    statsdict = []
    
    for n in range(num_planets+1):
        pdict = collections.OrderedDict()
        if verbose:
            print "Testing %d planet model.\n\n" % n

        for par in post.params.keys():
            try:
                num = int(par[-1])
                if num > n:
                    if par.startswith('k') or par.startswith('logk'):
                        post.params[par] = 0.0
                    post.vary[par] = False
                else:
                    post.params[par] = ipost.params[par]
                    post.vary[par] = ipost.vary[par]
            except (ValueError, KeyError):
                pass

        post = maxlike_fitting(post, verbose=False)

        ndata = len(post.likelihood.y)
        nfree = len(post.get_vary_params())
        chi = np.sum((post.likelihood.residuals()/post.likelihood.yerr)**2)
        chi_red = chi / (ndata - nfree)

        if verbose:
            print post
            print "N_free = %d" % nfree
            print "RMS = %4.2f" % np.std(post.likelihood.residuals())
            print "logprob = %4.2f" % post.logprob()
            print "chi (no jitter) = %4.2f" % chi
            print "chi_red (no jitter) = %4.2f" % chi_red
            print "BIC = %4.2f" % post.bic()
        
        pdict['$N_{\\rm data}$'] = (ndata, 'number of measurements')
        pdict['$N_{\\rm free}$'] = (nfree, 'number of free parameters')
        pdict['RMS'] = (
            np.round(np.std(post.likelihood.residuals()), 2), 
            'RMS of residuals in m s$^{-1}$'
        )
        pdict['$\\chi^{2}$'] = (np.round(chi,2), "assuming no jitter")
        pdict['$\\chi^{2}_{\\nu}$'] = (
            np.round(chi_red,2), "assuming no jitter"
        )
        pdict['$\\ln{\\mathcal{L}}$'] = (
            np.round(post.logprob(),2), "natural log of the likelihood"
        )
        pdict['BIC'] = (
            np.round(post.bic(),2), 
            'Bayesian information criterion'
        )
        
        statsdict.append(pdict)

    return statsdict
