# background.py: Module containing background fitting and model classes
#
#     Authors: William Cleveland (USRA),
#              Adam Goldstein (USRA) and
#              Daniel Kocevski (NASA)
#
#     Portions of the code are Copyright 2020 William Cleveland and
#     Adam Goldstein, Universities Space Research Association
#     All rights reserved.
#
#     Written for the Fermi Gamma-ray Burst Monitor (Fermi-GBM)
#
#     This program is free software: you can redistribute it and/or modify
#     it under the terms of the GNU General Public License as published by
#     the Free Software Foundation, either version 3 of the License, or
#     (at your option) any later version.
#
#     This program is distributed in the hope that it will be useful,
#     but WITHOUT ANY WARRANTY; without even the implied warranty of
#     MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
#     GNU General Public License for more details.
#
#     You should have received a copy of the GNU General Public License
#     along with this program.  If not, see <https://www.gnu.org/licenses/>.
#
import numpy as np
from ..data.primitives import EventList, TimeEnergyBins, EnergyBins
from ..data import PHAII, TTE


class BackgroundFitter:
    """Class for fitting a background, given a fitting algorithm,
    to time-energy data (e.g. :class:`~gbm.data.phaii.PHAII`, 
    :class:`~gbm.data.TTE`).
    
    When a BackgroundFitter is created, an algorithm must be specified.  In 
    particular, the algorithm must be a class that has two public methods:
    ``fit()`` and ``interpolate()``. 
    
    For PHAII data, the class, upon initialization, must take the following as 
    arguments: 
    
    1. a 2D array of counts with shape (``numtimes``, ``numchans``);
    2. a 1D array of time bin start times, shape (``numtimes``,); 
    3. a 1D array of time bin end times, shape (``numtimes``,);
    4. a 1D array of exposures for each time bin, shape (``numtimes``,).

    While for TTE data, the class, upon initialization, must take the following 
    as an argument: 

    1. a list, of length ``numchans``, where each item is a np.array
       of event times.
    
    The ``fit()`` method takes no arguments, hoewever, parameters that are 
    required for the algorithm may be specified as keywords.
    
    The interpolate() method must take in the following as arguments:
    
    1. a 1D array of time bin start times, shape (``numtimes``,);
    2. a 1D array of time bin end times, shape (``numtimes``,).
    
    Any additional parameters required can be specified as keywords. 
    The ``interpolate()`` method must return:
    
    1. a 2D rates array, shape (``numtimes``, ``numchans``);
    2. a 2D rate uncertainty array, shape (``numtimes``, ``numchans``).
    
    Additionally, the class can provide the following public attributes that 
    will be exposed by BackgroundFitter:
   
    - ``dof``: The degrees of freedom of the fits, array shape (``numchans``,)
    - ``statistic``: The fit statistic for each fit, array shape (``numchans``,)
    - ``statistic_name``: A string of the fit statistic used    
    
    
    Attributes:
        dof (np.array): If available, the degrees-of-freedom of the fit for 
                        each energy channel
        livetime (float): The total livetime of the data used for the background
        method (str): The name of the fitting algorithm class
        parameters (dict): All parameters passed to the fitting algorithm
        statistic (np.array): If available, the fit statistic for each energy 
                              channel
        statistic_name (str): If available, the name of the fit statistic
        type (str): The type of background algorithm, either 'binned' or 'unbinned'
    """

    def __init__(self):
        self._data_obj = None
        self._method = None
        self._type = None
        self._statistic = None
        self._dof = None
        self._livetime = None
        self._parameters = None

    @property
    def method(self):
        return self._method.__class__.__name__

    @property
    def type(self):
        return self._type

    @property
    def statistic_name(self):
        return getattr(self._method, 'statistic_name', None)

    @property
    def statistic(self):
        return getattr(self._method, 'statistic', None)

    @property
    def dof(self):
        return getattr(self._method, 'dof', None)

    @property
    def livetime(self):
        return self._livetime

    @property
    def parameters(self):
        return self._parameters

    @classmethod
    def from_phaii(cls, phaii, method, time_ranges=None):
        """Create a background fitter from a PHAII object
        
        Args:
            phaii (:class:`~gbm.data.phaii.PHAII`): A PHAII data object
            method (<class>): A background fitting/estimation class for binned data 
            time_ranges ([(float, float), ...]): 
                The time range or time ranges over which to fit the background. 
                If omitted, uses the full time range of the data 
        
        Returns:        
            :class:`BackgroundFitter`: The initialized background fitting object
        """
        if not isinstance(phaii, PHAII):
            raise TypeError('Input data must be a PHAII object')

        obj = cls()
        obj._data_obj = phaii
        obj._validate_method(method)
        time_ranges = obj._validate_time_ranges(time_ranges)

        # Slice the PHAII data and merge if multiple slices
        data = [phaii.data.slice_time(trange[0], trange[1]) for trange in
                time_ranges]
        data = TimeEnergyBins.merge_time(data)
        obj._method = method(data.counts, data.tstart, data.tstop,
                             data.exposure)
        obj._type = 'binned'
        obj._livetime = np.sum(data.exposure)
        return obj

    @classmethod
    def from_tte(cls, tte, method, time_ranges=None):
        """Create a background fitter from a TTE object
        
        Args:
            tte (:class:`~gbm.data.TTE`): A TTE data object
            method (<class>): A background fitting/estimation class for unbinned data 
            time_ranges ([(float, float), ...]): 
                The time range or time ranges over which to fit the background. 
                If omitted, uses the full time range of the data 
        
        Returns:        
            :class:`BackgroundFitter`: The initialized background fitting object
        """
        if not isinstance(tte, TTE):
            raise TypeError('Input data must be a TTE object')

        obj = cls()
        obj._data_obj = tte
        obj._validate_method(method)
        time_ranges = obj._validate_time_ranges(time_ranges)

        # Slice the TTE data and merge if multiple slices
        data = [tte.data.time_slice(trange[0], trange[1]) for trange in
                time_ranges]
        data = EventList.merge(data)
        data.sort('TIME')
        # pull out the events in each channel
        events = [data.channel_slice(i, i).time for i in
                  range(tte.numchans)]

        obj._method = method(events)
        obj._type = 'unbinned'
        obj._livetime = data.get_exposure(time_ranges=time_ranges)
        return obj

    def fit(self, **kwargs):
        """Perform a background fit of the data
        
        Args:
            **kwargs: Options to be passed as parameters to the fitting class
        """
        self._parameters = kwargs
        self._method.fit(**kwargs)

    def interpolate_bins(self, tstart, tstop, **kwargs):
        """Interpolate the fitted background model over a set of bins.
        The exposure is calculated for each bin of the background model 
        in case the background model counts is needed.
        
        Args:
            tstart (np.array): The starting times
            tstop (np.array): The ending times
            **kwargs: Options to be passed as parameters to the interpolation method 
        
        Returns:
            :class:`BackgroundRates`: The background rates/uncertainty object
        """
        tstart_, tstop_ = (tstart.copy(), tstop.copy())
        # do the interpolation
        rate, rate_uncert = self._method.interpolate(tstart_, tstop_, **kwargs)
        # get the exposure
        numtimes = tstart_.shape[0]
        exposure = np.array([self._data_obj.get_exposure((tstart_[i], tstop_[i])) \
                             for i in range(numtimes)])
        # create the rates object
        det = self._data_obj.detector
        dtype = self._data_obj.datatype
        trigtime = self._data_obj.trigtime
        rates = BackgroundRates(rate, rate_uncert, tstart_, tstop_,
                                self._data_obj.data.emin.copy(),
                                self._data_obj.data.emax.copy(), 
                                exposure=exposure,
                                detector=det, datatype=dtype,
                                trigtime=trigtime)

        return rates

    def interpolate_times(self, times, **kwargs):
        """Interpolate the fitted background model over a set of times.
        Does not calculate an exposure since this returns a set of point
        estimates of the background rates.
        
        Args:
            tstart (np.array): The sampling times
            **kwargs: Options to be passed as parameters to the interpolation method 
        
        Returns:
            :class:`BackgroundRates`: The background rates/uncertainty object
        """
        # do the interpolation
        rate, rate_uncert = self._method.interpolate(times, times, **kwargs)

        # create the rates object
        det = self._data_obj.detector
        dtype = self._data_obj.datatype
        trigtime = self._data_obj.trigtime
        rates = BackgroundRates(rate, rate_uncert, times, times,
                                self._data_obj.data.emin.copy(),
                                self._data_obj.data.emax.copy(), detector=det,
                                datatype=dtype, trigtime=trigtime)

        return rates

    def _validate_method(self, method):
        try:
            method
        except:
            raise NameError('Input method is not a known function')

        has_fit = callable(method.fit)
        has_interp = callable(method.interpolate)
        if (not has_fit) or (not has_interp):
            raise NotImplementedError(
                "User-defined Background class must have "
                "both fit() and an interpolate() methods")

    def _validate_time_ranges(self, time_ranges):
        if time_ranges is None:
            time_ranges = [self._data_obj.time_range]
        try:
            iter(time_ranges[0])
        except:
            raise TypeError('time_ranges must be a list of tuples')
        return time_ranges


