#!/usr/bin/env python
# -*- coding: utf-8 -*-
"""
A Metropolis-Hastings and Gibbs samplers for a state-ful model.

"""

import numpy as np
from emcee import autocorr
from emcee.sampler import Sampler
import logging
import h5py
from Starfish import constants as C

class StateSampler(Sampler):
    """
    The most basic possible Metropolis-Hastings style MCMC sampler, but with
    optional callbacks for acceptance and rejection.

    :param cov:
        The covariance matrix to use for the proposal distribution.

    :param dim:
        Number of dimensions in the parameter space.

    :param lnpostfn:
        A function that takes a vector in the parameter space as input and
        returns the natural logarithm of the posterior probability for that
        position.

    :param query_lnprob:
        A function which queries the model for the current lnprob value. This is
        because an intermediate sampler might have intervened.

    :param rejectfn: (optional)
        A function to be executed if the proposal is rejected. Generally this
        should be a function which reverts the model to the previous parameter
        values. Might need to be a closure.

    :param acceptfn: (optional)
        A function to be executed if the proposal is accepted.

    :param args: (optional)
        A list of extra positional arguments for ``lnpostfn``. ``lnpostfn``
        will be called with the sequence ``lnpostfn(p, *args, **kwargs)``.

    :param kwargs: (optional)
        A list of extra keyword arguments for ``lnpostfn``. ``lnpostfn``
        will be called with the sequence ``lnpostfn(p, *args, **kwargs)``.

    """
    def __init__(self, lnprob, p0, cov, query_lnprob=None, rejectfn=None,
        acceptfn=None, debug=False, outdir="", *args, **kwargs):
        dim = len(p0)
        super().__init__(dim, lnprob, *args, **kwargs)
        self.cov = cov
        self.p0 = p0
        self.query_lnprob = query_lnprob
        self.rejectfn = rejectfn
        self.acceptfn = acceptfn
        self.logger = logging.getLogger(self.__class__.__name__)
        self.outdir = outdir
        self.debug = debug
        if self.debug:
            self.logger.setLevel(logging.DEBUG)
        else:
            self.logger.setLevel(logging.INFO)

    def reset(self):
        super().reset()
        self._chain = np.empty((0, self.dim))
        self._lnprob = np.empty(0)

    def sample(self, p0, lnprob0=None, randomstate=None, thin=1,
               storechain=True, iterations=1, incremental_save=0, **kwargs):
        """
        Advances the chain ``iterations`` steps as an iterator

        :param p0:
            The initial position vector.

        :param lnprob0: (optional)
            The log posterior probability at position ``p0``. If ``lnprob``
            is not provided, the initial value is calculated.

        :param rstate0: (optional)
            The state of the random number generator. See the
            :func:`random_state` property for details.

        :param iterations: (optional)
            The number of steps to run.

        :param thin: (optional)
            If you only want to store and yield every ``thin`` samples in the
            chain, set thin to an integer greater than 1.

        :param storechain: (optional)
            By default, the sampler stores (in memory) the positions and
            log-probabilities of the samples in the chain. If you are
            using another method to store the samples to a file or if you
            don't need to analyse the samples after the fact (for burn-in
            for example) set ``storechain`` to ``False``.

        At each iteration, this generator yields:

        * ``pos`` - The current positions of the chain in the parameter
          space.

        * ``lnprob`` - The value of the log posterior at ``pos`` .

        * ``rstate`` - The current state of the random number generator.

        """

        self.random_state = randomstate

        p = np.array(p0)
        if lnprob0 is None:
            # See if there's something there
            lnprob0 = self.query_lnprob()

            # If not, we're on the first iteration
            if lnprob0 is None:
                lnprob0 = self.get_lnprob(p)

        # Resize the chain in advance.
        if storechain:
            N = int(iterations / thin)
            self._chain = np.concatenate((self._chain,
                                          np.zeros((N, self.dim))), axis=0)
            self._lnprob = np.append(self._lnprob, np.zeros(N))

        i0 = self.iterations

        # Use range instead of xrange for python 3 compatability
        for i in range(int(iterations)):
            self.iterations += 1

            # Since the independent nuisance sampling may have changed parameters,
            # query each process for the current lnprob
            lnprob0 = self.query_lnprob()
            self.logger.debug("Queried lnprob: {}".format(lnprob0))

            # Calculate the proposal distribution.
            if self.dim == 1:
                q = self._random.normal(loc=p[0], scale=self.cov[0], size=(1,))
            else:
                q = self._random.multivariate_normal(p, self.cov)

            newlnprob = self.get_lnprob(q)
            diff = newlnprob - lnprob0
            self.logger.debug("old lnprob: {}".format(lnprob0))
            self.logger.debug("proposed lnprob: {}".format(newlnprob))

            # M-H acceptance ratio
            if diff < 0:
                diff = np.exp(diff) - self._random.rand()
                if diff < 0:
                    #Reject the proposal and revert the state of the model
                    self.logger.debug("Proposal rejected")
                    if self.rejectfn is not None:
                        self.rejectfn()

            if diff > 0:
                #Accept the new proposal
                self.logger.debug("Proposal accepted")
                p = q
                lnprob0 = newlnprob
                self.naccepted += 1
                if self.acceptfn is not None:
                    self.acceptfn()

            if storechain and i % thin == 0:
                ind = i0 + int(i / thin)
                self._chain[ind, :] = p
                self._lnprob[ind] = lnprob0

            # The default of 0 evaluates to False
            if incremental_save:
                if (((i+1) % incremental_save) == 0) & (i > 0):
                    np.save('chain_backup.npy', self._chain)

            # Heavy duty iterator action going on right here...
            yield p, lnprob0, self.random_state

    @property
    def acor(self):
        """
        An estimate of the autocorrelation time for each parameter (length:
        ``dim``).

        """
        return self.get_autocorr_time()

    def get_autocorr_time(self, window=50):
        """
        Compute an estimate of the autocorrelation time for each parameter
        (length: ``dim``).

        :param window: (optional)
            The size of the windowing function. This is equivalent to the
            maximum number of lags to use. (default: 50)

        """
        return autocorr.integrated_time(self.chain, axis=0, window=window)

    def write(self, fname="mc.hdf5"):
        '''
        Write the samples to an HDF file.

        flatchain
        acceptance fraction

        Everything can be saved in the dataset self.fname

        '''

        filename = self.outdir + fname
        self.logger.debug("Opening {} for writing HDF5 flatchains".format(filename))
        hdf5 = h5py.File(filename, "w")
        samples = self.flatchain

        dset = hdf5.create_dataset("samples", samples.shape, compression='gzip', compression_opts=9)
        dset[:] = samples
        dset.attrs["acceptance"] = "{}".format(self.acceptance_fraction)
        dset.attrs["commit"] = "{}".format(C.get_git_commit())
        hdf5.close()

