from __future__ import division, print_function, absolute_import

from warnings import warn

import celerite
from . import numbers as nb
import numpy as np
from . import ranges
from scipy.optimize import minimize


class QuiescenceModel(celerite.GP):
    """
    An extension of the `celerite` Gassian Process model object for use in modeling the quiescence variations of a
    stellar lightcurve as correlated noise.
    """

    def __init__(self, t, f, e, tau_min=100., tau_logprior=None, params=None, mask=None, white_limit=2.0):
        """
        Create a QuiscenceModel object that models lightcurve fluctuations as correlated noise with a Gaussian Process
        (GP) where correlations decay as exp(-t/tau) where tau is a constant. Specifically, elements of the covariance
        matrix are parameterized as a*exp(-c*t), where ln(a), ln(c), and the mean data value, mu, are parameters of
        the model.

        An important deviation from a typical GP is that the QuiescenceModel enforces a penalty for "jagged" models
        to promote  smoothness. Specifically, the penalty is based on the power of the model at frequencies of 0.1 Hz.

        Parameters
        ----------
        t, f, e : arrays
            Time, flux, and error of data points.
        tau_min : float
            Minimum correlation time to consider in GP fit. Default
            is 100 s. For no limit, simply set to 0.
        tau_logprior : function, optional
            Prior on the correlation time. Should accept a single
            value for tau as input an return the natural
            logarithm of the likelihood of that value.
        params : list
            Parameters of the GP covariance model, a * exp(-c*t). 
            These are log(a), log(c), and the (constant) mean, mu. 
        mask : bool array, optional
            mask of data points where a value of True indicates a
            point that should be excluded from fitting (e.g., because
            it is in a flare or suspect region). If None, assume
            all data should be included.
        white_limit : float
            Likelihood ratio below which a white noise model should
            be used instead of a correlated noise model. Default is
            2.0 to give preference for the simpler model even if a
            GP gives a slightly better fit.
        """
        # set up celerite GP model as a
        terms = celerite.terms
        kernel = terms.RealTerm(log_a=np.log(np.var(f)), log_c=-7.)
        super(QuiescenceModel, self).__init__(kernel, mean=np.median(f), fit_mean=True)

        # initialize parameters of GP if desired
        if params is not None:
            self.set_parameter_vector(params)

        # set data mask
        if mask is not None:
            self.mask = mask
        else:
            self.mask = np.ones(len(t), bool)

        # store other attributes
        self.tau_min = tau_min
        self.tau_logprior = tau_logprior
        self.t, self.f, self.e = t, f, e
        self.n = len(self.t)
        self.fit_params = None
        self.white_limit = white_limit
        self.white = (self.get_parameter('kernel:log_c') == np.inf)
        self.quick_compute()

    def tau_loglike(self, tau):
        """
        Log likeliood of tau values based on lower limit and any user-defined prio.
        """
        if tau < self.tau_min:
            return -np.inf
        if self.tau_logprior is None:
            return 0.0
        else:
            return self.tau_logprior(tau)

    def smoothness_penalty(self, params):
        """
        Cost function for power at 0.1 Hz frequencies. Exact form was set by trial and error.
        
        Parameters
        ----------
        params : list
            Parameters of the GP covariance model, a * exp(-c*t). 
            These are log(a), log(c), and the (constant) mean, mu. 
        """
        self.set_parameter_vector(params)
        hi_freq_power = self.kernel.get_psd(2*np.pi/10)
        if hi_freq_power == 0:
            return 0.0
        if hi_freq_power == np.inf or np.isnan(hi_freq_power):
            return -np.inf
        else:
            return -5*np.sum(self.mask)*np.sqrt(hi_freq_power)/np.std(self.f[self.mask])

    def quick_compute(self):
        """
        Quickly recompute GP intermediate values (i.e. following a change in the parameters).

        """
        super(QuiescenceModel, self).compute(self.t[self.mask], self.e[self.mask])

    def _get_set_mask(self, mask=None):
        if mask is None:
            return self.mask
        else:
            self.mask = mask
            return mask

    def log_likelihood(self, params):
        """
        Log likelihood of the GP model based on the data likelihood and prior on tau. *Smoothness penalty is not applied
        here* so this can be used for sampling the posterior of parameters without that associated bias.

        Parameters
        ----------
        params : list
            Parameters of the GP covariance model, a * exp(-c*t). 
            These are log(a), log(c), and the (constant) mean, mu. 
        
        Returns
        -------
        Natural log of the likelihood of the data given the model and priors on parameters.
        """
        self.set_parameter_vector(params)
        if self.white:
            return self.log_likelihood_white_noise(params[[0,2]])
        else:
            data_loglike = super(QuiescenceModel, self).log_likelihood(self.f[self.mask])
            return data_loglike + self.tau_loglike(self.tau())

    def cost(self, params):
        """
        Log likelihood of the GP model including data likelihood, prior on tau, and smoothness penalty. This is
        what is actually used for the fitting.

        Parameters
        ----------
        params : list
            Parameters of the GP covariance model, a * exp(-c*t). 
            These are log(a), log(c), and the (constant) mean, mu.
            
        Returns
        -------
        Natural log of the likelihood of the data given the model, priors on parameters, and smoothness penalty.
        """
        return -(self.log_likelihood(params) + self.smoothness_penalty(params))
        # return -self.log_likelihood(params)

    def log_likelihood_white_noise(self, log_sig2_and_mu):
        """
        Log likelihood of a constant mean + white noise model.
        
        
        Parameters
        ----------
        log_sig2_and_mu : list
            Natural log of the "extra" variance and the constant mean.

        Returns
        -------
        Natural log of the likelihood of the data given the model.
        """
        self.set_parameter_vector([log_sig2_and_mu[0], np.inf, log_sig2_and_mu[1]])
        return super(QuiescenceModel, self).log_likelihood(self.f[self.mask])

    def log_likelihood_no_noise(self, mu):
        """
        Log likelihood of a model  with no noise and constant mean.
        
        Parameters
        ----------
        mu : float
            value of constant mean

        Returns
        -------
        Natural log of the likelihood of the data given the model.
        """
        self.set_parameter_vector([-np.inf, np.inf, mu])
        return super(QuiescenceModel, self).log_likelihood(self.f[self.mask])

    def fit(self, mask=None, method='Nelder-Mead'):
        """
        Perform a max-likelihood fit of the QuiescenceModel to the data.
        
        Parameters
        ----------
        mask : bool array
            Mask identifying data to be used in the fit.
        method : str
            numerical method to use in minimizing the likelihood function 
            (of those allowable by scipy.optimize.minimize)

        Returns
        -------
        None. Fit is performed in-place and parameters of the QuiescenceModel object set to the best-fit values.
        """
        self._get_set_mask(mask)
        self.quick_compute()
        guess = self.get_parameter_vector()
        if self.white:
            guess[1] = -7
        soln = minimize(self.cost, guess, method=method)
        soln_white = minimize(lambda params: -self.log_likelihood_white_noise(params), guess[[0,2]], method=method)
        if not (soln.success and soln_white.success):
            raise ValueError('Gaussian process fit to quiescence did not converge. Perhaps try a different minimize '
                             'method or different initial parameters.')
        if soln_white.fun + self.log_likelihood(soln.x) < np.log(self.white_limit):
            self.fit_params = np.array([soln_white.x[0], np.inf, soln_white.x[1]])
            self.white = True
        else:
            self.fit_params = soln.x
            self.white = False
        self.set_to_best_fit()

    def fit_with_frozen_tau(self, mask=None, method='Nelder-Mead'):
        """
        Perform a max-likelihood fit of the QuiescenceModel to the data holding the time constant, tau, constant.

        Parameters
        ----------
        mask : bool array
            Mask identifying data to be used in the fit.
        method : str
            numerical method to use in minimizing the likelihood function
            (of those allowable by scipy.optimize.minimize)

        Returns
        -------
        None. Fit is performed in-place and parameters of the QuiescenceModel object set to the best-fit values.
        """
        self._get_set_mask(mask)
        self.quick_compute()
        logc = self.get_parameter('kernel:log_c')
        if self.white:
            def cost(params):
                return -self.log_likelihood_white_noise([params[0], params[1]])
        else:
            def cost(params):
                return self.cost([params[0], logc, params[1]])
        soln = minimize(cost, [np.log(np.var(self.f)), np.median(self.f)], method=method)
        assert soln.success
        self.fit_params = np.array([soln.x[0], logc, soln.x[1]])
        self.set_to_best_fit()

    def _get_params_boiler(self, params):
        """Boilerplate code for getting the current model parameters as an array."""
        if params is None:
            params = self.get_parameter_vector()
        return np.reshape(params, [-1,3])

    def tau(self, params=None):
        """
        Value of the decay constant tau in the covariance model sigma**2 * exp(-t/tau). If the model is set
        to white noise, this  will be 0.

        Parameters
        ----------
        params : list
            Parameters of the GP covariance model, a * exp(-c*t).
            These are log(a), log(c), and the (constant) mean, mu.
            If None, the parameters currently set in the QuiescenceModel
            are used.
        """
        params = self._get_params_boiler(params)
        return np.squeeze(np.exp(-params[:,1]))

    def sigma(self, params=None):
        """
        Value of the standard deviation in the covariance model sigma**2 * exp(-t/tau). If the model is set
        to white noise, then technically this is sigma in sigma**2 * delta(t) where delta is the Dirac delta function.

        Parameters
        ----------
        params : list
            Parameters of the GP covariance model, a * exp(-c*t).
            These are log(a), log(c), and the (constant) mean, mu.
            If None, the parameters currently set in the QuiescenceModel
            are used.
        """
        params = self._get_params_boiler(params)
        return np.squeeze(np.exp(params[:,0]/2))

    def mu(self, params=None):
        """
        Value of model mean value.

        Parameters
        ----------
        params : list
            Parameters of the GP covariance model, a * exp(-c*t).
            These are log(a), log(c), and the (constant) mean, mu.
            If None, the parameters currently set in the QuiescenceModel
            are used.
        """
        params = self._get_params_boiler(params)
        return np.squeeze(params[:,2])

    def sigma_rel(self, params=None):
        """
        Value of the standard deviation in the covariance model sigma**2 * exp(-t/tau), normalized by the
        mean.

        Parameters
        ----------
        params : list
            Parameters of the GP covariance model, a * exp(-c*t).
            These are log(a), log(c), and the (constant) mean, mu.
            If None, the parameters currently set in the QuiescenceModel
            are used.
        """
        return np.squeeze(self.sigma(params)/self.mu(params))

    def sigma_relative_at_tbin(self, tbin, params=None):
        """
        Current value of the expected standard deviation of the model if values were binned to tbin.

        Parameters
        ----------
        tbin : float
            Width of time bins over which model values are assumed to be averaged.
        params : list
            Parameters of the GP covariance model, a * exp(-c*t).
            These are log(a), log(c), and the (constant) mean, mu.
            If None, the parameters currently set in the QuiescenceModel
            are used.
        """
        params = self._get_params_boiler(params)
        loga, logc, mu = params.T
        sig, c = np.exp(loga/2), np.exp(logc)
        with np.errstate(invalid='ignore', over='ignore'):
            x = c*tbin
            sig_dt = np.sqrt(2*(sig/x)**2 * (x + np.exp(-x) - 1))
            uncorrelated = (x == np.inf)
            if np.any(uncorrelated):
                dt = np.median(np.diff(self.t))
                sig_dt[uncorrelated] = sig[uncorrelated]/np.sqrt(tbin/dt)
        return np.squeeze(sig_dt/mu)

    def set_to_best_fit(self):
        """Return the model parameters to their best-fit values, if they have been computed.
        Else you will get an error."""
        self.set_parameter_vector(self.fit_params)
        self.quick_compute()

    def curve(self, t):
        """Compute the lightcurve of the quiescence  model at t."""
        if self.white:
            return self.mu()*np.ones_like(t), self.sigma()**2*np.ones_like(t)
        else:
            return self.predict(self.f[self.mask], t, return_var=True)


