# Copyright 2016 Valentine Svensson, James Hensman, alexggmatthews, Alexis Boukouvalas
# Copyright 2017 Artem Artemev @awav
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""
Likelihoods are another core component of GPflow. This describes how likely the
data is under the assumptions made about the underlying latent functions
p(Y|F). Different likelihoods make different
assumptions about the distribution of the data, as such different data-types
(continuous, binary, ordinal, count) are better modelled with different
likelihood assumptions.

Use of any likelihood other than Gaussian typically introduces the need to use
an approximation to perform inference, if one isn't already needed. A
variational inference and MCMC models are included in GPflow and allow
approximate inference with non-Gaussian likelihoods. An introduction to these
models can be found :ref:`here <implemented_models>`. Specific notebooks
illustrating non-Gaussian likelihood regressions are available for
`classification <notebooks/classification.html>`_ (binary data), `ordinal
<notebooks/ordinal.html>`_ and `multiclass <notebooks/multiclass.html>`_.

Creating new likelihoods
----------
Likelihoods are defined by their
log-likelihood. When creating new likelihoods, the
:func:`logp <gpflow.likelihoods.Likelihood.logp>` method (log p(Y|F)), the
:func:`conditional_mean <gpflow.likelihoods.Likelihood.conditional_mean>`,
:func:`conditional_variance
<gpflow.likelihoods.Likelihood.conditional_variance>`.

In order to perform variational inference with non-Gaussian likelihoods a term
called ``variational expectations``, ∫ q(F) log p(Y|F) dF, needs to
be computed under a Gaussian distribution q(F) ~ N(μ, Σ).