# class PSampler(ParallelSampler):
#     '''
#     Subclasses the GibbsSampler in emcee
#
#     :param cov:
#     :param starting_param_dict: the dictionary of starting parameters
#     :param cov: the MH proposal
#     :param revertfn:
#     :param acceptfn:
#     :param debug:
#
#     '''
#
#     def __init__(self, **kwargs):
#         self.dim = len(self.param_tuple)
#         #p0 = np.empty((self.dim,))
#         #starting_param_dict = kwargs.get("starting_param_dict")
#         #for i,param in enumerate(self.param_tuple):
#         #    p0[i] = starting_param_dict[param]
#
#         kwargs.update({"dim":self.dim})
#         #self.spectra_list = kwargs.get("spectra_list", [0])
#
#         super(PSampler, self).__init__(**kwargs)
#
#         #Each subclass will have to overwrite how it parses the param_dict into the correct order
#         #and sets the param_tuple
#
#         #SUBCLASS here and define self.param_tuple
#         #SUBCLASS here and define self.lnprob
#         #SUBCLASS here and do self.revertfn
#         #then do super().__init__() to call the following code
#
#         self.outdir = kwargs.get("outdir", "")
#
#     def startdict_to_tuple(self, startdict):
#         raise NotImplementedError("To be implemented by a subclass!")
#
#     def zip_p(self, p):
#         return dict(zip(self.param_tuple, p))
#
#     def lnprob(self):
#         raise NotImplementedError("To be implemented by a subclass!")
#
#     def revertfn(self):
#         raise NotImplementedError("To be implemented by a subclass!")
#
#     def acceptfn(self):
#         raise NotImplementedError("To be implemented by a subclass!")
#
#     def write(self):
#         '''
#         Write all of the relevant sample output to an HDF file.
#
#         Write the lnprobability to an HDF file.
#
#         flatchain
#         acceptance fraction
#         tuple parameters as an attribute in the header from self.param_tuple
#
#         The actual HDF5 file is structured as follows
#
#         /
#         stellar parameters.flatchain
#         00/
#         ...
#         22/
#         23/
#             global_cov.flatchain
#             regions/
#                 region1.flatchain
#
#         Everything can be saved in the dataset self.fname
#
#         '''
#
#         filename = self.outdir + "flatchains.hdf5"
#         self.logger.debug("Opening {} for writing HDF5 flatchains".format(filename))
#         hdf5 = h5py.File(filename, "w")
#         samples = self.flatchain
#
#         self.logger.debug("Creating dataset with fname:{}".format(self.fname))
#         dset = hdf5.create_dataset(self.fname, samples.shape, compression='gzip', compression_opts=9)
#         self.logger.debug("Storing samples and header attributes.")
#         dset[:] = samples
#         dset.attrs["parameters"] = "{}".format(self.param_tuple)
#         dset.attrs["acceptance"] = "{}".format(self.acceptance_fraction)
#         dset.attrs["acor"] = "{}".format(self.acor)
#         dset.attrs["commit"] = "{}".format(C.get_git_commit())
#         hdf5.close()
#
#         #lnprobability is the lnprob at each sample
#         filename = self.outdir + "lnprobs.hdf5"
#         self.logger.debug("Opening {} for writing HDF5 lnprobs".format(filename))
#         hdf5 = h5py.File(filename, "w") #creates if doesn't exist, otherwise read/write
#         lnprobs = self.lnprobability
#
#         dset = hdf5.create_dataset(self.fname, samples.shape[:1], compression='gzip', compression_opts=9)
#         dset[:] = lnprobs
#         dset.attrs["commit"] = "{}".format(C.get_git_commit())
#         hdf5.close()
#
#     def plot(self, triangle_plot=False):
#         '''
#         Generate the relevant plots once the sampling is done.
#         '''
#         samples = self.flatchain
#
#         plot_walkers(self.outdir + self.fname + "_chain_pos.png", samples, labels=self.param_tuple)
#
#         if triangle_plot:
#             import triangle
#             figure = triangle.corner(samples, labels=self.param_tuple, quantiles=[0.16, 0.5, 0.84],
#                                      show_titles=True, title_args={"fontsize": 12})
#             figure.savefig(self.outdir + self.fname + "_triangle.png")
#
#             plt.close(figure)
#
# class StellarSampler(PSampler):
#     """
#     Subclasses the Sampler specifically for stellar parameters
#
#     """
#     def __init__(self, **kwargs):
#         '''
#         :param pconns: Collection of parent ends of the PIPEs
#         :type pconns: dict
#
#         :param starting_param_dict:
#             the dictionary of starting parameters
#
#         :param cov:
#             the MH proposal
#
#         :param fix_logg:
#             fix logg? If so, to what value?
#
#         :param debug:
#
#         :param args: []
#         '''
#
#         self.fix_logg = kwargs.get("fix_logg", False)
#         starting_pram_dict = kwargs.get("starting_param_dict")
#         self.param_tuple = self.startdict_to_tuple(starting_pram_dict)
#         print("param_tuple is {}".format(self.param_tuple))
#         self.p0 = np.array([starting_pram_dict[key] for key in self.param_tuple])
#
#         kwargs.update({"p0":self.p0, "revertfn":self.revertfn, "acceptfn": self.acceptfn, "lnprobfn":self.lnprob})
#         super(StellarSampler, self).__init__(**kwargs)
#
#         #self.pconns is a dictionary of parent connections to each PIPE connecting to the child processes.
#         self.spectrum_ids = sorted(self.pconns.keys())
#         self.fname = "stellar"
#
#     def startdict_to_tuple(self, startdict):
#         tup = ()
#         for param in C.stellar_parameters:
#             #check if param is in keys, if so, add to the tuple
#             if param in startdict:
#                 tup += (param,)
#         return tup
#
#     def reset(self):
#         super(StellarSampler, self).reset()
#
#     def revertfn(self):
#         '''
#         Revert the model to the previous state of parameters, in the case of a rejected MH proposal.
#         '''
#         self.logger.debug("reverting stellar parameters")
#         self.prior = self.prior_last
#
#         #Decide we don't want these stellar params. Tell the children to reject the proposal.
#         for pconn in self.pconns.values():
#             pconn.send(("DECIDE", False))
#
#     def acceptfn(self):
#         '''
#         Execute this if the MH proposal is accepted.
#         '''
#         self.logger.debug("accepting stellar parameters")
#         #Decide we do want to keep these stellar params. Tell the children to accept the proposal.
#         for pconn in self.pconns.values():
#             pconn.send(("DECIDE", True))
#
#     def lnprob(self, p):
#         # We want to send the same stellar parameters to each model,
#         # but also send the different vz and logOmega parameters
#         # to the separate spectra, based upon spectrum_id.
#         #self.logger.debug("StellarSampler lnprob p is {}".format(p))
#
#         #Extract only the temp, logg, Z, vsini parameters
#         if not self.fix_logg:
#             params = self.zip_p(p[:4])
#             others = p[4:]
#         else:
#             #Coming in as temp, Z, vsini, vz, logOmega...
#             params = self.zip_p(p[:3])
#             others = p[3:]
#             params.update({"logg": self.fix_logg})
#
#         # Prior
#         self.prior_last = self.prior
#
#         logg = params["logg"]
#         self.prior = -0.5 * (logg - 5.0)**2/(0.05)**2
#
#         #others should now be either [vz, logOmega] or [vz0, logOmega0, vz1, logOmega1, ...] etc. Always div by 2.
#         #split p up into [vz, logOmega], [vz, logOmega] pairs that update the other parameters.
#         #mparams is now a list of parameter dictionaries
#
#         #Now, pack up mparams into a dictionary to send the right stellar parameters to the right subprocesses
#         mparams = {}
#         for (spectrum_id, order_id), (vz, logOmega) in zip(self.spectrum_ids, grouper(others, 2)):
#             p = params.copy()
#             p.update({"vz":vz, "logOmega":logOmega})
#             mparams[spectrum_id] = p
#
#         self.logger.debug("updated lnprob params: {}".format(mparams))
#
#         lnps = np.empty((self.nprocs,))
#
#         #Distribute the calculation to each process
#         self.logger.debug("Distributing params to children")
#         for ((spectrum_id, order_id), pconn) in self.pconns.items():
#             #Parse the parameters into what needs to be sent to each Model here.
#             pconn.send(("LNPROB", mparams[spectrum_id]))
#
#         #Collect the answer from each process
#         self.logger.debug("Collecting params from children")
#         for i, pconn in enumerate(self.pconns.values()):
#             lnps[i] = pconn.recv()
#
#         self.logger.debug("lnps : {}".format(lnps))
#         s = np.sum(lnps)
#         self.logger.debug("sum lnps {}".format(s))
#         return s + self.prior
#
#
#
#
# # From emcee
# class GibbsSampler(Sampler):
#     """
#     The most basic possible Metropolis-Hastings style MCMC sampler, designed to be
#     encapsulated as part of a state-ful Gibbs sampler.
#
#     :param cov:
#         The covariance matrix to use for the proposal distribution.
#
#     :param dim:
#         Number of dimensions in the parameter space.
#
#     :param lnpostfn:
#         A function that takes a vector in the parameter space as input and
#         returns the natural logarithm of the posterior probability for that
#         position.
#
#     :param revertfn:
#         A function which reverts the model to the previous parameter values, in the
#         case that the proposal parameters are rejected.
#
#     :param acceptfn:
#         A function to be executed if the proposal is accepted.
#
#     :param resample:
#         Redraw the previous lnprob at these parameters? Designed to work with the fact that an
#         Emulator is non-deterministic and prevent it from getting caught at a maximum.
#
#     :param args: (optional)
#         A list of extra positional arguments for ``lnpostfn``. ``lnpostfn``
#         will be called with the sequence ``lnpostfn(p, *args, **kwargs)``.
#
#     :param kwargs: (optional)
#         A list of extra keyword arguments for ``lnpostfn``. ``lnpostfn``
#         will be called with the sequence ``lnpostfn(p, *args, **kwargs)``.
#
#     """
#     def __init__(self, cov, p0, *args, **kwargs):
#         super(GibbsSampler, self).__init__(*args, **kwargs)
#         self.cov = cov
#         self.p0 = p0
#         self.revertfn = kwargs.get("revertfn", None)
#         self.acceptfn = kwargs.get("acceptfn", None)
#         self.logger = logging.getLogger(self.__class__.__name__)
#         self.debug = kwargs.get("debug", False)
#         if self.debug:
#             self.logger.setLevel(logging.DEBUG)
#         else:
#             self.logger.setLevel(logging.INFO)
#
#     def reset(self):
#         super(GibbsSampler, self).reset()
#         self._chain = np.empty((0, self.dim))
#         self._lnprob = np.empty(0)
#
#     def sample(self, p0, lnprob0=None, randomstate=None, thin=1,
#                storechain=True, iterations=1, **kwargs):
#         """
#         Advances the chain ``iterations`` steps as an iterator
#
#         :param p0:
#             The initial position vector.
#
#         :param lnprob0: (optional)
#             The log posterior probability at position ``p0``. If ``lnprob``
#             is not provided, the initial value is calculated.
#
#         :param rstate0: (optional)
#             The state of the random number generator. See the
#             :func:`random_state` property for details.
#
#         :param iterations: (optional)
#             The number of steps to run.
#
#         :param thin: (optional)
#             If you only want to store and yield every ``thin`` samples in the
#             chain, set thin to an integer greater than 1.
#
#         :param storechain: (optional)
#             By default, the sampler stores (in memory) the positions and
#             log-probabilities of the samples in the chain. If you are
#             using another method to store the samples to a file or if you
#             don't need to analyse the samples after the fact (for burn-in
#             for example) set ``storechain`` to ``False``.
#
#         At each iteration, this generator yields:
#
#         * ``pos`` - The current positions of the chain in the parameter
#           space.
#
#         * ``lnprob`` - The value of the log posterior at ``pos`` .
#
#         * ``rstate`` - The current state of the random number generator.
#
#         """
#
#         self.random_state = randomstate
#
#         p = np.array(p0)
#         if lnprob0 is None:
#             lnprob0 = self.get_lnprob(p)
#
#         # Resize the chain in advance.
#         if storechain:
#             N = int(iterations / thin)
#             self._chain = np.concatenate((self._chain,
#                                           np.zeros((N, self.dim))), axis=0)
#             self._lnprob = np.append(self._lnprob, np.zeros(N))
#
#         i0 = self.iterations
#         # Use range instead of xrange for python 3 compatability
#         for i in range(int(iterations)):
#             self.iterations += 1
#
#             # Calculate the proposal distribution.
#             if self.dim == 1:
#                 q = self._random.normal(loc=p[0], scale=self.cov[0], size=(1,))
#             else:
#                 q = self._random.multivariate_normal(p, self.cov)
#
#             newlnprob = self.get_lnprob(q)
#             diff = newlnprob - lnprob0
#             self.logger.debug("old lnprob: {}".format(lnprob0))
#             self.logger.debug("proposed lnprob: {}".format(newlnprob))
#
#             # M-H acceptance ratio
#             if diff < 0:
#                 diff = np.exp(diff) - self._random.rand()
#                 if diff < 0:
#                     #Reject the proposal and revert the state of the model
#                     self.logger.debug("Proposal rejected")
#                     if self.revertfn is not None:
#                         self.revertfn()
#
#             if diff > 0:
#                 #Accept the new proposal
#                 self.logger.debug("Proposal accepted")
#                 p = q
#                 lnprob0 = newlnprob
#                 self.naccepted += 1
#                 if self.acceptfn is not None:
#                     self.acceptfn()
#
#             if storechain and i % thin == 0:
#                 ind = i0 + int(i / thin)
#                 self._chain[ind, :] = p
#                 self._lnprob[ind] = lnprob0
#
#             # Heavy duty iterator action going on right here...
#             yield p, lnprob0, self.random_state
#
#     @property
#     def acor(self):
#         """
#         An estimate of the autocorrelation time for each parameter (length:
#         ``dim``).
#
#         """
#         return self.get_autocorr_time()
#
#     def get_autocorr_time(self, window=50):
#         """
#         Compute an estimate of the autocorrelation time for each parameter
#         (length: ``dim``).
#
#         :param window: (optional)
#             The size of the windowing function. This is equivalent to the
#             maximum number of lags to use. (default: 50)
#
#         """
#         return autocorr.integrated_time(self.chain, axis=0, window=window)
#
# class SSampler(GibbsSampler):
#     '''
#     Subclasses the GibbsSampler in emcee
#
#     :param cov:
#     :param starting_param_dict: the dictionary of starting parameters
#     :param cov: the MH proposal
#     :param revertfn:
#     :param acceptfn:
#     :param debug:
#
#     '''
#
#     def __init__(self, **kwargs):
#         self.dim = len(self.param_tuple)
#         #p0 = np.empty((self.dim,))
#         #starting_param_dict = kwargs.get("starting_param_dict")
#         #for i,param in enumerate(self.param_tuple):
#         #    p0[i] = starting_param_dict[param]
#
#         kwargs.update({"dim":self.dim})
#         #self.spectra_list = kwargs.get("spectra_list", [0])
#
#         super(Sampler, self).__init__(**kwargs)
#
#         #Each subclass will have to overwrite how it parses the param_dict into the correct order
#         #and sets the param_tuple
#
#         #SUBCLASS here and define self.param_tuple
#         #SUBCLASS here and define self.lnprob
#         #SUBCLASS here and do self.revertfn
#         #then do super().__init__() to call the following code
#
#         self.outdir = kwargs.get("outdir", "")
#
#     def startdict_to_tuple(self, startdict):
#         raise NotImplementedError("To be implemented by a subclass!")
#
#     def zip_p(self, p):
#         return dict(zip(self.param_tuple, p))
#
#     def lnprob(self):
#         raise NotImplementedError("To be implemented by a subclass!")
#
#     def revertfn(self):
#         raise NotImplementedError("To be implemented by a subclass!")
#
#     def acceptfn(self):
#         raise NotImplementedError("To be implemented by a subclass!")
#
#     def write(self):
#         '''
#         Write all of the relevant sample output to an HDF file.
#
#         Write the lnprobability to an HDF file.
#
#         flatchain
#         acceptance fraction
#         tuple parameters as an attribute in the header from self.param_tuple
#
#         The actual HDF5 file is structured as follows
#
#         /
#         stellar parameters.flatchain
#         00/
#         ...
#         22/
#         23/
#             global_cov.flatchain
#             regions/
#                 region1.flatchain
#
#         Everything can be saved in the dataset self.fname
#
#         '''
#
#         filename = self.outdir + "flatchains.hdf5"
#         self.logger.debug("Opening {} for writing HDF5 flatchains".format(filename))
#         hdf5 = h5py.File(filename, "w")
#         samples = self.flatchain
#
#         self.logger.debug("Creating dataset with fname:{}".format(self.fname))
#         dset = hdf5.create_dataset(self.fname, samples.shape, compression='gzip', compression_opts=9)
#         self.logger.debug("Storing samples and header attributes.")
#         dset[:] = samples
#         dset.attrs["parameters"] = "{}".format(self.param_tuple)
#         dset.attrs["acceptance"] = "{}".format(self.acceptance_fraction)
#         dset.attrs["acor"] = "{}".format(self.acor)
#         dset.attrs["commit"] = "{}".format(C.get_git_commit())
#         hdf5.close()
#
#         #lnprobability is the lnprob at each sample
#         filename = self.outdir + "lnprobs.hdf5"
#         self.logger.debug("Opening {} for writing HDF5 lnprobs".format(filename))
#         hdf5 = h5py.File(filename, "w")
#         lnprobs = self.lnprobability
#
#         dset = hdf5.create_dataset(self.fname, samples.shape[:1], compression='gzip', compression_opts=9)
#         dset[:] = lnprobs
#         dset.attrs["commit"] = "{}".format(C.get_git_commit())
#         hdf5.close()
#
#     def plot(self, triangle_plot=False):
#         '''
#         Generate the relevant plots once the sampling is done.
#         '''
#         samples = self.flatchain
#
#         plot_walkers(self.outdir + self.fname + "_chain_pos.png", samples, labels=self.param_tuple)
#
#         if triangle_plot:
#             import triangle
#             figure = triangle.corner(samples, labels=self.param_tuple, quantiles=[0.16, 0.5, 0.84],
#                                      show_titles=True, title_args={"fontsize": 12})
#             figure.savefig(self.outdir + self.fname + "_triangle.png")
#
#             plt.close(figure)
#
# class NuisanceSampler(SSampler):
#     def __init__(self, **kwargs):
#         '''
#
#         :param OrderModel: the parallel.OrderModel instance
#
#         :param starting_param_dict: the dictionary of starting parameters
#
#         :param cov:
#             the MH proposal
#
#         :param debug:
#
#         :param args: []
#
#         '''
#
#         starting_param_dict = kwargs.get("starting_param_dict")
#         self.param_tuple = self.startdict_to_tuple(starting_param_dict)
#         print("param_tuple is {}".format(self.param_tuple))
#         #print("param_tuple length {}".format(len(self.param_tuple)))
#
#         chebs = [starting_param_dict["cheb"][key] for key in self.cheb_tup]
#         covs = [starting_param_dict["cov"][key] for key in self.cov_tup]
#         regions = starting_param_dict["regions"]
#         #print("initializing {}".format(regions))
#         regs = [regions[id][kk] for id in sorted(regions) for kk in C.cov_region_parameters]
#         #print("regs {}".format(regs))
#
#         self.p0 = np.array(chebs + covs + regs)
#
#         kwargs.update({"p0":self.p0, "revertfn":self.revertfn, "lnprobfn":self.lnprob})
#         super(NuisanceSampler, self).__init__(**kwargs)
#
#         self.model = kwargs.get("OrderModel")
#         spectrum_id, order_id = self.model.id
#         order = kwargs.get("order", order_id)
#         #self.fname = "{}/{}/{}".format(spectrum_id, order, "nuisance")
#         self.fname = "nuisance"
#         self.params = None
#         self.prior_params = kwargs.get("prior_params", None)
#         if self.prior_params:
#             self.sigma0 = self.prior_params["regions"]["sigma0"]
#             self.mus = self.prior_params["regions"]["mus"]
#             self.mu_width = self.prior_params["regions"]["mu_width"]
#             self.sigma_knee = self.prior_params["regions"]["sigma_knee"]
#             self.frac_global = self.prior_params["regions"]["frac_global"]
#
#     def startdict_to_tuple(self, startdict):
#         #This is a little more tricky than the stellar parameters.
#         #How are the keys stored and passed in the dictionary?
#         #{"cheb": [c0, c1, c2, ..., cn], "cov": [sigAmp, logAmp, l],
#         #        "regions":{0: [logAmp, ], 1: [], N:[] }}
#
#         #Serialize the cheb parameters
#         self.ncheb = len(startdict["cheb"])
#         self.cheb_tup = ("logc0",) + tuple(["c{}".format(i) for i in range(1, self.ncheb)])
#
#         #Serialize the covariance parameters
#         self.ncov = 3
#         cov_tup = ()
#         for param in C.cov_global_parameters:
#             #check if param is in keys, if so, add to the tuple
#             if param in startdict["cov"]:
#                 cov_tup += (param,)
#         self.cov_tup = cov_tup
#
#         regions_tup = ()
#         self.regions = startdict.get("regions", None)
#         if self.regions:
#             self.nregions = len(self.regions)
#             for key in sorted(self.regions.keys()):
#                 for kk in C.cov_region_parameters:
#                     regions_tup += ("r{:0>2}-{}".format(key,kk),)
#             self.regions_tup = regions_tup
#         else:
#             self.nregions = 0
#             self.regions_tup = ()
#
#
#         tup = self.cheb_tup + self.cov_tup + self.regions_tup
#         #This should look like
#         #tup = ("c0", "c1", ..., "cn", "sigAmp", "logAmp", "l", "r00_logAmp", "r00_mu", "r00_sigma",
#         # "r01_logAmp", ..., "rNN_sigma")
#         return tup
#
#     def zip_p(self, p):
#         '''
#         Convert the vector to a dictionary
#         '''
#         cheb = dict(zip(self.cheb_tup, p[:self.ncheb]))
#         cov = dict(zip(self.cov_tup, p[self.ncheb:self.ncheb+self.ncov]))
#         regions = p[-self.nregions*3:]
#         rdict = {}
#         for i in range(self.nregions):
#             rdict[i] = dict(zip(("logAmp", "mu", "sigma"), regions[i*3:3*(i+1)]))
#
#         params = {"cheb":cheb, "cov":cov, "regions":rdict}
#         return params
#
#     def revertfn(self):
#         self.logger.debug("reverting model")
#         self.model.prior = self.prior_last
#         self.params = self.params_last
#         self.model.revert_nuisance()
#
#     def lnprob(self, p):
#         self.params_last = self.params
#         params = self.zip_p(p)
#         self.params = params
#         self.logger.debug("Updating nuisance params {}".format(params))
#
#         # Nuisance parameter priors implemented here
#         self.prior_last = self.model.prior
#         # Region parameter priors implemented here
#         if self.nregions > 0:
#             regions = params["regions"]
#             keys = sorted(regions)
#
#             #Unpack the region parameters into a vector of mus, amps, and sigmas
#             amps = 10**np.array([regions[key]["logAmp"] for key in keys])
#             cov_amp = 10**params["cov"]["logAmp"]
#
#             #First check to make sure that amplitude can't be some factor less than the global covariance
#             if np.any(amps < (cov_amp * self.frac_global)):
#                 return -np.inf
#
#             mus = np.array([regions[key]["mu"] for key in keys])
#             sigmas = np.array([regions[key]["sigma"] for key in keys])
#
#             #Make sure the region hasn't strayed too far from the original specification
#             if np.any(np.abs(mus - self.mus) > self.sigma0):
#                 # The region has strayed too far from the original specification
#                 return -np.inf
#
#             #Use a Gaussian prior on mu, that it keeps the region within the original setting.
#             # 1/(sqrt(2pi) * sigma) exp(-0.5 (mu-x)^2/sigma^2)
#             #-ln(sigma * sqrt(2 pi)) - 0.5 (mu - x)^2 / sigma^2
#             #width = 0.04
#             lnGauss = -0.5 * np.sum(np.abs(mus - self.mus)**2/self.mu_width**2 -
#                                     np.log(self.mu_width * np.sqrt(2. * np.pi)))
#
#             # Use a ln(logistic) function on sigma, that is flat before the knee and dies off for anything
#             # greater, to prevent dilution into global cov kernel
#             lnLogistic = np.sum(np.log(-1./(1. + np.exp(self.sigma_knee - sigmas)) + 1.))
#
#             self.model.prior = lnLogistic + lnGauss
#
#         try:
#             self.model.update_nuisance(params)
#             lnp = self.model.evaluate() # also sets OrderModel.lnprob to proposed value. Includes self.model.prior
#             return lnp
#         except C.ModelError:
#             return -np.inf