class BackgroundRates(TimeEnergyBins):
    """Class containing the background rate data.

    Parameters:
        rates (np.array): The array of background rates in each bin
        rate_uncertainty (np.array): The array of background rate uncertainties 
                                     in each bin
        tstart (np.array): The low-value edges of the time bins
        tstop (np.array): The high-value edges of the time bins
        emin (np.array): The low-value edges of the energy bins
        emax (np.array): The high-value edges of the energy bins
        exposure (np.array, optional): The exposure of each bin
        detector (str, optional): The associated detector
        datatype (str, optional): The datatype associated with the background
        trigtime (float, optional): The trigger time
    
    Attributes:
        chanwidths (np.array): The bin widths along the energy axis
        count_uncertainty (np.array): The counts uncertainty in each bin
        energy_centroids (np.array): The bin centroids along the energy axis
        energy_range (float, float): The range of the data along the energy axis
        numchans (int): The number of energy channels along the energy axis
        numtimes (int): The number of bins along the time axis
        rates (np.array): The rates in each Time-Energy Bin
        rates_per_kev (np.array): The differential rates in units of counts/s/keV
        rate_uncertainty (np.array): The rate uncertainty in each bin
        rate_uncertainty_per_kev (np.array):
            The differential rate uncertainty in units of counts/s/keV
        size (int, int): The number of bins along both axes (numtimes, numchans)
        time_centroids (np.array): The bin centroids along the time axis
        time_range (float, float): The range of the data along the time axis
        time_widths (np.array): The bin widths along the time axis    
    """

    def __init__(self, rates, rate_uncertainty, tstart, tstop, emin, emax,
                 exposure=None, detector=None, datatype=None, trigtime=None):
        if exposure is None:
            exposure = np.zeros_like(tstart)

        counts = np.squeeze(rates * exposure[:, np.newaxis])
        super().__init__(counts, tstart, tstop, exposure, emin, emax)
        self._count_uncertainty = np.squeeze(rate_uncertainty * exposure[:, np.newaxis])
        self._rates = rates.squeeze()
        self._rate_uncertainty = rate_uncertainty.squeeze()
        self._detector = detector
        self._datatype = datatype
        self._trigtime = trigtime

    @property
    def count_uncertainty(self):
        return self._count_uncertainty

    @property
    def rates(self):
        return self._rates

    @property
    def rate_uncertainty(self):
        return self._rate_uncertainty

    def integrate_energy(self, emin=None, emax=None):
        """Integrate the over the energy axis.
        Limits on the integration smaller than the full range can be set.
        
        Args:
            emin (float, optional): The low end of the integration range. If not 
                               set, uses the lowest energy edge of the histogram
            emax (float, optional): The high end of the integration range. If not 
                              set, uses the highest energy edge of the histogram
        
        Returns:
            :class:`gbm.data.primitives.TimeBins`: A TimeBins object containing \
                                                   the lightcurve histogram
        """
        if emin is None:
            emin = self.energy_range[0]
        if emax is None:
            emax = self.energy_range[1]

        mask = self._slice_energy_mask(emin, emax)
        emin = self.emin[mask][0]
        emax = self.emax[mask][-1]
        rates = np.nansum(self.rates[:, mask], axis=1).reshape(-1,1)
        rate_uncert = np.sqrt(
            np.nansum(self.rate_uncertainty[:, mask] ** 2, axis=1)).reshape(-1,1)

        obj = BackgroundRates(rates, rate_uncert, self.tstart.copy(), 
                              self.tstop.copy(), np.array([emin]), 
                              np.array([emax]), exposure=self.exposure.copy())
        return obj

    def integrate_time(self, tstart=None, tstop=None):
        """Integrate the background over the time axis (producing a count rate
        spectrum). Limits on the integration smaller than the full range can 
        be set.
        
        Args:
            tstart (float, optional): The low end of the integration range.  
                          If not set, uses the lowest time edge of the histogram
            tstop (float, optional): The high end of the integration range. 
                         If not set, uses the highest time edge of the histogram
        
        Returns:
            :class:`BackgroundSpectrum`: A BackgroundSpectrum object containing \
                                         the count rate spectrum
        """
        if tstart is None:
            tstart = self.time_range[0]
        if tstop is None:
            tstop = self.time_range[1]

        mask = self._slice_time_mask(tstart, tstop)
        exposure = np.nansum(self.exposure[mask])
        rates = np.nansum(self.counts[mask, :], axis=0) / exposure
        rate_uncert = np.sqrt(np.nansum(self.count_uncertainty[mask, :] ** 2,
                                        axis=0)) / exposure
        exposure = np.full(rates.size, exposure)

        obj = BackgroundSpectrum(rates, rate_uncert, self.emin.copy(), 
                                 self.emax.copy(), exposure)
        return obj

    def to_bak(self, time_range=None, **kwargs):
        """Integrate over the time axis and produce a BAK object
        
        Args:
            time_range ((float, float), optional): 
                The time range to integrate over
            **kwargs: Options to pass to BAK.from_data()
        
        Returns:
            :class:`gbm.data.BAK`: The background BAK object
        """
        from gbm.data.pha import BAK

        if time_range is None:
            time_range = self.time_range
        back_spec = self.integrate_time(*time_range)
        dtype = self._datatype.lower()
        bak = BAK.from_data(back_spec, *time_range, datatype=dtype,
                            detector=self._detector, trigtime=self._trigtime,
                            **kwargs)
        return bak

    def rebin_time(self, method, *args, tstart=None, tstop=None):
        """Not Implemented      
        """
        raise NotImplementedError

    def rebin_energy(self, method, *args, emin=None, emax=None):
        """Not Implemented      
        """
        raise NotImplementedError

    @classmethod
    def merge_time(cls, histos):
        """Merge multiple BackroundRates together along the time axis.
        
        Args:
            histos (list of :class:`BackgroundRates`): 
                A list containing the BackgroundRates to be merged
        
        Returns:
            :class:`BackgroundRates`: A new BackgroundRates object containing \
                                      the merged BackgroundRates
        """
        rates = np.vstack((histo.rates for histo in histos))
        rate_uncertainty = np.vstack(
            (histo.rate_uncertainty for histo in histos))
        bins = TimeEnergyBins.merge_time(histos)
        obj = cls(rates, rate_uncertainty, bins.tstart, bins.tstop,
                  bins.emin, bins.emax, exposure=bins.exposure)
        return obj

    @classmethod
    def sum_time(cls, bkgds):
        """Sum multiple TimeBins together if they have the same time range.  
        Example use would be summing two backgrounds from two detectors.
        
        Args:
            bkgds (list of :class:`BackgroundRates):
                A list containing the BackgroundRates to be summed
        
        Returns:
            :class:`BackgroundRates`: A new BackgroundRates object containing \
                                      the summed BackgroundRates
        """
        rates = np.zeros_like(bkgds[0].rates)
        rates_var = np.zeros_like(bkgds[0].rates)
        for bkgd in bkgds:
            assert bkgd.numtimes == bkgds[0].numtimes, \
                "The backgrounds must all have the same support"
            rates += bkgd.rates
            rates_var += bkgd.rate_uncertainty ** 2

        # averaged exposure, sampling times
        exposure = np.mean([bkgd.exposure for bkgd in bkgds], axis=0)
        tstart = np.mean([bkgd.tstart for bkgd in bkgds], axis=0)
        tstop = np.mean([bkgd.tstop for bkgd in bkgds], axis=0)
        emin = np.array([np.min([bkgd.emin for bkgd in bkgds])])
        emax = np.array([np.min([bkgd.emax for bkgd in bkgds])])

        sum_bkgd = cls(rates, np.sqrt(rates_var), tstart, tstop, emin, emax,
                       exposure=exposure, datatype=bkgds[0]._datatype,
                       trigtime=bkgds[0]._trigtime)
        return sum_bkgd