The :func:`variational_expectations <gpflow.likelihoods.Likelihood.variational_expectations>`
method can be overriden if this can be computed in closed form, otherwise; if
the new likelihood inherits
:class:`Likelihood <gpflow.likelihoods.Likelihood>` the default will use
Gauss-Hermite numerical integration (works well when F is 1D
or 2D), if the new likelihood inherits from
:class:`MonteCarloLikelihood <gpflow.likelihoods.MonteCarloLikelihood>` the
integration is done by sampling (can be more suitable when F is higher dimensional).
"""

import numpy as np
import tensorflow as tf
import abc
import warnings

from .. import logdensities
from ..base import Module, Parameter
from ..config import default_float
from ..quadrature import hermgauss, ndiag_mc, ndiagquad
from ..utilities import positive, to_default_int
from .robustmax import RobustMax

def inv_probit(x):
    jitter = 1e-3  # ensures output is strictly between 0 and 1
    return 0.5 * (1.0 + tf.math.erf(x / np.sqrt(2.0))) * (1 - 2 * jitter) + jitter

class Likelihood(Module, metaclass=abc.ABCMeta):
    def __init__(self, latent_dim: int, observation_dim: int):
        """
        A base class for likelihoods, which specifies an observation model
        connecting the latent functions ('F') to the data ('Y').

        All of the members of this class are expected to obey some shape conventions, as specified
        by latent_dim and observation_dim.

        If we're operating on an array of function values 'F', then the last dimension represents
        multiple functions (preceding dimensions could represent different data points, or
        different random samples, for example). Similarly, the last dimension of Y represents a
        single data point. We check that the dimensions are as this object expects.

        The return shapes of all functions in this class is the broadcasted shape of the arguments,
        excluding the last dimension of each argument.

        :param latent_dim: the dimension of the vector F of latent functions for a single data point
        :param observation_dim: the dimension of the observation vector Y for a single data point
        """
        super().__init__()
        self.latent_dim = latent_dim
        self.observation_dim = observation_dim

    def _check_last_dims_valid(self, F, Y):
        """
        Assert that the dimensions of the latent functions F and the data Y are compatible.

        :param F: function evaluation Tensor, with shape [..., latent_dim]
        :param Y: observation Tensor, with shape [..., observation_dim]
        """
        self._check_latent_dims(F)
        self._check_data_dims(Y)

    def _check_return_shape(self, result, F, Y):
        """
        Check that the shape of a computed statistic of the data
        is the broadcasted shape from F and Y.

        :param result: result Tensor, with shape [...]
        :param F: function evaluation Tensor, with shape [..., latent_dim]
        :param Y: observation Tensor, with shape [..., observation_dim]
        """
        expected_shape = tf.broadcast_dynamic_shape(tf.shape(F)[:-1], tf.shape(Y)[:-1])
        tf.debugging.assert_equal(tf.shape(result), expected_shape)

    def _check_latent_dims(self, F):
        """
        Ensure that a tensor of latent functions F has latent_dim as right-most dimension.

        :param F: function evaluation Tensor, with shape [..., latent_dim]
        """
        tf.debugging.assert_shapes([(F, (..., self.latent_dim))])

    def _check_data_dims(self, Y):
        """
        Ensure that a tensor of data Y has observation_dim as right-most dimension.

        :param Y: observation Tensor, with shape [..., observation_dim]
        """
        tf.debugging.assert_shapes([(Y, (..., self.observation_dim))])

    def log_prob(self, F, Y):
        """
        The log probability density log p(Y|F)

        :param F: function evaluation Tensor, with shape [..., latent_dim]
        :param Y: observation Tensor, with shape [..., observation_dim]:
        :returns: log pdf, with shape [...]
        """
        self._check_last_dims_valid(F, Y)
        res = self._log_prob(F, Y)
        self._check_return_shape(res, F, Y)
        return res

    @abc.abstractmethod
    def _log_prob(self, F, Y):
        raise NotImplementedError

    def conditional_mean(self, F):
        """
        The conditional mean of Y|F: [E[Y₁|F], ..., E[Yₖ|F]]
        where K = observation_dim

        :param F: function evaluation Tensor, with shape [..., latent_dim]
        :returns: mean [..., observation_dim]
        """
        self._check_latent_dims(F)
        expected_Y = self._conditional_mean(F)
        self._check_data_dims(expected_Y)
        return expected_Y

    def _conditional_mean(self, F):
        raise NotImplementedError

    def conditional_variance(self, F):
        """
        The conditional marginal variance of Y|F: [var(Y₁|F), ..., var(Yₖ|F)]
        where K = observation_dim

        :param F: function evaluation Tensor, with shape [..., latent_dim]
        :returns: variance [..., observation_dim]
        """
        self._check_latent_dims(F)
        var_Y = self._conditional_variance(F)
        self._check_data_dims(var_Y)
        return var_Y

    def _conditional_variance(self, F):
        raise NotImplementedError

    def predict_mean_and_var(self, Fmu, Fvar):
        """
        Given a Normal distribution for the latent function,
        return the mean and marginal variance of Y,

        i.e. if
            q(f) = N(Fmu, Fvar)

        and this object represents

            p(y|f)

        then this method computes the predictive mean

           ∫∫ y p(y|f)q(f) df dy

        and the predictive variance

           ∫∫ y² p(y|f)q(f) df dy  - [ ∫∫ y p(y|f)q(f) df dy ]²

        :param Fmu: mean function evaluation Tensor, with shape [..., latent_dim]
        :param Fvar: variance of function evaluation Tensor, with shape [..., latent_dim]
        :returns: mean and variance, both with shape [..., observation_dim]
        """
        self._check_latent_dims(Fmu)
        self._check_latent_dims(Fvar)
        mu, var = self._predict_mean_and_var(Fmu, Fvar)
        self._check_data_dims(mu)
        self._check_data_dims(var)
        return mu, var

    @abc.abstractmethod
    def _predict_mean_and_var(self, Fmu, Fvar):
        raise NotImplementedError

    def predict_log_density(self, Fmu, Fvar, Y):
        r"""
        Given a Normal distribution for the latent function, and a datum Y,
        compute the log predictive density of Y,

        i.e. if
            q(F) = N(Fmu, Fvar)

        and this object represents

            p(y|F)

        then this method computes the predictive density

            log ∫ p(y=Y|F)q(F) df

        :param Fmu: mean function evaluation Tensor, with shape [..., latent_dim]
        :param Fvar: variance of function evaluation Tensor, with shape [..., latent_dim]
        :param Y: observation Tensor, with shape [..., observation_dim]:
        :returns: log predictive density, with shape [...]
        """
        tf.debugging.assert_equal(tf.shape(Fmu), tf.shape(Fvar))
        self._check_last_dims_valid(Fmu, Y)
        res = self._predict_log_density(Fmu, Fvar, Y)
        self._check_return_shape(res, Fmu, Y)
        return res

    @abc.abstractmethod
    def _predict_log_density(self, Fmu, Fvar, Y):
        raise NotImplementedError

    def predict_density(self, Fmu, Fvar, Y):
        """
        Deprecated: see `predict_log_density`
        """
        warnings.warn(
            "predict_density is deprecated and will be removed in GPflow 2.1, use predict_log_density instead",
            DeprecationWarning,
        )
        return self.predict_log_density(Fmu, Fvar, Y)

    def variational_expectations(self, Fmu, Fvar, Y):
        r"""
        Compute the expected log density of the data, given a Gaussian
        distribution for the function values,

        i.e. if
            q(f) = N(Fmu, Fvar)

        and this object represents

            p(y|f)

        then this method computes

           ∫ log(p(y=Y|f)) q(f) df.

        This only works if the broadcasting dimension of the statistics of q(f) (mean and variance)
        are broadcastable with that of the data Y.

        :param Fmu: mean function evaluation Tensor, with shape [..., latent_dim]
        :param Fvar: variance of function evaluation Tensor, with shape [..., latent_dim]
        :param Y: observation Tensor, with shape [..., observation_dim]:
        :returns: expected log density of the data given q(F), with shape [...]
        """
        tf.debugging.assert_equal(tf.shape(Fmu), tf.shape(Fvar))
        # returns an error if Y[:-1] and Fmu[:-1] do not broadcast together
        _ = tf.broadcast_dynamic_shape(tf.shape(Fmu)[:-1], tf.shape(Y)[:-1])
        self._check_last_dims_valid(Fmu, Y)
        ret = self._variational_expectations(Fmu, Fvar, Y)
        self._check_return_shape(ret, Fmu, Y)
        return ret

    @abc.abstractmethod
    def _variational_expectations(self, Fmu, Fvar, Y):
        raise NotImplementedError

class ScalarLikelihood(Likelihood):
    """
    A likelihood class that helps with scalar likelihood functions: likelihoods where
    each scalar latent function is associated with a single scalar observation variable.

    If there are multiple latent functions, then there must be a corresponding number of data: we
    check for this.

    The `Likelihood` class contains methods to compute marginal statistics of functions
    of the latents and the data ϕ(y,f):
     * variational_expectations:  ϕ(y,f) = log p(y|f)
     * predict_log_density: ϕ(y,f) = p(y|f)
    Those statistics are computed after having first marginalized the latent processes f
    under a multivariate normal distribution q(f) that is fully factorized.

    Some univariate integrals can be done by quadrature: we implement quadrature routines for 1D
    integrals in this class, though they may be overwritten by inheriting classes where those
    integrals are available in closed form.
    """

    def __init__(self, **kwargs):
        super().__init__(latent_dim=None, observation_dim=None, **kwargs)
        self.num_gauss_hermite_points = 20

    def _check_last_dims_valid(self, F, Y):
        """
        Assert that the dimensions of the latent functions and the data are compatible
        :param F: function evaluation Tensor, with shape [..., latent_dim]
        :param Y: observation Tensor, with shape [..., latent_dim]
        """
        tf.debugging.assert_shapes([(F, (..., "num_latent")), (Y, (..., "num_latent"))])

    def _log_prob(self, F, Y):
        r"""
        Compute log p(Y|F), where by convention we sum out the last axis as it represented
        independent latent functions and observations.
        :param F: function evaluation Tensor, with shape [..., latent_dim]
        :param Y: observation Tensor, with shape [..., latent_dim]
        """
        return tf.reduce_sum(self._scalar_log_density(F, Y), axis=-1)

    def _scalar_log_density(self, F, Y):
        raise NotImplementedError

    def _variational_expectations(self, Fmu, Fvar, Y):
        r"""
        Here, we implement a default Gauss-Hermite quadrature routine, but some
        likelihoods (Gaussian, Poisson) will implement specific cases.
        :param Fmu: mean function evaluation Tensor, with shape [..., latent_dim]
        :param Fvar: variance of function evaluation Tensor, with shape [..., latent_dim]
        :param Y: observation Tensor, with shape [..., latent_dim]:
        :returns: variational expectations, with shape [...]
        """
        return tf.reduce_sum(
            ndiagquad(self._scalar_log_density, self.num_gauss_hermite_points, Fmu, Fvar, Y=Y),
            axis=-1,
        )

    def _predict_log_density(self, Fmu, Fvar, Y):
        r"""
        Here, we implement a default Gauss-Hermite quadrature routine, but some
        likelihoods (Gaussian, Poisson) will implement specific cases.
        :param Fmu: mean function evaluation Tensor, with shape [..., latent_dim]
        :param Fvar: variance of function evaluation Tensor, with shape [..., latent_dim]
        :param Y: observation Tensor, with shape [..., latent_dim]:
        :returns: log predictive density, with shape [...]
        """
        return tf.reduce_sum(
            ndiagquad(
                self._scalar_log_density,
                self.num_gauss_hermite_points,
                Fmu,
                Fvar,
                logspace=True,
                Y=Y,
            ),
            axis=-1,
        )

    def _predict_mean_and_var(self, Fmu, Fvar):
        r"""
        Here, we implement a default Gauss-Hermite quadrature routine, but some
        likelihoods (e.g. Gaussian) will implement specific cases.

        :param Fmu: mean function evaluation Tensor, with shape [..., latent_dim]
        :param Fvar: variance of function evaluation Tensor, with shape [..., latent_dim]
        :returns: mean and variance, both with shape [..., observation_dim]
        """

        def integrand(*X):
            return self.conditional_variance(*X) + self.conditional_mean(*X) ** 2

        integrands = [self.conditional_mean, integrand]
        E_y, E_y2 = ndiagquad(integrands, self.num_gauss_hermite_points, Fmu, Fvar)
        V_y = E_y2 - E_y ** 2
        return E_y, V_y

class Gaussian(ScalarLikelihood):
    r"""
    The Gaussian likelihood is appropriate where uncertainties associated with the data are
    believed to follow a normal distribution, with constant variance.

    Very small uncertainties can lead to numerical instability during the
    optimization process. A lower bound of 1e-6 is therefore imposed on the likelihood variance
    by default.
    """

    def __init__(self, variance=1.0, variance_lower_bound=1e-6, **kwargs):
        super().__init__(**kwargs)
        self.variance = Parameter(variance, transform=positive(lower=variance_lower_bound))

    def _scalar_log_density(self, F, Y):
        return logdensities.gaussian(Y, F, self.variance)

    def _conditional_mean(self, F):  # pylint: disable=R0201
        return tf.identity(F)

    def _conditional_variance(self, F):
        return tf.fill(tf.shape(F), tf.squeeze(self.variance))

    def _predict_mean_and_var(self, Fmu, Fvar):
        return tf.identity(Fmu), Fvar + self.variance

    def _predict_log_density(self, Fmu, Fvar, Y):
        return tf.reduce_sum(logdensities.gaussian(Y, Fmu, Fvar + self.variance), axis=-1)

    def _variational_expectations(self, Fmu, Fvar, Y):
        return tf.reduce_sum(
            -0.5 * np.log(2 * np.pi)
            - 0.5 * tf.math.log(self.variance)
            - 0.5 * ((Y - Fmu) ** 2 + Fvar) / self.variance,
            axis=-1,
        )

class Poisson(ScalarLikelihood):
    r"""
    Poisson likelihood for use with count data, where the rate is given by the (transformed) GP.

    let g(.) be the inverse-link function, then this likelihood represents

    p(yᵢ | fᵢ) = Poisson(yᵢ | g(fᵢ) * binsize)

    Note:binsize
    For use in a Log Gaussian Cox process (doubly stochastic model) where the
    rate function of an inhomogeneous Poisson process is given by a GP.  The
    intractable likelihood can be approximated via a Riemann sum (with bins
    of size 'binsize') and using this Poisson likelihood.
    """

    def __init__(self, invlink=tf.exp, binsize=1.0, **kwargs):
        super().__init__(**kwargs)
        self.invlink = invlink
        self.binsize = np.array(binsize, dtype=default_float())

    def _scalar_log_density(self, F, Y):
        return logdensities.poisson(Y, self.invlink(F) * self.binsize)

    def _conditional_variance(self, F):
        return self.invlink(F) * self.binsize

    def _conditional_mean(self, F):
        return self.invlink(F) * self.binsize

    def _variational_expectations(self, Fmu, Fvar, Y):
        if self.invlink is tf.exp:
            return tf.reduce_sum(
                Y * Fmu
                - tf.exp(Fmu + Fvar / 2) * self.binsize
                - tf.math.lgamma(Y + 1)
                + Y * tf.math.log(self.binsize),
                axis=-1,
            )
        return super()._variational_expectations(Fmu, Fvar, Y)

class Exponential(ScalarLikelihood):
    def __init__(self, invlink=tf.exp, **kwargs):
        super().__init__(**kwargs)
        self.invlink = invlink

    def _scalar_log_density(self, F, Y):
        return logdensities.exponential(Y, self.invlink(F))

    def _conditional_mean(self, F):
        return self.invlink(F)

    def _conditional_variance(self, F):
        return tf.square(self.invlink(F))

    def _variational_expectations(self, Fmu, Fvar, Y):
        if self.invlink is tf.exp:
            return tf.reduce_sum(-tf.exp(-Fmu + Fvar / 2) * Y - Fmu, axis=-1)
        return super()._variational_expectations(Fmu, Fvar, Y)

class StudentT(ScalarLikelihood):
    def __init__(self, scale=1.0, df=3.0, **kwargs):
        """
        :param scale float: scale parameter
        :param df float: degrees of freedom
        """
        super().__init__(**kwargs)
        self.df = df
        self.scale = Parameter(scale, transform=positive())

    def _scalar_log_density(self, F, Y):
        return logdensities.student_t(Y, F, self.scale, self.df)

    def _conditional_mean(self, F):
        return F

    def _conditional_variance(self, F):
        var = (self.scale ** 2) * (self.df / (self.df - 2.0))
        return tf.fill(tf.shape(F), tf.squeeze(var))

class Bernoulli(ScalarLikelihood):
    def __init__(self, invlink=inv_probit, **kwargs):
        super().__init__(**kwargs)
        self.invlink = invlink

    def _scalar_log_density(self, F, Y):
        return logdensities.bernoulli(Y, self.invlink(F))

    def _predict_mean_and_var(self, Fmu, Fvar):
        if self.invlink is inv_probit:
            p = inv_probit(Fmu / tf.sqrt(1 + Fvar))
            return p, p - tf.square(p)
        else:
            # for other invlink, use quadrature
            return super()._predict_mean_and_var(Fmu, Fvar)

    def _predict_log_density(self, Fmu, Fvar, Y):
        p = self.predict_mean_and_var(Fmu, Fvar)[0]
        return logdensities.bernoulli(Y, p)

    def _conditional_mean(self, F):
        return self.invlink(F)

    def _conditional_variance(self, F):
        p = self.conditional_mean(F)
        return p - (p ** 2)

class Gamma(ScalarLikelihood):
    """
    Use the transformed GP to give the *scale* (inverse rate) of the Gamma
    """

    def __init__(self, invlink=tf.exp, **kwargs):
        super().__init__(**kwargs)
        self.invlink = invlink
        self.shape = Parameter(1.0, transform=positive())

    def _scalar_log_density(self, F, Y):
        return logdensities.gamma(Y, self.shape, self.invlink(F))

    def _conditional_mean(self, F):
        return self.shape * self.invlink(F)

    def _conditional_variance(self, F):
        scale = self.invlink(F)
        return self.shape * (scale ** 2)

    def _variational_expectations(self, Fmu, Fvar, Y):
        if self.invlink is tf.exp:
            return tf.reduce_sum(
                -self.shape * Fmu
                - tf.math.lgamma(self.shape)
                + (self.shape - 1.0) * tf.math.log(Y)
                - Y * tf.exp(-Fmu + Fvar / 2.0),
                axis=-1,
            )
        else:
            return super()._variational_expectations(Fmu, Fvar, Y)

class Beta(ScalarLikelihood):
    """
    This uses a reparameterisation of the Beta density. We have the mean of the
    Beta distribution given by the transformed process:

        m = invlink(f)

    and a scale parameter. The familiar α, β parameters are given by

        m     = α / (α + β)
        scale = α + β

    so:
        α = scale * m
        β  = scale * (1-m)
    """

    def __init__(self, invlink=inv_probit, scale=1.0, **kwargs):
        super().__init__(**kwargs)
        self.scale = Parameter(scale, transform=positive())
        self.invlink = invlink

    def _scalar_log_density(self, F, Y):
        mean = self.invlink(F)
        alpha = mean * self.scale
        beta = self.scale - alpha
        return logdensities.beta(Y, alpha, beta)

    def _conditional_mean(self, F):
        return self.invlink(F)

    def _conditional_variance(self, F):
        mean = self.invlink(F)
        return (mean - tf.square(mean)) / (self.scale + 1.0)

class MultiClass(Likelihood):
    def __init__(self, num_classes, invlink=None, **kwargs):
        """
        A likelihood for multi-way classification.  Currently the only valid
        choice of inverse-link function (invlink) is an instance of RobustMax.

        For most problems, the stochastic `Softmax` likelihood may be more
        appropriate (note that you then cannot use Scipy optimizer).
        """
        super().__init__(latent_dim=num_classes, observation_dim=None, **kwargs)
        self.num_classes = num_classes
        self.num_gauss_hermite_points = 20

        if invlink is None:
            invlink = RobustMax(self.num_classes)

        if not isinstance(invlink, RobustMax):
            raise NotImplementedError

        self.invlink = invlink

    def _log_prob(self, F, Y):
        hits = tf.equal(tf.expand_dims(tf.argmax(F, 1), 1), tf.cast(Y, tf.int64))
        yes = tf.ones(tf.shape(Y), dtype=default_float()) - self.invlink.epsilon
        no = tf.zeros(tf.shape(Y), dtype=default_float()) + self.invlink.eps_k1
        p = tf.where(hits, yes, no)
        return tf.reduce_sum(tf.math.log(p), axis=-1)

    def _variational_expectations(self, Fmu, Fvar, Y):
        gh_x, gh_w = hermgauss(self.num_gauss_hermite_points)
        p = self.invlink.prob_is_largest(Y, Fmu, Fvar, gh_x, gh_w)
        ve = p * tf.math.log(1.0 - self.invlink.epsilon) + (1.0 - p) * tf.math.log(
            self.invlink.eps_k1
        )
        return tf.reduce_sum(ve, axis=-1)

    def _predict_mean_and_var(self, Fmu, Fvar):
        possible_outputs = [
            tf.fill(tf.stack([tf.shape(Fmu)[0], 1]), np.array(i, dtype=np.int64))
            for i in range(self.num_classes)
        ]
        ps = [self._predict_non_logged_density(Fmu, Fvar, po) for po in possible_outputs]
        ps = tf.transpose(tf.stack([tf.reshape(p, (-1,)) for p in ps]))
        return ps, ps - tf.square(ps)

    def _predict_log_density(self, Fmu, Fvar, Y):
        return tf.reduce_sum(tf.math.log(self._predict_non_logged_density(Fmu, Fvar, Y)), axis=-1)

    def _predict_non_logged_density(self, Fmu, Fvar, Y):
        gh_x, gh_w = hermgauss(self.num_gauss_hermite_points)
        p = self.invlink.prob_is_largest(Y, Fmu, Fvar, gh_x, gh_w)
        den = p * (1.0 - self.invlink.epsilon) + (1.0 - p) * (self.invlink.eps_k1)
        return den

    def _conditional_mean(self, F):
        return self.invlink(F)

    def _conditional_variance(self, F):
        p = self.conditional_mean(F)
        return p - tf.square(p)

class SwitchedLikelihood(ScalarLikelihood):
    def __init__(self, likelihood_list, **kwargs):
        """
        In this likelihood, we assume at extra column of Y, which contains
        integers that specify a likelihood from the list of likelihoods.
        """
        super().__init__(**kwargs)
        for l in likelihood_list:
            assert isinstance(l, ScalarLikelihood)
        self.likelihoods = likelihood_list

    def _partition_and_stitch(self, args, func_name):
        """
        args is a list of tensors, to be passed to self.likelihoods.<func_name>

        args[-1] is the 'Y' argument, which contains the indexes to self.likelihoods.

        This function splits up the args using dynamic_partition, calls the
        relevant function on the likelihoods, and re-combines the result.
        """
        # get the index from Y
        Y = args[-1]
        ind = Y[..., -1]
        ind = tf.cast(ind, tf.int32)
        Y = Y[..., :-1]
        args[-1] = Y

        # split up the arguments into chunks corresponding to the relevant likelihoods
        args = zip(*[tf.dynamic_partition(X, ind, len(self.likelihoods)) for X in args])

        # apply the likelihood-function to each section of the data
        funcs = [getattr(lik, func_name) for lik in self.likelihoods]
        results = [f(*args_i) for f, args_i in zip(funcs, args)]

        # stitch the results back together
        partitions = tf.dynamic_partition(tf.range(0, tf.size(ind)), ind, len(self.likelihoods))
        results = tf.dynamic_stitch(partitions, results)

        return results

    def _check_last_dims_valid(self, F, Y):
        tf.assert_equal(tf.shape(F)[-1], tf.shape(Y)[-1] - 1)

    def _scalar_log_density(self, F, Y):
        return self._partition_and_stitch([F, Y], "_scalar_log_density")

    def _predict_log_density(self, Fmu, Fvar, Y):
        return self._partition_and_stitch([Fmu, Fvar, Y], "predict_log_density")

    def _variational_expectations(self, Fmu, Fvar, Y):
        return self._partition_and_stitch([Fmu, Fvar, Y], "variational_expectations")

    def _predict_mean_and_var(self, Fmu, Fvar):
        mvs = [lik.predict_mean_and_var(Fmu, Fvar) for lik in self.likelihoods]
        mu_list, var_list = zip(*mvs)
        mu = tf.concat(mu_list, 1)
        var = tf.concat(var_list, 1)
        return mu, var

    def _conditional_mean(self, F):
        raise NotImplementedError

    def _conditional_variance(self, F):
        raise NotImplementedError

class Ordinal(ScalarLikelihood):
    """
    A likelihood for doing ordinal regression.

    The data are integer values from 0 to k, and the user must specify (k-1)
    'bin edges' which define the points at which the labels switch. Let the bin
    edges be [a₀, a₁, ... aₖ₋₁], then the likelihood is

    p(Y=0|F) = ɸ((a₀ - F) / σ)
    p(Y=1|F) = ɸ((a₁ - F) / σ) - ɸ((a₀ - F) / σ)
    p(Y=2|F) = ɸ((a₂ - F) / σ) - ɸ((a₁ - F) / σ)
    ...
    p(Y=K|F) = 1 - ɸ((aₖ₋₁ - F) / σ)

    where ɸ is the cumulative density function of a Gaussian (the inverse probit
    function) and σ is a parameter to be learned. A reference is:

    @article{chu2005gaussian,
      title={Gaussian processes for ordinal regression},
      author={Chu, Wei and Ghahramani, Zoubin},
      journal={Journal of Machine Learning Research},
      volume={6},
      number={Jul},
      pages={1019--1041},
      year={2005}
    }
    """

    def __init__(self, bin_edges, **kwargs):
        """
        bin_edges is a numpy array specifying at which function value the
        output label should switch. If the possible Y values are 0...K, then
        the size of bin_edges should be (K-1).
        """
        super().__init__(**kwargs)
        self.bin_edges = bin_edges
        self.num_bins = bin_edges.size + 1
        self.sigma = Parameter(1.0, transform=positive())

    def _scalar_log_density(self, F, Y):
        Y = to_default_int(Y)
        scaled_bins_left = tf.concat([self.bin_edges / self.sigma, np.array([np.inf])], 0)
        scaled_bins_right = tf.concat([np.array([-np.inf]), self.bin_edges / self.sigma], 0)
        selected_bins_left = tf.gather(scaled_bins_left, Y)
        selected_bins_right = tf.gather(scaled_bins_right, Y)

        return tf.math.log(
            inv_probit(selected_bins_left - F / self.sigma)
            - inv_probit(selected_bins_right - F / self.sigma)
            + 1e-6
        )

    def _make_phi(self, F):
        """
        A helper function for making predictions. Constructs a probability
        matrix where each row output the probability of the corresponding
        label, and the rows match the entries of F.

        Note that a matrix of F values is flattened.
        """
        scaled_bins_left = tf.concat([self.bin_edges / self.sigma, np.array([np.inf])], 0)
        scaled_bins_right = tf.concat([np.array([-np.inf]), self.bin_edges / self.sigma], 0)
        return inv_probit(scaled_bins_left - tf.reshape(F, (-1, 1)) / self.sigma) - inv_probit(
            scaled_bins_right - tf.reshape(F, (-1, 1)) / self.sigma
        )

    def _conditional_mean(self, F):
        phi = self._make_phi(F)
        Ys = tf.reshape(np.arange(self.num_bins, dtype=default_float()), (-1, 1))
        return tf.reshape(tf.linalg.matmul(phi, Ys), tf.shape(F))

    def _conditional_variance(self, F):
        phi = self._make_phi(F)
        Ys = tf.reshape(np.arange(self.num_bins, dtype=default_float()), (-1, 1))
        E_y = phi @ Ys
        E_y2 = phi @ (Ys ** 2)
        return tf.reshape(E_y2 - E_y ** 2, tf.shape(F))

class MonteCarloLikelihood(Likelihood):
    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self.num_monte_carlo_points = 100

    def _mc_quadrature(self, funcs, Fmu, Fvar, logspace: bool = False, epsilon=None, **Ys):
        return ndiag_mc(funcs, self.num_monte_carlo_points, Fmu, Fvar, logspace, epsilon, **Ys)

    def _predict_mean_and_var(self, Fmu, Fvar, epsilon=None):
        r"""
        Given a Normal distribution for the latent function,
        return the mean of Y

        if
            q(f) = N(Fmu, Fvar)

        and this object represents

            p(y|f)

        then this method computes the predictive mean

           ∫∫ y p(y|f)q(f) df dy

        and the predictive variance

           ∫∫ y² p(y|f)q(f) df dy  - [ ∫∫ y p(y|f)q(f) df dy ]²

        Here, we implement a default Monte Carlo routine.
        """
        integrand2 = lambda *X: self.conditional_variance(*X) + tf.square(self.conditional_mean(*X))
        E_y, E_y2 = self._mc_quadrature(
            [self.conditional_mean, integrand2], Fmu, Fvar, epsilon=epsilon
        )
        V_y = E_y2 - tf.square(E_y)
        return E_y, V_y  # [N, D]

    def _predict_log_density(self, Fmu, Fvar, Y, epsilon=None):
        r"""
        Given a Normal distribution for the latent function, and a datum Y,
        compute the log predictive density of Y.

        i.e. if
            q(f) = N(Fmu, Fvar)

        and this object represents

            p(y|f)

        then this method computes the predictive density

            log ∫ p(y=Y|f)q(f) df

        Here, we implement a default Monte Carlo routine.
        """
        return tf.reduce_sum(
            self._mc_quadrature(self.log_prob, Fmu, Fvar, Y=Y, logspace=True, epsilon=epsilon),
            axis=-1,
        )

    def _variational_expectations(self, Fmu, Fvar, Y, epsilon=None):
        r"""
        Compute the expected log density of the data, given a Gaussian
        distribution for the function values.

        if
            q(f) = N(Fmu, Fvar)  - Fmu: [N, D]  Fvar: [N, D]

        and this object represents

            p(y|f)  - Y: [N, 1]

        then this method computes

           ∫ (log p(y|f)) q(f) df.

        Here, we implement a default Monte Carlo quadrature routine.
        """
        return tf.reduce_sum(
            self._mc_quadrature(self.log_prob, Fmu, Fvar, Y=Y, epsilon=epsilon), axis=-1
        )

class GaussianMC(MonteCarloLikelihood, Gaussian):
    """
    Stochastic version of Gaussian likelihood for comparison.
    """

    pass

class Softmax(MonteCarloLikelihood):
    """
    The soft-max multi-class likelihood.  It can only provide a stochastic
    Monte-Carlo estimate of the variational expectations term, but this
    added variance tends to be small compared to that due to mini-batching
    (when using the SVGP model).
    """

    def __init__(self, num_classes, **kwargs):
        super().__init__(latent_dim=num_classes, observation_dim=None, **kwargs)
        self.num_classes = self.latent_dim

    def _log_prob(self, F, Y):
        return -tf.nn.sparse_softmax_cross_entropy_with_logits(logits=F, labels=Y[:, 0])

    def _conditional_mean(self, F):
        return tf.nn.softmax(F)

    def _conditional_variance(self, F):
        p = self.conditional_mean(F)
        return p - p ** 2


""" hat_model.py models the environment and components for heliotropic hat
    in 3D.

    Author: Jonathon Sather
    Last updated: 4/15/2017