#!/usr/bin/python

'''
Use synthetic partial columns from GEOS-Chem to obtain cloud-sliced
NO2 in the upper troposphere and compare this to UT NO2 obtained if
simply average the NO2 mixing ratios from the model over the same
pressure range (the "truth"). Both are obtained by Gaussian
weighting toward the pressure center.

GEOS-Chem partial columns are obtained over Europe, North America,
and China at the GEOS-FP meteorology native resolution (0.25x0.3125)
(latxlon) for June-August 2016-2017.

Input options to process the data include the region, the horizontal
resolution, and model simulation years.

.. code-block:: bash

    usage: ut_no2_gc_test.py [-h] [--gc_dir GC_DIR] [--out_dir OUT_DIR]
                         [--resolution RESOLUTION] [--region REGION]
                         [--strat_filter_threshold STRAT_FILTER_THRESHOLD]
                         [--start_date START_DATE]
                         [--end_date END_DATE] [-p PLOT]
                         [--do_temp_correct DO_TEMP_CORRECT]
                         [--apply_cld_frac_filter APPLY_CLD_FRAC_FILTER]
                         [--do_cld_hght_test DO_CLD_HGHT_TEST]

    optional arguments:
      -h, --help            show this help message and exit
      --gc_dir GC_DIR
      --out_dir OUT_DIR
      --resolution RESOLUTION
                            Can be 8x10, 4x5, 2x25 or 1x1
      --region REGION       Can be EU, NA, or CH
      --strat_filter_threshold STRAT_FILTER_THRESHOLD
      --start_date START_DATE
      --end_date END_DATE
      -p PLOT, --plot PLOT
      --do_temp_correct DO_TEMP_CORRECT
      --apply_cld_frac_filter APPLY_CLD_FRAC_FILTER
      --do_cld_hght_test DO_CLD_HGHT_TEST

'''