def lightcurve_fill(t, f, e, qmodel, fill_ranges):
    """
    Replace flare times with simulated data based on the qmodel with appropriately correlated noise.

    Parameters
    ----------
    t, f, e : arrays
        Lightcurve points -- time, flux, energy.
    qmodel : QuiescenceModel
        Gassian Process model for quiescent variations in lightcurve.
    fill_ranges : Nx2 array
        Start and end time of each flare.

    Returns
    -------
    f_filled, e_filled : arrays
        Flux and error arrays where regions within flares  have been filled with simulated data.
    """

    f_filled, e_filled = list(map(np.copy, [f, e]))
    if len(fill_ranges) > 0:
        # pull random draws to fill where flares were until no false positives occur
        flare = ranges.inranges(t, fill_ranges)
        if not np.any(flare):
            warn("Flare ranges were supplied, yet no points were within these ranges.")
            return f, e

        isort = np.argsort(f) # comes in handy in a sec

        # estimate white noise error in flare range
        mean_flare, _ = qmodel.curve(t[flare])
        e_sim = np.interp(mean_flare, f[isort], e[isort])

        # draw random fluxes and estimate what the uncertainty estimate would have been for those points
        f_fill = nb.conditional_qmodel_draw(qmodel, t[~flare], f[~flare], t[flare])
        f_fill += np.random.randn(np.sum(flare))*e_sim
        e_fill = np.interp(f_fill, f[isort], e[isort])
        e_filled[flare] = e_fill
        f_filled[flare] = f_fill
        return f_filled, e_filled
    else:
        return f, e