# Copyright 2016 The TensorFlow Authors. All Rights Reserved.
#
# 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.
# ==============================================================================
"""A Transformed Distribution class."""
from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

import numpy as np

# Bijectors must be directly imported because `remove_undocumented` prevents
# individual file imports.
from tensorflow.python.framework import constant_op
from tensorflow.python.framework import dtypes
from tensorflow.python.framework import ops
from tensorflow.python.framework import tensor_util
from tensorflow.python.ops import array_ops
from tensorflow.python.ops import check_ops
from tensorflow.python.ops import control_flow_ops
from tensorflow.python.ops import math_ops
from tensorflow.python.ops.distributions import distribution as distribution_lib
from tensorflow.python.ops.distributions import identity_bijector
from tensorflow.python.ops.distributions import util as distribution_util

__all__ = [
    "TransformedDistribution",
]

# The following helper functions attempt to statically perform a TF operation.
# These functions make debugging easier since we can do more validation during
# graph construction.

def _static_value(x):
  """Returns the static value of a `Tensor` or `None`."""
  return tensor_util.constant_value(ops.convert_to_tensor(x))

def _logical_and(*args):
  """Convenience function which attempts to statically `reduce_all`."""
  args_ = [_static_value(x) for x in args]
  if any(x is not None and not bool(x) for x in args_):
    return constant_op.constant(False)
  if all(x is not None and bool(x) for x in args_):
    return constant_op.constant(True)
  if len(args) == 2:
    return math_ops.logical_and(*args)
  return math_ops.reduce_all(args)

def _logical_equal(x, y):
  """Convenience function which attempts to statically compute `x == y`."""
  x_ = _static_value(x)
  y_ = _static_value(y)
  if x_ is None or y_ is None:
    return math_ops.equal(x, y)
  return constant_op.constant(np.array_equal(x_, y_))

def _logical_not(x):
  """Convenience function which attempts to statically apply `logical_not`."""
  x_ = _static_value(x)
  if x_ is None:
    return math_ops.logical_not(x)
  return constant_op.constant(np.logical_not(x_))

def _concat_vectors(*args):
  """Convenience function which concatenates input vectors."""
  args_ = [_static_value(x) for x in args]
  if any(x_ is None for x_ in args_):
    return array_ops.concat(args, 0)
  return constant_op.constant([x_ for vec_ in args_ for x_ in vec_])

def _pick_scalar_condition(pred, cond_true, cond_false):
  """Convenience function which chooses the condition based on the predicate."""
  # Note: This function is only valid if all of pred, cond_true, and cond_false
  # are scalars. This means its semantics are arguably more like tf.cond than
  # tf.select even though we use tf.select to implement it.
  pred_ = _static_value(pred)
  if pred_ is None:
    return array_ops.where(pred, cond_true, cond_false)
  return cond_true if pred_ else cond_false

def _ones_like(x):
  """Convenience function attempts to statically construct `ones_like`."""
  # Should only be used for small vectors.
  if x.get_shape().is_fully_defined():
    return array_ops.ones(x.get_shape().as_list(), dtype=x.dtype)
  return array_ops.ones_like(x)

def _ndims_from_shape(shape):
  """Returns `Tensor`'s `rank` implied by a `Tensor` shape."""
  if shape.get_shape().ndims not in (None, 1):
    raise ValueError("input is not a valid shape: not 1D")
  if not shape.dtype.is_integer:
    raise TypeError("input is not a valid shape: wrong dtype")
  if shape.get_shape().is_fully_defined():
    return constant_op.constant(shape.get_shape().as_list()[0])
  return array_ops.shape(shape)[0]

def _is_scalar_from_shape(shape):
  """Returns `True` `Tensor` if `Tensor` shape implies a scalar."""
  return _logical_equal(_ndims_from_shape(shape), 0)