# Note for the future; most of this can probably be replaced with a modification of the GridAggregator class from
# tropomi_ut_no2

# Import relevant packages:
import numpy as np
from netCDF4 import Dataset
from scipy import stats
import matplotlib.pyplot as plt
from mpl_toolkits.basemap import Basemap
import argparse
from sklearn.linear_model import LinearRegression
import sys
import os
import datetime as dt
from dateutil import rrule as rr

# Import hack
sys.path.append(
    os.path.join(
        os.path.dirname(os.path.abspath(__file__)),
        '..'))

from uptrop.date_file_utils import get_gc_file_list
from uptrop.constants import AVOGADRO
from uptrop.constants import G
from uptrop.constants import MW_AIR
from uptrop.bootstrap import rma
from uptrop.cloud_slice_no2 import cldslice, CLOUD_SLICE_ERROR_ENUM
from uptrop.height_pressure_converter import alt2pres, pres2alt

# Turn off warnings:
np.warnings.filterwarnings('ignore')

# Define pressure range:
# These are fixed, so can be constants, as opposed to inputs.
P_MIN=180
P_MAX=450

class ProcessingException(Exception):
    pass

class CloudSliceException(Exception):
    pass

class InvalidRegionException(Exception):
    pass

class InvalidResolutionException(Exception):
    pass

class DomainIssueException(Exception):
    pass