class BackgroundSpectrum(EnergyBins):
    """A class defining a Background Spectrum.

    Parameters:
        rates (np.array): The array of background rates in each bin
        rate_uncertainty (np.array): The array of background rate uncertainties 
                                     in each bin
        lo_edges (np.array): The low-value edges of the bins
        hi_edges (np.array): The high-value edges of the bins
        exposure (np.array): The exposure of each bin

    Attributes:
        centroids (np.array): The geometric centroids of the bins
        counts (np.array): The counts in each bin
        count_uncertainty (np.array): The count uncertainty in each bin
        exposure (np.array): The exposure of each bin
        hi_edges (np.array): The high-value edges of the bins
        lo_edges (np.array): The low-value edges of the bins
        range (float, float): The range of the bin edges
        rates (np.array): count rate of each bin
        rate_uncertainty (np.array): The count rate uncertainty of each bin
        size (int): Number of bins
        widths (np.array): The widths of the bins
    """

    def __init__(self, rates, rate_uncertainty, lo_edges, hi_edges, exposure):
        counts = rates * exposure
        super().__init__(counts, lo_edges, hi_edges, exposure)
        self._count_uncertainty = rate_uncertainty * exposure
        self._rates = rates
        self._rate_uncertainty = rate_uncertainty

    @property
    def count_uncertainty(self):
        return self._count_uncertainty

    @property
    def rates(self):
        return self._rates

    @property
    def rate_uncertainty(self):
        return self._rate_uncertainty

    @classmethod
    def sum(cls, histos):
        """Sum multiple BackgroundSpectrums together if they have the same 
        energy range (support).
        
        Args:
            histos (list of :class:`BackgroundSpectrum`):  
                A list containing the background spectra to be summed
        
        Returns:        
            :class:`BackgroundSpectrum`: A new object containing the \
                                         summed BackgroundSpectrum
        """
        counts = np.zeros(histos[0].size)
        count_variance = np.zeros(histos[0].size)
        exposure = 0.0
        for histo in histos:
            assert histo.size == histos[0].size, \
                "The histograms must all have the same size"
            assert np.all(histo.lo_edges == histos[0].lo_edges), \
                "The histograms must all have the same support"
            counts += histo.counts
            count_variance += histo.count_uncertainty**2
            exposure += histo.exposure
        
        rates = counts/exposure
        rate_uncertainty = np.sqrt(count_variance)/exposure
        
        sum_bins = cls(rates, rate_uncertainty, histos[0].lo_edges, 
                       histos[0].hi_edges, exposure)
        return sum_bins

    def slice(self, emin, emax):
        """Perform a slice over an energy range and return a new 
        BackgroundSpectrum object. Note that the emin and emax values that fall 
        inside a bin will result in that bin being included.

        Args:
            emin (float): The low energy edge of the slice
            emax (float): The high energy of the slice
        
        Returns:
            :class:`BackgroundSpectrum`: A new object containing the energy slice
        """
        emin_snap = self.closest_edge(emin, which='low')
        emax_snap = self.closest_edge(emax, which='high')

        mask = (self.lo_edges < emax_snap) & (self.hi_edges > emin_snap)
        obj = self.__class__(self.rates[mask], self.rate_uncertainty[mask],
                             self.lo_edges[mask], self.hi_edges[mask], 
                             self.exposure[mask])
        return obj