class TransformedDistribution(distribution_lib.Distribution):
  """A Transformed Distribution.

  A `TransformedDistribution` models `p(y)` given a base distribution `p(x)`,
  and a deterministic, invertible, differentiable transform, `Y = g(X)`. The
  transform is typically an instance of the `Bijector` class and the base
  distribution is typically an instance of the `Distribution` class.

  A `Bijector` is expected to implement the following functions:
  - `forward`,
  - `inverse`,
  - `inverse_log_det_jacobian`.
  The semantics of these functions are outlined in the `Bijector` documentation.

  We now describe how a `TransformedDistribution` alters the input/outputs of a
  `Distribution` associated with a random variable (rv) `X`.

  Write `cdf(Y=y)` for an absolutely continuous cumulative distribution function
  of random variable `Y`; write the probability density function `pdf(Y=y) :=
  d^k / (dy_1,...,dy_k) cdf(Y=y)` for its derivative wrt to `Y` evaluated at
  `y`. Assume that `Y = g(X)` where `g` is a deterministic diffeomorphism,
  i.e., a non-random, continuous, differentiable, and invertible function.
  Write the inverse of `g` as `X = g^{-1}(Y)` and `(J o g)(x)` for the Jacobian
  of `g` evaluated at `x`.

  A `TransformedDistribution` implements the following operations:

    * `sample`
      Mathematically:   `Y = g(X)`
      Programmatically: `bijector.forward(distribution.sample(...))`

    * `log_prob`
      Mathematically:   `(log o pdf)(Y=y) = (log o pdf o g^{-1})(y)
                         + (log o abs o det o J o g^{-1})(y)`
      Programmatically: `(distribution.log_prob(bijector.inverse(y))
                         + bijector.inverse_log_det_jacobian(y))`

    * `log_cdf`
      Mathematically:   `(log o cdf)(Y=y) = (log o cdf o g^{-1})(y)`
      Programmatically: `distribution.log_cdf(bijector.inverse(x))`

    * and similarly for: `cdf`, `prob`, `log_survival_function`,
     `survival_function`.

  A simple example constructing a Log-Normal distribution from a Normal
  distribution:

  ```python
  ds = tf.contrib.distributions
  log_normal = ds.TransformedDistribution(
    distribution=ds.Normal(loc=0., scale=1.),
    bijector=ds.bijectors.Exp(),
    name="LogNormalTransformedDistribution")
  ```

  A `LogNormal` made from callables:

  ```python
  ds = tf.contrib.distributions
  log_normal = ds.TransformedDistribution(
    distribution=ds.Normal(loc=0., scale=1.),
    bijector=ds.bijectors.Inline(
      forward_fn=tf.exp,
      inverse_fn=tf.log,
      inverse_log_det_jacobian_fn=(
        lambda y: -tf.reduce_sum(tf.log(y), axis=-1)),
    name="LogNormalTransformedDistribution")
  ```

  Another example constructing a Normal from a StandardNormal:

  ```python
  ds = tf.contrib.distributions
  normal = ds.TransformedDistribution(
    distribution=ds.Normal(loc=0., scale=1.),
    bijector=ds.bijectors.Affine(
      shift=-1.,
      scale_identity_multiplier=2.)
    name="NormalTransformedDistribution")
  ```

  A `TransformedDistribution`'s batch- and event-shape are implied by the base
  distribution unless explicitly overridden by `batch_shape` or `event_shape`
  arguments. Specifying an overriding `batch_shape` (`event_shape`) is
  permitted only if the base distribution has scalar batch-shape (event-shape).
  The bijector is applied to the distribution as if the distribution possessed
  the overridden shape(s). The following example demonstrates how to construct a
  multivariate Normal as a `TransformedDistribution`.

  ```python
  ds = tf.contrib.distributions
  # We will create two MVNs with batch_shape = event_shape = 2.
  mean = [[-1., 0],      # batch:0
          [0., 1]]       # batch:1
  chol_cov = [[[1., 0],
               [0, 1]],  # batch:0
              [[1, 0],
               [2, 2]]]  # batch:1
  mvn1 = ds.TransformedDistribution(
      distribution=ds.Normal(loc=0., scale=1.),
      bijector=ds.bijectors.Affine(shift=mean, scale_tril=chol_cov),
      batch_shape=[2],  # Valid because base_distribution.batch_shape == [].
      event_shape=[2])  # Valid because base_distribution.event_shape == [].
  mvn2 = ds.MultivariateNormalTriL(loc=mean, scale_tril=chol_cov)
  # mvn1.log_prob(x) == mvn2.log_prob(x)
  ```

  """

  def __init__(self,
               distribution,
               bijector=None,
               batch_shape=None,
               event_shape=None,
               validate_args=False,
               name=None):
    """Construct a Transformed Distribution.

    Args:
      distribution: The base distribution instance to transform. Typically an
        instance of `Distribution`.
      bijector: The object responsible for calculating the transformation.
        Typically an instance of `Bijector`. `None` means `Identity()`.
      batch_shape: `integer` vector `Tensor` which overrides `distribution`
        `batch_shape`; valid only if `distribution.is_scalar_batch()`.
      event_shape: `integer` vector `Tensor` which overrides `distribution`
        `event_shape`; valid only if `distribution.is_scalar_event()`.
      validate_args: Python `bool`, default `False`. When `True` distribution
        parameters are checked for validity despite possibly degrading runtime
        performance. When `False` invalid inputs may silently render incorrect
        outputs.
      name: Python `str` name prefixed to Ops created by this class. Default:
        `bijector.name + distribution.name`.
    """
    parameters = locals()
    name = name or (("" if bijector is None else bijector.name) +
                    distribution.name)
    with ops.name_scope(name, values=[event_shape, batch_shape]) as name:
      # For convenience we define some handy constants.
      self._zero = constant_op.constant(0, dtype=dtypes.int32, name="zero")
      self._empty = constant_op.constant([], dtype=dtypes.int32, name="empty")

      if bijector is None:
        bijector = identity_bijector.Identity(validate_args=validate_args)

      # We will keep track of a static and dynamic version of
      # self._is_{batch,event}_override. This way we can do more prior to graph
      # execution, including possibly raising Python exceptions.

      self._override_batch_shape = self._maybe_validate_shape_override(
          batch_shape, distribution.is_scalar_batch(), validate_args,
          "batch_shape")
      self._is_batch_override = _logical_not(_logical_equal(
          _ndims_from_shape(self._override_batch_shape), self._zero))
      self._is_maybe_batch_override = bool(
          tensor_util.constant_value(self._override_batch_shape) is None or
          tensor_util.constant_value(self._override_batch_shape).size != 0)

      self._override_event_shape = self._maybe_validate_shape_override(
          event_shape, distribution.is_scalar_event(), validate_args,
          "event_shape")
      self._is_event_override = _logical_not(_logical_equal(
          _ndims_from_shape(self._override_event_shape), self._zero))
      self._is_maybe_event_override = bool(
          tensor_util.constant_value(self._override_event_shape) is None or
          tensor_util.constant_value(self._override_event_shape).size != 0)

      # To convert a scalar distribution into a multivariate distribution we
      # will draw dims from the sample dims, which are otherwise iid. This is
      # easy to do except in the case that the base distribution has batch dims
      # and we're overriding event shape. When that case happens the event dims
      # will incorrectly be to the left of the batch dims. In this case we'll
      # cyclically permute left the new dims.
      self._needs_rotation = _logical_and(
          self._is_event_override,
          _logical_not(self._is_batch_override),
          _logical_not(distribution.is_scalar_batch()))
      override_event_ndims = _ndims_from_shape(self._override_event_shape)
      self._rotate_ndims = _pick_scalar_condition(
          self._needs_rotation, override_event_ndims, 0)
      # We'll be reducing the head dims (if at all), i.e., this will be []
      # if we don't need to reduce.
      self._reduce_event_indices = math_ops.range(
          self._rotate_ndims - override_event_ndims, self._rotate_ndims)

    self._distribution = distribution
    self._bijector = bijector
    super(TransformedDistribution, self).__init__(
        dtype=self._distribution.dtype,
        reparameterization_type=self._distribution.reparameterization_type,
        validate_args=validate_args,
        allow_nan_stats=self._distribution.allow_nan_stats,
        parameters=parameters,
        # We let TransformedDistribution access _graph_parents since this class
        # is more like a baseclass than derived.
        graph_parents=(distribution._graph_parents +  # pylint: disable=protected-access
                       bijector.graph_parents),
        name=name)

  @property
  def distribution(self):
    """Base distribution, p(x)."""
    return self._distribution

  @property
  def bijector(self):
    """Function transforming x => y."""
    return self._bijector

  def _event_shape_tensor(self):
    return self.bijector.forward_event_shape_tensor(
        distribution_util.pick_vector(
            self._is_event_override,
            self._override_event_shape,
            self.distribution.event_shape_tensor()))

  def _event_shape(self):
    # If there's a chance that the event_shape has been overridden, we return
    # what we statically know about the `event_shape_override`. This works
    # because: `_is_maybe_event_override` means `static_override` is `None` or a
    # non-empty list, i.e., we don't statically know the `event_shape` or we do.
    #
    # Since the `bijector` may change the `event_shape`, we then forward what we
    # know to the bijector. This allows the `bijector` to have final say in the
    # `event_shape`.
    static_override = tensor_util.constant_value_as_shape(
        self._override_event_shape)
    return self.bijector.forward_event_shape(
        static_override
        if self._is_maybe_event_override
        else self.distribution.event_shape)

  def _batch_shape_tensor(self):
    return distribution_util.pick_vector(
        self._is_batch_override,
        self._override_batch_shape,
        self.distribution.batch_shape_tensor())

  def _batch_shape(self):
    # If there's a chance that the batch_shape has been overridden, we return
    # what we statically know about the `batch_shape_override`. This works
    # because: `_is_maybe_batch_override` means `static_override` is `None` or a
    # non-empty list, i.e., we don't statically know the `batch_shape` or we do.
    #
    # Notice that this implementation parallels the `_event_shape` except that
    # the `bijector` doesn't get to alter the `batch_shape`. Recall that
    # `batch_shape` is a property of a distribution while `event_shape` is
    # shared between both the `distribution` instance and the `bijector`.
    static_override = tensor_util.constant_value_as_shape(
        self._override_batch_shape)
    return (static_override
            if self._is_maybe_batch_override
            else self.distribution.batch_shape)

  def _sample_n(self, n, seed=None):
    sample_shape = _concat_vectors(
        distribution_util.pick_vector(self._needs_rotation, self._empty, [n]),
        self._override_batch_shape,
        self._override_event_shape,
        distribution_util.pick_vector(self._needs_rotation, [n], self._empty))
    x = self.distribution.sample(sample_shape=sample_shape, seed=seed)
    x = self._maybe_rotate_dims(x)
    # We'll apply the bijector in the `_call_sample_n` function.
    return x

  def _call_sample_n(self, sample_shape, seed, name, **kwargs):
    # We override `_call_sample_n` rather than `_sample_n` so we can ensure that
    # the result of `self.bijector.forward` is not modified (and thus caching
    # works).
    with self._name_scope(name, values=[sample_shape]):
      sample_shape = ops.convert_to_tensor(
          sample_shape, dtype=dtypes.int32, name="sample_shape")
      sample_shape, n = self._expand_sample_shape_to_vector(
          sample_shape, "sample_shape")

      # First, generate samples. We will possibly generate extra samples in the
      # event that we need to reinterpret the samples as part of the
      # event_shape.
      x = self._sample_n(n, seed, **kwargs)

      # Next, we reshape `x` into its final form. We do this prior to the call
      # to the bijector to ensure that the bijector caching works.
      batch_event_shape = array_ops.shape(x)[1:]
      final_shape = array_ops.concat([sample_shape, batch_event_shape], 0)
      x = array_ops.reshape(x, final_shape)

      # Finally, we apply the bijector's forward transformation. For caching to
      # work, it is imperative that this is the last modification to the
      # returned result.
      y = self.bijector.forward(x, **kwargs)
      y = self._set_sample_static_shape(y, sample_shape)

      return y

  def _log_prob(self, y):
    # For caching to work, it is imperative that the bijector is the first to
    # modify the input.
    x = self.bijector.inverse(y)
    event_ndims = self._maybe_get_event_ndims_statically()

    ildj = self.bijector.inverse_log_det_jacobian(y, event_ndims=event_ndims)
    if self.bijector._is_injective:  # pylint: disable=protected-access
      return self._finish_log_prob_for_one_fiber(y, x, ildj, event_ndims)

    lp_on_fibers = [
        self._finish_log_prob_for_one_fiber(y, x_i, ildj_i, event_ndims)
        for x_i, ildj_i in zip(x, ildj)]
    return math_ops.reduce_logsumexp(array_ops.stack(lp_on_fibers), axis=0)

  def _finish_log_prob_for_one_fiber(self, y, x, ildj, event_ndims):
    """Finish computation of log_prob on one element of the inverse image."""
    x = self._maybe_rotate_dims(x, rotate_right=True)
    log_prob = self.distribution.log_prob(x)
    if self._is_maybe_event_override:
      log_prob = math_ops.reduce_sum(log_prob, self._reduce_event_indices)
    log_prob += math_ops.cast(ildj, log_prob.dtype)
    if self._is_maybe_event_override and isinstance(event_ndims, int):
      log_prob.set_shape(array_ops.broadcast_static_shape(
          x.get_shape().with_rank_at_least(1)[:-event_ndims], self.batch_shape))
    return log_prob

  def _prob(self, y):
    x = self.bijector.inverse(y)
    event_ndims = self._maybe_get_event_ndims_statically()
    ildj = self.bijector.inverse_log_det_jacobian(y, event_ndims=event_ndims)
    if self.bijector._is_injective:  # pylint: disable=protected-access
      return self._finish_prob_for_one_fiber(y, x, ildj, event_ndims)

    prob_on_fibers = [
        self._finish_prob_for_one_fiber(y, x_i, ildj_i, event_ndims)
        for x_i, ildj_i in zip(x, ildj)]
    return sum(prob_on_fibers)

  def _finish_prob_for_one_fiber(self, y, x, ildj, event_ndims):
    """Finish computation of prob on one element of the inverse image."""
    x = self._maybe_rotate_dims(x, rotate_right=True)
    prob = self.distribution.prob(x)
    if self._is_maybe_event_override:
      prob = math_ops.reduce_prod(prob, self._reduce_event_indices)
    prob *= math_ops.exp(math_ops.cast(ildj, prob.dtype))
    if self._is_maybe_event_override and isinstance(event_ndims, int):
      prob.set_shape(array_ops.broadcast_static_shape(
          y.get_shape().with_rank_at_least(1)[:-event_ndims], self.batch_shape))
    return prob

  def _log_cdf(self, y):
    if self._is_maybe_event_override:
      raise NotImplementedError("log_cdf is not implemented when overriding "
                                "event_shape")
    if not self.bijector._is_injective:  # pylint: disable=protected-access
      raise NotImplementedError("log_cdf is not implemented when "
                                "bijector is not injective.")
    x = self.bijector.inverse(y)
    return self.distribution.log_cdf(x)

  def _cdf(self, y):
    if self._is_maybe_event_override:
      raise NotImplementedError("cdf is not implemented when overriding "
                                "event_shape")
    if not self.bijector._is_injective:  # pylint: disable=protected-access
      raise NotImplementedError("cdf is not implemented when "
                                "bijector is not injective.")
    x = self.bijector.inverse(y)
    return self.distribution.cdf(x)

  def _log_survival_function(self, y):
    if self._is_maybe_event_override:
      raise NotImplementedError("log_survival_function is not implemented when "
                                "overriding event_shape")
    if not self.bijector._is_injective:  # pylint: disable=protected-access
      raise NotImplementedError("log_survival_function is not implemented when "
                                "bijector is not injective.")
    x = self.bijector.inverse(y)
    return self.distribution.log_survival_function(x)

  def _survival_function(self, y):
    if self._is_maybe_event_override:
      raise NotImplementedError("survival_function is not implemented when "
                                "overriding event_shape")
    if not self.bijector._is_injective:  # pylint: disable=protected-access
      raise NotImplementedError("survival_function is not implemented when "
                                "bijector is not injective.")
    x = self.bijector.inverse(y)
    return self.distribution.survival_function(x)

  def _quantile(self, value):
    if self._is_maybe_event_override:
      raise NotImplementedError("quantile is not implemented when overriding "
                                "event_shape")
    if not self.bijector._is_injective:  # pylint: disable=protected-access
      raise NotImplementedError("quantile is not implemented when "
                                "bijector is not injective.")
    # x_q is the "qth quantile" of X iff q = P[X <= x_q].  Now, since X =
    # g^{-1}(Y), q = P[X <= x_q] = P[g^{-1}(Y) <= x_q] = P[Y <= g(x_q)],
    # implies the qth quantile of Y is g(x_q).
    inv_cdf = self.distribution.quantile(value)
    return self.bijector.forward(inv_cdf)

  def _entropy(self):
    if not self.bijector.is_constant_jacobian:
      raise NotImplementedError("entropy is not implemented")
    if not self.bijector._is_injective:  # pylint: disable=protected-access
      raise NotImplementedError("entropy is not implemented when "
                                "bijector is not injective.")
    # Suppose Y = g(X) where g is a diffeomorphism and X is a continuous rv. It
    # can be shown that:
    #   H[Y] = H[X] + E_X[(log o abs o det o J o g)(X)].
    # If is_constant_jacobian then:
    #   E_X[(log o abs o det o J o g)(X)] = (log o abs o det o J o g)(c)
    # where c can by anything.
    entropy = self.distribution.entropy()
    if self._is_maybe_event_override:
      # H[X] = sum_i H[X_i] if X_i are mutually independent.
      # This means that a reduce_sum is a simple rescaling.
      entropy *= math_ops.cast(math_ops.reduce_prod(self._override_event_shape),
                               dtype=entropy.dtype.base_dtype)
    if self._is_maybe_batch_override:
      new_shape = array_ops.concat([
          _ones_like(self._override_batch_shape),
          self.distribution.batch_shape_tensor()
      ], 0)
      entropy = array_ops.reshape(entropy, new_shape)
      multiples = array_ops.concat([
          self._override_batch_shape,
          _ones_like(self.distribution.batch_shape_tensor())
      ], 0)
      entropy = array_ops.tile(entropy, multiples)
    dummy = array_ops.zeros(
        shape=array_ops.concat(
            [self.batch_shape_tensor(), self.event_shape_tensor()],
            0),
        dtype=self.dtype)
    event_ndims = (self.event_shape.ndims if self.event_shape.ndims is not None
                   else array_ops.size(self.event_shape_tensor()))
    ildj = self.bijector.inverse_log_det_jacobian(
        dummy, event_ndims=event_ndims)

    entropy -= math_ops.cast(ildj, entropy.dtype)
    entropy.set_shape(self.batch_shape)
    return entropy

  def _maybe_validate_shape_override(self, override_shape, base_is_scalar,
                                     validate_args, name):
    """Helper to __init__ which ensures override batch/event_shape are valid."""
    if override_shape is None:
      override_shape = []

    override_shape = ops.convert_to_tensor(override_shape, dtype=dtypes.int32,
                                           name=name)

    if not override_shape.dtype.is_integer:
      raise TypeError("shape override must be an integer")

    override_is_scalar = _is_scalar_from_shape(override_shape)
    if tensor_util.constant_value(override_is_scalar):
      return self._empty

    dynamic_assertions = []

    if override_shape.get_shape().ndims is not None:
      if override_shape.get_shape().ndims != 1:
        raise ValueError("shape override must be a vector")
    elif validate_args:
      dynamic_assertions += [check_ops.assert_rank(
          override_shape, 1,
          message="shape override must be a vector")]

    if tensor_util.constant_value(override_shape) is not None:
      if any(s <= 0 for s in tensor_util.constant_value(override_shape)):
        raise ValueError("shape override must have positive elements")
    elif validate_args:
      dynamic_assertions += [check_ops.assert_positive(
          override_shape,
          message="shape override must have positive elements")]

    is_both_nonscalar = _logical_and(_logical_not(base_is_scalar),
                                     _logical_not(override_is_scalar))
    if tensor_util.constant_value(is_both_nonscalar) is not None:
      if tensor_util.constant_value(is_both_nonscalar):
        raise ValueError("base distribution not scalar")
    elif validate_args:
      dynamic_assertions += [check_ops.assert_equal(
          is_both_nonscalar, False,
          message="base distribution not scalar")]

    if not dynamic_assertions:
      return override_shape
    return control_flow_ops.with_dependencies(
        dynamic_assertions, override_shape)

  def _maybe_rotate_dims(self, x, rotate_right=False):
    """Helper which rolls left event_dims left or right event_dims right."""
    needs_rotation_const = tensor_util.constant_value(self._needs_rotation)
    if needs_rotation_const is not None and not needs_rotation_const:
      return x
    ndims = array_ops.rank(x)
    n = (ndims - self._rotate_ndims) if rotate_right else self._rotate_ndims
    return array_ops.transpose(
        x, _concat_vectors(math_ops.range(n, ndims), math_ops.range(0, n)))

  def _maybe_get_event_ndims_statically(self):
    if self.event_shape.ndims is not None:
      return self.event_shape.ndims

    event_ndims = array_ops.size(self.event_shape_tensor())

    static_event_ndims = tensor_util.constant_value(event_ndims)

    if static_event_ndims is not None:
      return static_event_ndims

    return event_ndims