class ProcessedData:
    """A class for comparing geoschem and tropomi data"""
    # Note for anyone reading this; this performs very similar functions to the GridAggregator class, but in a different
    # way. I prefer the method in GridAggregator, but this is fine for now. -John Roberts
    def __init__(self, region, str_res, strat_thld, do_temperature_correction=False, do_cld_frac_filter=False, do_cld_hght_test=False):
        """Creates and returns an instance of ProcessedData for a given region

        This class aggregates geoschem data over a grid defined by the :ref:`define_grid method`.

        :param region: The region to aggregate geoschem over; can be NA, EU or CH.
        :type region: str
        :param str_res: The resolution of the grid around teh region to analyse; can be 8x10, 4x5, 2x25 or 1x1
        :type str_res: str
        :param do_temperature_correction: Whether to perform the temperature correction step
        :type do_temperature_correction: bool
        :param do_cld_frac_filter: Whether to perform fractional filtering of clouds
        :type do_cld_frac_filter: bool

        :returns: A ProcessedData isntance ready to recieve data
        :rtype: ProcessedData
        """

        self.temperature_correction = do_temperature_correction
        self.cld_frac_filter = do_cld_frac_filter
        self.cloud_height_test = do_cld_hght_test

        # Filtering threshold for stratosphere:
        self.strat_filt = strat_thld

        self.define_grid(region, str_res)
        grid_shape = self.X.shape

        # Define output arrays:
        self.g_no2_vcd = np.zeros(grid_shape)
        self.g_no2_vmr = np.zeros(grid_shape)
        self.g_cld_fr = np.zeros(grid_shape)
        self.g_cld_p = np.zeros(grid_shape)
        self.g_slope_err = np.zeros(grid_shape)
        self.g_gaus_wgt = np.zeros(grid_shape)
        self.true_o3 = np.zeros(grid_shape)  # ozone mixing ratio coincident with cloud-sliced ut self
        self.true_no2 = np.zeros(grid_shape)  # "true" cloudy UT NO2
        self.g_askut_no2 = np.zeros(grid_shape)  # "true" all-sky UT NO2
        self.g_ask_gaus_wgt = np.zeros(grid_shape)  # "true" all-sky UT NO2 weights
        self.g_as_cnt = np.zeros(grid_shape)  # Count all-sky
        self.g_cnt = np.zeros(grid_shape)

        # Type of data loss in the cloud-slicing retrieval:
        self.loss_count = {
            "too_few_points": 0,
            "low_cloud_height_range": 0,
            "low_cloud_height_std": 0,
            "large_error": 0,
            "sig_diff_from_zero": 0,
            "no2_outlier": 0,
            "non_uni_strat": 0,
        }

        # Initialize:
        self.cloud_slice_count = 0
        self.maxcnt = 0
        self.grad_retain = 0
        self.grad_remove = 0
        # Print out min and max cloud-sliced UT NO2 error size:
        self.minerr = 1   # Set to maximum possible error (100%)
        self.maxerr = 0

        # Define string to represent the layer range:
        self.prange = str(P_MIN) + '-' + str(P_MAX)

        # Define factor to convert slope of NO2 mixing ratio versus pressure
        # to VMR:
        self.den2mr = np.divide((np.multiply(G, MW_AIR)), AVOGADRO)

    def define_grid(self, region, str_res):
        """Defines a grid based on a set of regional bounding boxes and a resolution

        :param region: The region to aggregate geoschem over; can be NA, EU or CH.
        :type region: str
        :param str_res: The resolution of the grid around the region to analyse; can be 8x10, 4x5, 2x25 or 1x1
        :type str_res: str
        """
        # Define target grid:
        if region == 'NA':
            self.minlat = 2.
            self.maxlat = 62.
            self.minlon = -140.
            self.maxlon = -55.
        elif region == 'EU':
            self.minlat = 26.
            self.maxlat = 66.
            self.minlon = -25.
            self.maxlon = 45.
        elif region == 'CH':
            self.minlat = 6.
            self.maxlat = 62.
            self.minlon = 60.
            self.maxlon = 140.
        else:
            print("Invalid region; valid regions are 'NA','EU','CH'.")
            raise InvalidRegionException
        # NOTE FOR LATER: This snippet turns up in fresco_cld_err; a candidate for the library.
        # Define grid information:
        if str_res == '8x10':
            self.dellat, self.dellon = 8, 10
        elif str_res == '4x5':
            self.dellat, self.dellon = 4, 5
        elif str_res == '2x25':
            self.dellat, self.dellon = 2, 2.5
        elif str_res == '1x1':
            self.dellat, self.dellon = 1, 1
        else:
            print("Invalid resolution: valid resolutions are 8x10, 4x5, 2x25 (two pt 5) and 1x1")
            raise InvalidResolutionException
        self.out_lon = np.arange(self.minlon, self.maxlon + self.dellon, self.dellon)
        self.out_lat = np.arange(self.minlat, self.maxlat + self.dellat, self.dellat)
        # Convert output lats and long to 2D:
        self.X, self.Y = np.meshgrid(self.out_lon, self.out_lat, indexing='ij')
        # Dimensions of output data:
        self.xdim = len(self.out_lon)
        self.ydim = len(self.out_lat)
        # Get maximum possible number of native resolution model gridsquares
        # in the coarser grid (to use in process_grid_square to check that
        # don't exceed this):
        self.max_limit=(self.dellon/0.3125)*(self.dellat/0.25)

    # ----Processing methods----
    def process_geoschem_day(self, file_path):
        """Aggregates geoschem data into the processing grid for a file on a given day.

        Calls :ref:`regrid_and_process` for each pixel in the geoschem file, then calls
        :ref:`process_grid_square` for each square of the grid

        :param file_path: Path to the geoschem file to load into the processing grid
        :type file_path: str"""
        # Define output data for this day:
        out_shape = (self.xdim, self.ydim)  # This feel gross. A 3-d list of appendable lists.
        self.g_no2 = [[[] for n in range(self.ydim)] for m in range(self.xdim)]
        self.g_grad = [[[] for n in range(self.ydim)] for m in range(self.xdim)]
        self.g_o3 = [[[] for n in range(self.ydim)] for m in range(self.xdim)]
        self.all_cld_fr = [[[] for n in range(self.ydim)] for m in range(self.xdim)]
        self.strat_no2 = [[[] for n in range(self.ydim)] for m in range(self.xdim)]
        self.all_cld_p = [[[] for n in range(self.ydim)] for m in range(self.xdim)]
        self.g_true_no2 = [[[] for n in range(self.ydim)] for m in range(self.xdim)]

        this_geoschem_day = GeosChemDay(file_path,
                                        temperature_correction=self.temperature_correction,
                                        cloud_height_test=self.cloud_height_test)
        # Get column values:
        for y in range(len(this_geoschem_day.t_lat)):
            for x in range(len(this_geoschem_day.t_lon)):
                this_geoschem_day.prepare_no2_pixel(x, y)
                self.regrid_and_process(x, y, this_geoschem_day)
        for i in range(self.xdim):
            for j in range(self.ydim):
                if any(value != 0 for value in self.g_no2[i][j]):
                    self.process_grid_square(i, j)

    def regrid_and_process(self, x, y, gc_data):
        """Inserts data from pixel x,y from geoschem data into the appropriate grid cell

        :param x: X index of geoschem grid cell
        :type x: int
        :param y: Y index of geoschem grid cell
        :type y: int
        :param gc_data: An instance of GeosChemDay data
        :type gc_data: GeosChemDay
        """
        # If no valid data, skip:
        if ( len(gc_data.askind) == 0):
            return

        # Find nearest gridsquare in output grid:
        lon = int(np.argmin(abs(self.out_lon - gc_data.t_lon[x])))
        lat = int(np.argmin(abs(self.out_lat - gc_data.t_lat[y])))

        self.g_ask_gaus_wgt[lon, lat] += np.sum(gc_data.twgt)

        self.g_askut_no2[lon, lat] += np.sum(gc_data.t_gc_no2[gc_data.askind, y, x] * gc_data.twgt * 1e3)
        self.g_as_cnt[lon, lat] += 1.0

        # Add relevant data to array if there are clouds at 450-180 hPa:
        if (gc_data.lcld < gc_data.level_min) or (gc_data.lcld > gc_data.level_max):
            # Check for clouds between min and max pressure. If none, move to next pixel.
            #print("Cloud top outside pressure range in pixel {},{}".format(x,y))
            return

        # Find where cloud fraction in UT exceeds 0.7 after calculating
        # true all-sky NO2:
        # (Keep for testing effect of thick clouds on cloud-sliced UT NO2):
        if ( self.cld_frac_filter ):
            if np.sum(gc_data.t_cld_fr[gc_data.level_min:gc_data.level_max + 1, y, x]) < 0.7:
                return

        self.g_no2[lon][lat].append(gc_data.no2_2d)
        self.g_grad[lon][lat].append(gc_data.grad)
        self.strat_no2[lon][lat].append(gc_data.strat_col)
        self.all_cld_p[lon][lat].append(gc_data.t_cld_hgt[y, x])
        self.g_o3[lon][lat].append(
            np.mean(gc_data.t_gc_o3[gc_data.level_min:gc_data.level_max + 1, y, x]))
        self.g_true_no2[lon][lat].append(
            np.mean(gc_data.t_gc_no2[gc_data.level_min:gc_data.level_max + 1, y, x]))
        self.all_cld_fr[lon][lat].append(
            np.sum(gc_data.t_cld_fr[gc_data.level_min:gc_data.level_max + 1, y, x]))
        pass

    def process_grid_square(self, i, j):

        # Define vectors of relevant data:
        # These should all have the same length
        t_col_no2 = np.asarray(self.g_no2[i][j],dtype=np.float)
        t_strat_no2 = np.asarray(self.strat_no2[i][j],dtype=np.float)
        t_fr_c = np.asarray(self.all_cld_fr[i][j],dtype=np.float)
        t_cld = np.asarray(self.all_cld_p[i][j],dtype=np.float)
        t_mr_no2 = np.asarray(self.g_true_no2[i][j],dtype=np.float)
        t_o3 = np.asarray(self.g_o3[i][j],dtype=np.float)
        t_grad_no2 = np.asarray(self.g_grad[i][j],dtype=np.float)

        # Skip if fewer than 10 points:
        if len(t_col_no2) < 10:
            #print("Grid square {}, {} failed with condition too_few_points".format(i, j))
            self.loss_count["too_few_points"] += 1
            return

        # Remove non-uniform stratosphere:
        if (np.std(t_strat_no2) / np.mean(t_strat_no2)) > self.strat_filt:
            self.loss_count["non_uni_strat"] += 1
            return  # continue

        # Get number of points:
        n_pnts = len(t_cld)
        if n_pnts > self.maxcnt:
            self.maxcnt = n_pnts
            print('Max no. of points in a cluster: ', self.maxcnt, flush=True)

        # Check if number of gridsquares exceeds max possible
        # (indicates error in coarser grid domain):
        if ( self.maxcnt > self.max_limit ):
            print('No. of model grids exceeds max possible')
            raise DomainIssueException

        # Use cloud_slice_ut_no2 function to get cloud-sliced UT NO2.
        # Treat data differently depending on whether there are 10-99 points:
        if n_pnts >= 10 and n_pnts <100:
            self.add_slice(i,j,t_cld,t_col_no2, t_mr_no2, t_fr_c, t_grad_no2, t_o3)
        # Or at least 100 points:
        elif n_pnts >= 100:
            num_slices = 40
            stride = round(n_pnts / num_slices)
            nloop = list(range(stride))
            for w in nloop:
                subset_t_col_no2 = t_col_no2[w::stride]
                subset_t_cld = t_cld[w::stride]
                subset_t_mr_no2 = t_mr_no2[w::stride]
                subset_t_fr_c = t_fr_c[w::stride]
                subset_t_grad_no2 = t_grad_no2[w::stride]
                subset_t_o3 = t_o3[w::stride]
                self.add_slice(i, j, subset_t_cld, subset_t_col_no2, subset_t_mr_no2, subset_t_fr_c, subset_t_grad_no2, subset_t_o3)

    def add_slice(self, i, j, t_cld, t_col_no2, t_mr_no2, t_fr_c, t_grad_no2, t_o3):
        """Extracts the upper troposphere gc_data, gc_data error, ozone data, ozone error,
         and mean cloud pressure for grid square [i,j]

        This method uses the cloud-slicing function :ref:`uptrop.cloud_slice_ut_no2.cldslice`
        Once calculated, the a weighting is derived from cloud pressure.
        The weighted upper tropospheric gc_data and error is added to the rolling total for this season.
        If the cloud slicing fails, then the reason is added to loss_count for the end report.

        :param i: X-index of grid square
        :type i: int
        :param j: Y-index of grid square
        :type j: int
        :param t_cld: Tropospheric cloud
        :type t_cld: float
        :param t_col_no2: Total (stratosphere+troposphere) column no2
        :type t_col_no2: float
        :param t_mr_no2: NO2 mixing ratio between pmax and pmin
        :type t_mr_no2: float
        :param t_fr_c: Cloud fraction between pmax and pmin
        :type t_fr_c: float
        :param t_grad_no2: NO2 gradient between pmax and pmin in pptv/hPa
        :type t_grad_no2: float
        :param t_o3: O3 mixing ratio to test influence from stratosphere
        :type t_o3: float
        """
        utmrno2, utmrno2err, stage_reached, mean_cld_pres = cldslice(t_col_no2, t_cld,140)

        # Skip if approach didn't work (i.e. cloud-sliced UT NO2 is NaN):
        # Drop out after the reason for data loss is added to loss_count.
        if np.isnan(utmrno2) or np.isnan(utmrno2err):
            # Cloud-slicing led to nan, but due to rma regression rather
            # than data filtering (these are rare):
            if ( stage_reached==0 ):
                print("Cloud-sliced NO2 NAN for pixel i:{} j:{}".format(i, j))
                return
            self.loss_count[CLOUD_SLICE_ERROR_ENUM[stage_reached]] += 1
            # Track non-uniform NO2 removed:
            grad_ind=np.where(np.abs(t_grad_no2) >= 0.33)[0]
            self.grad_remove += len(grad_ind)
            #print("Cloud-slice exception {} in pixel i:{} j:{}".format(
            #    CLOUD_SLICE_ERROR_ENUM[stage_reached], i, j))
        else:
            # Calculate weights:
            #err_wgt = 1.0 / (utmrno2err ** 2)  # not use, but preserve anyway
            gaus_wgt = np.exp((-(mean_cld_pres - 315) ** 2) / (2 * 135 ** 2))

            # Get and track error range:
            nerr = utmrno2err / utmrno2
            if nerr > self.maxerr:
                self.maxerr = nerr
                print('Max rel. cloud-sliced NO2 error: ', self.maxerr, flush=True)
            if nerr < self.minerr:
                self.minerr = nerr
                print('Min rel. cloud-sliced NO2 error: ', self.minerr, flush=True)

            # Track non-uniform NO2 retained:
            grad_ind=np.where(np.abs(t_grad_no2) >= 0.33)[0]
            self.grad_retain += len(grad_ind)

            # Weighted mean for each pass of cldslice:
            weight_mean = np.mean(t_mr_no2 * 1e3) * gaus_wgt
            self.true_no2[i, j] += weight_mean
            self.true_o3[i, j] += (np.mean(t_o3) * gaus_wgt)
            self.g_no2_vmr[i, j] += utmrno2 * gaus_wgt
            self.g_slope_err[i, j] += utmrno2err * gaus_wgt
            self.g_gaus_wgt[i, j] += gaus_wgt
            self.g_cnt[i, j] += 1
            self.g_cld_fr[i, j] += np.mean(t_fr_c)
            self.cloud_slice_count += 1

    def get_weighted_mean(self):
        """Applies weighting to the aggregated means.
        """
        # Mean physical parameters:
        self.g_no2_vmr = np.divide(self.g_no2_vmr, self.g_gaus_wgt, where=self.g_cnt != 0)
        self.g_cld_fr = np.divide(self.g_cld_fr, self.g_cnt, where=self.g_cnt != 0)
        self.true_no2 = np.divide(self.true_no2, self.g_gaus_wgt, where=self.g_cnt != 0)
        self.g_cld_p = np.divide(self.g_cld_p, self.g_cnt, where=self.g_cnt != 0)
        self.true_o3 = np.divide(self.true_o3, self.g_gaus_wgt, where=self.g_cnt != 0)
        self.g_askut_no2 = np.divide(self.g_askut_no2, self.g_ask_gaus_wgt, where=self.g_as_cnt != 0)
        # Mean cloud-sliced UT NO2 error:
        self.g_slope_err = np.divide(self.g_slope_err, self.g_gaus_wgt, where=self.g_cnt != 0)
        # Mean weights:
        self.g_gaus_wgt = np.divide(self.g_gaus_wgt, self.g_cnt, where=self.g_cnt != 0)
        # No data (nan):
        self.true_no2[self.g_cnt == 0] = np.nan
        self.true_o3[self.g_cnt == 0] = np.nan
        self.g_slope_err[self.g_cnt == 0] = np.nan
        self.g_gaus_wgt[self.g_cnt == 0] = np.nan
        self.g_askut_no2[self.g_as_cnt == 0] = np.nan
        self.g_no2_vmr[self.g_cnt == 0] = np.nan
        self.g_cld_fr[self.g_cnt == 0] = np.nan
        self.g_cld_p[self.g_cnt == 0] = np.nan

    # ----Reporting and saving methods----
    def print_data_report(self):
        """Prints a set of reasons for missing data
        """
        # Output code diagnostics:
        # No. of data points:
        # print('No. of valid data points: ',cloud_slice_count,flush=True)
        print('Max no. of data points in a gridsquare: ', np.amax(self.g_cnt), flush=True)
        # Track reasons for data loss:
        print('(1) Too few points: ', self.loss_count["too_few_points"], flush=True)
        print('(2) Low cloud height range: ', self.loss_count["low_cloud_height_range"], flush=True)
        print('(3) Low cloud height std dev: ', self.loss_count["low_cloud_height_std"], flush=True)
        print('(4) Large error: ', self.loss_count["large_error"], flush=True)
        print('(5) Significantly less than zero: ', self.loss_count["sig_diff_from_zero"], flush=True)
        print('(6) Outlier (NO2 > 200 pptv): ', self.loss_count["no2_outlier"], flush=True)
        print('(7) Non-uniform stratosphere: ', self.loss_count["non_uni_strat"], flush=True)
        print('(8) Successful retrievals: ', self.cloud_slice_count, flush=True)
        print('(9) Total possible points: ', (sum(self.loss_count.values()) + self.cloud_slice_count), flush=True)
        print('Non-uniform NO2 retained = ', self.grad_retain)
        print('Non-uniform NO2 removed = ', self.grad_remove)
        print('Error range of cloud-sliced UT NO2: ', np.nanmin(self.g_slope_err), np.nanmax(self.g_slope_err))

    def plot_data(self):
        """Plots the present state of the gridded data
        """
        # ===> FUTURE UPDATE : change from basemap to cartopy <===
        # Plot the data:
        m = Basemap(resolution='l', projection='merc',
                    lat_0=0, lon_0=0, llcrnrlon=self.minlon,
                    llcrnrlat=self.minlat, urcrnrlon=self.maxlon, urcrnrlat=self.maxlat)
        xi, yi = m(self.X, self.Y)
        plt.subplot(2, 3, 1)
        cs = m.pcolor(xi, yi, np.squeeze(self.g_no2_vmr), vmin=0, vmax=80, cmap='jet')
        m.drawparallels(np.arange(-80., 81., 45.), labels=[1, 0, 0, 0], fontsize=8)
        m.drawmeridians(np.arange(-180., 181., 45.), labels=[0, 0, 0, 1], fontsize=8)
        m.drawcoastlines()
        m.drawcountries()
        cbar = m.colorbar(cs, location='bottom', pad="10%")
        plt.title('Cloud-sliced NO2 VMRs')
        plt.subplot(2, 3, 2)
        cs = m.pcolor(xi, yi, np.squeeze(self.true_no2), vmin=0.0, vmax=80, cmap='jet')
        m.drawparallels(np.arange(-80., 81., 45.), labels=[1, 0, 0, 0], fontsize=8)
        m.drawmeridians(np.arange(-180., 181., 45.), labels=[0, 0, 0, 1], fontsize=8)
        m.drawcoastlines()
        m.drawcountries()
        cbar = m.colorbar(cs, location='bottom', pad="10%")
        plt.title('True cloudy NO2')
        plt.subplot(2, 3, 3)
        cs = m.pcolor(xi, yi, np.squeeze(self.g_askut_no2), vmin=0, vmax=80, cmap='jet')
        m.drawparallels(np.arange(-80., 81., 45.), labels=[1, 0, 0, 0], fontsize=8)
        m.drawmeridians(np.arange(-180., 181., 45.), labels=[0, 0, 0, 1], fontsize=8)
        m.drawcoastlines()
        m.drawcountries()
        cbar = m.colorbar(cs, location='bottom', pad="10%")
        plt.title('True NO2 VMRs under all-sky conditions')
        plt.subplot(2, 3, 4)
        plt.plot(self.true_no2, self.g_no2_vmr, 'o', color='black', markersize=6)
        r = stats.pearsonr(self.true_no2[~np.isnan(self.g_no2_vmr)],
                           self.g_no2_vmr[~np.isnan(self.g_no2_vmr)])
        # print('Correlation = ', r[0])
        result = rma(self.true_no2[~np.isnan(self.g_no2_vmr)], self.g_no2_vmr[~np.isnan(self.g_no2_vmr)],
                     len(self.true_no2[~np.isnan(self.g_no2_vmr)]), 1000)
        print(result, flush=True)
        xvals = np.arange(0, 100, 5)
        yvals = result[1] + xvals * result[0]
        plt.plot(xvals, yvals, '-')
        plt.xlim(-4, 80)
        plt.ylim(-4, 80)
        plt.xlabel('True NO2 (cloudy)')
        plt.ylabel('Cloud-sliced NO2')
        print('===== True (cloudy) vs cloud-sliced UT NO2 ====')
        print('R = ', r[0], flush=True)
        print('Slope = ', result[0])
        print('Slope Err = ', result[2], flush=True)
        print('Intercept = ', result[1], flush=True)
        print('Intercept Err = ', result[3], flush=True)
        add2plt = ("y = {a:6.2f}x + {b:6.3f}".format(a=result[0], b=result[1]))
        plt.text(2, 75, add2plt, fontsize=8,
                 ha='left', va='center')  # , transform=ax.transAxes)
        add2plt = ("R = {a:6.2f}".format(a=r[0]))
        plt.text(2, 65, add2plt, fontsize=8,
                 ha='left', va='center')  # , transform=ax.transAxes)
        plt.subplot(2, 3, 5)
        plt.plot(self.g_askut_no2, self.g_no2_vmr, 'o', color='black', markersize=6)
        r = stats.pearsonr(self.g_askut_no2[(~np.isnan(self.g_no2_vmr)) & (~np.isnan(self.g_askut_no2))],
                           self.g_no2_vmr[(~np.isnan(self.g_no2_vmr)) & (~np.isnan(self.g_askut_no2))])
        result = rma(self.g_askut_no2[(~np.isnan(self.g_no2_vmr)) & (~np.isnan(self.g_askut_no2))],
                     self.g_no2_vmr[(~np.isnan(self.g_no2_vmr)) & (~np.isnan(self.g_askut_no2))],
                     len(self.g_askut_no2[(~np.isnan(self.g_no2_vmr)) & (~np.isnan(self.g_askut_no2))]),
                     1000)
        xvals = np.arange(0, 100, 5)
        yvals = result[1] + xvals * result[0]
        plt.plot(xvals, yvals, '-')
        plt.xlim(-4, 80)
        plt.ylim(-4, 80)
        plt.xlabel('True NO2 (all-sky)')
        plt.ylabel('Cloud-sliced NO2')
        add2plt = ("y = {a:6.2f}x + {b:6.3f}". \
                   format(a=result[0], b=result[1]))
        plt.text(2, 75, add2plt, fontsize=8, \
                 ha='left', va='center')  # , transform=ax.transAxes)
        add2plt = ("R = {a:6.2f}".format(a=r[0]))
        plt.text(2, 65, add2plt, fontsize=8, \
                 ha='left', va='center')  # , transform=ax.transAxes)
        plt.subplot(2, 3, 6)
        plt.plot(self.g_askut_no2, self.true_no2, 'o', color='black', markersize=6)
        r = stats.pearsonr(self.g_askut_no2[(~np.isnan(self.true_no2)) & (~np.isnan(self.g_askut_no2))],
                           self.true_no2[(~np.isnan(self.true_no2)) & (~np.isnan(self.g_askut_no2))])
        result = rma(self.g_askut_no2[(~np.isnan(self.true_no2)) & (~np.isnan(self.g_askut_no2))],
                     self.true_no2[(~np.isnan(self.true_no2)) & (~np.isnan(self.g_askut_no2))],
                     len(self.g_askut_no2[(~np.isnan(self.true_no2)) & (~np.isnan(self.g_askut_no2))]),
                     1000)
        xvals = np.arange(0, 100, 5)
        yvals = result[1] + xvals * result[0]
        plt.plot(xvals, yvals, '-')
        plt.xlim(-4, 80)
        plt.ylim(-4, 80)
        plt.xlabel('True NO2 (all-sky)')
        plt.ylabel('True NO2 (cloudy)')
        add2plt = ("y = {a:6.2f}x + {b:6.3f}".format(a=result[0], b=result[1]))
        plt.text(2, 75, add2plt, fontsize=8, ha='left', va='center')
        add2plt = ("R = {a:6.2f}".format(a=r[0]))
        plt.text(2, 65, add2plt, fontsize=8, ha='left', va='center')
        plt.show()

    def save_to_netcdf(self, out_path):
        """Saves the present state of the grid to a netCDF4 file

        :param out_path: Path to the output file
        :type out_path: str
        """
        # Save the data to NetCDF:
        ncout = Dataset(out_path, mode='w', format='NETCDF4')
        # Create data dimensions:
        ncout.createDimension('lat', self.ydim)
        ncout.createDimension('lon', self.xdim)
        # create lon axis:
        lon = ncout.createVariable('lon', np.float32, ('lon',))
        lon.units = 'degrees_east'
        lon.long_name = 'longitude'
        lon[:] = self.out_lon
        # create lat axis:
        lat = ncout.createVariable('lat', np.float32, ('lat',))
        lat.units = 'degrees_north'
        lat.long_name = 'latgitude'
        lat[:] = self.out_lat
        # create data axes:
        # (1) Cloud-sliced NO2:
        csutno2 = ncout.createVariable('csutno2', np.float32, ('lon', 'lat'))
        csutno2.units = 'pptv'
        csutno2.long_name = 'UT NO2 mixing ratio (180-450 hPa) obtained using cloud-slicing'
        csutno2[:] = self.g_no2_vmr
        # (2a) Double-weighting error:
        utdblerr = ncout.createVariable('utdblerr', np.float32, ('lon', 'lat'))
        utdblerr.units = 'pptv'
        utdblerr.long_name = 'Standard error of the NO2 mixing ratios in the UT (180-450 hPa) obtained using cloud-slicing'
        utdblerr[:] = self.g_slope_err
        # (2b) Gaussian-weighting error:
        utgauserr = ncout.createVariable('utgauserr', np.float32, ('lon', 'lat'))
        utgauserr.units = 'pptv'
        utgauserr.long_name = 'Standard error of the NO2 mixing ratios in the UT (180-450 hPa) obtained using cloud-slicing'
        utgauserr[:] = self.g_gaus_wgt
        # (3) Number of observations in each gridsquare:
        nobs = ncout.createVariable('nobs', np.float32, ('lon', 'lat'))
        nobs.units = 'unitless'
        nobs.long_name = 'Number of observations in each gridsquare used to obtain cloud-sliced UT NO2 mixing ratios'
        nobs[:] = self.g_cnt
        # (4) Mean cloud pressure for season between 450-180 hPa:
        utcld = ncout.createVariable('utcld', np.float32, ('lon', 'lat'))
        utcld.units = 'hPa'
        utcld.long_name = 'Mean cloud pressure between 450 and 180 hPa'
        utcld[:] = self.g_cld_p
        # (5) Mean NO2 mixing ratio at 450-180 hPa for scenes with clouds:
        cldutno2 = ncout.createVariable('cldutno2', np.float32, ('lon', 'lat'))
        cldutno2.units = 'pptv'
        cldutno2.long_name = 'UT NO2 mixing ratio (180-450 hPa) obtained if clouds are present'
        cldutno2[:] = self.true_no2
        # (6) Mean NO2 mixing ratio at 450-180 hPa under all conditions (all-sky):
        askutno2 = ncout.createVariable('askutno2', np.float32, ('lon', 'lat'))
        askutno2.units = 'pptv'
        askutno2.long_name = 'UT NO2 mixing ratio (180-450 hPa) obtained under all conditions (all-sky)'
        askutno2[:] = self.g_askut_no2
        # (7) Cloud fraction:
        utcldfrc = ncout.createVariable('utcldfrc', np.float32, ('lon', 'lat'))
        utcldfrc.units = 'unitless'
        utcldfrc.long_name = 'GEOS-FP cloud fraction obtained as sum of 3D cloud fractions across range of interest (180-450 hPa)'
        utcldfrc[:] = self.g_cld_fr
        # (8) O3 sampled coincident with cloud-slicing retrieval:
        uto3 = ncout.createVariable('uto3', np.float32, ('lon', 'lat'))
        uto3.units = 'ppbv'
        uto3.long_name = 'GEOS-Chem ozone obtained coincident with cloud-sliced NO2'
        uto3[:] = self.true_o3
        # Close the file:
        ncout.close()