"""

import numpy as np
import pdb

class DirectionalSource:
    """ Directional light source w/ direction and intensity. """

    def __init__(self, (dirX, dirY, dirZ)):
        self.direction = np.array([dirX, dirY, dirZ])
        self.intensity = np.linalg.norm(self.direction)

class PointSource:
    """ Point light source w/ position and intensity. """

    def __init__(self, (x, y, z), intensity):
        self.position = np.array([x, y, z])
        self.intensity = intensity

class LightSensor:
    """ Light sensor w/ value proportional to light intensity. """

    def __init__(self, (x, y, z), (dirX, dirY, dirZ)):
        self.position = np.array([x, y, z])
        self.orientation = (np.array([dirX, dirY, dirZ]) /
         np.linalg.norm(np.array([dirX, dirY, dirZ]))) # Ensure unit magnitude
        self.value = 0
        self.color = np.array([255, 255, 0])

class Hat:
    """ A standard baseball hat with spacial orientation and bill position.
        Bill is represented as circular w/ offset from center. Hat class also
        stores light sensors attached to hat as a list of light sensor objects.
    """

    def __init__(self, (x, y, z), theta, scale, restricted=True):
        self.billDiam = 1.5 * scale
        self.billOffset = 0.9 * scale
        self.color = np.array([0, 0, 255])
        self.diam = 2 * scale
        self.lightSensors = []
        self.position = np.array([x, y, z])
        self.restricted = restricted    # Restricted bill rotation
        self.theta = theta
        self.thetaDot = 2 * np.pi / 100 # Constant rotational speed

    def getBillCenter(self):
        """ Method to find the center of the bill based on angular position
            and hat geometry.
        """

        center = self.position + np.array([np.cos(self.theta) *
                 self.billOffset, np.sin(self.theta) * self.billOffset, 0])
        return center

    def getLSValues(self, mode='no_noise'):
        """ Fetches values of light sensors currently attached to hat. """

        num_ls = len(self.lightSensors)
        ls_vals = np.empty((1, num_ls))

        for ls in range(num_ls):
            ls_vals[0,ls] = self.lightSensors[ls].value

        if mode == 'add_noise':
            ls_vals += np.random.normal(scale=0.00001,size=(1,num_ls))

        return ls_vals

    def get_state(self, pot=True, actions=None):
        """ Returns current state of hat (light sensor values + bill angle in
            numpy array).
        """

        if pot: # State w/ potentiometer
            state = np.transpose(np.hstack((self.getLSValues(),
                                 np.array([[self.theta]]))))
        else:   # State w/ last actions
            state = np.transpose(np.hstack((self.getLSValues(),
                                 np.array([actions]))))
        return state

    def includeLightSensors(self, quantity):
        """ Adds <quantity> light sensors evenly spaced around hat base.
            Note that this method gets rid of any previously included
            light sensors.
        """

        self.lightSensors = []
        sectorAngle = 2 * np.pi / quantity

        for i in range(quantity):
            (dx, dy, dz) = (np.cos(sectorAngle * i), np.sin(sectorAngle *
             i), 0)
            (x, y, z) = (self.position[0] + dx, self.position[1] + dy,
                         self.position[2] + dz)

            self.lightSensors.append(LightSensor((x, y, z), (dx, dy, dz)))

    def rotateBill(self, direction):
        """ Rotates bill based on input direction and angular velocity. """

        update = self.theta + direction * self.thetaDot
        if self.restricted:
            if update < 0 or update > (2 * np.pi):  # Boundary at 0/2*pi
                pass
            else:                                   # No boundary issues
                self.theta = update
        else:
            self.theta = update % (2 * np.pi)

    def rotateBillCCW(self):
        """ Rotates bill counter-clockwise. Includes stop at 0 radians by
            default. Return 1 if non-restricted rotation. 0 if restricted.
        """

        update = self.theta - self.thetaDot
        rotation = 1

        if self.restricted:
            if update < 0:
                rotation = 0
            else:
                self.theta = update
        else:
            self.theta = self.theta % (2 * np.pi)

        return rotation

    def rotateBillCW(self):
        """ Rotates bill clockwise. Includes stop at 0 radians by default.
            Return 1 if non-restricted rotation. 0 if restricted.
        """

        update = self.theta + self.thetaDot
        rotation = 1

        if self.restricted:
            if update > (2 * np.pi):   # Overflow pi to -pi
                rotation = 0
            else:
                self.theta = update
        else:
            self.theta = update % (2 * np.pi)

        return rotation

    def updateSpeed(self, thetaDot):
        """ Method to update self.thetadot. """

        self.thetaDot = thetaDot

class Environment:
    """ Environment contains data and regarding the hat environment. """

    def __init__(self, (width, height)):
        self.boundary = np.array([width, height])
        self.color = np.array([255, 255, 255])

        self.hat = None
        self.directionalSources = []
        self.pointSources = []
        self.ambient = 0

    def addHat(self, hat):
        """ Method for adding or replacing hat. """

        self.hat = hat

    def updateColor(self, (r, g, b)):
        """ Updates the color of the environment. """

        self.color = np.array([r, g, b])

    def updateAmbient(self, ambient):
        """ Updates the ambient light value of the environment. """

        self.ambient = ambient

    def addDirectionalSource(self, (dirX, dirY, dirZ)):
        """ Adds a directional source to the environment. """

        self.directionalSources.append(DirectionalSource((dirX, dirY, dirZ)))

    def addPointSource(self, (x, y, z), intensity):
        """ Adds a point light source to the environment. """

        self.pointSources.append(PointSource((x, y, z), intensity))

    def clearSources(self):
        """ Clears the ambient, point and directional light sources from
            the environment.
        """

        self.directionalSources = []
        self.pointSources = []
        self.ambient = 0

    def notObstructedDirectional(self, hat, lightSensor, directionalSource):
        """ Returns 0 if light sensor obstructed from directional source by
            hat bill or base.
        """
        distance = np.linalg.norm(lightSensor.position - hat.getBillCenter())

        # Only account for bill if light directed downwards.
        if np.dot(directionalSource.direction, np.array([0,0,1])) < 0:
            direction_down = 1
        else:
            direction_down = 0

        if (((distance < hat.billDiam / 2) and direction_down) or
         (np.dot(lightSensor.orientation, directionalSource.direction) > 0)):
            return 0

        return 1

    def notObstructedPoint(self, hat, lightSensor, pointSource):
        """ Returns 0 if light sensor obstructed from point source by hat
            bill or base.
        """

        distance = np.linalg.norm(lightSensor.position - hat.getBillCenter())

        dir = lightSensor.position - pointSource.position

        # Only account for bill if light directed downwards.
        if np.dot(dir, np.array([0, 0, 1])) < 0:
            direction_down = 1
        else:
            direction_down = 0

        if (((distance < hat.billDiam / 2) and direction_down) or
         (np.dot(lightSensor.orientation, dir) > 0)):
            return 0

        return 1

    def updateLightSensor(self, lightSensor):
        """ Updates a single light sensor's value, considering ambient, point,
            and directional components.
        """

        # Ambient
        lightSensor.value = self.ambient

        # Point
        for p in self.pointSources:
            #pdb.set_trace()
            inv_square = (1 /
             (np.linalg.norm(p.position - lightSensor.position) ** 2))

            lightSensor.value +=  (
             inv_square * (self.notObstructedPoint(self.hat,
             lightSensor, p) * np.dot(lightSensor.orientation, -
             (lightSensor.position - p.position) * p.intensity /
             np.linalg.norm(lightSensor.position - p.position))))

        # Directional
        for d in self.directionalSources:
            lightSensor.value += (self.notObstructedDirectional(self.hat,
             lightSensor, d) * np.dot(lightSensor.orientation, - d.direction))

    def update(self):
        """ Updates all light sensors on the hat. """

        for lightSensor in self.hat.lightSensors:
            self.updateLightSensor(lightSensor)

if __name__ == '__main__':
    pass


#! /usr/bin/env python
import os
import numpy as np
from PythonPhot import getpsf, aper
from astropy.io import fits as pyfits

def getcamera(imfile):
    """ Determine the camera name (e.g. ACS-WFC or WFC3-IR)
    from the header of the given image.  The input imfile
    may be a string giving a fits filename, a pyfits hdu object,
    or a pyfits header.
    :param imfile:
    :return:
    """
    import hstphot
    hdr = hstphot.getheader(imfile)
    instrument, detector = '', ''
    if 'CAMERA' in hdr:
        instrument = hdr['CAMERA']
        detector = ''
    elif 'INSTRUME' in hdr:
        instrument = hdr['INSTRUME']
    if 'DETECTOR' in hdr:
        detector = hdr['DETECTOR']
    camera = '-'.join([instrument, detector]).rstrip('-')
    return camera

def mkpsfmodel_stdstar(psfimdir="./", psfrad=0.3, fitrad=0.15,
                       starname='g191b2b',
                       pixscale=0.03,
                       bandlist=['f435w', 'f606w', 'f814w'],
                       verbose=True):
    """ Construct a set of psf models from drizzled HST images of a
    standard star, with the star in the center of the image.

    :param psfimdir: psf image directory (where the input std star images are
       and also where the output psf model images will be put)
    :param psfrad: the scalar radius, in arcsec, of the circular area within
                which the PSF will be defined.  This should be slightly larger
                than the radius of the brightest star that one will be
                interested in.
    :param fitrad: the scalar radius, in arcsec, of the circular area used in
               the least-square star fits.  Stetson suggest that fitrad should
               approximately equal to the FWHM, slightly less for crowded
               fields.  (fitrad must be smaller than psfrad.)

    :return : None
    """
    # radius of aperture and sky annulus for first-pass aperture photometry,
    # used to scale the psf
    aparcsec = 2.5
    skyradarcsec = np.array([2.5, 3.5])

    for band in bandlist:

        # convert user-supplied radii from arcsec to pixels
        appix = aparcsec / pixscale
        skyradpix = skyradarcsec / pixscale
        psfradpix = psfrad / pixscale
        fitradpix = fitrad / pixscale

        inputfile = os.path.join(
            psfimdir, '%s.e00/%s_%s_e00_reg_drz_sci.fits' % (
                starname, starname, band))
        outputfile = os.path.join(
            psfimdir, '%s_%s_%2imas_psf_model.fits' % (
                starname, band, int(pixscale * 1000)))

        hdulist = pyfits.open(inputfile)
        hdr = hdulist[0].header
        imdat = hdulist[0].data
        if 'FILTER1' in hdr:
            if 'CLEAR' in hdr['FILTER1']:
                filtname = hdr['FILTER2']
            else:
                filtname = hdr['FILTER1']
        else:
            filtname = hdr['FILTER']
        camera = getcamera(hdr)

        # Define the conversion factor from the values in this image
        # to photons : photons per ADU.
        if 'BUNIT' not in hdr:
            if camera == 'WFC3-IR' and 'EXPTIME' in hdr:
                phpadu = hdr['EXPTIME']
            else:
                phpadu = 1
        elif hdr['BUNIT'].lower() in ['cps', 'electrons/s']:
            phpadu = hdr['EXPTIME']
        elif hdr['BUNIT'].lower() in ['counts', 'electrons']:
            phpadu = 1
        assert (phpadu is not None), \
            "Can't determine units from the image header."

        rdnoise = 0
        if 'READNSEA' in hdr:
            rdnoise = np.mean([hdr[key] for key in hdr.keys()
                               if key.startswith('READNSE')])

        xmid = np.array([hdr['NAXIS1'] / 2.]) - 1
        ymid = np.array([hdr['NAXIS2'] / 2.]) - 1
        xpos, ypos = hstphot.getxycenter(
            inputfile, xmid, ymid, ext=0, radec=False,
            fitsconvention=False, verbose=True)
        if verbose:
            print("PSF recentering : (%.2f,%.2f) ==> (%.2f,%.2f)" % (
                xmid, ymid, xpos, ypos))
        xpos = np.array([xpos])
        ypos = np.array([ypos])

        idpsf = np.arange(len(xpos))
        image = pyfits.getdata(inputfile)

        # run aper to get mags and sky values for specified coords
        mag, magerr, flux, fluxerr, sky, skyerr, badflag, outstr = \
            aper.aper(image, xpos, ypos, phpadu=phpadu, apr=appix,
                      zeropoint=25,
                      skyrad=skyradpix, badpix=[-12000, 60000], exact=True)

        # use the star at those coords to generate a PSF model
        gauss, psf, psfmag = getpsf.getpsf(image, xpos, ypos, mag,
                                           np.asfarray([sky]), rdnoise,
                                           phpadu,
                                           idpsf, psfradpix, fitradpix,
                                           outputfile, zeropoint=25,
                                           debug=False)
        if verbose:
            print("PSF image written to %s" % outputfile)
    return

def bin_image_data(arr, binfactor):
    # bin an array by reshaping it into a higher-order array and taking
    # the sum over the new dimension
    old_shape = arr.shape
    new_shape = np.array(old_shape) / binfactor
    shape = (new_shape[0], arr.shape[0] // new_shape[0],
             new_shape[1], arr.shape[1] // new_shape[1])
    return arr.reshape(shape).sum(-1).sum(1)

def mkpsfmodel(psfimage, psfrad=0.6, fitrad=0.3, pixscale=0.03, binning=None,
               phpadu=1, rdnoise=0, mag=25, zeropoint=25, sky=0,
               verbose=True):
    """ Construct a psf model from drizzled HST images of a
    standard star or composite star, with the star in the center of the image.

    :param psfimage: filename of the .fits file with the star at the center in
               image array 0.
    :param psfrad: the scalar radius, in arcsec, of the circular area within
                which the PSF will be defined.  This should be slightly larger
                than the radius of the brightest star that one will be
                interested in measuring photometry for.
    :param fitrad: the scalar radius, in arcsec, of the circular area used in
               the least-square star fits.  Stetson suggest that fitrad should
               approximately equal to the FWHM, slightly less for crowded
               fields.  (fitrad must be smaller than psfrad.)
    :param pixscale: the pixel scale of the input image [arcseconds per pixel]
    :param phpadu:  the "gain" of the input image in photons per ADU  [?]
    :param rdnoise: the readnoise of the input image, in ADU [?]
    :param mag: the magnitude of the star in the input image
    :param zeropoint: the zero point of the input image (magnitude that
             produces a flux of 1 [ADU per second?]
    :param sky: the sky brightness of the input image, in ADU per second [?]
    :return : None
    """
    # TODO : currently the units in the __doc__ text are guesses. Need to
    # review the code  in the PythonPhot source and check
    import hstphot

    # convert user-supplied radii from arcsec to pixels
    psfradpix = psfrad / pixscale
    fitradpix = fitrad / pixscale

    outputfile = psfimage.replace('.fits', '_model.fits')
    hdulist = pyfits.open(psfimage)
    hdr = hdulist[0].header
    imdat = hdulist[0].data

    xmid = np.array([hdr['NAXIS1'] / 2.]) - 1
    ymid = np.array([hdr['NAXIS2'] / 2.]) - 1
    xpos, ypos = hstphot.getxycenter(
        psfimage, xmid, ymid, ext=0, radec=False,
        fitsconvention=False, verbose=True)
    if verbose:
        print("PSF recentering : (%.2f,%.2f) ==> (%.2f,%.2f)" % (
            xmid, ymid, xpos, ypos))
    xpos = np.array([xpos])
    ypos = np.array([ypos])
    mag = np.array([mag])
    idpsf = np.arange(len(xpos))

    if binning is not None:
        assert isinstance(binning, int)
        assert not (imdat.shape[0] % binning)
        assert not (imdat.shape[1] % binning)
        imdat = bin_image_data(imdat, binning)
        xpos /= binning
        ypos /= binning
        psfradpix /= binning
        fitradpix /= binning
        print("Binning input image %i x %i" % (binning, binning))

    # use the star at those coords to generate a PSF model
    gauss, psf, psfmag = getpsf.getpsf(
        imdat, xpos, ypos, mag,  np.asfarray([sky]), rdnoise,
        phpadu, idpsf, psfradpix, fitradpix,  outputfile, zeropoint=zeropoint,
        debug=False)
    if verbose:
        print("PSF image written to %s" % outputfile)
    return outputfile, gauss, psf, psfmag

def main():
    import argparse

    parser = argparse.ArgumentParser(
        description="Make a PythonPhot psf model from the a single-star image")

    # Required positional argument
    parser.add_argument('inputimage',
                        help='FITS file with a single star at the center.')

    # optional keyword arguments
    parser.add_argument('--pixscale', type=float,
                        default=0.03,
                        help="arcseconds per pixel.")

    argv = parser.parse_args()

    mkpsfmodel(argv.inputimage, pixscale=argv.pixscale)

if __name__ == '__main__':
    main()


# /mpcutilities/mpcutilities/leapsec.py
"""

--------------------------------------------------------------

Nov 2018

Payne

Original code by Sonia Keys

Provides leap second information by year.  It requires a file
of leap second data in the format prepared by Judah Levine at NIST.
See https://www.ietf.org/timezones/data/leap-seconds.list for example.

*** Payne: I believe this code to be unnecessary / unused ***

*** Payne: I am keeping it for archival purposes ***

--------------------------------------------------------------

"""

# Import third-party packages
# --------------------------------------------------------------
import numpy as np
from pkg_resources import resource_filename

__all__ = ["LeapSeconds"]

class LeapSeconds:
    """loads leap second data from a file.

    Parameters
    ----------
    fn : string
        File name of leap second data, by default leap-seconds.list.
    jd_ref_utc : float
        Reference date, JD at which to start counting leap seconds.
        The default value of 2415020.5 is year 1900.0.
    """

    def __init__(self, fn=resource_filename('mpcutilities','data/leap-seconds.list'), jd_ref_utc=2415020.5):
        # Create two numpy arrays, secSinceArray and leapSecArray,
        # that I can use to make a look-up table.
        # This is just reading Judah Levine's file.
        self.jd_ref_utc = jd_ref_utc
        secSinceList = []
        leapSecList = []
        with open(fn) as f:
            for line in f:
                if not line.startswith('#'):
                    secondsSince, leapSeconds = line.rsplit('#')[0].split()
                    secSinceList.append(int(secondsSince))
                    leapSecList.append(int(leapSeconds))
        self.secSinceArray = np.array(secSinceList)
        self.leapSecArray = np.array(leapSecList)

    def getLeapSeconds(self, jd_utc):
        """
        gets number of leap seconds since the reference Julian date.

        Parameters
        ----------
        jd_utc : float
            JD of end of period since reference date.

        Returns
        -------
        int
            number of leap seconds since reference date.
        """
        # Given a Julian Date in UTC, determine the number of seconds
        # that have elapsed since the reference time.  Then find the
        # correspoding number of leap seconds.
        max_idx = self.secSinceArray.shape[0] - 1
        secsSince = self.secondsSinceRef(jd_utc)
        idx = np.searchsorted(self.secSinceArray, int(secsSince),
            side='right') - 1
        if idx < 0:
            return 0
        if idx > max_idx:
            return self.leapSecArray[max_idx]
        else:
            return self.leapSecArray[idx]

    def secondsSinceRef(self, jd_utc):
        """Determines the seconds (not just leap seconds) that have elapsed
        since the reference Julian date.

        Parameters
        ----------
        jd_utc : float
            JD of end of period since reference date.

        Returns
        -------
        float
            number of seconds since reference date.
        """
        secondsSince = (jd_utc - self.jd_ref_utc) * 24.0 * 60 * 60
        return secondsSince


#!/usr/bin/env python3
import configparser
import os, h5py, sys
import numpy as np
from matplotlib import pyplot as plt
thisDir = os.path.dirname(os.path.abspath(__file__))
parentDir = os.path.dirname(thisDir)
sys.path.insert(0,parentDir)
from pyCRTM import pyCRTM, profilesCreate

def main(coefficientPath, sensor_id):
    thisDir = os.path.dirname(os.path.abspath(__file__))
    cases = os.listdir( os.path.join(thisDir,'data') )
    cases.sort()
    # create 4 profiles for each of the 4 cases
    profiles = profilesCreate( 4, 92 )
    storedTb = []
    storedEmis = []
    # populate the cases, and previously calculated Tb from crtm test program.
    for i,c in enumerate(cases):
        h5 = h5py.File(os.path.join(thisDir,'data',c) , 'r')
        profiles.Angles[i,0] = h5['zenithAngle'][()]
        profiles.Angles[i,1] = 999.9
        profiles.Angles[i,2] = 100.0  # 100 degrees zenith below horizon.
        profiles.Angles[i,3] = 0.0 # zero solar azimuth
        profiles.Angles[i,4] = h5['scanAngle'][()]
        profiles.DateTimes[i,0] = 2001
        profiles.DateTimes[i,1] = 1
        profiles.DateTimes[i,2] = 1
        profiles.Pi[i,:] = np.asarray(h5['pressureLevels'] )
        profiles.P[i,:] = np.asarray(h5['pressureLayers'][()])
        profiles.T[i,:] = np.asarray(h5['temperatureLayers'])
        profiles.Q[i,:] = np.asarray(h5['humidityLayers'])
        profiles.O3[i,:] = np.asarray(h5['ozoneConcLayers'])
        profiles.clouds[i,:,0,0] = np.asarray(h5['cloudConcentration'])
        profiles.clouds[i,:,0,1] = np.asarray(h5['cloudEffectiveRadius'])
        profiles.aerosols[i,:,0,0] = np.asarray(h5['aerosolConcentration'])
        profiles.aerosols[i,:,0,1] = np.asarray(h5['aerosolEffectiveRadius'])
        profiles.aerosolType[i] = h5['aerosolType'][()]
        profiles.cloudType[i] = h5['cloudType'][()]
        profiles.cloudFraction[i,:] = h5['cloudFraction'][()]
        profiles.climatology[i] = h5['climatology'][()]
        profiles.surfaceFractions[i,:] = h5['surfaceFractions']
        profiles.surfaceTemperatures[i,:] = h5['surfaceTemperatures']
        profiles.S2m[i,1] = 33.0 # just use salinity out of S2m for the moment.
        profiles.windSpeed10m[i] = 5.0
        profiles.LAI[i] = h5['LAI'][()]
        profiles.windDirection10m[i] = h5['windDirection10m'][()]
        # land, soil, veg, water, snow, ice
        profiles.surfaceTypes[i,0] = h5['landType'][()]
        profiles.surfaceTypes[i,1] = h5['soilType'][()]
        profiles.surfaceTypes[i,2] = h5['vegType'][()]
        profiles.surfaceTypes[i,3] = h5['waterType'][()]
        profiles.surfaceTypes[i,4] = h5['snowType'][()]
        profiles.surfaceTypes[i,5] = h5['iceType'][()]
        storedTb.append(np.asarray(h5['Tb']))
        storedEmis.append(np.asarray(h5['emissivity_atms']))
        h5.close()

    crtmOb = pyCRTM()
    crtmOb.profiles = profiles
    crtmOb.coefficientPath = pathInfo['CRTM']['coeffs_dir']
    crtmOb.sensor_id = sensor_id
    crtmOb.nThreads = 4

    crtmOb.loadInst()

    crtmOb.runDirect()
    forwardTb = crtmOb.Bt
    forwardEmissivity = crtmOb.surfEmisRefl[0,:]
    crtmOb.surfEmisRefl = []

    crtmOb.runK()
    kTb = crtmOb.Bt
    kEmissivity = crtmOb.surfEmisRefl[0,:]

    if ( all( np.abs( forwardTb.flatten() - np.asarray(storedTb).flatten() ) <= 1e-5)  and all( np.abs( kTb.flatten() - np.asarray(storedTb).flatten() ) <= 1e-5) ):
        print("Yay! all values are close enough to what CRTM test program produced!")
    else:
        print("Boo! something failed. Look at cris plots")
        wavenumbers = np.zeros([4,1305])
        wavenumbers[0:4,:] = np.linspace(1,1306,1305)
        plt.figure()
        plt.plot(wavenumbers.T,forwardTb.T-np.asarray(storedTb).T )
        plt.legend(['1','2','3','4'])
        plt.savefig(os.path.join(thisDir,'cris'+'_spectrum_forward.png'))
        plt.figure()
        plt.plot(wavenumbers.T,forwardEmissivity.T-np.asarray(storedEmis).T)
        plt.savefig(os.path.join(thisDir,'cris'+'_emissivity_forward.png'))

        plt.figure()
        plt.plot(wavenumbers.T,kTb.T-np.asarray(storedTb).T)
        plt.savefig(os.path.join(thisDir,'cris'+'_spectrum_k.png'))
        plt.figure()
        plt.plot(wavenumbers.T,kEmissivity.T-np.asarray(storedEmis).T)
        plt.savefig(os.path.join(thisDir,'cris'+'_emissivity_k.png'))
        sys.exit("Boo! didn't pass tolerance with CRTM test program.")

if __name__ == "__main__":
    pathInfo = configparser.ConfigParser()
    pathInfo.read( os.path.join(parentDir,'crtm.cfg') )
    coefficientPath = pathInfo['CRTM']['coeffs_dir']
    sensor_id = 'cris_npp'
    main(coefficientPath, sensor_id)


import ctypes
import os
import numpy as np
import time

if os.name == 'nt':
    import msvcrt
    def getch():
        return msvcrt.getch().decode()
else:
    import sys, tty, termios
    fd = sys.stdin.fileno()
    old_settings = termios.tcgetattr(fd)
    def getch():
        try:
            tty.setraw(sys.stdin.fileno())
            ch = sys.stdin.read(1)
        finally:
            termios.tcsetattr(fd, termios.TCSADRAIN, old_settings)
        return ch

from dynamixel_sdk import *                    # Uses Dynamixel SDK library

# Length of the arms
a0 				= 10
a1 				= 5
a2 				= 20
# Angles of the joints
t0 				= 0
t1				= 0
t2				= 0

# Angles in radians
t0 		= t0/180*np.pi
t1 		= t1/180*np.pi
t2 		= t2/180*np.pi

# Control Table address
ADDR_TORQUE_ENABLE                 = 64
ADDR_PROFILE_VELOCITY              = 112
ADDR_PROFILE_ACCELERATION          = 108
ADDR_GOAL_POSITION                 = 116
ADDR_PRESENT_POSITION              = 132
ADDR_DRIVE_MODE                    = 10
ADDR_MOVING_STATUS                 = 123

# Data Byte Length
LEN_GOAL_POSITION       = 4
LEN_PRESENT_POSITION    = 4

# Protocol version
PROTOCOL_VERSION            = 2.0               # See which protocol version is used in the Dynamixel

# Default setting
DXL1_ID                     = 71                 # Dynamixel#1 ID : 1
DXL2_ID                     = 72                 # Dynamixel#1 ID : 2
BAUDRATE                    = 1000000             # Dynamixel default baudrate : 57600
DEVICENAME                  = '/dev/ttyUSB0'    # Check which port is being used on your controller
                                                # ex) Windows: "COM1"   Linux: "/dev/ttyUSB0" Mac: "/dev/tty.usbserial-*"

TORQUE_ENABLE               = 1                 # Value for enabling the torque
TORQUE_DISABLE              = 0                 # Value for disabling the torque
#DXL1_MIN_POSITION_VALUE     = -1000			# Dynamixel 1 will rotate between this value
#DXL1_MAX_POSITION_VALUE     = 1000		# and this value
#DXL2_MIN_POSITION_VALUE     = -2000 		# Dynamixel 2 will rotate between this value
#DXL2_MAX_POSITION_VALUE     = 2000		# and this value
DXL_MOVING_STATUS_THRESHOLD = 20 		# Dynamixel moving status threshold

index = 0
dxl1_goal_position = [DXL1_MIN_POSITION_VALUE, DXL1_MAX_POSITION_VALUE]         # Goal position of dynamixel 1
dxl2_goal_position = [DXL2_MIN_POSITION_VALUE, DXL2_MAX_POSITION_VALUE]         # Goal position of dynamixel 2

# Initialize PortHandler instance
# Set the port path
# Get methods and members of PortHandlerLinux or PortHandlerWindows
portHandler = PortHandler(DEVICENAME)

# Initialize PacketHandler instance
# Set the protocol version
# Get methods and members of Protocol1PacketHandler or Protocol2PacketHandler
packetHandler = PacketHandler(PROTOCOL_VERSION)

# Initialize GroupBulkWrite instance
groupBulkWrite = GroupBulkWrite(portHandler, packetHandler)

# Initialize GroupBulkRead instace for Present Position
groupBulkRead = GroupBulkRead(portHandler, packetHandler)

# Open port
if portHandler.openPort():
    print("Succeeded to open the port")
else:
    print("Failed to open the port")
    print("Press any key to terminate...")
    getch()
    quit()

# Set port baudrate
if portHandler.setBaudRate(BAUDRATE):
    print("Succeeded to change the baudrate")
else:
    print("Failed to change the baudrate")
    print("Press any key to terminate...")
    getch()
    quit()

# Enable Dynamixel Torque DXL1
dxl_comm_result, dxl_error = packetHandler.write1ByteTxRx(portHandler, DXL1_ID, ADDR_TORQUE_ENABLE, TORQUE_ENABLE)
if dxl_comm_result != COMM_SUCCESS:
    print("%s" % packetHandler.getTxRxResult(dxl_comm_result))
elif dxl_error != 0:
    print("%s" % packetHandler.getRxPacketError(dxl_error))
else:
    print("Dynamixel 1 has been successfully connected")

# Enable Dynamixel Torque DXL2
dxl_comm_result, dxl_error = packetHandler.write1ByteTxRx(portHandler, DXL2_ID, ADDR_TORQUE_ENABLE, TORQUE_ENABLE)
if dxl_comm_result != COMM_SUCCESS:
    print("%s" % packetHandler.getTxRxResult(dxl_comm_result))
elif dxl_error != 0:
    print("%s" % packetHandler.getRxPacketError(dxl_error))
else:
    print("Dynamixel 2 has been successfully connected")

# Add parameter storage for Dynamixel#1 present position
dxl_addparam_result = groupBulkRead.addParam(DXL1_ID, ADDR_PRESENT_POSITION, LEN_PRESENT_POSITION)
if dxl_addparam_result != True:
    print("[ID:%03d] groupBulkRead addparam failed" % DXL1_ID)
    quit()

# Add parameter storage for Dynamixel#2 present position
dxl_addparam_result = groupBulkRead.addParam(DXL2_ID, ADDR_PRESENT_POSITION, LEN_PRESENT_POSITION)
if dxl_addparam_result != True:
    print("[ID:%03d] groupBulkRead addparam failed" % DXL2_ID)
    quit()

def end_eff_pos(t0,t1,te):
    # Parameter table with columns " theta, alpha, r or a, d

    param_tab = [[t0 ,      0         ,  a0   ,  0  ],
                 [t1 ,      0         ,  a1   ,  0  ],
                 [t2 ,      0         ,  a2   ,  0  ]]

    # Transformation matrix

    i=0
    T0_1 = [[np.cos(param_tab[i][0]) , -np.cos(param_tab[i][1])*np.sin(param_tab[i][0]), np.sin(param_tab[i][1])*np.sin(param_tab[i][0])  , param_tab[i][2]*np.cos(param_tab[i][0])],
            [np.sin(param_tab[i][0]) , np.cos(param_tab[i][1])*np.cos(param_tab[i][0]) , -np.sin(param_tab[i][1])*np.cos(param_tab[i][0]) , param_tab[i][2]*np.sin(param_tab[i][0])],
            [           0            ,            np.sin(param_tab[i][1])              ,       np.cos(param_tab[i][1])                    ,     param_tab[i][3]                    ],
            [0  ,  0  ,  0  ,  1 ]]

    i=1
    T1_2 = [[np.cos(param_tab[i][0]) , -np.cos(param_tab[i][1])*np.sin(param_tab[i][0]), np.sin(param_tab[i][1])*np.sin(param_tab[i][0])  , param_tab[i][2]*np.cos(param_tab[i][0])],
            [np.sin(param_tab[i][0]) , np.cos(param_tab[i][1])*np.cos(param_tab[i][0]) , -np.sin(param_tab[i][1])*np.cos(param_tab[i][0]) , param_tab[i][2]*np.sin(param_tab[i][0])],
            [           0            ,            np.sin(param_tab[i][1])              ,       np.cos(param_tab[i][1])                    ,     param_tab[i][3]                    ],
            [0  ,  0  ,  0  ,  1 ]]

    i=2
    T2_3 = [[np.cos(param_tab[i][0]) , -np.cos(param_tab[i][1])*np.sin(param_tab[i][0]), np.sin(param_tab[i][1])*np.sin(param_tab[i][0])  , param_tab[i][2]*np.cos(param_tab[i][0])],
            [np.sin(param_tab[i][0]) , np.cos(param_tab[i][1])*np.cos(param_tab[i][0]) , -np.sin(param_tab[i][1])*np.cos(param_tab[i][0]) , param_tab[i][2]*np.sin(param_tab[i][0])],
            [           0            ,            np.sin(param_tab[i][1])              ,       np.cos(param_tab[i][1])                    ,     param_tab[i][3]                    ],
            [0  ,  0  ,  0  ,  1 ]]

while 1:
    print("Press any key to continue! (or press ESC to quit!)")
    if getch() == chr(0x1b):
        break

    # Allocate goal position value into byte array
    param_goal_position_dxl1 = [DXL_LOBYTE(DXL_LOWORD(dxl1_goal_position[index])), DXL_HIBYTE(DXL_LOWORD(dxl1_goal_position[index])), DXL_LOBYTE(DXL_HIWORD(dxl1_goal_position[index])), DXL_HIBYTE(DXL_HIWORD(dxl1_goal_position[index]))]
    param_goal_position_dxl2 = [DXL_LOBYTE(DXL_LOWORD(dxl2_goal_position[index])), DXL_HIBYTE(DXL_LOWORD(dxl2_goal_position[index])), DXL_LOBYTE(DXL_HIWORD(dxl2_goal_position[index])), DXL_HIBYTE(DXL_HIWORD(dxl2_goal_position[index]))]

    # Add Dynamixel#1 goal position value to the Bulkwrite parameter storage
    dxl_addparam_result = groupBulkWrite.addParam(DXL1_ID, ADDR_GOAL_POSITION, LEN_GOAL_POSITION, param_goal_position_dxl1)
    if dxl_addparam_result != True:
        print("[ID:%03d] groupBulkWrite addparam failed" % DXL1_ID)
        quit()
    # Add Dynamixel#2 goal position value to the Bulkwrite parameter storage
    dxl_addparam_result = groupBulkWrite.addParam(DXL2_ID, ADDR_GOAL_POSITION, LEN_GOAL_POSITION, param_goal_position_dxl2)
    if dxl_addparam_result != True:
        print("[ID:%03d] groupBulkWrite addparam failed" % DXL2_ID)
        quit()

    # Bulkwrite goal position values for both the dynamixels
    dxl_comm_result = groupBulkWrite.txPacket()
    if dxl_comm_result != COMM_SUCCESS:
        print("%s" % packetHandler.getTxRxResult(dxl_comm_result))

    # Clear bulkwrite parameter storage
    groupBulkWrite.clearParam()

    while 1:
        # Bulkread present positions of both dynamixels
        dxl_comm_result = groupBulkRead.txRxPacket()
        if dxl_comm_result != COMM_SUCCESS:
            print("%s" % packetHandler.getTxRxResult(dxl_comm_result))

        # Check if groupbulkread data of Dynamixel#1 is available
        dxl_getdata_result = groupBulkRead.isAvailable(DXL1_ID, ADDR_PRESENT_POSITION, LEN_PRESENT_POSITION)
        if dxl_getdata_result != True:
            print("[ID:%03d] groupBulkRead getdata failed" % DXL1_ID)
            quit()

        # Check if groupbulkread data of Dynamixel#2 is available
        dxl_getdata_result = groupBulkRead.isAvailable(DXL2_ID, ADDR_PRESENT_POSITION, LEN_PRESENT_POSITION)
        if dxl_getdata_result != True:
            print("[ID:%03d] groupBulkRead getdata failed" % DXL2_ID)
            quit()

	# Get present position value
        dxl1_present_position = groupBulkRead.getData(DXL1_ID, ADDR_PRESENT_POSITION, LEN_PRESENT_POSITION)
        dxl2_present_position = groupBulkRead.getData(DXL2_ID, ADDR_PRESENT_POSITION, LEN_PRESENT_POSITION)

	number = dxl1_present_position & 0xFFFFFFFF
        dxl1_present_position = ctypes.c_long(number).value

	number = dxl2_present_position & 0xFFFFFFFF
        dxl2_present_position = ctypes.c_long(number).value

	print("[ID:%03d] Present Position : %d \t [ID:%03d] LED Value: %d" % (DXL1_ID, dxl1_present_position, DXL2_ID, dxl2_present_position))

	if ((abs(dxl1_goal_position[index] - dxl1_present_position) > DXL_MOVING_STATUS_THRESHOLD) and (abs(dx2_goal_position[index] - dxl2_present_position) > DXL_MOVING_STATUS_THRESHOLD)):
	    break

    # Change goal position
    if index == 0:
        index = 1
    else:
        index = 0

# Clear bulkread parameter storage
groupBulkRead.clearParam()

# Disable Dynamixel Torque DXL1
dxl_comm_result, dxl_error = packetHandler.write1ByteTxRx(portHandler, DXL1_ID, ADDR_TORQUE_ENABLE, TORQUE_DISABLE)
if dxl_comm_result != COMM_SUCCESS:
    print("%s" % packetHandler.getTxRxResult(dxl_comm_result))
elif dxl_error != 0:
    print("%s" % packetHandler.getRxPacketError(dxl_error))

# Disable Dynamixel Torque DXL2
dxl_comm_result, dxl_error = packetHandler.write1ByteTxRx(portHandler, DXL2_ID, ADDR_TORQUE_ENABLE, TORQUE_DISABLE)
if dxl_comm_result != COMM_SUCCESS:
    print("%s" % packetHandler.getTxRxResult(dxl_comm_result))
elif dxl_error != 0:
    print("%s" % packetHandler.getRxPacketError(dxl_error))

# Close port
portHandler.closePort()


import os

import numpy as np

from astropy import units as u
from astropy import wcs
from astropy.coordinates import EarthLocation
from astropy.coordinates import FK5
from astropy.coordinates import SkyCoord
from astropy.io import fits
from astropy.time import Time
from ccdproc import CCDData
from collections import namedtuple
from skimage.feature import register_translation
from skimage.util import view_as_blocks

from pocs import PanBase
from pocs.utils import images as img_utils

PointingError = namedtuple('PointingError', ['delta_ra', 'delta_dec', 'magnitude'])

class Image(PanBase):

    def __init__(self, fits_file, wcs_file=None):
        """Object to represent a single image from a PANOPTES camera.

        Instantiate the object by providing a .cr2 (or .dng) file.

        Args:
            fits_file (str): Name of FITS file to be read (can be .fz)
            wcs_file (str, optional): Name of FITS file to use for WCS
        """
        super().__init__()
        assert os.path.exists(fits_file), self.logger.warning('File does not exist: {}'.format(fits_file))

        if fits_file.endswith('.fz'):
            fits_file = img_utils.fpack(fits_file, unpack=True)

        assert fits_file.lower().endswith(('.fits')), self.logger.warning('File must end with .fits')

        self.wcs = None
        self._wcs_file = None
        self.fits_file = fits_file

        if wcs_file is not None:
            self.wcs_file = wcs_file
        else:
            self.wcs_file = fits_file

        with fits.open(self.fits_file, 'readonly') as hdu:
            self.header = hdu[0].header
            self.data = hdu[0].data

        assert 'DATE-OBS' in self.header, self.logger.warning('FITS file must contain the DATE-OBS keyword')
        assert 'EXPTIME' in self.header, self.logger.warning('FITS file must contain the EXPTIME keyword')

        self.RGGB = CCDData(data=self.data, unit='adu',
                            meta=self.header,
                            mask=np.zeros(self.data.shape))

        # Location Information
        cfg_loc = self.config['location']
        self.loc = EarthLocation(lat=cfg_loc['latitude'],
                                 lon=cfg_loc['longitude'],
                                 height=cfg_loc['elevation'],
                                 )
        # Time Information
        self.starttime = Time(self.header['DATE-OBS'], location=self.loc)
        self.exptime = float(self.header['EXPTIME']) * u.second
        self.midtime = self.starttime + (self.exptime / 2.0)
        self.sidereal = self.midtime.sidereal_time('apparent')
        self.FK5_Jnow = FK5(equinox=self.midtime)

        # Coordinates from header keywords
        self.header_pointing = None
        self.header_RA = None
        self.header_Dec = None
        self.header_HA = None

        # Coordinates from WCS
        self.pointing = None
        self.RA = None
        self.Dec = None
        self.HA = None

        self.get_header_pointing()
        self.get_wcs_pointing()

        self._luminance = None
        self._pointing = None
        self._pointing_error = None

    @property
    def wcs_file(self):
        """WCS file name

        When setting the WCS file name, the WCS information will be read,
        setting the `wcs` property.
        """
        return self._wcs_file

    @wcs_file.setter
    def wcs_file(self, filename):
        if filename is not None:
            try:
                w = wcs.WCS(filename)
                assert w.is_celestial

                self.wcs = w
                self._wcs_file = filename
            except Exception:
                self.logger.warn("Can't get WCS from FITS file (try solve_field)")

    @property
    def luminance(self):
        """Luminance for the image

        Bin the image 2x2 combining each RGGB set of pixels in to a single
        luminance value.
        """
        if self._luminance is None:
            block_size = (2, 2)
            image_out = view_as_blocks(self.RGGB.data, block_size)

            for i in range(len(image_out.shape) // 2):
                image_out = np.average(image_out, axis=-1)

            self._luminance = image_out

        return self._luminance

    @property
    def pointing_error(self):
        """Pointing error namedtuple (delta_ra, delta_dec, magnitude)

        Returns pointing error information. The first time this is accessed
        this will solve the field if not previously solved.

        Returns:
            namedtuple: Pointing error information
        """
        if self._pointing_error is None:
            assert self.pointing is not None, self.logger.warn("No WCS, can't get pointing_error")
            assert self.header_pointing is not None

            if self.wcs is None:
                self.solve_field()

            mag = self.pointing.separation(self.header_pointing)
            dDec = self.pointing.dec - self.header_pointing.dec
            dRA = self.pointing.ra - self.header_pointing.ra

            self._pointing_error = PointingError(dRA.to(u.degree), dDec.to(u.degree), mag)

        return self._pointing_error

    def get_header_pointing(self):
        """Get the pointing information from the header

        The header should contain the `RA-MNT` and `DEC-MNT` keywords, from which
        the header pointing coordinates are built.
        """
        try:
            self.header_pointing = SkyCoord(ra=float(self.header['RA-MNT']) * u.degree,
                                            dec=float(self.header['DEC-MNT']) * u.degree)

            self.header_RA = self.header_pointing.ra.to(u.hourangle)
            self.header_Dec = self.header_pointing.dec.to(u.degree)

            # Precess to the current equinox otherwise the RA - LST method will be off.
            self.header_HA = self.header_pointing.transform_to(self.FK5_Jnow).ra.to(u.hourangle) - self.sidereal
        except Exception as e:
            self.logger.warning('Cannot get header pointing information: {}'.format(e))

    def get_wcs_pointing(self):
        """Get the pointing information from the WCS

        Builds the pointing coordinates from the plate-solved WCS. These will be
        compared with the coordinates stored in the header.
        """
        if self.wcs is not None:
            ny, nx = self.RGGB.data.shape
            decimals = self.wcs.all_pix2world(nx // 2, ny // 2, 1)

            self.pointing = SkyCoord(ra=decimals[0] * u.degree,
                                     dec=decimals[1] * u.degree)

            self.RA = self.pointing.ra.to(u.hourangle)
            self.Dec = self.pointing.dec.to(u.degree)

            # Precess to the current equinox otherwise the RA - LST method will be off.
            self.HA = self.pointing.transform_to(self.FK5_Jnow).ra.to(u.hourangle) - self.sidereal

    def solve_field(self, **kwargs):
        """ Solve field and populate WCS information

        Args:
            **kwargs (dict): Options to be passed to `get_solve_field`
        """
        solve_info = img_utils.get_solve_field(self.fits_file,
                                               ra=self.header_pointing.ra.value,
                                               dec=self.header_pointing.dec.value,
                                               **kwargs)

        self.wcs_file = solve_info['solved_fits_file']
        self.get_wcs_pointing()

        return solve_info

    def compute_offset(self, ref, units='arcsec', rotation=True):
        """Offset information between this image and a reference

        Args:
            ref (str): Refernce image, either another `Image` instance or a
                filename that will be read
            units (str, optional): Can be either `arcsec` or `pixel`
            rotation (bool, optional): If rotation information should be included,
                defaults to True

        Returns:
            dict: Offset information in key/value pairs
        """
        if isinstance(units, (u.Unit, u.Quantity, u.IrreducibleUnit)):
            units = units.name
        assert units in ['pix', 'pixel', 'arcsec']

        if isinstance(ref, str):
            assert os.path.exists(ref)
            ref = Image(ref)
        assert isinstance(ref, Image)

        offset_pix = compute_offset_rotation(ref.luminance, self.luminance)
        offset_pix['X'] *= 2
        offset_pix['Y'] *= 2

        if self.HA:
            selfHA = self.HA
        else:
            selfHA = self.header_HA
        if self.Dec:
            selfDec = self.Dec
        else:
            selfDec = self.header_Dec
        if ref.HA:
            refHA = ref.HA
        else:
            stime_diff = (self.midtime.sidereal_time('apparent') - ref.midtime.sidereal_time('apparent'))
            refHA = selfHA - stime_diff.to(u.hourangle)

        time_diff = (self.midtime - ref.midtime)

        info = {'image': self.fits_file,
                'time': self.midtime.to_datetime().isoformat(),
                'HA': selfHA.to(u.hourangle).value,
                'HA unit': 'hours',
                'Dec': selfDec.to(u.degree).value,
                'Dec unit': 'deg',

                'refimage': ref.fits_file,
                'reftime': ref.midtime.to_datetime().isoformat(),
                'refHA': refHA.to(u.hourangle).value,

                'dt': time_diff.to(u.second).value,
                'dt unit': 'seconds',
                'angle': offset_pix['angle'].to(u.degree).value,
                'angle unit': 'deg',
                'offset units': units,
                }

        if units in ['pix', 'pixel']:
            info['offsetX'] = offset_pix['X'].to(u.pixel).value
            info['offsetY'] = offset_pix['Y'].to(u.pixel).value
        elif units == 'arcsec':
            deltapix = [offset_pix['X'].to(u.pixel).value,
                        offset_pix['Y'].to(u.pixel).value]
            offset_deg = self.wcs.pixel_scale_matrix.dot(deltapix)
            info['offsetX'] = (offset_deg[0] * u.degree).to(u.arcsecond).value
            info['offsetY'] = (offset_deg[1] * u.degree).to(u.arcsecond).value
        return info

##################################################################################################
# Private Methods
##################################################################################################

    def __str__(self):
        return "{}: {}".format(self.fits_file, self.header_pointing)

def compute_offset_rotation(im, imref, upsample_factor=20, subframe_size=200, corners=True):
    """Determine rotation information between two images

    Detremine the rotation information for the center and, if `corner`, the
    four corner boxes, each of `subframe_size` pixels.

    Args:
        im (numpy.array): Image data
        imref (numpy.array): Comparison image data
        upsample_factor (int, optional): Subpixel fraction to compute
        subframe_size (int, optional): Box size
        corners (bool, optional): If corner boxes should be included, defaults
            to True

    Returns:
        dict: Rotation offset in `X`, `Y`, and `angle`
    """
    assert im.shape == imref.shape
    ny, nx = im.shape

    subframe_half = int(subframe_size / 2)

    # Create the center point for each of our regions
    regions = {'center': (int(nx / 2), int(ny / 2)), }
    offsets = {'center': None, }

    if corners:
        regions.update({
            'upper_right': (int(nx - subframe_half), int(ny - subframe_half)),
            'upper_left': (int(subframe_half), int(ny - subframe_half)),
            'lower_right': (int(nx - subframe_half), int(subframe_half)),
            'lower_left': (int(subframe_half), int(subframe_half)),
        })

        offsets.update({
            'upper_right': None,
            'upper_left': None,
            'lower_right': None,
            'lower_left': None,
        })

    # Get im/imref offsets for each region
    for region, midpoint in regions.items():
        imarr = img_utils.crop_data(im, center=midpoint, box_width=subframe_size)
        imrefarr = img_utils.crop_data(imref, center=midpoint, box_width=subframe_size)

        shifts, err, h = register_translation(imrefarr, imarr, upsample_factor=upsample_factor)
        offsets[region] = shifts

    # Rotate the offsets according to region
    angles = []
    for region in regions.keys():
        if region != 'center':
            offsets[region] -= offsets['center']

            relpos = (regions[region][0] - regions['center'][0],
                      regions[region][1] - regions['center'][1])

            theta1 = np.arctan(relpos[1] / relpos[0])
            theta2 = np.arctan((relpos[1] + offsets[region][1]) / (relpos[0] + offsets[region][0]))
            angles.append(theta2 - theta1)

    angle = np.mean(angles)

    result = {'X': offsets['center'][0] * u.pix,
              'Y': offsets['center'][1] * u.pix,
              'angle': (angle * u.radian).to(u.degree)}

    return result


#!/usr/bin/env python

# A library of stuff used for getting Balrog data from the databse,
# HEALPixifying it, and general munging.

import desdb
import numpy as np
import esutil
import sys
import healpy as hp
import os
import numpy.lib.recfunctions as rf

def GetDepthMap(depth_file):
    map = hp.read_map(depth_file, nest=True)
    nside = hp.npix2nside(map.size)
    return map, nside

def GetPhi(ra):
    return ra * np.pi / 180.0

def GetRa(phi):
    return phi*180.0/np.pi

def GetTheta(dec):
    return (90.0 - dec) * np.pi / 180.0

def GetDec(theta):
    return 90.0 - theta*180.0/np.pi

def GetRaDec(theta, phi):
    return [GetRa(phi), GetDec(theta)]

def GetPix(nside, ra, dec, nest=True):
    phi = GetPhi(ra)
    theta = GetTheta(dec)
    pix = hp.ang2pix(nside, theta, phi, nest=nest)
    return pix

def GetDepthCut(map, nside, ra, dec, depth = 0.0):
    pix = GetPix(nside, ra, dec)
    depths = map[pix]
    ok_depths =  (depths > depth)
    return ok_depths

def ValidDepth(map, nside, arr, rakey='ra', deckey='dec', depth = 0.0):
    ok_depths = GetDepthCut(map, nside, arr[rakey], arr[deckey], depth = depth)
    arr = arr[ok_depths]
    return arr

def InTile(data, ura, udec, rakey='ra', deckey='dec'):
    inside = (data[rakey] > ura[0]) & (data[rakey] < ura[1]) & (data[deckey] > udec[0]) & (data[deckey] < udec[1])
    return inside

def RemoveTileOverlap(tilestuff, data, col='tilename', rakey='ra', deckey='dec'):
    datatile = data[col]
    tiles = np.unique(datatile)
    keep = np.zeros( len(data), dtype=np.bool_)
    for tile in tiles:
        cut = (datatile==tile)
        entry = tilestuff[tile]
        ura = (entry['urall'], entry['uraur'])
        udec = (entry['udecll'], entry['udecur'])
        u = InTile(data[cut], ura, udec, rakey=rakey, deckey=deckey)
        keep[cut] =  u
    return data[keep]

def NoSimFields(band='i'):
    q = """
    SELECT
        balrog_index,
        mag_auto,
        flags
    FROM
        SUCHYTA1.balrog_sva1v2_nosim_%s
    """ %(band)
    return q

def SimFields(band='i',table='sva1v2'):
    q = """
    SELECT
        t.tilename as tilename,
        m.xwin_image as xwin_image,
        m.ywin_image as ywin_image,
        m.xmin_image as xmin_image,
        m.ymin_image as ymin_image,
        m.xmax_image as xmax_image,
        m.ymax_image as ymax_image,
        m.balrog_index as balrog_index,
        m.alphawin_j2000 as ra,
        m.deltawin_j2000 as dec,
        m.mag_auto as mag_auto,
        m.spread_model as spread_model,
        m.spreaderr_model as spreaderr_model,
        m.class_star as class_star,
        m.mu_max as mu_max,
        t.sersicindex_0 as sersic_index,
        m.disk_scale_image as disk_scale,
        m.flux_radius as flux_radius,
        m.mag_disk as mag_disk,
        m.mag_psf as mag_psf,
        t.mag as truth_mag_auto,
        m.flags as flags
    FROM
        SUCHYTA1.balrog_%s_sim_%s m
        JOIN SUCHYTA1.balrog_%s_truth_%s t ON t.balrog_index = m.balrog_index
    """ %(table, band, table, band)
    return q

def DESFields(tilestuff, band='i'):
    q = """
        SELECT
           tilename,
           coadd_objects_id,
           mag_auto_%s as mag_auto,
           alphawin_j2000_%s as ra,
           deltawin_j2000_%s as dec,
           spread_model_%s as spread_model,
           spreaderr_model_%s as spreaderr_model,
           class_star_%s as class_star,
           mag_psf_%s as mag_psf,
           flux_radius_%s as flux_radius,
           flags_%s as flags
        FROM
           sva1_coadd_objects
        WHERE
           tilename in %s
        """ % (band,band,band,band,band,band,band,band,band,str(tuple(np.unique(tilestuff['tilename']))))
    return q

def TruthFields(band='i', table = 'sva1v2'):
    q = """
    SELECT
        balrog_index,
        tilename,
        ra,
        dec,
        objtype,
        HALFLIGHTRADIUS_0 as radius,
        mag,
        z,
        sersicindex_0
    FROM
        SUCHYTA1.balrog_%s_truth_%s
    """%(table,band)
    return q

def GetDESCat( depthmap, nside, tilestuff, tileinfo, band='i',depth = 0.0):
    cur = desdb.connect()
    q = DESFields(tileinfo, band=band)
    detcat = cur.quick(q, array=True)
    detcat = ValidDepth(depthmap, nside, detcat, rakey='ra', deckey='dec',depth = depth)
    detcat = RemoveTileOverlap(tilestuff, detcat, col='tilename', rakey='ra', deckey='dec')
    return detcat

def getTileInfo(catalog, HealConfig=None):
    if HealConfig is None:
        HealConfig = getHealConfig()

    tiles = np.unique(catalog['tilename'])
    cur = desdb.connect()
    q = "SELECT tilename, udecll, udecur, urall, uraur FROM coaddtile"
    tileinfo = cur.quick(q, array=True)
    tilestuff = {}
    for i in range(len(tileinfo)):
        tilestuff[ tileinfo[i]['tilename'] ] = tileinfo[i]
    max = np.power(map_nside/float(HealConfig['out_nside']), 2.0)
    depthmap, nside = GetDepthMap(HealConfig['depthfile'])
    return depthmap, nside

def cleanCatalog(catalog, tag='mag_auto'):
    # We should get rid of obviously wrong things.
    keep = np.where( (catalog[tag] > 15. ) & (catalog[tag] < 30.) & (catalog['flags'] < 2) )
    return catalog[keep]

def removeBadTilesFromTruthCatalog(truth, tag='mag_auto', goodfrac = 0.8):
    tileList = np.unique(truth['tilename'])
    number = np.zeros(tileList.size)
    for tile, i in zip(tileList,xrange(number.size)):
        number[i] = np.sum(truth['tilename'] == tile)
    tileList = tileList[number > goodfrac*np.max(number)]
    keep = np.in1d( truth['tilename'], tileList )
    return truth[keep]

def mergeCatalogsUsingPandas(sim=None, truth=None, key='balrog_index', suffixes = ['_sim','']):
    import pandas as pd
    simData = pd.DataFrame(sim)
    truthData = pd.DataFrame(truth)
    matched = pd.merge(simData, truthData, on=key, suffixes = suffixes)
    matched_arr = matched.to_records(index=False)
    # This last step is necessary because Pandas converts strings to Objects when eating structured arrays.
    # And np.recfunctions flips out when it has one.
    oldDtype = matched_arr.dtype.descr
    newDtype = oldDtype
    for thisOldType,i in zip(oldDtype, xrange(len(oldDtype) )):
        if 'O' in thisOldType[1]:
            newDtype[i] = (thisOldType[0], 'S12')
    matched_arr = np.array(matched_arr,dtype=newDtype)
    return matched_arr

def GetFromDB( band='i', depth = 0.0,tables =['sva1v2','sva1v3','sva1v3_2','sva1v3_3']):
    depthfile = '../../Data/sva1_gold_1.0.2-4_nside4096_nest_i_auto_weights.fits'

    cur = desdb.connect()
    q = "SELECT tilename, udecll, udecur, urall, uraur FROM coaddtile"
    tileinfo = cur.quick(q, array=True)
    tilestuff = {}
    for i in range(len(tileinfo)):
        tilestuff[ tileinfo[i]['tilename'] ] = tileinfo[i]
    depthmap, nside = GetDepthMap(depthfile)
    truths = []
    sims = []
    truthMatcheds = []

    for tableName in tables:
        q = TruthFields(band=band,table=tableName)
        truth = cur.quick(q, array=True)

        truth = removeBadTilesFromTruthCatalog(truth)
        truth = ValidDepth(depthmap, nside, truth, depth = depth)
        truth = RemoveTileOverlap(tilestuff, truth)
        #truth = cleanCatalog(truth,tag='mag')
        unique_binds, unique_inds = np.unique(truth['balrog_index'],return_index=True)
        truth = truth[unique_inds]

        q = SimFields(band=band, table=tableName)
        sim = cur.quick(q, array=True)
        sim = cleanCatalog(sim,tag='mag_auto')
        unique_binds, unique_inds = np.unique(sim['balrog_index'],return_index=True)
        sim = sim[unique_inds]

        truthMatched = mergeCatalogsUsingPandas(sim=sim,truth=truth)

        sim = sim[np.in1d(sim['balrog_index'],truthMatched['balrog_index'])]
        sim.sort(order='balrog_index')
        truthMatched.sort(order='balrog_index')

        truthMatcheds.append(truthMatched)
        truths.append(truth)
        sims.append(sim)

    sim = np.hstack(sims)
    truth = np.hstack(truths)
    truthMatched = np.hstack(truthMatcheds)

    des = GetDESCat(depthmap, nside, tilestuff, sim, band=band,depth = depth)
    des = cleanCatalog(des, tag='mag_auto')

    return des, sim, truthMatched, truth, tileinfo

def hpRaDecToHEALPixel(ra, dec, nside=  4096, nest= True):
    phi = ra * np.pi / 180.0
    theta = (90.0 - dec) * np.pi / 180.0
    hpInd = hp.ang2pix(nside, theta, phi, nest= nest)
    return hpInd

def convertThetaPhiToRaDec(theta, phi):
    ra = phi*180.0/np.pi
    dec = 90.0 - theta*180.0/np.pi
    return ra,dec

def convertRaDecToThetaPhi(ra, dec):
    theta = (90.0 - dec) * np.pi / 180.0
    phi =  ra * np.pi / 180.0
    return theta, phi

def HealPixifyCatalogs(catalog=None, healConfig=None, ratag='ra', dectag = 'dec'):
    HealInds = hpRaDecToHEALPixel( catalog[ratag],catalog[dectag], nside= healConfig['out_nside'], nest= healConfig['nest'])
    if 'HEALIndex' in catalog.dtype.fields:
        healCat = catalog.copy()
        healCat['HEALIndex'] = HealInds
    else:
        healCat = rf.append_fields(catalog,'HEALIndex',HealInds,dtypes=HealInds.dtype)
    return healCat

def getHealConfig(map_nside = 4096, out_nside = 128, depthfile = '../../Data/sva1_gold_1.0.2-4_nside4096_nest_i_auto_weights.fits'):
    HealConfig = {}
    HealConfig['map_nside'] = map_nside
    HealConfig['out_nside'] = out_nside
    HealConfig['finer_nside'] = map_nside
    HealConfig['depthfile'] = depthfile
    HealConfig['nest'] = True
    return HealConfig

def getGoodRegionIndices(catalog=None, badHPInds=None, nside=4096,band=None, raTag = 'ra', decTag = 'dec'):
    hpInd = hpRaDecToHEALPixel(catalog[raTag], catalog[decTag], nside=nside, nest= True)
    keep = ~np.in1d(hpInd, badHPInds)
    return keep

def excludeBadRegions(des,balrogObs, balrogTruthMatched, balrogTruth, band=None):
    eliMap = hp.read_map("sva1_gold_1.0.4_goodregions_04_equ_nest_4096.fits", nest=True)
    nside = hp.npix2nside(eliMap.size)
    maskIndices = np.arange(eliMap.size)
    badIndices = maskIndices[eliMap == 1]
    if band is not None:
        raTag = 'ra_'+band
        decTag = 'dec_'+band
    else:
        raTag = 'ra'
        decTag = 'dec'
    obsKeepIndices = getGoodRegionIndices(catalog=balrogObs, badHPInds=badIndices, nside=nside, raTag = raTag, decTag = decTag)
    truthKeepIndices = getGoodRegionIndices(catalog=balrogTruth, badHPInds=badIndices, nside=nside, raTag = raTag, decTag = decTag)
    desKeepIndices = getGoodRegionIndices(catalog=des, badHPInds=badIndices, nside=nside, raTag = raTag, decTag = decTag)

    balrogObs = balrogObs[obsKeepIndices]
    balrogTruthMatched = balrogTruthMatched[obsKeepIndices]
    balrogTruth = balrogTruth[truthKeepIndices]
    des = des[desKeepIndices]

    return des,balrogObs, balrogTruthMatched, balrogTruth

def removeNeighbors(thing1, thing2, radius= 2./3600,
                    raTag1 = 'ra', decTag1 = 'dec', raTag2 = 'ra', decTag2 = 'dec'):
    # Returns the elements of thing 1 that are outside of the matching radius from thing 2

    depth=10
    h = esutil.htm.HTM(depth)
    m1, m2, d12 = h.match(thing1[raTag1],thing1[decTag1],thing2[raTag2],thing2[decTag2],radius,maxmatch=0)
    keep = ~np.in1d(thing1['balrog_index'],thing1['balrog_index'][m1])
    return keep

def getCatalogs(reload=False,band='i', path = '../../Data/'):
    # Check to see whether the catalog files exist.  If they do, then
    # use the files. If at least one does not, then get what we need
    # from the database

    fileNames = ['desCatalogFile-'+band+'.fits','BalrogObsFile-'+band+'.fits',
                 'BalrogTruthFile-'+band+'.fits', 'BalrogTruthMatchedFile-'+band+'.fits',
                 'BalrogTileInfo.fits']
    exists = True
    for thisFile in fileNames:
        print "Checking for existence of: "+path+thisFile
        if not os.path.isfile(path+thisFile): exists = False
    if exists and not reload:

        desCat = esutil.io.read(path+fileNames[0])
        BalrogObs = esutil.io.read(path+fileNames[1])
        BalrogTruth = esutil.io.read(path+fileNames[2])
        BalrogTruthMatched = esutil.io.read(path+fileNames[3])
        BalrogTileInfo = esutil.io.read(path+fileNames[4])
    else:
        print "Cannot find files, or have been asked to reload. Getting data from DESDB."
        desCat, BalrogObs, BalrogTruthMatched, BalrogTruth, BalrogTileInfo = GetFromDB(band=band)
        esutil.io.write( path+fileNames[0], desCat , clobber=True)
        esutil.io.write( path+fileNames[1], BalrogObs , clobber=True)
        esutil.io.write( path+fileNames[2], BalrogTruth , clobber=True)
        esutil.io.write( path+fileNames[3], BalrogTruthMatched , clobber=True)
        esutil.io.write( path+fileNames[4], BalrogTileInfo, clobber=True)

    return desCat, BalrogObs, BalrogTruthMatched, BalrogTruth, BalrogTileInfo

def modestify(data):
    modest = np.zeros(len(data), dtype=np.int32)

    galcut = (data['flags'] <=3) & -( ((data['class_star'] > 0.3) & (data['mag_auto'] < 18.0)) | ((data['spread_model'] + 3*data['spreaderr_model']) < 0.003) | ((data['mag_psf'] > 30.0) & (data['mag_auto'] < 21.0)))
    modest[galcut] = 1

    starcut = (data['flags'] <=3) & ((data['class_star'] > 0.3) & (data['mag_auto'] < 18.0) & (data['mag_psf'] < 30.0) | (((data['spread_model'] + 3*data['spreaderr_model']) < 0.003) & ((data['spread_model'] +3*data['spreaderr_model']) > -0.003)))
    modest[starcut] = 3

    neither = -(galcut | starcut)
    modest[neither] = 5

    data = rf.append_fields(data, 'modtype', modest)
    print len(data), np.sum(galcut), np.sum(starcut), np.sum(neither)
    return data

def getCleanCatalogs( reload = False, band=None, isolated=False, nside=None):

    if band is None:
        print "Please specify a filter using the 'band' keyword argument."
        print "Choose one of [grizY]."

    des, balrogObs, balrogTruthMatched, balrogTruth, balrogTileInfo = getCatalogs(reload=reload, band=band)

    # Remove things in regions that are officially masked.
    des, balrogObs, balrogTruthMatched, balrogTruth = excludeBadRegions(des,balrogObs, balrogTruthMatched, balrogTruth)

    # if the isolated keyword is set, exclude things within some
    # separation from a pre-existing DES detection.
    if isolated is not False:
        print "isolated keyword should be set to a number, in arcseconds."
        keep = removeNeighbors(balrogTruthMatched, des, radius= float(isolated)/3600)
        balrogTruthMatched = balrogTruthMatched[keep]
        balrogObs = balrogObs[keep]

    # nside is set, then create a HEALPixel configuration and
    # assign pixel indices to each object.
    if nside is not None:
        HEALConfig = getHealConfig(map_nside = 4096, out_nside = nside )
        des = HealPixifyCatalogs(catalog=des, healConfig=HEALConfig)
        balrogObs = HealPixifyCatalogs(catalog=balrogObs, healConfig=HEALConfig)
        balrogTruth = HealPixifyCatalogs(catalog=balrogTruth, healConfig=HEALConfig)
        balrogTruthMatched = HealPixifyCatalogs(catalog=balrogTruthMatched, healConfig=HEALConfig)
        return des, balrogObs, balrogTruthMatched, balrogTruth,HEALConfig
    else:
        return des, balrogObs, balrogTruthMatched, balrogTruth

def matchCatalogs(cat1, cat2 ,tag = 'balrog_index'):
    ind1, ind2 = esutil.numpy_util.match(cat1[tag], cat2[tag])
    return cat1[ind1], cat2[ind2]

def getStellarityCatalogs( reload = False, band = None, nside = None):

    des, balrogObs, balrogTruthMatched, balrogTruth = getCleanCatalogs(reload=reload,band=band)
    des = modestify(des)
    balrogObs = modestify(balrogObs)

    return des, balrogObs, balrogTruthMatched, balrogTruth