"""Module containing main Neural Style Transfer class and methods"""
from keras import backend as K
from keras.models import Sequential
from keras.layers import AveragePooling2D, InputLayer
from keras.applications.vgg19 import preprocess_input, VGG19
from .utils import Image
import tensorflow as tf
from typing import Optional, Tuple
import numpy as np
from copy import deepcopy
import os
from tqdm import tqdm
from streamlit.DeltaGenerator import DeltaGenerator
import datetime as dt

class NeuralStyleTransfer:
    """Neural style transfer class that implements the whole workflow for
    transfering artistic style to pictures, as described in the NST paper by
    Gatys et al. (2015) (https://arxiv.org/abs/1508.06576),
    and using a Keras pre-trained VGG19 neural network and the Adam optimizer.
    """
    style_layers = {'block1_conv1': 0.2,
                    'block2_conv1': 0.2,
                    'block3_conv1': 0.2,
                    'block4_conv1': 0.2,
                    'block5_conv1': 0.2}

    def __init__(self, alpha: float = 5., beta: float = 100.,
                 gamma: float = 1e-3, dir_path: Optional[str] = None):
        """Initialize with the default loss weights and the directory
        path that will be used to write the outputs.

        Parameters
        ----------
        alpha
            Weight associated with content image loss (default = 5)
        beta
            Weight associated with style image loss (default = 100)
        gamma
            Weight associated with image variation loss (default = 1e-3)
        dir_path
            Directory where generated images will be written
            (default = current working directory)
        """
        self.dir_path = dir_path or os.getcwd()
        self.run_dir = None
        self.default_alpha = alpha
        self.default_beta = beta
        self.default_gamma = gamma
        # Image parameters
        self.img_height = None
        self.img_width = None
        self.color_channels = None
        self.input_shape = None
        # Initialize images
        self.img_content = None
        self.img_style = None
        self.input_content = None
        self.input_style = None
        self.input_generated = None
        # Tensorflow stuff
        self.session = None
        self.kmodel = None
        # Cost functions
        self.j_style = None
        self.j_content = None
        self.j_var = None
        self.j_total = None

    def load_content_image(self, fpath: str, content_size: int = 512,
                           interpolation: str = 'bilinear'):
        """Load content image, and scale it. Scaling the style image will
        affect the artistic features transfered to the generated image.

        Parameters
        ----------
        fpath
            Path of style image
        content_size
            Height to be used for scaling the content image
            (default = 512)
        interpolation
            Interpolation scheme to use when rescaling image
            (default = bilinear)
        """
        self.img_content = Image.from_fp(fpath, label='content')
        # Resize content image
        n_h, n_w, n_c = self.img_content.data.shape
        n_h_scaled = content_size
        n_w_scaled = int(n_w * content_size / n_h)
        self.img_content.resize(n_w_scaled, n_h_scaled,
                                resample=interpolation)
        model_input = np.expand_dims(deepcopy(self.img_content.data), axis=0)
        self.input_content = preprocess_input(model_input)
        # Save the content image shape for loading model
        self.img_height, self.img_width, self.color_channels = \
            self.img_content.data.shape
        self.input_shape = (1, self.img_height, self.img_width,
                            self.color_channels)

    def load_style_image(self, fpath: str, scale: float = 1.0,
                         interpolation: str = 'bilinear'):
        """Load style image, and scale it. Scaling the style image will
        affect the artistic features transfered to the generated image.

        Parameters
        ----------
        fpath
            Path of style image
        scale
            Scaling factor of style image (default = 1., no scaling)
        interpolation
            Interpolation scheme to use when rescaling image
            (default = bilinear)
        """
        self.img_style = Image.from_fp(fpath, label='content')
        if scale != 1.0:
            n_h, n_w, n_c = self.img_style.data.shape
            n_h_scaled = int(scale * n_h)
            n_w_scaled = int(scale * n_w)
            self.img_style.resize(n_w_scaled, n_h_scaled,
                                  resample=interpolation)
        model_input = np.expand_dims(deepcopy(self.img_style.data), axis=0)
        self.input_style = preprocess_input(model_input)

    def run(self, num_iterations: int = 1000, learning_rate: float = 2.0,
            alpha: Optional[float] = None, beta: Optional[float] = None,
            gamma: Optional[float] = None,
            noise_low: float = -20, noise_high: float = 20,
            noise_ratio: float = 0.6, write_steps: int = 20,
            style_layers: Optional[dict] = None,
            progress_bar: DeltaGenerator = None):
        """Run style transfer optimization that will generate the new image
        build from content and style images.

        Parameters
        ----------
        num_iterations
            Number of optimization steps to run (default = 1000)
        learning_rate
            Learning rate to use for the Adam optimizer (default = 2.0)
        alpha
            Weight associated with content image loss (default = default_alpha)
        beta
            Weight associated with style image loss (default = default_beta)
        gamma
            Weight associated with image variation loss
            (default = default_gamma)
        noise_low
            Min of random values to use when first creating the generated image
            from noise (default = -20)
        noise_high
            Max of random values to use when first creating the generated image
            from noise (default = 20)
        noise_ratio
            When first creating the generated image, ratio of amount coming
            from noise to amount coming from content image (default = 0.6)
        write_steps
            A generated image will be written every time after this number
            of iterations (default = 20)
        style_layers
            The weight to give to each VGG19 layer used in the style loss
            calculation (default = all equal weights)
        progress_bar
            Streamlit progress bar to be updated when training
            (default = None)
        """
        # Create a directory specific to this run
        self.run_dir = os.path.join(
            self.dir_path, dt.datetime.now().strftime("%Y%m%d-%H%M%S"))
        if not os.path.exists(self.run_dir):
            os.makedirs(self.run_dir)

        # Initialize Keras model to be used
        print('...creating Keras model to be trained')
        self.kmodel = self._init_model(self.img_content.data.shape)
        # For name formatting
        n_digits = len(str(abs(num_iterations)))
        # Input weights
        alpha = alpha or self.default_alpha
        beta = beta or self.default_beta
        gamma = self.default_gamma if gamma is None else gamma
        # Session: need to get session from Keras, because it contains
        # the pre-trained weight initialization for the model variables
        self.session = K.get_session()

        # get parameters for style cost
        style_layers = style_layers or self.style_layers
        # build total cost function
        print('...building cost function')
        self._build_cost_function(self.session, alpha, beta, gamma,
                                  style_layers)
        # Define optimizer
        optimizer = tf.train.AdamOptimizer(learning_rate)
        # Define training step, and specify input variable explicitely,
        # otherwise will train all the variables in the graph
        train_step = optimizer.minimize(self.j_total,
                                        var_list=[self.kmodel.input])
        # All variables need to be initialized because we rely on tensorflow,
        # but we don't want to lose the pretrained weights
        self._custom_global_variable_initialization(self.session)

        # Assign generated image to variable
        self._generate_noise_image(low=noise_low, high=noise_high,
                                   noise_ratio=noise_ratio)
        K.set_value(self.kmodel.input, self.input_generated)

        print('...training')
        # Start progress bar
        if progress_bar is not None:
            progress_bar.progress(0)
        for i in tqdm(range(num_iterations)):
            # Update progress bar if any
            if progress_bar is not None:
                progress_bar.progress((i + 1) / num_iterations)
            # Print every `write_steps` iteration.
            if i % write_steps == 0:
                jt, jc, js, jv = self.session.run(
                    [self.j_total, self.j_content, self.j_style, self.j_var])
                print("Iteration: {}".format(i))
                print("total cost = {}".format(jt))
                print("content cost = {}".format(jc))
                print("style cost = {}".format(js))
                print("variation cost = {}".format(jv))

                # save current generated image in `dir_path`
                fname = str(i).zfill(n_digits) + '.png'
                fpath = os.path.join(self.run_dir, fname)
                img = self._model_input_to_image(self.input_generated, label=i)
                img.save(fpath)

            # Run minimization step
            self.session.run(train_step)
            # Get generated image
            self.input_generated = self.session.run(self.kmodel.input)

        print('...done training')
        # save last generated image
        fpath = os.path.join(self.run_dir, 'image_generated.png')
        img = self._model_input_to_image(self.input_generated, label='final')
        img.save(fpath)

    @property
    def img_generated(self) -> Image:
        """Deprocess the generated input into the generated image"""
        return self._model_input_to_image(self.input_generated)

    def _init_model(self, input_shape: Tuple) -> Sequential:
        """Create keras model used for neural style transfer
        - import VGG19 model with pre-trained weights
        - remove end layers if existing (no top)
        - change max pool layers into avg pool layers
        - set weights as non trainable (doesn't really matter since we'll use
        tensorflow directly for training)
        - use variable as input
        - use input shape from given image

        Check out original model build here:
        https://github.com/keras-team/keras-applications/blob/master/keras_applications/vgg19.py

        Parameters
        ----------
        input_shape
            Input shape of the image that will be used as input to the model:
            (input height, input width, number of channels)

        Returns
        -------
        model
            Keras sequential model
        """
        # Get pretrained keras model: if using docker image, the weights
        # will be saved inside the mounted volume (since it takes a while
        # to download)
        kmodel = VGG19(include_top=False, weights='imagenet',
                       input_tensor=None, input_shape=input_shape)
        # Build new model from pretrained keras model
        new_kmodel = Sequential(name='VGG19_nst')
        # Create variable input tensor
        var = K.zeros(shape=(1,) + input_shape)
        input_var = InputLayer(input_tensor=var)
        input_var.trainable = True
        # Add input layer
        new_kmodel.add(input_var)
        # Add pretrained layers
        i_pool = 1
        for layer in kmodel.layers[1:]:
            if 'pool' in layer.name:
                name = 'block%i_avgpool' % i_pool
                layer = AveragePooling2D((2, 2), strides=(2, 2),
                                         padding='same', name=name)
                i_pool += 1
            else:
                # clear all the nodes from the layer before adding it to model
                layer._inbound_nodes = []
                layer._outbound_nodes = []
            # make all layers non trainable
            layer.trainable = False
            new_kmodel.add(layer)

        # build: need to rebuild model wo input layer since already exist
        new_kmodel.build(None)

        return new_kmodel

    def _custom_global_variable_initialization(self, session: tf.Session):
        """Custom variable initialization:
        the problem is that tensorflow requires all variables to
        be initialized, but that would erase all the pre-trained values
        obtained with Keras. So here we save their values before
        initialization, and then reassign it afterwards.

        Parameters
        ----------
        session
            Tensorflow session where graph resides
        """
        # There are other hidden variables in the model that need to
        # be initialized, and we can't get their values
        all_variables = tf.global_variables()
        model_variables = self.kmodel.weights
        # save variable values
        list_values = []
        for var in model_variables:
            list_values.append(session.run(var))
        # run initialization on that variable only
        init_op = tf.initialize_variables(all_variables)
        session.run(init_op)
        # then assign previous value again
        for idx, var in enumerate(model_variables):
            K.set_value(var, list_values[idx])

    def _generate_noise_image(self, low: int = -20, high: int = 20,
                              noise_ratio: float = 0.6):
        """Generate noise image from combination of random noise and
        content image pixel values.

        Parameters
        ----------
        noise_low
            Min of random values to use when first creating the generated image
            from noise (default = -20)
        noise_high
            Max of random values to use when first creating the generated image
            from noise (default = 20)
        noise_ratio
            When first creating the generated image, ratio of amount coming
            from noise to amount coming from content image (default = 0.6)
        """
        shape_img = (1, self.img_height, self.img_width, self.color_channels)
        noise_data = np.random.uniform(low, high, shape_img).astype('float32')
        input_data = (noise_data * noise_ratio
                      + self.input_content * (1 - noise_ratio))
        self.input_generated = input_data

    def _build_cost_function(self, session: tf.Session, alpha: float,
                             beta: float, gamma: float, style_layers: dict):
        """Build total cost function, which is the weighted sum of the
        content loss, the style loss, and the variation loss functions.
        This will create the loss functions and sum them up.

        Parameters
        ----------
        session
            Tensorflow session where the graph resides
        alpha
            Weight associated with content image loss
        beta
            Weight associated with style image loss
        gamma
            Weight associated with image variation loss
        style_layers
            The weight to give to each VGG19 layer used in the style loss
            calculation
        """
        # Calculate costs
        self.j_content = self._compute_content_cost(self.input_content,
                                                    session=session)
        self.j_style = self._compute_style_cost(self.input_style,
                                                style_layers, session=session)
        self.j_var = self._compute_variation_cost(session)
        # Get total weighted cost
        self.j_total = self._total_cost(self.j_content, self.j_style,
                                        self.j_var, alpha, beta, gamma)

    def _model_input_to_image(self, input_data: np.ndarray,
                              label: Optional[str] = None) -> Image:
        """Undo all the preprocessing done for modeling and return image
        Reversing this:
        https://github.com/keras-team/keras-applications/blob/master/keras_applications/imagenet_utils.py#L60-L61

        Parameters
        ----------
        input_data
            Image input data used for model, of shape
            (1, image height, image width, number of channels)
        label
            Optional label for created image (default = None)
        """
        x = deepcopy(input_data)
        # Add what was subtracted
        mean = [103.939, 116.779, 123.68]
        x[..., 0] += mean[0]
        x[..., 1] += mean[1]
        x[..., 2] += mean[2]
        # BGR -> RGB
        x = x[..., ::-1]
        # Images need integer
        x = np.clip(x[0, ...], 0, 255).astype('uint8')
        return Image(x, label)

    @staticmethod
    def _total_cost(j_content: tf.Tensor, j_style: tf.Tensor, j_var: tf.Tensor,
                    alpha: float, beta: float, gamma: float) -> tf.Tensor:
        """Compute total weighted averaged cost of the content, style, and
        variation losses.

        Parameters
        ----------
        j_content
            Content loss
        j_style
            Style loss
        j_var
            Variation loss
        alpha
            Weight associated with content image loss
        beta
            Weight associated with style image loss
        gamma
            Weight associated with image variation loss

        Returns
        -------
        total loss
            Weighted sum of all losses
        """
        return alpha * j_content + beta * j_style + gamma * j_var

    def _compute_content_cost(self, input_content: np.ndarray,
                              session: Optional[tf.Session] = None
                              ) -> tf.Tensor:
        """
        Compute content loss

        Parameters
        ----------
        input_content
            Model input of content image
            Shape = (1, image height, image width, number of channels)
        session
            Session where graph resides (default = NST object session)

        Returns
        -------
        content loss
        """

        session = session or self.session
        # get model outputs
        output = self.kmodel.get_layer('block4_conv2').output
        # Get activations from content inputs
        K.set_value(self.kmodel.input, input_content)
        a_c = session.run(output)
        # Get symbolic activation for future updates of generated image
        a_g = output
        # Get shapes of activation
        _, n_h, n_w, n_c = a_c.shape
        # compute content loss
        j_content = 1. / (4 * n_h * n_w * n_c) * K.sum(K.square(a_c - a_g))

        return j_content

    def _compute_style_cost(self, input_style: np.ndarray, style_layers: dict,
                            session: tf.Session = None) -> tf.Tensor:
        """
        Compute the total style loss

        Parameters
        ----------
        input_style
            Model input for style image
        style_layers
            The weight to give to each VGG19 layer used in the style loss
            calculation
        session
            Session where graph resides (default = NST object session)

        Returns
        -------
        total style loss
        """

        session = session or self.session
        # initialize the overall style cost
        j_style = 0
        # Create temporary model to calculate activations for style image
        kmodel_style = self._init_model(self.img_style.data.shape)
        # Sum up losses for each layer
        for layer_name, coeff in style_layers.items():
            # Set input variable to style image inputs
            K.set_value(kmodel_style.input, input_style)
            # Set a_s to be the hidden layer activation from the layer we
            # have selected, by running the session on layer.output
            a_s = session.run(kmodel_style.get_layer(layer_name).output)
            # Get symbolic a_g
            a_g = self.kmodel.get_layer(layer_name).output
            # Compute style_cost for the current layer
            j_style_layer = self._compute_layer_style_cost(a_s, a_g, session)
            # Add coeff * J_style_layer of this layer to overall style cost
            j_style += coeff * j_style_layer

        return j_style

    @staticmethod
    def _compute_layer_style_cost(a_s: np.ndarray, a_g: tf.Tensor,
                                  session: tf.Session):
        """
        Calculate layer-specific style loss using style image and generated
        image activations. Style image and generated image number of channels
        must be equal.

        Parameters
        ----------
        a_s
            Activations of style image for layer
            Shape = (1, style img height, style img width, n channels)
        a_g
            Activations of generated image for layer
            Shape = (1, generated img height, generated img width, n channels)
        session
            Session where graph resides

        Returns
        -------
        Layer-specific style loss
        """

        # Retrieve dimensions from a_g (≈1 line)
        _, n_h, n_w, n_c = K.shape(a_g).eval(session=session)

        # Reshape the images to have them of shape (n_c, n_h*n_w) (≈2 lines)
        a_s = K.reshape(K.permute_dimensions(a_s, (3, 0, 1, 2)), (n_c, -1))
        a_g = K.reshape(K.permute_dimensions(a_g, (3, 0, 1, 2)), (n_c, -1))

        # Computing gram_matrices for both images S and G (≈2 lines)
        gs = NeuralStyleTransfer._gram_matrix(a_s)
        gg = NeuralStyleTransfer._gram_matrix(a_g)

        # Computing the loss (≈1 line)
        j_style_layer = 1 / (2 * n_c * n_w * n_h)**2 * K.sum(K.square(gs - gg))

        return j_style_layer

    def _compute_variation_cost(self, session: tf.Session) -> tf.Tensor:
        """Compute variation loss of Keras model input to keep the image locally
        coherent.

        Inspired by:
        https://github.com/keras-team/keras/blob/master/examples/neural_style_transfer.py

        Parameters
        ----------
        session
            Session where graph resides

        Returns
        -------
        Variation loss
        """
        # Get shape of input var
        input_var = self.kmodel.input
        _, n_h, n_w, n_c = K.shape(input_var).eval(session=session)
        a = K.square(
            input_var[:, :n_h - 1, :n_w - 1, :]
            - input_var[:, 1:, :n_w - 1, :])
        b = K.square(
            input_var[:, :n_h - 1, :n_w - 1, :]
            - input_var[:, :n_h - 1, 1:, :])
        return K.sum(K.pow(a + b, 1.25))

    @staticmethod
    def _gram_matrix(a: tf.Tensor) -> tf.Tensor:
        """
        Calculate gram matrix of given tensor.

        Parameters
        ----------
        a
            Input tensor for which to calculate gram matrix

        Returns
        -------
        ga
            Gram matrix of input tensor
        """
        ga = K.dot(a, K.transpose(a))
        return ga