class GeosChemDay:
    """A class for reading, preprocessing and accessing Geoschem data on a given day
    """
    def __init__(self, file_path, temperature_correction=False, cloud_height_test=False):
        """Reads the data at file_path and returns a GeosChemDay object containing that data

        :param file_path: Path to the netcdf4 file containing the GeosChem data
        :type file_path: str
        :param temperature_correction: Whether to apply temperature correction
        :type temperature_correction: bool
        :param cloud_height_test: Whether to test effect of systematic underestimate in cloud height
        :type cloud_height_test: bool

        :returns: A GeosChemDay class
        :rtype: GeosChemDay
        """
        print('File path: ',file_path, flush=True)

        self.temperature_correction = temperature_correction
        self.cloud_height_test = cloud_height_test

        # Read dataset:
        fh = Dataset(file_path, mode='r')
        # Extract data of interest:
        tlon, tlat, tgcno2, tcldfr, tcldhgt, tadn, tbxhgt, tpedge, tpause, tgco3, tdegk = \
            fh.variables['LON'], fh.variables['LAT'], \
            fh.variables['IJ-AVG-S__NO2'], fh.variables['TIME-SER__CF'], \
            fh.variables['TIME-SER__CThgt'], fh.variables['TIME-SER__AIRDEN'], \
            fh.variables['BXHGHT-S__BXHEIGHT'], fh.variables['PEDGE-S__PSURF'], \
            fh.variables['TR-PAUSE__TP-PRESS'], fh.variables['IJ-AVG-S__O3'], \
            fh.variables['DAO-3D-S__TMPU']
        self.t_lon = tlon[:]
        self.t_lat = tlat[:]
        self.t_gc_no2 = tgcno2[:]
        self.t_cld_fr = tcldfr[:]
        self.t_cld_hgt = tcldhgt[0, :, :]
        self.t_adn = tadn[:]  # in molec/cm3
        self.t_bx_hgt = tbxhgt[:]
        self.t_p_edge = tpedge[:]
        self.t_pause = tpause[0, :, :]
        self.t_gc_o3 = tgco3[:]
        self.t_deg_k = tdegk[:]
        # Convert box height from m to cm:
        self.t_bx_hgt = self.t_bx_hgt * 1e2

        if self.cloud_height_test:
            # Lower cloud heights by 1 km to roughly mimic lower altitude
            # clouds retrieved for TROPOMI assuming clouds are reflective
            # boundaries with uniform reflectivity:
            # Calculate cloud top height in m:
            t_cld_hgt = pres2alt(self.t_cld_hgt*1e2)
            # Lower the clouds by 1 km (1000 m) (this won't work for low-altitude
            # clouds):
            t_cld_hgt = t_cld_hgt - 1e3
            # Convert back to Pa and convert that to hPa:
            self.t_cld_hgt = alt2pres(t_cld_hgt)*1e-2

        #Get outputs ready here for tidyness:
        self.no2_2d = None
        self.trop_col = None
        self.strat_col = None
        self.gcutno2 = None
        self.gascnt = None
        self.grad = None

        self.level_min = None
        self.level_max = None
        self.askind = None

    def prepare_no2_pixel(self, x, y):
        """Extracts preprocesed no2 from the geoschem pixel at x,y

        :param x: The x index of the pixel
        :type x: int
        :param y: The y index of the pixel
        :type y: int
        """
        # Calculate corresponding mid-pressure values:
        tp_mid = np.zeros(len(self.t_p_edge[:, y, x]))
        # Get mid-pressure values, except for highest layer:
        for k in range(len(self.t_p_edge[:, y, x]) - 1):
            tp_mid[k] = np.multiply(0.5, (self.t_p_edge[k, y, x] + self.t_p_edge[k + 1, y, x]))
        # Estimate mid-pressure for highest value (doesn't need to
        # be accurate, as surpassing the range of interset):
        # Data output from the model includes 47 vertical layers. This means that only 46 pressure centres can be calculated as the calculation requires pressure edges.
        tp_mid[46] = np.multiply(0.5, (self.t_p_edge[46, y, x] + (self.t_p_edge[46, y, x] - 0.1)))
        # Get model layer of tropopause:
        tppind = np.argmin(abs(tp_mid - self.t_pause[y, x]))
        # Get indices that fall between 450 and 180 hPa for estimating
        # "true' all-sky UT NO2 and partial columns:
        lind = np.where((tp_mid >= P_MIN) & (tp_mid <= P_MAX))[0]
        # Get UT NO2 under "true" all-sky conditions:
        # Make sure this is below the tropopause:
        # If below tropopause, use full extent (180-450 hPa):
        if lind[len(lind) - 1] <= tppind:
            self.askind = lind
        # If above tropopause, trim to tropopause-450 hPa:
        if lind[len(lind) - 1] > tppind:
            self.askind = lind[np.where(lind <= tppind)[0]]
        # If tropopause below 450 hPa, skip entirely:
        if self.t_pause[y, x] > P_MAX:
            #print("Tropopause less than P_MAX in geoschem pixel x:{}, y:{}".format(x,y))
            return  # continue

        # Get Guassian weights that allocate higher weights to points
        # closest to the pressure centre (315 hPa):
        # Equation is:
        #   w = exp(-(p-315)^2/2*135^2 ) where 315 hPa is the centre and
        #         135 hPa is the standard deviation.
        self.twgt = np.exp((-(tp_mid[self.askind] - 315) ** 2) / (2 * 135 ** 2))

        # Get model level of cloud top height closest to lowest
        # pressure-altitude of interest (P_MIN):
        self.lcld = np.argmin(abs(self.t_cld_hgt[y, x] - tp_mid))
        # Skip if cloud top height ouside pressure range of interest:
        self.level_min, self.level_max = np.amin(lind), np.amax(lind)

        if (self.temperature_correction):

            # Equation is from the TROPOMI product ATBD (p. 32, Eqn 18)
            # (product document abbrevation: S5P-KNMI-L2-0005-RP)
            self.temp_corr = 1 - (3.16e-3 * (self.t_deg_k[self.level_min:, y, x] - 220.)) + \
                        (3.39e-6 * ((self.t_deg_k[self.level_min:, y, x] - 220) ** 2))
        else:
            # Set to 1 so that no scaling is applied:
            # (might be a more eloquent way to do this)
            self.temp_corr = np.ones(len(self.t_gc_no2[self.level_min:, y, x]))

        # Calculate NO2 gradient:
        regr = LinearRegression()
        # Pressure (hPa)
        x_vals=tp_mid[self.level_min:self.level_max].reshape(-1,1)
        # NO2 (ppbv)
        y_vals=self.t_gc_no2[self.level_min:self.level_max, y, x].reshape(-1,1)
        # NO2 (pptv)
        y_vals=y_vals*1e3
        # Perform regression:
        regr.fit(x_vals,y_vals)
        # Define gradient from regression slope (pptv/hPa):
        self.grad = regr.coef_

        # Get partial NO2 column in molec/m2 from cloud top height
        # to highest model level (output up to level 47):
        # print(t_gc_no2[self.level_min:tppind,y,x])
        # print(t_gc_no2[self.level_min:tppind,y,x]*1.5)
        self.no2_2d = np.sum(self.t_gc_no2[self.level_min:, y, x]
                             * 1e-5
                             * self.temp_corr
                             * self.t_adn[self.level_min:, y, x]
                             * self.t_bx_hgt[self.level_min:, y, x])
        # Get stratospheric column from 180 hPa aloft:
        # Previous approach (remove when model simulations done):
        # tppind=np.where(tpmid<180.)[0]
        self.strat_col = np.sum(self.t_gc_no2[tppind:, y, x]
                                * 1e-5 * self.t_adn[tppind:, y, x]
                                * self.t_bx_hgt[tppind:, y, x])