import os
import sys
import numpy as np
from scipy.interpolate import interp1d
import matplotlib.pyplot as plt
from matplotlib import gridspec
import matplotlib.lines as mlines
import warnings

from MulensModel.utils import Utils, MAG_ZEROPOINT
from MulensModel import Event, Trajectory
from MulensModel import __version__ as mm_version

from MCPM import utils

K2_MAG_ZEROPOINT = 25.

# Multiple cpm_source-s to be done:
#     def set_pixel_coeffs_from_models
#     def satellite_maximum

# Issues with flux uncertsinteis read from TPF files:
#  - not yet properly used in chi2_fun()
#  - plot functions here and in cpmfitsource.py
#  - also CpmFitPixel.target_masked needs _err equivalent

class Minimizer(object):
    """
    An object to link an Event to the functions necessary to minimize chi2.

    Arguments :
        event: *MulensModel.Event*
            ...
            It is assumed that the last datasets are the satellite ones.

        parameters_to_fit: *list* of *str*
            Parameters that will be fitted. Except
            the *MulensModel.ModelParameters* parameters one can use satellite
            source fluxes: 'f_s_sat', 'f_s_sat_over_u_0',
            and 'f_b_sat' (I suggest not to use it).

        cpm_sources: *CpmFitSource* or *list* of them

    To force periodic flush of file with all models set n_flush to
    100 or 1000 etc.

    """
    def __init__(self, event, parameters_to_fit, cpm_sources):
        self.event = event
        self._MM = isinstance(event, Event)
        self.n_datasets = len(self.event.datasets)
        self.parameters_to_fit = parameters_to_fit
        self.n_parameters = len(self.parameters_to_fit)
        # self.n_parameters += 1
        if not isinstance(cpm_sources, list):
            cpm_sources = [cpm_sources]
        self.cpm_sources = cpm_sources
        self.n_sat = len(self.cpm_sources)
        self.reset_min_chi2()
        self._chi2_0 = None

        self._prior_min_values = None
        self._prior_max_values = None
        self._prior_gaussian = dict()
        self._prior_tabulated = dict()

        self._n_calls = 0
        self._color_constraint = None
        self.model_masks = [None] * self.n_sat
        self.fit_blending = [True] * self.n_datasets

        self._file_all_models_name = None
        self._file_all_models = None
        self.save_fluxes = True

        self._sat_masks = None
        self._sat_times = None
        self._sat_models = None
        self._sat_magnifications = None
        self._sat_source_flux = None
        self._sat_blending_flux = 0.

        self._coeffs_cache = None
        self.n_flush = None

        self.other_constraints = dict()

        self.sigma_scale = 1.

    def close_file_all_models(self):
        """closes the file to which all models are saved"""
        self._file_all_models.close()
        self._file_all_models = None

    @property
    def file_all_models(self):
        """name of the file to save all the models"""
        return self._file_all_models_name

    @file_all_models.setter
    def file_all_models(self, file_name):
        if self._file_all_models_name is not None:
            self.close_file_all_models()
        self._file_all_models_name = file_name
        if self._file_all_models_name is not None:
            self._file_all_models = open(self._file_all_models_name, 'w')

    def reset_min_chi2(self):
        """reset best model and corresponding parameters"""
        self._best_ln_prob = None
        self._best_ln_prob_theta = None

    def print_min_chi2(self):
        """Print best model (including prior) and corresponding chi2 value"""
        # fmt = " ".join(["{:.4f}"] * self.n_parameters)
        fmt = " ".join(["{:}"] * self.n_parameters)
        parameters = fmt.format(*list(self._best_ln_prob_theta))
        print("{:.3f}  {:}".format(self._min_chi2, parameters))

    def set_satellite_source_flux(self, sat_source_flux):
        """
        Provide value of satellite source flux. It's useful only if f_s_sat
        is not a fitting parameter.

        Parameters :
            sat_source_flux: *float*
                Satellite source flux
        """
        self._sat_source_flux = sat_source_flux

    def set_parameters(self, theta):
        """
        for given event set attributes from parameters_to_fit (list of str)
        to values from theta list
        """
        if len(self.parameters_to_fit) != len(theta):
            raise ValueError('wrong number of parameters {:} vs {:}'.format(
                    len(self.parameters_to_fit), len(theta)))
        combined = dict(zip(self.parameters_to_fit, theta))
        if 't_0_pl' in self.parameters_to_fit:
            combined = utils.get_standard_parameters(combined)
        for (param, value) in combined.items():
            if param == 'f_s_sat':
                self.set_satellite_source_flux(value)
            elif param == 'f_s_sat_over_u_0':
                try:
                    u_0 = self.parameters_to_fit.index('u_0')
                except Exception:
                    raise ValueError(
                        'This case is not yet coded: f_s_sat_over_u_0 is '
                        'fitted, but u_0 is not')
                self.set_satellite_source_flux(value*u_0)
            elif param == 'f_b_sat':
                self._sat_blending_flux = value
            elif param == 'q_f':
                self.event.model.set_source_flux_ratio(value)
            elif param[:4] == 'q_f_':
                self.event.model.set_source_flux_ratio_for_band(
                    param[4:], value)
            elif param == 'log_q_f':
                self.event.model.set_source_flux_ratio(10**value)
            elif param[:8] == 'log_q_f_':
                self.event.model.set_source_flux_ratio_for_band(
                    param[8:], 10**value)
            else:
                setattr(self.event.model.parameters, param, value)

    def _run_cpm(self, theta):
        """set the satellite light curve and run CPM"""
        self.set_parameters(theta)
        n_0 = self.n_datasets - self.n_sat

        if self._sat_masks is None:
            self._sat_masks = [
                cpm_source.residuals_mask for cpm_source in self.cpm_sources]
            self._sat_times = [
                self.cpm_sources[i].pixel_time[self._sat_masks[i]] + 2450000.
                for i in range(self.n_sat)]
            self._sat_models = [
                np.zeros(len(cpm_source.pixel_time))
                for cpm_source in self.cpm_sources]
            self._sat_magnifications = [None] * self.n_sat

        data = self.event.datasets
        for i in range(self.n_sat):
            # Here we prepare the satellite lightcurves:
            kwargs = {
                'time': self._sat_times[i],
                'satellite_skycoord': data[n_0+i].satellite_skycoord}
            if hasattr(self.event.model, "bandpasses"):
                if 'K2' in self.event.model.bandpasses:
                    kwargs['gamma'] = (
                        self.event.model.get_limb_coeff_gamma('K2'))
            if self.event.model.n_sources == 2:
                if not self._MM:
                    raise NotImplementedError('not yet coded in pixel_lensing')
                kwargs['flux_ratio_constraint'] = 'K2'
            if self._MM:
                if mm_version[0] in ['0', '1']:
                    out = self.event.model.magnification(**kwargs)
                else:
                    out = self.event.model.get_magnification(**kwargs)
                self._sat_magnifications[i] = out
                model = self._magnification_to_sat_flux(out)
            else:
                model = self.event.model.flux_difference(self._sat_times[i])
                # XXX satellite_skycoord is ignored above
            self._sat_models[i][self._sat_masks[i]] = model
            self.cpm_sources[i].run_cpm(
                self._sat_models[i], model_mask=self.model_masks[i])

    def _sat_flux_to_magnification(self, sat_flux):
        """
        translates satellite flux to magnification
        """
        if self._sat_blending_flux != 0.:
            warnings.warn(
                "self._sat_blending_flux is not 0. - not sure if this works")

        magnification = (
            (sat_flux - self._sat_blending_flux) / self._sat_source_flux + 1.)
        return magnification

    def _magnification_to_sat_flux(self, magnification):
        """
        translates magnification to satellite flux
        """
        if self._sat_blending_flux != 0.:
            warnings.warn(
                "self._sat_blending_flux is not 0. - not sure if this works")

        flux = (magnification - 1.) * self._sat_source_flux
        flux += self._sat_blending_flux
        return flux

    def _magnitude_to_sat_flux(self, magnitude):
        """
        translates magnitude in reference frame (OGLE I-band in most cases)
        to satellite flux scale
        """
        if self._sat_blending_flux != 0.:
            warnings.warn(
                "self._sat_blending_flux is not 0. - not sure if this works")

        flux = Utils.get_flux_from_mag(magnitude)
        (fs, fb) = self.event.model.get_ref_fluxes()
        magnification = (flux - fb) / fs[0]
        flux_sat = self._magnification_to_sat_flux(magnification)
        return flux_sat

    def set_satellite_data(self, theta):
        """set satellite dataset magnitudes and fluxes"""
        self._run_cpm(theta)
        n_0 = self.n_datasets - self.n_sat
        for i in range(self.n_sat):
            ii = n_0 + i
            # OLD:
            sat_residuals = self.cpm_sources[i].residuals[self._sat_masks[i]]
            flux = self._sat_models[i][self._sat_masks[i]] + sat_residuals
            # NEW - IF ACCEPTED!!! :
            # sat_residuals = self.cpm_sources[i].residuals[
            #     self.cpm_sources[i].residuals_mask]
            # flux = self._sat_models[i][self.cpm_sources[i].residuals_mask]
            # flux += sat_residuals
            #
            # Below we use private properties from MulensModel - this
            # should be corrected (maybe make new MulensData object and
            # delete the old one).
            self.event.datasets[ii]._flux = flux
            mag_and_err = Utils.get_mag_and_err_from_flux(
                flux,
                self.event.datasets[ii].err_flux, zeropoint=K2_MAG_ZEROPOINT)
            self.event.datasets[ii]._mag = mag_and_err[0]
            self.event.datasets[ii]._err_mag = mag_and_err[1]

    def add_full_color_constraint(
            self,
            ref_dataset_0, ref_dataset_1, ref_dataset_2,
            polynomial_2, sigma, ref_zero_point_0=MAG_ZEROPOINT,
            ref_zero_point_1=MAG_ZEROPOINT, ref_zero_point_2=MAG_ZEROPOINT):
        """
        Specify parameters that are used to constrain the source flux in
        satellite band:
                Kp-m0 = polynomial_value(m1-m2)
        In common case (e.g., Zhu+17 method: Kp-I = f(V-I)) m0 can be equal to
        m1 or m2.
            ref_dataset_0 (int) - dataset to calculate satellite color
            ref_dataset_1 (int) - first dataset for ground-based color
            ref_dataset_2 (int) - second dataset for ground-based color
            polynomial (np.array of floats) - color polynomial coefficients:
                    a0, a1, a2, that will be translated to
                    a0 + a1*(m1-m2) + a2*(m1-m2)**2
            sigma (float) - scatter or color for the constraint
            ref_zero_point_0 (float) - defines magnitude scale for 0-th dataset
            ref_zero_point_1 (float) - defines magnitude scale for 1-st dataset
            ref_zero_point_2 (float) - defines magnitude scale for 2-nd dataset
        """
        self._color_constraint = [ref_dataset_0, ref_dataset_1, ref_dataset_2,
                                  ref_zero_point_0, ref_zero_point_1,
                                  ref_zero_point_2, polynomial_2, sigma]

    def add_color_constraint(self, ref_dataset, ref_zero_point,
                             color, sigma_color):
        """
        Specify parameters that are used to constrain the source flux in
        satellite band:
            ref_dataset (int) - reference dataset
            ref_zero_point (float) - magnitude zeropoint of reference dataset
                                    probably MulensModel.utils.MAG_ZEROPOINT
            color (float) - (satellite-ref_dataset) color (NOTE ORDER)
            sigma_color (float) - sigma of color
        """
        self._color_constraint = [
            ref_dataset, ref_zero_point, color, sigma_color]

    def _get_source_flux(self, index):
        """
        Return source flux for dataset index.
        It uses self.fit_blending properly.
        index - int
        """
        if not self._MM:
            raise NotImplementedError('not yet coded in pixel_lensing')
        if self.fit_blending[index]:
            return self.event.get_ref_fluxes(index)[0][0]

        if self.event.model.n_lenses > 1:
            warnings.warn('The code is not optimized for binary lens and ' +
                          'no blending flux. It can be very slow...')
        fit_before = self.event.fit
        self.event.get_chi2_for_dataset(index, fit_blending=False)
        f_s = self.event.fit.flux_of_sources(self.event.datasets[index])[0]
        self.event.fit = fit_before
        return f_s

    def _chi2_for_color_constraint(self, satellite_flux):
        """calculate chi2 for flux constraint"""
        if not self._MM:
            raise NotImplementedError('not yet coded in pixel_lensing')
        before_ref = self.event.data_ref
        if len(self._color_constraint) == 4:
            (ref_dataset, ref_zero_point) = self._color_constraint[:2]
            (color, sigma_color) = self._color_constraint[2:]
        elif len(self._color_constraint) == 8:
            (ref_dataset, ref_dataset_1) = self._color_constraint[:2]
            (ref_dataset_2, ref_zero_point) = self._color_constraint[2:4]
            (ref_zero_point_1, ref_zero_point_2) = self._color_constraint[4:6]
            (polynomial, sigma_color) = self._color_constraint[6:]
            flux_1 = self._get_source_flux(ref_dataset_1)
            flux_2 = self._get_source_flux(ref_dataset_2)
            mag_1 = ref_zero_point_1 - 2.5 * np.log10(flux_1)
            mag_2 = ref_zero_point_2 - 2.5 * np.log10(flux_2)
            c = mag_1 - mag_2
            x = 1.
            color = 0.
            for p in polynomial:
                color += x * p
                x *= c
        else:
            raise ValueError('wrong size of internal variable')

        flux_ref = self._get_source_flux(ref_dataset)
        self.event.get_ref_fluxes(before_ref)

        mag_ref = ref_zero_point - 2.5 * np.log10(flux_ref)
        mag_sat = K2_MAG_ZEROPOINT - 2.5 * np.log10(satellite_flux)
        color_value = mag_sat - mag_ref
        out = ((color_value - color) / sigma_color)**2
        return out

    def chi2_fun(self, theta):
        """for a given set of parameters (theta), return the chi2"""
        self._run_cpm(theta)
        n = self.n_datasets - self.n_sat
        chi2_sat = []
        for source in self.cpm_sources:
            residuals = source.residuals_prf()[source.residuals_mask]
            # OLD:
            # residuals = source.residuals[source.residuals_mask]
            sigma = source.all_pixels_flux_err[source.residuals_mask]
            sigma *= self.sigma_scale
            chi2_sat.append(np.sum((residuals/sigma)**2))
        # Correct the line below.
        # chi2_sat = [
        #     np.sum(self._sat_masks[i]) *
        #     (self.cpm_sources[i].residuals_rms/np.mean(
        #     self.event.datasets[n+i].err_flux))**2
        #     for i in range(self.n_sat)]
        # We also tried:
        # chi2_sat = 0.
        # for i in range(self.n_sat):
            # ii = n + i
            # rms = self.cpm_sources[i].residuals_rms_prf_photometry(
            #     self._sat_models[i])
            # rms /= np.mean(self.event.datasets[n+i].err_flux)
            # chi2_sat += np.sum(self._sat_masks[i]) * rms**2

        # self.chi2 = [
        #     self.event.get_chi2_for_dataset(
        #         i, fit_blending=self.fit_blending[i])
        #     for i in range(n)]
        if self._MM:
            temp_chi2 = self.event.get_chi2_per_point()  # XXX - this
            # ignores self.fit_blending
        else:
            if n > 0:
                raise NotImplementedError('not yet coded in pixel_lensing')
            temp_chi2 = []
        self.chi2 = [np.sum(temp_chi2[i]) for i in range(n)]
        self.chi2 += chi2_sat
        if self._color_constraint is not None:
            self.chi2.append(
                self._chi2_for_color_constraint(self._sat_source_flux))
        chi2 = sum(self.chi2)
        self._last_chi2 = chi2
        self._n_calls += 1
        if self.save_fluxes:
            self.fluxes = 2 * n * [0.]
            if not self._MM:
                if n > 0:
                    raise NotImplementedError('not yet coded in pixel_lensing')
            else:
                self.event.get_chi2_per_point()
                for i in range(n):
                    d = self.event.datasets[i]
                    if self.fit_blending[i]:
                        self.fluxes[2*i] = self.event.fit.flux_of_sources(d)[0]
                        self.fluxes[2*i+1] = self.event.fit.blending_flux(d)
                    else:
                        self.fluxes[2*i] = self._get_source_flux(i)
        if self._file_all_models is not None:
            text = " ".join([repr(chi2)] + [repr(ll) for ll in theta])
            if self.save_fluxes:
                text += " " + " ".join(
                    ["{:.5f}".format(f) for f in self.fluxes])
            self._file_all_models.write(text + '\n')
            if self.n_flush is not None and self._n_calls % self.n_flush == 0:
                self._file_all_models.flush()
                os.fsync(self._file_all_models.fileno())
        if self._coeffs_cache is not None:
            coeffs = []
            for i in range(self.n_sat):
                n_pixels = self.cpm_sources[i].n_pixels
                source = self.cpm_sources[i]
                c = [source.pixel_coeffs(j).flatten() for j in range(n_pixels)]
                coeffs.append(np.array(c))
            self._coeffs_cache[tuple(theta.tolist())] = coeffs

        return chi2

    def set_chi2_0(self, chi2_0=None):
        """set reference value of chi2"""
        if chi2_0 is None:
            chi2_0 = np.sum([np.sum(d.good) for d in self.event.datasets])
        self._chi2_0 = chi2_0

    def set_pixel_coeffs_from_samples(self, samples):
        """
        Provide a matrix samples[n_models, n_params] and for each
        model there get the cached coeffs (caching MUST be turned ON)
        and set pixel coeffs to the mean of these cached coeffs.
        You may want to run stop_coeffs_cache() afterward.
        """
        weights = dict()
        coeffs = dict()
        for sample in samples:
            key = tuple(sample.tolist())
            if key in weights:
                weights[key] += 1
            else:
                weights[key] = 1
                coeffs[key] = self.get_cached_coeffs(key)
        self.set_pixel_coeffs_from_dicts(coeffs, weights)

    def set_pixel_coeffs_from_dicts(self, coeffs, weights=None):
        """
        Take coeffs, average them, and set pixel coeffs to the averages.

        Arguments :
            coeffs: *dict*
                Dictionary of pixel coeffs. Each value specifies a list
                (length same as number of satellite datasets) of coeffs for all
                pixels i.e., coeffs[key][i][j] is for
                j-th pixel in i-th cpm_source.
                The keys can be whatever, but most probably you want
                tuple(list(model_parameters)) to be the keys.
            weights: *dict*, optional
                Dictionary of weights - uses the same keys as coeffs.
        """
        keys = list(coeffs.keys())
        if weights is None:
            weights_ = None
        else:
            weights_ = [weights[key] for key in keys]

        for i in range(self.n_sat):
            for j in range(self.cpm_sources[i].n_pixels):
                data = [coeffs[key][i][j] for key in keys]
                average = np.average(np.array(data), 0, weights=weights_)
                self.cpm_sources[i].set_pixel_coeffs(
                    j, average.reshape((-1, 1)))

    def set_pixel_coeffs_from_models(self, models, weights=None):
        """run a set of models, remember the coeffs for every pixel,
        then average them (using weights) and remember

        NOTE: this version may be not very stable numerically. Try using
        e.g., set_pixel_coeffs_from_dicts()
        """
        if self.n_sat > 1:
            raise ValueError(
                "set_pixel_coeffs_from_models() doesn't allow " +
                "multiple cpm_sources")
        n_models = len(models)
        shape = (n_models, self.cpm_source.predictor_matrix.shape[1])
        coeffs = [np.zeros(shape) for i in range(self.cpm_source.n_pixels)]

        for i in range(n_models):
            self._run_cpm(models[i])
            for j in range(self.cpm_source.n_pixels):
                coeffs[j][i] = self.cpm_source.pixel_coeffs(j).reshape(
                    shape[1])

        for j in range(self.cpm_source.n_pixels):
            average = np.average(coeffs[j], 0, weights=weights)
            self.cpm_source.set_pixel_coeffs(j, average.reshape((-1, 1)))

    def start_coeffs_cache(self):
        """
        Start internally remembering coeffs; also resets cache if caching was
        working.
        """
        self._coeffs_cache = dict()

    def get_cached_coeffs(self, theta):
        """
        Get pixel coeffs for model defined by theta; note that
        theta = tuple(list(model_parameters))
        """
        if self._coeffs_cache is None:
            raise ValueError(
                "You want to get cached values and you haven't " +
                "turned on caching (see start_coeffs_cache())? Strange...")
        if not isinstance(theta, tuple):
            raise TypeError(
                'wrong type of get_cached_coeffs() input: \n' +
                'got {:}, expected tuple'.format(type(theta)))
        return self._coeffs_cache[theta]

    def stop_coeffs_cache(self):
        """turn off internally remembering coeffs"""
        self._coeffs_cache = None

    def save_coeffs_to_fits(self, files):
        """saves coeffs to fits files"""
        for (file_, cpm_source) in zip(files, self.cpm_sources):
            cpm_source.save_coeffs_to_fits(file_)

    def read_coeffs_from_fits(self, files):
        """read coeffs from fits files"""
        for (file_, cpm_source) in zip(files, self.cpm_sources):
            cpm_source.read_coeffs_from_fits(file_)

    def set_prior_boundaries(
            self, parameters_min_values, parameters_max_values):
        """
        remembers 2 dictionaries that set minimum and maximum values of
        parameters
        """
        self._prior_min_values = parameters_min_values
        self._prior_max_values = parameters_max_values

    def set_prior_gaussian(self, settings):
        """
        Set gaussian priors. The keywords *settings* is a *dict* with keys
        being names of parameters (e.g., 't_E') and values being lists of 2
        *floats* - mean and sigma (e.g., [10., 2.]).
        """
        self._prior_gaussian = settings

    def set_prior_tabulated(self, parameter, file_name, outside_factor=1.e-4):
        """
        Add settings to calculate prior based on tabulated data.

        Parameters :
            parameter: *str*
                Name of the parameter to be used for prior, e.g. 't_0'

            file_name: *str*
                Name of 2-column text file with a histogram. Columns are bin
                centers and counts (the latter are normalized by this function)

            outside_factor: *str*
                Ratio of largest to smallest value of prior - used both to
                places where histogram = 0 and values that are beyond the prior
        """
        (bin_centers, counts) = np.loadtxt(file_name, unpack=True)
        counts = counts / float(np.sum(counts))
        min_counts = outside_factor * np.max(counts)
        counts[counts < min_counts] = min_counts
        self._prior_tabulated[parameter] = interp1d(
            bin_centers, counts, fill_value=min_counts, bounds_error=False)

    def ln_prior(self, theta):
        """return 0 in most cases, or -np.inf if beyond ranges provided"""
        out = 0.
        outside = -np.inf

        if self._prior_min_values is not None:
            for (parameter, value) in self._prior_min_values.items():
                index = self.parameters_to_fit.index(parameter)
                if theta[index] < value:
                    return outside

        if self._prior_max_values is not None:
            for (parameter, value) in self._prior_max_values.items():
                index = self.parameters_to_fit.index(parameter)
                if theta[index] > value:
                    return outside

        for (key, value) in self.other_constraints.items():
            if key == 't_0':
                t_0_1 = theta[self.parameters_to_fit.index('t_0_1')]
                t_0_2 = theta[self.parameters_to_fit.index('t_0_2')]
                if value == 't_0_1 < t_0_2':
                    if t_0_1 >= t_0_2:
                        return outside
                elif value == 't_0_1 > t_0_2':
                    if t_0_2 >= t_0_1:
                        return outside
                else:
                    raise ValueError('urecognized value: {:}'.format(value))
            elif key == 'min_blending_flux':
                (data, limit) = value
                index = self.event.datasets.index(data)
                if not self.fit_blending[index]:
                    raise NotImplementedError(
                        'min_blending_flux and no blending flux make no sense')
                if not self._MM:
                    raise NotImplementedError('not yet coded in pixel_lensing')
                self.event.get_chi2_for_dataset(index)
                if self.event.fit.blending_flux(data) <= limit:
                    return outside
            else:
                raise KeyError('unkown constraint: {:}'.format(key))

        for (parameter, gauss) in self._prior_gaussian.items():
            value = theta[self.parameters_to_fit.index(parameter)]
            out -= ((value - gauss[0]) / gauss[1])**2

        for (parameter, function) in self._prior_tabulated.items():
            value = theta[self.parameters_to_fit.index(parameter)]
            out += np.log(function(value))

        return out

    def ln_like(self, theta):
        """logarithm of likelihood"""
        chi2 = self.chi2_fun(theta)

        ln_likelihood = -0.5 * (chi2 - self._chi2_0)

        return ln_likelihood

    def ln_prob(self, theta):
        """combines prior and likelihood"""
        ln_prior = self.ln_prior(theta)
        if not np.isfinite(ln_prior):
            if self.save_fluxes:
                return (-np.inf, self.fluxes)
            else:
                return -np.inf

        ln_like = self.ln_like(theta)
        if np.isnan(ln_like):
            if self.save_fluxes:
                return (-np.inf, self.fluxes)
            else:
                return -np.inf

        ln_probability = ln_prior + ln_like

        if self._best_ln_prob is None or ln_probability > self._best_ln_prob:
            self._min_chi2 = self._last_chi2
            self._best_ln_prob = ln_probability
            self._best_ln_prob_theta = theta

        if self.save_fluxes:
            return (ln_probability, self.fluxes)
        else:
            return ln_probability

    def set_MN_cube(self, min_values, max_values):
        """
        remembers how to transform unit cube to physical parameters for MN
        """
        self._MN_cube = [(min_values[i], (max_values[i]-min_values[i]))
                         for i in range(self.n_parameters)]

    def transform_MN_cube(self, cube):
        """transform unit cube to physical parameters"""
        out = []
        for i in range(len(cube)):
            (zero_point, range_) = self._MN_cube[i]
            out.append(zero_point + range_ * cube[i])
        return np.array(out)

    def satellite_maximum(self):
        """
        return time of maximum magnification, its value, and corresponding
        flux for the satellite dataset; takes into account the epochs when
        satellite data exist

        NOTE: This function is not yet fully tested.
        """
        if self._MM:
            return self._satellite_maximum_MM()
        else:
            return (self.event.model.parameters.t_0,
                    self.event.model.parameters.f_s_sat_over_beta)

    def _satellite_maximum_MM(self):
        """
        satellite_maximum() for MulensModel
        """
        if not self._MM:
            raise NotImplementedError('not yet coded in pixel_lensing')
        if self.n_sat > 1:
            raise ValueError(
                "satellite_maximum() doesn't allow " +
                "multiple cpm_sources")
        index = np.argmax(self._sat_magnifications[0])
        magnification = self._sat_magnifications[0][index]
        u_0 = (2*magnification*(magnification**2-1.)**-.5 - 2.)**.5
        if self.event.model.n_sources == 1:
            trajectory = Trajectory(
                self._sat_times[-1], parameters=self.event.model.parameters,
                parallax=self.event.model._parallax, coords=self.event.coords,
                satellite_skycoord=self.event.datasets[-1].satellite_skycoord)
            if trajectory.y[index] < 0.:
                u_0 = -u_0
        else:
            warnings.warn("binary source model - code is not yet ready " +
                          "to provide signed u_0 for K2")
        return (self._sat_times[0][index], u_0, magnification)

    def plot_sat_magnitudes(self, **kwargs):
        """
        Plot satellite model in reference magnitude system
        """
        if not self._MM:
            raise NotImplementedError('not yet coded in pixel_lensing')
        data_ref = self.event.model.data_ref
        (fs, fb) = self.event.model.get_ref_fluxes()
        n = self.n_datasets - self.n_sat
        if 'zorder' not in kwargs:
            kwargs['zorder'] = np.inf

        for i in range(self.n_sat):
            times = self._sat_times[i] - 2450000.
            flux = self._sat_magnifications[i] * fs[0] + fb
            plt.plot(
                times, Utils.get_mag_from_flux(flux), **kwargs)
        self.event.model.data_ref = data_ref

    def _legend_standard_plot(self, legend_order, legend_kwargs, color_list,
                              label_list, alphas):
        """
        plots legend for standard plots
        """
        # legend_kwargs['ncol'] = 2 # XXX
        if legend_order is not None:
            (handles, labels) = plt.gca().get_legend_handles_labels()
            if isinstance(legend_order, tuple):
                for (i, l_o) in enumerate(legend_order):
                    handles_ = [handles[idx] for idx in l_o]
                    labels_ = [labels[idx] for idx in l_o]
                    if i == 0:
                        first_legend = plt.legend(handles_, labels_,
                                                  loc='upper left')
                        plt.gca().add_artist(first_legend)
                    else:
                        plt.legend(handles_, labels_, loc='upper right',
                                   **legend_kwargs)
            else:
                handles_ = [handles[idx] for idx in legend_order]
                labels_ = [labels[idx] for idx in legend_order]
                plt.legend(handles_, labels_, **legend_kwargs)
        elif color_list is not None and label_list is not None:
            plt.legend(loc='best', **legend_kwargs)
        else:  # Prepare legend "manually":
            if self.n_sat == 0:
                plt.legend(loc='best', **legend_kwargs)
            else:
                black_line = mlines.Line2D(
                    [], [], color='black', marker='o',
                    lw=0, markersize=5, label='ground-based', alpha=alphas[0])
                red_line = mlines.Line2D(
                    [], [], color='red', marker='o',
                    lw=0, markersize=5, label='K2C9 data')
                blue_line = mlines.Line2D(
                    [], [], color='orange', lw=2, markersize=5,
                    label='K2C9 model')  # alpha=0.75,
                handles_ = [red_line, blue_line, black_line]
                plt.legend(handles=handles_, loc='best', **legend_kwargs)

    def standard_plot(self, t_start, t_stop, ylim, title=None,
                      label_list=None, color_list=None, line_width=1.5,
                      legend_order=None, separate_residuals=False,
                      model_line_width=4., legend_kwargs=None,
                      fluxes_y_axis=None, ground_model_zorder=None,
                      sat_model_zorder=None):
        """
        Make plot of the event and residuals.

        Parameters :
            XXX

            ylim: [*float*, *float*]
                A list o two values that are used to set y axis limits
                (in mag).

            fluxes_y_axis: *list* or *np.ndarray* of *floats*
                K2 fluxes which will be marked on right side of Y axis.

            ground_model_zorder: *float*
                Passed to pyplot to control if ground model is plotted
                at the top or bottom.

            sat_model_zorder: *float*
                Passed to pyplot to control if satellite model is
                plotted at the top or bottom. Defualts to *np.inf*.
        """
        if not self._MM:
            raise NotImplementedError('not yet coded in pixel_lensing')
        if (label_list is None) != (color_list is None):
            raise ValueError('wrong input in standard_plot')
        if not separate_residuals:
            grid_spec = gridspec.GridSpec(2, 1, height_ratios=[5, 1],
                                          hspace=0.12)
        else:
            grid_spec = gridspec.GridSpec(3, 1, height_ratios=[5, 1, 1],
                                          hspace=0.13)
        plt.figure()
        plt.subplot(grid_spec[0])
        if title is not None:
            plt.title(title)
        alphas = [0.5] * self.n_datasets
        for i in range(self.n_sat):
            alphas[-(i+1)] = 1.

        self.event.plot_model(
            color='black', subtract_2450000=True,
            t_start=t_start+2450000., t_stop=t_stop+2450000.,
            label="ground-based model", lw=model_line_width,
            zorder=ground_model_zorder)
        self.plot_sat_magnitudes(color='orange', lw=2,
                                 label="K2 model", zorder=sat_model_zorder)

        if color_list is not None:
            color_list_ = color_list
        else:
            if self.n_sat == 0:
                color_list_ = None
            else:
                color_list_ = ['black'] * (self.n_datasets-self.n_sat)
                color_list_ += ['red']*self.n_sat
        zorder_list = np.arange(self.n_datasets, 0, -1)
        zorder_list[1] = self.n_datasets + 1

        self.event.plot_data(  # alpha_list=alphas,
            zorder_list=zorder_list, mfc='none', lw=line_width, mew=line_width,
            marker='o', markersize=6, subtract_2450000=True,
            color_list=color_list_, label_list=label_list)
        if ylim is not None:
            plt.ylim(ylim[0], ylim[1])
        else:
            ylim = plt.ylim()
        plt.xlim(t_start, t_stop)
        plt.gca().tick_params(top=True, direction='in')

        if legend_kwargs is None:
            legend_kwargs = dict()
        self._legend_standard_plot(legend_order, legend_kwargs, color_list,
                                   label_list, alphas)

        y_K2_max = self._magnitude_to_sat_flux(np.min(ylim))
        y_K2_min = self._magnitude_to_sat_flux(np.max(ylim))
        print("Y-axis limits:")
        print("   mag:      {:.3f} {:.3f}".format(*ylim))
        print("   K2 flux:  {:.2f} {:.2f}".format(y_K2_min, y_K2_max))

        if fluxes_y_axis is not None:
            y_color = 'red'
            y_label = r'K2 differential counts [e$^-$s$^{-1}$]'

            min_ = np.min(fluxes_y_axis)
            if min_ < y_K2_min or np.max(fluxes_y_axis) > y_K2_max:
                raise ValueError('ylim incompatible with fluxes_y_axis')

            (fs, fb) = self.event.model.get_ref_fluxes()
            mags = self._sat_flux_to_magnification(np.array(fluxes_y_axis))
            mags_fluxes = Utils.get_mag_from_flux(fb + fs[0] * mags)
            ax2 = plt.gca().twinx()
            ax2.set_ylabel(y_label).set_color(y_color)
            ax2.spines['right'].set_color(y_color)
            ax2.set_ylim(ylim[0], ylim[1])
            ax2.tick_params(axis='y', colors=y_color)
            plt.yticks(mags_fluxes.tolist(), fluxes_y_axis, color=y_color)

        plt.subplot(grid_spec[1])
        kwargs_ = dict(mfc='none', lw=line_width, mew=line_width)
        if not separate_residuals:
            self.event.plot_residuals(subtract_2450000=True, **kwargs_)
            plt.xlim(t_start, t_stop)
        else:
            plt.plot([0., 3000000.], [0., 0.], color='black')
            self.event.datasets[-1].plot(
                phot_fmt='mag', show_errorbars=True, subtract_2450000=True,
                model=self.event.model, plot_residuals=True, **kwargs_)
            plt.ylim(0.29, -0.29)  # XXX
            plt.ylabel('K2 residuals')
            plt.xlim(t_start, t_stop)
            plt.gca().tick_params(top=True, direction='in')
            plt.gca().tick_params(right=True, direction='in')

            plt.subplot(grid_spec[2])
            plt.plot([0., 3000000.], [0., 0.], color='black')
            for data in self.event.datasets[:-1]:
                data.plot(
                    phot_fmt='mag', show_errorbars=True, subtract_2450000=True,
                    model=self.event.model, plot_residuals=True, **kwargs_)
            plt.ylabel('Residuals')
            plt.xlim(t_start, t_stop)
        plt.gca().tick_params(top=True, direction='in')
        plt.gca().tick_params(right=True, direction='in')

    def very_standard_plot(self, t_start, t_stop, ylim, title=None):
        """Make plot of the event and residuals. """
        if not self._MM:
            raise NotImplementedError('not yet coded in pixel_lensing')
        grid_spec = gridspec.GridSpec(2, 1, height_ratios=[5, 1], hspace=0.1)
        plt.figure()
        plt.subplot(grid_spec[0])
        if title is not None:
            plt.title(title)
        alphas = [0.95] * self.n_datasets
        for i in range(self.n_sat):
            alphas[-(i+1)] = 1.

        constraint = None
        if self.event.model.n_sources > 1:
            constraint = 'I'  # XXX
        self.event.plot_model(
            color='black', subtract_2450000=True,
            t_start=t_start+2450000., t_stop=t_stop+2450000.,
            flux_ratio_constraint=constraint)
        self.plot_sat_magnitudes(color='blue', lw=3.5, alpha=0.75)

        if self.n_sat == 0:
            colors = None
        else:
            colors = ['black'] * (self.n_datasets - self.n_sat)
            colors += ['red'] * self.n_sat

        self.event.plot_data(
            alpha_list=alphas,
            zorder_list=np.arange(self.n_datasets, 0, -1),
            marker='o', markersize=5, subtract_2450000=True,
            color_list=colors)
        if ylim is not None:
            plt.ylim(ylim[0], ylim[1])
        plt.xlim(t_start, t_stop)

        # Prepare legend "manually":
        if self.n_sat == 0:
            plt.legend(fontsize='small')
        else:
            black_line = mlines.Line2D(
                [], [], color='black', marker='o', lw=0,
                markersize=5, label='ground-based', alpha=alphas[0])
            red_line = mlines.Line2D(
                [], [], color='red', marker='o', lw=0,
                markersize=5, label='K2C9 data')
            blue_line = mlines.Line2D(
                [], [], color='blue', lw=3.5, alpha=0.75,
                markersize=5, label='K2C9 model')
            plt.legend(handles=[red_line, blue_line, black_line], loc='best')

        plt.subplot(grid_spec[1])
        self.event.plot_residuals(subtract_2450000=True)
        plt.xlim(t_start, t_stop)