if __name__ == "__main__":

    parser = argparse.ArgumentParser()
    # Shorten directory name to up to "GC/", then define the subdirectory
    # as 'geosfp' + dirreg + 'iccw/' in get_file_list.
    # This is now done in get_gc_file_list
    parser.add_argument("--gc_dir")
    parser.add_argument("--out_dir")
    parser.add_argument("--resolution", default="4x5", help="Can be 8x10, 4x5, 2x25 or 1x1")
    parser.add_argument("--region", default="EU", help="Can be EU, NA, or CH")
    parser.add_argument("--strat_filter_threshold", default="002", help="")
    #parser.add_argument("--start_date", default="2016-06-01")
    #parser.add_argument("--end_date", default="2017-08-31")
    parser.add_argument("-p", "--plot", type=bool, default=False)
    parser.add_argument("--do_temp_correct", type=bool)
    parser.add_argument("--apply_cld_frac_filter", type=bool)
    parser.add_argument("--do_cld_hght_test", type=bool)
    args = parser.parse_args()

    # Get files:
    gc_dir = args.gc_dir
    STR_RES = args.resolution
    REGION = args.region

    # Hard code for GEOS-Chem synthetic experiment:
    yrrange = '2016-2017'

    # Define output file depending on input arguments:
    out_file = os.path.join(args.out_dir, 'gc-v12-1-0-ut-no2'
                            + '-' + args.region
                            + '-' + args.resolution
                            + '-jja-' + yrrange
                            + '-' + "gaus-wgt" )

    if args.do_cld_hght_test:
        out_file += "-cldtop"
    if args.strat_filter_threshold != "002":
        out_file += '-' + args.strat_filter_threshold +'strat'
    if args.do_temp_correct:
        out_file += "-temp-corr"
    if args.apply_cld_frac_filter:
        out_file += "-cld-filt"

    out_file += "-v1.nc"

    strat_filter_threshold = float(args.strat_filter_threshold)/100
    files = get_gc_file_list(gc_dir, args.region)
    print('Number of files:', len(files), flush=True)

    rolling_total = ProcessedData(REGION, STR_RES, strat_filter_threshold,
                                  do_temperature_correction=args.do_temp_correct,
                                  do_cld_frac_filter=args.apply_cld_frac_filter,
                                  do_cld_hght_test=args.do_cld_hght_test)

    for file_path in files:
        rolling_total.process_geoschem_day(file_path)

    rolling_total.get_weighted_mean()
    rolling_total.print_data_report()
    if args.plot==True: rolling_total.plot_data()
    rolling_total.save_to_netcdf(out_file)