# -*- coding: utf-8 -*-
from . import config
from ._constants import *
from .maps import MapBase, RVBase, ReflectedBase
from ._core import OpsSystem, math
import numpy as np
from astropy import units
from inspect import getmro
import matplotlib.pyplot as plt
from matplotlib.animation import FuncAnimation
from IPython.display import HTML
import os
import logging

logger = logging.getLogger("starry.maps")

__all__ = ["Primary", "Secondary", "System"]

class Body(object):
    """A generic body. Must be subclassed."""

    def __init__(
        self,
        map,
        r=1.0,
        m=1.0,
        prot=1.0,
        t0=0.0,
        theta0=0.0,
        length_unit=units.Rsun,
        mass_unit=units.Msun,
        time_unit=units.day,
        angle_unit=units.degree,
        **kwargs,
    ):
        # Surface map
        self._lazy = map._lazy
        self._math = map._math
        self.map = map

        # Units
        self.length_unit = length_unit
        self.mass_unit = mass_unit
        self.time_unit = time_unit
        self.angle_unit = angle_unit

        # Attributes
        self.r = r
        self.m = m
        self.prot = prot
        self.t0 = t0
        self.theta0 = theta0

    @property
    def length_unit(self):
        """An ``astropy.units`` unit defining the length metric for this body."""
        return self._length_unit

    @length_unit.setter
    def length_unit(self, value):
        assert value.physical_type == "length"
        self._length_unit = value
        self._length_factor = value.in_units(units.Rsun)

    @property
    def mass_unit(self):
        """An ``astropy.units`` unit defining the mass metric for this body."""
        return self._mass_unit

    @mass_unit.setter
    def mass_unit(self, value):
        assert value.physical_type == "mass"
        self._mass_unit = value
        self._mass_factor = value.in_units(units.Msun)

    @property
    def time_unit(self):
        """An ``astropy.units`` unit defining the time metric for this body."""
        return self._time_unit

    @time_unit.setter
    def time_unit(self, value):
        assert value.physical_type == "time"
        self._time_unit = value
        self._time_factor = value.in_units(units.day)

    @property
    def angle_unit(self):
        """An ``astropy.units`` unit defining the angle metric for this body."""
        return self._angle_unit

    @angle_unit.setter
    def angle_unit(self, value):
        assert value.physical_type == "angle"
        self._angle_unit = value
        self._angle_factor = value.in_units(units.radian)

    @property
    def _angle_unit(self):
        return self._map._angle_unit

    @_angle_unit.setter
    def _angle_unit(self, value):
        self._map._angle_unit = value

    @property
    def _angle_factor(self):
        return self._map._angle_factor

    @_angle_factor.setter
    def _angle_factor(self, value):
        self._map._angle_factor = value

    @property
    def map(self):
        """The surface map for this body."""
        return self._map

    @map.setter
    def map(self, value):
        assert MapBase in getmro(
            type(value)
        ), "The `map` attribute must be a `starry` map instance."
        assert (
            value._lazy == self._lazy
        ), "Map must have the same evaluation mode (lazy/greedy)."
        self._map = value

    @property
    def r(self):
        """The radius in units of :py:attr:`length_unit`."""
        return self._r / self._length_factor

    @r.setter
    def r(self, value):
        self._r = self._math.cast(value * self._length_factor)

    @property
    def m(self):
        """The mass in units of :py:attr:`mass_unit`."""
        return self._m / self._mass_factor

    @m.setter
    def m(self, value):
        self._m = self._math.cast(value * self._mass_factor)

    @property
    def prot(self):
        """The rotation period in units of :py:attr:`time_unit`."""
        return self._prot / self._time_factor

    @prot.setter
    def prot(self, value):
        self._prot = self._math.cast(value * self._time_factor)

    @property
    def t0(self):
        """A reference time in units of :py:attr:`time_unit`."""
        return self._t0 / self._time_factor

    @t0.setter
    def t0(self, value):
        self._t0 = self._math.cast(value * self._time_factor)

    @property
    def theta0(self):
        """The map rotational phase at time :py:attr:`t0`."""
        return self._theta0 / self._angle_factor

    @theta0.setter
    def theta0(self, value):
        self._theta0 = self._math.cast(value * self._angle_factor)

    def _check_kwargs(self, method, kwargs):
        if not config.quiet:
            for key in kwargs.keys():
                message = "Invalid keyword `{0}` in call to `{1}()`. Ignoring."
                message = message.format(key, method)
                logger.warning(message)

class Primary(Body):
    """A primary (central) body.

    Args:
        map: The surface map of this body. This should be an instance
            returned by :py:func:`starry.Map`.
        r (scalar, optional): The radius of the body in units of
            :py:attr:`length_unit`. Defaults to 1.0.
        m (scalar, optional): The mass of the body in units of
            :py:attr:`mass_unit`. Defaults to 1.0.
        prot (scalar, optional): The rotation period of the body in units of
            :py:attr:`time_unit`. Defaults to 1.0.
        t0 (scalar, optional): A reference time in units of
            :py:attr:`time_unit`. Defaults to 0.0.
        theta0 (scalar, optional): The rotational phase of the map at time
            :py:attr:`t0` in units of :py:attr:`angle_unit`. Defaults to 0.0.
        length_unit (optional): An ``astropy.units`` unit defining the
            distance metric for this object. Defaults to
            :py:attr:`astropy.units.Rsun.`
        mass_unit (optional): An ``astropy.units`` unit defining the
            mass metric for this object. Defaults to
            :py:attr:`astropy.units.Msun.`
        time_unit (optional): An ``astropy.units`` unit defining the
            time metric for this object. Defaults to
            :py:attr:`astropy.units.day.`
        angle_unit (optional): An ``astropy.units`` unit defining the
            angular metric for this object. Defaults to
            :py:attr:`astropy.units.degree.`
    """

    def __init__(self, map, **kwargs):
        # Initialize `Body`
        super(Primary, self).__init__(map, **kwargs)
        for kw in [
            "r",
            "m",
            "prot",
            "t0",
            "theta0",
            "length_unit",
            "mass_unit",
            "time_unit",
            "angle_unit",
        ]:
            kwargs.pop(kw, None)
        self._check_kwargs("Primary", kwargs)

class Secondary(Body):
    """A secondary (orbiting) body.

    Args:
        map: The surface map of this body. This should be an instance
            returned by :py:func:`starry.Map`.
        r (scalar, optional): The radius of the body in units of
            :py:attr:`length_unit`. Defaults to 1.0.
        m (scalar, optional): The mass of the body in units of
            :py:attr:`mass_unit`. Defaults to 1.0.
        a (scalar, optional): The semi-major axis of the body in units of
            :py:attr:`time_unit`. Defaults to 1.0. If :py:attr:`porb` is
            also provided, this value is ignored.
        porb (scalar, optional): The orbital period of the body in units of
            :py:attr:`time_unit`. Defaults to 1.0. Setting this value
            overrides :py:attr:`a`.
        prot (scalar, optional): The rotation period of the body in units of
            :py:attr:`time_unit`. Defaults to 1.0.
        t0 (scalar, optional): A reference time in units of
            :py:attr:`time_unit`. This is taken to be the time of a reference
            transit. Defaults to 0.0.
        ecc (scalar, optional): The orbital eccentricity of the body.
            Defaults to 0.
        w, omega (scalar, optional): The argument of pericenter of the body
            in units of :py:attr:`angle_unit`. Defaults to 90 degrees.
        Omega (scalar, optional): The longitude of ascending node of the
            body in units of :py:attr:`angle_unit`. Defaults to 0 degrees.
        inc (scalar, optional): The orbital inclination of the body in
            units of :py:attr:`angle_unit`. Defaults to 90 degrees.
        theta0 (scalar, optional): The rotational phase of the map at time
            :py:attr:`t0` in units of :py:attr:`angle_unit`. Defaults to
            0.0.
        length_unit (optional): An ``astropy.units`` unit defining the
            distance metric for this object. Defaults to
            :py:attr:`astropy.units.Rsun.`
        mass_unit (optional): An ``astropy.units`` unit defining the
            mass metric for this object. Defaults to
            :py:attr:`astropy.units.Msun.`
        time_unit (optional): An ``astropy.units`` unit defining the
            time metric for this object. Defaults to
            :py:attr:`astropy.units.day.`
        angle_unit (optional): An ``astropy.units`` unit defining the
            angular metric for this object. Defaults to
            :py:attr:`astropy.units.degree.`
    """

    def __init__(self, map, **kwargs):
        # Initialize `Body`
        super(Secondary, self).__init__(map, **kwargs)
        for kw in [
            "r",
            "m",
            "prot",
            "t0",
            "theta0",
            "length_unit",
            "mass_unit",
            "time_unit",
            "angle_unit",
        ]:
            kwargs.pop(kw, None)

        # Attributes
        if kwargs.get("porb", None) is not None:
            self.porb = kwargs.pop("porb", None)
        elif kwargs.get("a", None) is not None:
            self.a = kwargs.pop("a", None)
        else:
            raise ValueError("Must provide a value for either `porb` or `a`.")
        self.ecc = kwargs.pop("ecc", 0.0)
        self.w = kwargs.pop(
            "w", kwargs.pop("omega", 0.5 * np.pi / self._angle_factor)
        )
        self.Omega = kwargs.pop("Omega", 0.0)
        self.inc = kwargs.pop("inc", 0.5 * np.pi / self._angle_factor)
        self._check_kwargs("Secondary", kwargs)

    @property
    def porb(self):
        """The orbital period in units of :py:attr:`time_unit`.

        .. note::
            Setting this value overrides the value of :py:attr:`a`.
        """
        if self._porb == 0.0:
            return None
        else:
            return self._porb / self._time_factor

    @porb.setter
    def porb(self, value):
        self._porb = self._math.cast(value * self._time_factor)
        self._a = 0.0

    @property
    def a(self):
        """The semi-major axis in units of :py:attr:`length_unit`.

        .. note::
            Setting this value overrides the value of :py:attr:`porb`.
        """
        if self._a == 0.0:
            return None
        else:
            return self._a / self._length_factor

    @a.setter
    def a(self, value):
        self._a = self._math.cast(value * self._length_factor)
        self._porb = 0.0

    @property
    def ecc(self):
        """The orbital eccentricity."""
        return self._ecc

    @ecc.setter
    def ecc(self, value):
        self._ecc = value

    @property
    def w(self):
        """The longitude of pericenter in units of :py:attr:`angle_unit`."""
        return self._w / self._angle_factor

    @w.setter
    def w(self, value):
        self._w = self._math.cast(value * self._angle_factor)

    @property
    def omega(self):
        """Alias for the longitude of pericenter :py:attr:`w`."""
        return self.w

    @omega.setter
    def omega(self, value):
        self.w = value

    @property
    def Omega(self):
        """The longitude of ascending node in units of :py:attr:`angle_unit`."""
        return self._Omega / self._angle_factor

    @Omega.setter
    def Omega(self, value):
        self._Omega = self._math.cast(value * self._angle_factor)

    @property
    def inc(self):
        """The orbital inclination in units of :py:attr:`angle_unit`."""
        return self._inc / self._angle_factor

    @inc.setter
    def inc(self, value):
        self._inc = self._math.cast(value * self._angle_factor)

class System(object):
    """A system of bodies in Keplerian orbits about a central primary body.

    Args:
        primary (:py:class:`Primary`): The central body.
        secondaries (:py:class:`Secondary`): One or more secondary bodies
            in orbit about the primary.
        time_unit (optional): An ``astropy.units`` unit defining the
            time metric for this object. Defaults to
            :py:attr:`astropy.units.day.`
        light_delay (bool, optional): Account for the light travel time
            delay to the barycenter of the system? Default is False.
        texp (scalar): The exposure time of each observation. This can be a
            scalar or a tensor with the same shape as ``t``. If ``texp`` is
            provided, ``t`` is assumed to indicate the timestamp at the middle
            of an exposure of length ``texp``.
        oversample (int): The number of function evaluations to use when
            numerically integrating the exposure time.
        order (int): The order of the numerical integration scheme. This must
            be one of the following: ``0`` for a centered Riemann sum
            (equivalent to the "resampling" procedure suggested by Kipping 2010),
            ``1`` for the trapezoid rule, or ``2`` for Simpson’s rule.
    """

    def _no_spectral(self):
        if self._primary._map.nw is not None:  # pragma: no cover
            raise NotImplementedError(
                "Method not yet implemented for spectral maps."
            )

    def __init__(
        self,
        primary,
        *secondaries,
        time_unit=units.day,
        light_delay=False,
        texp=None,
        oversample=7,
        order=0,
    ):
        # Units
        self.time_unit = time_unit
        self._light_delay = bool(light_delay)
        if texp is None:
            self._texp = 0.0
        else:
            self._texp = texp
        assert self._texp >= 0.0, "Parameter `texp` must be >= 0."
        self._oversample = int(oversample)
        assert self._oversample > 0, "Parameter `oversample` must be > 0."
        self._order = int(order)
        assert self._order in [0, 1, 2], "Invalid value for parameter `order`."

        # Primary body
        assert (
            type(primary) is Primary
        ), "Argument `primary` must be an instance of `Primary`."
        assert (
            primary._map.__props__["reflected"] == False
        ), "Reflected light map not allowed for the primary body."
        self._primary = primary
        self._rv = primary._map.__props__["rv"]
        self._lazy = primary._lazy
        self._math = primary._math
        if self._lazy:
            self._linalg = math.lazy_linalg
        else:
            self._linalg = math.greedy_linalg

        # Secondary bodies
        assert len(secondaries) > 0, "There must be at least one secondary."
        for sec in secondaries:
            assert type(sec) is Secondary, (
                "Argument `*secondaries` must be a sequence of "
                "`Secondary` instances."
            )
            assert (
                sec._map.nw == self._primary._map.nw
            ), "All bodies must have the same number of wavelength bins `nw`."
            assert sec._map.__props__["rv"] == self._rv, (
                "Radial velocity must be enabled "
                "for either all or none of the bodies."
            )
            assert (
                sec._lazy == self._lazy
            ), "All bodies must have the same evaluation mode (lazy/greedy)."

        reflected = [sec._map.__props__["reflected"] for sec in secondaries]
        if np.all(reflected):
            self._reflected = True
        elif np.any(reflected):
            raise ValueError(
                "Reflected light must be enabled "
                "for either all or none of the secondaries."
            )
        else:
            self._reflected = False
        self._secondaries = secondaries

        # All bodies
        self._bodies = [self._primary] + list(self._secondaries)

        # Indices of each of the bodies in the design matrix
        Ny = [self._primary._map.Ny] + [
            sec._map.Ny for sec in self._secondaries
        ]
        self._inds = []
        cur = 0
        for N in Ny:
            self._inds.append(cur + np.arange(N))
            cur += N

        # Theano ops class
        self.ops = OpsSystem(
            self._primary,
            self._secondaries,
            reflected=self._reflected,
            rv=self._rv,
            light_delay=self._light_delay,
            texp=self._texp,
            oversample=self._oversample,
            order=self._order,
        )

        # Solve stuff
        self._flux = None
        self._C = None
        self._solution = None
        self._solved_bodies = []

    @property
    def light_delay(self):
        """Account for the light travel time delay? *Read-only*"""
        return self._light_delay

    @property
    def texp(self):
        """The exposure time in units of :py:attr:`time_unit`. *Read-only*"""

    @property
    def oversample(self):
        """Oversample factor when integrating over exposure time. *Read-only*"""
        return self._oversample

    @property
    def order(self):
        """The order of the numerical integration scheme. *Read-only*

        - ``0``: a centered Riemann sum
        - ``1``: trapezoid rule
        - ``2``: Simpson’s rule
        """
        return self._order

    @property
    def time_unit(self):
        """An ``astropy.units`` unit defining the time metric for the system."""
        return self._time_unit

    @time_unit.setter
    def time_unit(self, value):
        assert value.physical_type == "time"
        self._time_unit = value
        self._time_factor = value.in_units(units.day)

    @property
    def primary(self):
        """The primary (central) object in the Keplerian system."""
        return self._primary

    @property
    def secondaries(self):
        """A list of the secondary (orbiting) object(s) in the Keplerian system."""
        return self._secondaries

    @property
    def bodies(self):
        """A list of all objects in the Keplerian system."""
        return self._bodies

    @property
    def map_indices(self):
        """A list of the indices corresponding to each body in the design matrix."""
        return self._inds

    def show(
        self,
        t,
        cmap="plasma",
        res=300,
        interval=75,
        file=None,
        figsize=(3, 3),
        html5_video=True,
        window_pad=1.0,
    ):
        """Visualize the Keplerian system.

        Note that the body surface intensities are not normalized.

        Args:
            t (scalar or vector): The time(s) at which to evaluate the orbit and
                the map in units of :py:attr:`time_unit`.
            cmap (string or colormap instance, optional): The matplotlib colormap
                to use. Defaults to ``plasma``.
            res (int, optional): The resolution of the map in pixels on a
                side. Defaults to 300.
            figsize (tuple, optional): Figure size in inches. Default is
                (3, 3) for orthographic maps and (7, 3.5) for rectangular
                maps.
            interval (int, optional): Interval between frames in milliseconds
                (animated maps only). Defaults to 75.
            file (string, optional): The file name (including the extension)
                to save the animation to (animated maps only). Defaults to None.
            html5_video (bool, optional): If rendering in a Jupyter notebook,
                display as an HTML5 video? Default is True. If False, displays
                the animation using Javascript (file size will be larger.)
            window_pad (float, optional): Padding around the primary in units
                of the primary radius. Bodies outside of this window will be
                cropped. Default is 1.0.
        """
        # Not yet implemented
        if self._primary._map.nw is not None:  # pragma: no cover
            raise NotImplementedError(
                "Method not implemented for spectral maps."
            )

        # Render the maps & get the orbital positions
        if self._rv:
            self._primary.map._set_RV_filter()
            for sec in self._secondaries:
                sec.map._set_RV_filter()
        img_pri, img_sec, x, y, z = self.ops.render(
            self._math.reshape(self._math.to_array_or_tensor(t), [-1])
            * self._time_factor,
            res,
            self._primary._r,
            self._primary._m,
            self._primary._prot,
            self._primary._t0,
            self._primary._theta0,
            self._primary._map._inc,
            self._primary._map._obl,
            self._primary._map._y,
            self._primary._map._u,
            self._primary._map._f,
            self._primary._map._alpha,
            self._primary._map._tau,
            self._primary._map._delta,
            self._math.to_array_or_tensor(
                [sec._r for sec in self._secondaries]
            ),
            self._math.to_array_or_tensor(
                [sec._m for sec in self._secondaries]
            ),
            self._math.to_array_or_tensor(
                [sec._prot for sec in self._secondaries]
            ),
            self._math.to_array_or_tensor(
                [sec._t0 for sec in self._secondaries]
            ),
            self._math.to_array_or_tensor(
                [sec._theta0 for sec in self._secondaries]
            ),
            self._get_periods(),
            self._math.to_array_or_tensor(
                [sec._ecc for sec in self._secondaries]
            ),
            self._math.to_array_or_tensor(
                [sec._w for sec in self._secondaries]
            ),
            self._math.to_array_or_tensor(
                [sec._Omega for sec in self._secondaries]
            ),
            self._math.to_array_or_tensor(
                [sec._inc for sec in self._secondaries]
            ),
            self._math.to_array_or_tensor(
                [sec._map._inc for sec in self._secondaries]
            ),
            self._math.to_array_or_tensor(
                [sec._map._obl for sec in self._secondaries]
            ),
            self._math.to_array_or_tensor(
                [sec._map._y for sec in self._secondaries]
            ),
            self._math.to_array_or_tensor(
                [sec._map._u for sec in self._secondaries]
            ),
            self._math.to_array_or_tensor(
                [sec._map._f for sec in self._secondaries]
            ),
            self._math.to_array_or_tensor(
                [sec._map._alpha for sec in self._secondaries]
            ),
            self._math.to_array_or_tensor(
                [sec._map._tau for sec in self._secondaries]
            ),
            self._math.to_array_or_tensor(
                [sec._map._delta for sec in self._secondaries]
            ),
            self._math.to_array_or_tensor(
                [sec._map._sigr for sec in self._secondaries]
            ),
        )

        # Convert to units of the primary radiu
        x, y, z = (
            x / self._primary._r,
            y / self._primary._r,
            z / self._primary._r,
        )
        r = self._math.to_array_or_tensor(
            [sec._r for sec in self._secondaries]
        )
        r = r / self._primary._r

        # Evaluate if needed
        if config.lazy:
            img_pri = img_pri.eval()
            img_sec = img_sec.eval()
            x = x.eval()
            y = y.eval()
            z = z.eval()
            r = r.eval()

        # We need this to be of shape (nplanet, nframe)
        x = x.T
        y = y.T
        z = z.T

        # Ensure we have an array of frames
        if len(img_pri.shape) == 3:
            nframes = img_pri.shape[0]
        else:  # pragma: no cover
            nframes = 1
            img_pri = np.reshape(img_pri, (1,) + img_pri.shape)
            img_sec = np.reshape(img_sec, (1,) + img_sec.shape)
        animated = nframes > 1

        # Set up the plot
        fig, ax = plt.subplots(1, figsize=figsize)
        ax.axis("off")
        ax.set_xlim(-1.0 - window_pad, 1.0 + window_pad)
        ax.set_ylim(-1.0 - window_pad, 1.0 + window_pad)

        # Render the first frame
        img = [None for n in range(1 + len(self._secondaries))]
        circ = [None for n in range(1 + len(self._secondaries))]
        extent = np.array([-1.0, 1.0, -1.0, 1.0])
        img[0] = ax.imshow(
            img_pri[0],
            origin="lower",
            extent=extent,
            cmap=cmap,
            interpolation="none",
            vmin=np.nanmin(img_pri),
            vmax=np.nanmax(img_pri),
            animated=animated,
            zorder=0.0,
        )
        circ[0] = plt.Circle(
            (0, 0), 1, color="k", fill=False, zorder=1e-3, lw=2
        )
        ax.add_artist(circ[0])
        for i, _ in enumerate(self._secondaries):
            extent = np.array([x[i, 0], x[i, 0], y[i, 0], y[i, 0]]) + (
                r[i] * np.array([-1.0, 1.0, -1.0, 1.0])
            )
            img[i + 1] = ax.imshow(
                img_sec[i, 0],
                origin="lower",
                extent=extent,
                cmap=cmap,
                interpolation="none",
                vmin=np.nanmin(img_sec),
                vmax=np.nanmax(img_sec),
                animated=animated,
                zorder=z[i, 0],
            )
            circ[i] = plt.Circle(
                (x[i, 0], y[i, 0]),
                r[i],
                color="k",
                fill=False,
                zorder=z[i, 0] + 1e-3,
                lw=2,
            )
            ax.add_artist(circ[i])

        # Animation
        if animated:

            def updatefig(k):

                # Update Primary map
                img[0].set_array(img_pri[k])

                # Update Secondary maps & positions
                for i, _ in enumerate(self._secondaries):
                    extent = np.array([x[i, k], x[i, k], y[i, k], y[i, k]]) + (
                        r[i] * np.array([-1.0, 1.0, -1.0, 1.0])
                    )
                    if np.any(np.abs(extent) < 1.0 + window_pad):
                        img[i + 1].set_array(img_sec[i, k])
                        img[i + 1].set_extent(extent)
                        img[i + 1].set_zorder(z[i, k])
                        circ[i].center = (x[i, k], y[i, k])
                        circ[i].set_zorder(z[i, k] + 1e-3)

                return img + circ

            ani = FuncAnimation(
                fig, updatefig, interval=interval, blit=False, frames=nframes
            )

            # Business as usual
            if (file is not None) and (file != ""):
                if file.endswith(".mp4"):
                    ani.save(file, writer="ffmpeg")
                elif file.endswith(".gif"):
                    ani.save(file, writer="imagemagick")
                else:  # pragma: no cover
                    # Try and see what happens!
                    ani.save(file)
                plt.close()
            else:  # pragma: no cover
                try:
                    if "zmqshell" in str(type(get_ipython())):
                        plt.close()
                        if html5_video:
                            display(HTML(ani.to_html5_video()))
                        else:
                            display(HTML(ani.to_jshtml()))
                    else:
                        raise NameError("")
                except NameError:
                    plt.show()
                    plt.close()

            # Matplotlib generates an annoying empty
            # file when producing an animation. Delete it.
            try:
                os.remove("None0000000.png")
            except FileNotFoundError:
                pass

        else:

            if (file is not None) and (file != ""):
                fig.savefig(file)
                plt.close()
            else:  # pragma: no cover
                plt.show()

        if self._rv:
            self._primary.map._unset_RV_filter()
            for sec in self._secondaries:
                sec.map._unset_RV_filter()

    def design_matrix(self, t):
        """Compute the system flux design matrix at times ``t``.

        .. note::

            This is the *unweighted* design matrix, i.e., it does not
            include the scaling by the amplitude of each body's map.
            To perform this weighting, do

            .. code-block:: python

                X = sys.design_matrix(**kwargs)
                for i, body in zip(sys.map_indices, sys.bodies):
                    X[:, i] *= body.map.amp

        Args:
            t (scalar or vector): An array of times at which to evaluate
                the design matrix in units of :py:attr:`time_unit`.
        """
        return self.ops.X(
            self._math.reshape(self._math.to_array_or_tensor(t), [-1])
            * self._time_factor,
            self._primary._r,
            self._primary._m,
            self._primary._prot,
            self._primary._t0,
            self._primary._theta0,
            self._math.to_array_or_tensor(1.0),
            self._primary._map._inc,
            self._primary._map._obl,
            self._primary._map._u,
            self._primary._map._f,
            self._primary._map._alpha,
            self._primary._map._tau,
            self._primary._map._delta,
            self._math.to_array_or_tensor(
                [sec._r for sec in self._secondaries]
            ),
            self._math.to_array_or_tensor(
                [sec._m for sec in self._secondaries]
            ),
            self._math.to_array_or_tensor(
                [sec._prot for sec in self._secondaries]
            ),
            self._math.to_array_or_tensor(
                [sec._t0 for sec in self._secondaries]
            ),
            self._math.to_array_or_tensor(
                [sec._theta0 for sec in self._secondaries]
            ),
            self._get_periods(),
            self._math.to_array_or_tensor(
                [sec._ecc for sec in self._secondaries]
            ),
            self._math.to_array_or_tensor(
                [sec._w for sec in self._secondaries]
            ),
            self._math.to_array_or_tensor(
                [sec._Omega for sec in self._secondaries]
            ),
            self._math.to_array_or_tensor(
                [sec._inc for sec in self._secondaries]
            ),
            self._math.to_array_or_tensor(
                [
                    self._math.to_array_or_tensor(1.0)
                    for sec in self._secondaries
                ]
            ),
            self._math.to_array_or_tensor(
                [sec._map._inc for sec in self._secondaries]
            ),
            self._math.to_array_or_tensor(
                [sec._map._obl for sec in self._secondaries]
            ),
            self._math.to_array_or_tensor(
                [sec._map._u for sec in self._secondaries]
            ),
            self._math.to_array_or_tensor(
                [sec._map._f for sec in self._secondaries]
            ),
            self._math.to_array_or_tensor(
                [sec._map._alpha for sec in self._secondaries]
            ),
            self._math.to_array_or_tensor(
                [sec._map._tau for sec in self._secondaries]
            ),
            self._math.to_array_or_tensor(
                [sec._map._delta for sec in self._secondaries]
            ),
            self._math.to_array_or_tensor(
                [sec._map._sigr for sec in self._secondaries]
            ),
        )

    def flux(self, t, total=True):
        """Compute the system flux at times ``t``.

        Args:
            t (scalar or vector): An array of times at which to evaluate
                the flux in units of :py:attr:`time_unit`.
            total (bool, optional): Return the total system flux? Defaults to
                True. If False, returns arrays corresponding to the flux
                from each body.
        """
        X = self.design_matrix(t)

        # Weight the ylms by amplitude
        if self._reflected:
            # If we're doing reflected light, scale the amplitude of
            # each of the secondaries by the amplitude of the primary
            # (the illumination source).
            ay = [self._primary.map.amp * self._primary._map._y] + [
                self._primary.map.amp * body.map.amp * body._map._y
                for body in self._secondaries
            ]
        else:
            ay = [body.map.amp * body._map._y for body in self._bodies]

        if total:
            return self._math.dot(X, self._math.concatenate(ay))
        else:
            return [
                self._math.dot(X[:, idx], ay[i])
                for i, idx in enumerate(self._inds)
            ]

    def rv(self, t, keplerian=True, total=True):
        """Compute the observed radial velocity of the system at times ``t``.

        Args:
            t (scalar or vector): An array of times at which to evaluate
                the radial velocity in units of :py:attr:`time_unit`.
            keplerian (bool): Include the Keplerian component of the radial
                velocity of the primary? Default is True. If False, this
                method returns a model for only the radial velocity anomaly
                due to transits (the Rossiter-McLaughlin effect) and
                time-variable surface features (Doppler tomography) for all
                bodies in the system.
            total (bool, optional): Return the total system RV? Defaults to
                True. If False, returns arrays corresponding to the RV
                contribution from each body.

        """
        rv = self.ops.rv(
            self._math.reshape(self._math.to_array_or_tensor(t), [-1])
            * self._time_factor,
            self._primary._r,
            self._primary._m,
            self._primary._prot,
            self._primary._t0,
            self._primary._theta0,
            self._primary._map._amp,
            self._primary._map._inc,
            self._primary._map._obl,
            self._primary._map._y,
            self._primary._map._u,
            self._primary._map._alpha,
            self._primary._map._tau,
            self._primary._map._delta,
            self._primary._map._veq,
            self._math.to_array_or_tensor(
                [sec._r for sec in self._secondaries]
            ),
            self._math.to_array_or_tensor(
                [sec._m for sec in self._secondaries]
            ),
            self._math.to_array_or_tensor(
                [sec._prot for sec in self._secondaries]
            ),
            self._math.to_array_or_tensor(
                [sec._t0 for sec in self._secondaries]
            ),
            self._math.to_array_or_tensor(
                [sec._theta0 for sec in self._secondaries]
            ),
            self._get_periods(),
            self._math.to_array_or_tensor(
                [sec._ecc for sec in self._secondaries]
            ),
            self._math.to_array_or_tensor(
                [sec._w for sec in self._secondaries]
            ),
            self._math.to_array_or_tensor(
                [sec._Omega for sec in self._secondaries]
            ),
            self._math.to_array_or_tensor(
                [sec._inc for sec in self._secondaries]
            ),
            self._math.to_array_or_tensor(
                [sec._map._amp for sec in self._secondaries]
            ),
            self._math.to_array_or_tensor(
                [sec._map._inc for sec in self._secondaries]
            ),
            self._math.to_array_or_tensor(
                [sec._map._obl for sec in self._secondaries]
            ),
            self._math.to_array_or_tensor(
                [sec._map._y for sec in self._secondaries]
            ),
            self._math.to_array_or_tensor(
                [sec._map._u for sec in self._secondaries]
            ),
            self._math.to_array_or_tensor(
                [sec._map._alpha for sec in self._secondaries]
            ),
            self._math.to_array_or_tensor(
                [sec._map._tau for sec in self._secondaries]
            ),
            self._math.to_array_or_tensor(
                [sec._map._delta for sec in self._secondaries]
            ),
            self._math.to_array_or_tensor(
                [sec._map._sigr for sec in self._secondaries]
            ),
            self._math.to_array_or_tensor(
                [sec._map._veq for sec in self._secondaries]
            ),
            np.array(keplerian),
        )
        if total:
            return self._math.sum(rv, axis=0)
        else:
            return rv

    def position(self, t):
        """Compute the Cartesian positions of all bodies at times ``t``.

        Args:
            t (scalar or vector): An array of times at which to evaluate
                the position in units of :py:attr:`time_unit`.
        """
        x, y, z = self.ops.position(
            self._math.reshape(self._math.to_array_or_tensor(t), [-1])
            * self._time_factor,
            self._primary._m,
            self._primary._t0,
            self._math.to_array_or_tensor(
                [sec._m for sec in self._secondaries]
            ),
            self._math.to_array_or_tensor(
                [sec._t0 for sec in self._secondaries]
            ),
            self._get_periods(),
            self._math.to_array_or_tensor(
                [sec._ecc for sec in self._secondaries]
            ),
            self._math.to_array_or_tensor(
                [sec._w for sec in self._secondaries]
            ),
            self._math.to_array_or_tensor(
                [sec._Omega for sec in self._secondaries]
            ),
            self._math.to_array_or_tensor(
                [sec._inc for sec in self._secondaries]
            ),
        )
        fac = np.reshape(
            [self._primary._length_factor]
            + [sec._length_factor for sec in self._secondaries],
            [-1, 1],
        )
        return (x / fac, y / fac, z / fac)

    def _get_periods(self):
        periods = [None for sec in self._secondaries]
        for i, sec in enumerate(self._secondaries):
            if sec.porb:
                periods[i] = sec.porb
            else:
                periods[i] = (
                    (2 * np.pi)
                    * sec._a ** (3 / 2)
                    / (self._math.sqrt(G_grav * (self._primary._m + sec._m)))
                )
        return self._math.to_array_or_tensor(periods)

    def set_data(self, flux, C=None, cho_C=None):
        """Set the data vector and covariance matrix.

        This method is required by the :py:meth:`solve` method, which
        analytically computes the posterior over surface maps for all bodies
        in the system given a dataset and a prior, provided both are described
        as multivariate Gaussians.

        Args:
            flux (vector): The observed system light curve.
            C (scalar, vector, or matrix): The data covariance. This may be
                a scalar, in which case the noise is assumed to be
                homoscedastic, a vector, in which case the covariance
                is assumed to be diagonal, or a matrix specifying the full
                covariance of the dataset. Default is None. Either `C` or
                `cho_C` must be provided.
            cho_C (matrix): The lower Cholesky factorization of the data
                covariance matrix. Defaults to None. Either `C` or
                `cho_C` must be provided.
        """
        self._flux = self._math.cast(flux)
        self._C = self._linalg.Covariance(
            C=C, cho_C=cho_C, N=self._flux.shape[0]
        )

    def solve(self, *, design_matrix=None, t=None):
        """Solve the least-squares problem for the posterior over maps for all bodies.

        This method solves the generalized least squares problem given a system
        light curve and its covariance (set via the :py:meth:`set_data` method)
        and a Gaussian prior on the spherical harmonic coefficients
        (set via the :py:meth:`set_prior` method). The map amplitudes and
        coefficients of each of the bodies in the system are then set to the
        maximum a posteriori (MAP) solution.

        Args:
            design_matrix (matrix, optional): The flux design matrix, the
                quantity returned by :py:meth:`design_matrix`. Default is
                None, in which case this is computed based on ``kwargs``.
            t (vector, optional): The vector of times at which to evaluate
                :py:meth:`design_matrix`, if a design matrix is not provided.
                Default is None.

        Returns:
            The posterior mean for the spherical harmonic \
            coefficients `l > 0` and the Cholesky factorization of the \
            posterior covariance of all of the bodies in the system, \
            stacked in order (primary, followed by each of the secondaries \
            in the order they were provided.)

        .. note::
            Users may call the :py:meth:`draw` method of this class to draw
            from the posterior after calling :py:meth:`solve`.
        """
        # TODO: Implement for spectral maps?
        self._no_spectral()

        # Check that the data is set
        if self._flux is None or self._C is None:
            raise ValueError("Please provide a dataset with `set_data()`.")

        # Get the full design matrix
        if design_matrix is None:
            assert t is not None, "Please provide a time vector `t`."
            design_matrix = self.design_matrix(t)
        X = self._math.cast(design_matrix)

        # Get the data vector
        f = self._math.cast(self._flux)

        # Check for bodies whose priors are set
        self._solved_bodies = []
        inds = []
        dense_L = False
        for k, body in enumerate(self._bodies):

            if body.map._mu is None or body.map._L is None:

                # Subtract out this term from the data vector,
                # since it is fixed
                f -= body.map.amp * self._math.dot(
                    X[:, self._inds[k]], body.map.y
                )

            else:

                # Add to our list of indices/bodies to solve for
                inds.extend(self._inds[k])
                self._solved_bodies.append(body)
                if body.map._L.kind in ["matrix", "cholesky"]:
                    dense_L = True

        # Do we have at least one body?
        if len(self._solved_bodies) == 0:
            raise ValueError("Please provide a prior for at least one body.")

        # Keep only the terms we'll solve for
        X = X[:, inds]

        # Stack our priors
        mu = self._math.concatenate(
            [body.map._mu for body in self._solved_bodies]
        )

        if not dense_L:
            # We can just concatenate vectors
            LInv = self._math.concatenate(
                [
                    body.map._L.inverse * self._math.ones(body.map.Ny)
                    for body in self._solved_bodies
                ]
            )
        else:
            # FACT: The inverse of a block diagonal matrix
            # is the block diagonal matrix of the inverses.
            LInv = self._math.block_diag(
                *[
                    body.map._L.inverse * self._math.eye(body.map.Ny)
                    for body in self._solved_bodies
                ]
            )

        # Compute the MAP solution
        self._solution = self._linalg.solve(X, f, self._C.cholesky, mu, LInv)

        # Set all the map vectors
        x, cho_cov = self._solution
        n = 0
        for body in self._solved_bodies:
            inds = slice(n, n + body.map.Ny)
            body.map.amp = x[inds][0]
            body.map[1:, :] = x[inds][1:] / body.map.amp
            n += body.map.Ny

        # Return the mean and covariance
        self._solution = (x, cho_cov)
        return self._solution

    @property
    def solution(self):
        r"""The posterior probability distribution for the maps in the system.

        This is a tuple containing the mean and lower Cholesky factorization of the
        covariance of the amplitude-weighted spherical harmonic coefficient vectors,
        obtained by solving the regularized least-squares problem
        via the :py:meth:`solve` method.

        Note that to obtain the actual covariance matrix from the lower Cholesky
        factorization :math:`L`, simply compute :math:`L L^\top`.

        Note also that this is the posterior for the **amplitude-weighted**
        map vectors. Under this convention, the map amplitude is equal to the
        first term of the vector of each body and the spherical harmonic coefficients are
        equal to the vector normalized by the first term.
        """
        if self._solution is None:
            raise ValueError("Please call `solve()` first.")
        return self._solution

    def draw(self):
        """
        Draw a map from the posterior distribution and set
        the :py:attr:`y` map vector of each body.

        Users should call :py:meth:`solve` to enable this attribute.
        """
        if self._solution is None:
            raise ValueError("Please call `solve()` first.")

        # Number of coefficients
        N = np.sum([body.map.Ny for body in self._solved_bodies])

        # Fast multivariate sampling using the Cholesky factorization
        yhat, cho_ycov = self._solution
        u = self._math.cast(np.random.randn(N))
        x = yhat + self._math.dot(cho_ycov, u)

        # Set all the map vectors
        n = 0
        for body in self._solved_bodies:
            inds = slice(n, n + body.map.Ny)
            body.map.amp = x[inds][0]
            body.map[1:, :] = x[inds][1:] / body.map.amp
            n += body.map.Ny

    def lnlike(self, *, design_matrix=None, t=None, woodbury=True):
        """Returns the log marginal likelihood of the data given a design matrix.

        This method computes the marginal likelihood (marginalized over the
        spherical harmonic coefficients of all bodies) given a system
        light curve and its covariance (set via the :py:meth:`set_data` method)
        and a Gaussian prior on the spherical harmonic coefficients
        (set via the :py:meth:`set_prior` method).

        Args:
            design_matrix (matrix, optional): The flux design matrix, the
                quantity returned by :py:meth:`design_matrix`. Default is
                None, in which case this is computed based on ``kwargs``.
            t (vector, optional): The vector of times at which to evaluate
                :py:meth:`design_matrix`, if a design matrix is not provided.
                Default is None.
            woodbury (bool, optional): Solve the linear problem using the
                Woodbury identity? Default is True. The
                `Woodbury identity <https://en.wikipedia.org/wiki/Woodbury_matrix_identity>`_
                is used to speed up matrix operations in the case that the
                number of data points is much larger than the number of
                spherical harmonic coefficients. In this limit, it can
                speed up the code by more than an order of magnitude. Keep
                in mind that the numerical stability of the Woodbury identity
                is not great, so if you're getting strange results try
                disabling this. It's also a good idea to disable this in the
                limit of few data points and large spherical harmonic degree.

        Returns:
            lnlike: The log marginal likelihood.
        """
        # TODO: Implement for spectral maps?
        self._no_spectral()

        # Check that the data is set
        if self._flux is None or self._C is None:
            raise ValueError("Please provide a dataset with `set_data()`.")

        # Get the full design matrix
        if design_matrix is None:
            assert t is not None, "Please provide a time vector `t`."
            design_matrix = self.design_matrix(t)
        X = self._math.cast(design_matrix)

        # Get the data vector
        f = self._math.cast(self._flux)

        # Check for bodies whose priors are set
        self._solved_bodies = []
        inds = []
        dense_L = False
        for k, body in enumerate(self._bodies):

            if body.map._mu is None or body.map._L is None:

                # Subtract out this term from the data vector,
                # since it is fixed
                f -= body.map.amp * self._math.dot(
                    X[:, self._inds[k]], body.map.y
                )

            else:

                # Add to our list of indices/bodies to solve for
                inds.extend(self._inds[k])
                self._solved_bodies.append(body)
                if body.map._L.kind in ["matrix", "cholesky"]:
                    dense_L = True

        # Do we have at least one body?
        if len(self._solved_bodies) == 0:
            raise ValueError("Please provide a prior for at least one body.")

        # Keep only the terms we'll solve for
        X = X[:, inds]

        # Stack our priors
        mu = self._math.concatenate(
            [body.map._mu for body in self._solved_bodies]
        )

        # Compute the likelihood
        if woodbury:
            if not dense_L:
                # We can just concatenate vectors
                LInv = self._math.concatenate(
                    [
                        body.map._L.inverse * self._math.ones(body.map.Ny)
                        for body in self._solved_bodies
                    ]
                )
            else:
                LInv = self._math.block_diag(
                    *[
                        body.map._L.inverse * self._math.eye(body.map.Ny)
                        for body in self._solved_bodies
                    ]
                )
            lndetL = self._math.cast(
                [body.map._L.lndet for body in self._solved_bodies]
            )
            return self._linalg.lnlike_woodbury(
                X, f, self._C.inverse, mu, LInv, self._C.lndet, lndetL
            )
        else:
            if not dense_L:
                # We can just concatenate vectors
                L = self._math.concatenate(
                    [
                        body.map._L.value * self._math.ones(body.map.Ny)
                        for body in self._solved_bodies
                    ]
                )
            else:
                L = self._math.block_diag(
                    *[
                        body.map._L.value * self._math.eye(body.map.Ny)
                        for body in self._solved_bodies
                    ]
                )
            return self._linalg.lnlike(X, f, self._C.value, mu, L)


from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

import collections
# Dependency imports

import numpy as np
import tensorflow as tf

from tensorflow.python.ops import array_ops
from tensorflow.python.ops import control_flow_ops
from tensorflow.python.ops import tensor_array_ops
from tensorflow_probability.python import distributions
from tensorflow_probability.python.internal import distribution_util
from tensorflow_probability.python.mcmc import kernel as kernel_base
from tensorflow_probability.python.mcmc import metropolis_hastings
from tensorflow_probability.python.mcmc.internal import util as mcmc_util
from tensorflow.python.util import deprecation  # pylint: disable=g-direct-tensorflow-import

from inspect import signature

from GeneralLeapfrogIntegrator import GeneralLeapfrogIntegrator

UncalibratedRiemannManifoldHamiltonianMonteCarloKernelResults = collections.namedtuple(
    'UncalibratedRiemannManifoldHamiltonianMonteCarloKernelResults',
    [
        'log_acceptance_correction',
        'target_log_prob',        # For "next_state".
        'step_size',
        'num_leapfrog_steps',
    ])

RiemannManifoldHamiltonianMonteCarloExtraKernelResults = collections.namedtuple(
    'HamiltonianMonteCarloExtraKernelResults',
    [
        'step_size_assign',
    ])

class RiemannManifoldHamiltonianMonteCarlo(kernel_base.TransitionKernel):

  @deprecation.deprecated_args(
      '2019-05-22', 'The `step_size_update_fn` argument is deprecated. Use '
      '`tfp.mcmc.SimpleStepSizeAdaptation` instead.', 'step_size_update_fn')
  def __init__(self,
               target_log_prob_fn,
               step_size,
               num_leapfrog_steps,
               state_gradients_are_stopped=False,
               step_size_update_fn=None,
               seed=None,
               store_parameters_in_results=False,
               name=None):
    """Initializes this transition kernel.
    Args:
      target_log_prob_fn: Python callable which takes an argument like
        `current_state` (or `*current_state` if it's a list) and returns its
        (possibly unnormalized) log-density under the target distribution.
      step_size: `Tensor` or Python `list` of `Tensor`s representing the step
        size for the leapfrog integrator. Must broadcast with the shape of
        `current_state`. Larger step sizes lead to faster progress, but
        too-large step sizes make rejection exponentially more likely. When
        possible, it's often helpful to match per-variable step sizes to the
        standard deviations of the target distribution in each variable.
      num_leapfrog_steps: Integer number of steps to run the leapfrog integrator
        for. Total progress per HMC step is roughly proportional to
        `step_size * num_leapfrog_steps`.
      state_gradients_are_stopped: Python `bool` indicating that the proposed
        new state be run through `tf.stop_gradient`. This is particularly useful
        when combining optimization over samples from the HMC chain.
        Default value: `False` (i.e., do not apply `stop_gradient`).
      step_size_update_fn: Python `callable` taking current `step_size`
        (typically a `tf.Variable`) and `kernel_results` (typically
        `collections.namedtuple`) and returns updated step_size (`Tensor`s).
        Default value: `None` (i.e., do not update `step_size` automatically).
      seed: Python integer to seed the random number generator.
      store_parameters_in_results: If `True`, then `step_size` and
        `num_leapfrog_steps` are written to and read from eponymous fields in
        the kernel results objects returned from `one_step` and
        `bootstrap_results`. This allows wrapper kernels to adjust those
        parameters on the fly. This is incompatible with `step_size_update_fn`,
        which must be set to `None`.
      name: Python `str` name prefixed to Ops created by this function.
        Default value: `None` (i.e., 'hmc_kernel').
    """
    if step_size_update_fn and store_parameters_in_results:
      raise ValueError('It is invalid to simultaneously specify '
                       '`step_size_update_fn` and set '
                       '`store_parameters_in_results` to `True`.')

    impl = metropolis_hastings.MetropolisHastings(
        inner_kernel=UncalibratedRiemannManifoldHamiltonianMonteCarlo(
            target_log_prob_fn=target_log_prob_fn,
            step_size=step_size,
            num_leapfrog_steps=num_leapfrog_steps,
            state_gradients_are_stopped=state_gradients_are_stopped,
            seed=seed,
            name='rm_hmc_kernel' if name is None else name,
            store_parameters_in_results=store_parameters_in_results),
        seed=seed)
    parameters = impl.inner_kernel.parameters.copy()

    parameters['step_size_update_fn'] = step_size_update_fn
    self._impl = impl
    self._parameters = parameters

  @property
  def target_log_prob_fn(self):
    return self._impl.inner_kernel.target_log_prob_fn

  @property
  def step_size(self):
    """Returns the step_size parameter.
    If `store_parameters_in_results` argument to the initializer was set to
    `True`, this only returns the value of the `step_size` placed in the kernel
    results by the `bootstrap_results` method. The actual step size in that
    situation is governed by the `previous_kernel_results` argument to
    `one_step` method.
    Returns:
      step_size: A floating point `Tensor` or a list of such `Tensors`.
    """
    return self._impl.inner_kernel.step_size

  @property
  def num_leapfrog_steps(self):
    """Returns the num_leapfrog_steps parameter.
    If `store_parameters_in_results` argument to the initializer was set to
    `True`, this only returns the value of the `num_leapfrog_steps` placed in
    the kernel results by the `bootstrap_results` method. The actual
    `num_leapfrog_steps` in that situation is governed by the
    `previous_kernel_results` argument to `one_step` method.
    Returns:
      num_leapfrog_steps: An integer `Tensor`.
    """
    return self._impl.inner_kernel.num_leapfrog_steps

  @property
  def state_gradients_are_stopped(self):
    return self._impl.inner_kernel.state_gradients_are_stopped

  @property
  def step_size_update_fn(self):
    return self._parameters['step_size_update_fn']

  @property
  def seed(self):
    return self._impl.inner_kernel.seed

  @property
  def name(self):
    return self._impl.inner_kernel.name

  @property
  def parameters(self):
    """Return `dict` of ``__init__`` arguments and their values."""
    return self._parameters

  @property
  def is_calibrated(self):
    return True

  #@tf.function
  def one_step(self, current_state, previous_kernel_results):
    """Runs one iteration of Hamiltonian Monte Carlo.
    Args:
      current_state: `Tensor` or Python `list` of `Tensor`s representing the
        current state(s) of the Markov chain(s). The first `r` dimensions index
        independent chains, `r = tf.rank(target_log_prob_fn(*current_state))`.
      previous_kernel_results: `collections.namedtuple` containing `Tensor`s
        representing values from previous calls to this function (or from the
        `bootstrap_results` function.)
    Returns:
      next_state: Tensor or Python list of `Tensor`s representing the state(s)
        of the Markov chain(s) after taking exactly one step. Has same type and
        shape as `current_state`.
      kernel_results: `collections.namedtuple` of internal calculations used to
        advance the chain.
    Raises:
      ValueError: if there isn't one `step_size` or a list with same length as
        `current_state`.
    """
    previous_step_size_assign = (
        [] if self.step_size_update_fn is None
        else (previous_kernel_results.extra.step_size_assign
              if mcmc_util.is_list_like(
                  previous_kernel_results.extra.step_size_assign)
              else [previous_kernel_results.extra.step_size_assign]))

    with tf.control_dependencies(previous_step_size_assign):
      next_state, kernel_results = self._impl.one_step(
          current_state, previous_kernel_results)
      if self.step_size_update_fn is not None:
        step_size_assign = self.step_size_update_fn(  # pylint: disable=not-callable
            self.step_size, kernel_results)
        kernel_results = kernel_results._replace(
            extra=RiemannManifoldHamiltonianMonteCarloExtraKernelResults(
                step_size_assign=step_size_assign))

    return next_state, kernel_results

  def bootstrap_results(self, init_state):
    """Creates initial `previous_kernel_results` using a supplied `state`."""
    kernel_results = self._impl.bootstrap_results(init_state)
    if self.step_size_update_fn is not None:
      step_size_assign = self.step_size_update_fn(self.step_size, None)  # pylint: disable=not-callable
      kernel_results = kernel_results._replace(
          extra=RiemannManifoldHamiltonianMonteCarloExtraKernelResults(
              step_size_assign=step_size_assign))
    return kernel_results

class UncalibratedRiemannManifoldHamiltonianMonteCarlo(kernel_base.TransitionKernel):
  """Runs one step of Uncalibrated Hamiltonian Monte Carlo.
  Warning: this kernel will not result in a chain which converges to the
  `target_log_prob`. To get a convergent MCMC, use `HamiltonianMonteCarlo(...)`
  or `MetropolisHastings(UncalibratedHamiltonianMonteCarlo(...))`.
  For more details on `UncalibratedHamiltonianMonteCarlo`, see
  `HamiltonianMonteCarlo`.
  """

  def __init__(self,
               target_log_prob_fn,
               step_size,
               num_leapfrog_steps,
               state_gradients_are_stopped=False,
               seed=None,
               store_parameters_in_results=False,
               name=None):
    """Initializes this transition kernel.
    Args:
      target_log_prob_fn: Python callable which takes an argument like
        `current_state` (or `*current_state` if it's a list) and returns its
        (possibly unnormalized) log-density under the target distribution.
      step_size: `Tensor` or Python `list` of `Tensor`s representing the step
        size for the leapfrog integrator. Must broadcast with the shape of
        `current_state`. Larger step sizes lead to faster progress, but
        too-large step sizes make rejection exponentially more likely. When
        possible, it's often helpful to match per-variable step sizes to the
        standard deviations of the target distribution in each variable.
      num_leapfrog_steps: Integer number of steps to run the leapfrog integrator
        for. Total progress per HMC step is roughly proportional to
        `step_size * num_leapfrog_steps`.
      state_gradients_are_stopped: Python `bool` indicating that the proposed
        new state be run through `tf.stop_gradient`. This is particularly useful
        when combining optimization over samples from the HMC chain.
        Default value: `False` (i.e., do not apply `stop_gradient`).
      seed: Python integer to seed the random number generator.
      store_parameters_in_results: If `True`, then `step_size` and
        `num_leapfrog_steps` are written to and read from eponymous fields in
        the kernel results objects returned from `one_step` and
        `bootstrap_results`. This allows wrapper kernels to adjust those
        parameters on the fly.
      name: Python `str` name prefixed to Ops created by this function.
        Default value: `None` (i.e., 'hmc_kernel').
    """
    if seed is not None and tf.executing_eagerly():
      # TODO(b/68017812): Re-enable once TFE supports `tf.random_shuffle` seed.
      raise NotImplementedError('Specifying a `seed` when running eagerly is '
                                'not currently supported. To run in Eager '
                                'mode with a seed, use `tf.set_random_seed`.')
    self._seed_stream = distributions.SeedStream(seed, 'hmc_one_step')
    if not store_parameters_in_results:
      mcmc_util.warn_if_parameters_are_not_simple_tensors(
          dict(step_size=step_size, num_leapfrog_steps=num_leapfrog_steps))
    self._parameters = dict(
        target_log_prob_fn=target_log_prob_fn,
        step_size=step_size,
        num_leapfrog_steps=num_leapfrog_steps,
        state_gradients_are_stopped=state_gradients_are_stopped,
        seed=seed,
        name=name,
        store_parameters_in_results=store_parameters_in_results,
    )
    self._momentum_dtype = None

  @property
  def target_log_prob_fn(self):
    return self._parameters['target_log_prob_fn']

  @property
  def step_size(self):
    """Returns the step_size parameter.
    If `store_parameters_in_results` argument to the initializer was set to
    `True`, this only returns the value of the `step_size` placed in the kernel
    results by the `bootstrap_results` method. The actual step size in that
    situation is governed by the `previous_kernel_results` argument to
    `one_step` method.
    Returns:
      step_size: A floating point `Tensor` or a list of such `Tensors`.
    """
    return self._parameters['step_size']

  @property
  def num_leapfrog_steps(self):
    """Returns the num_leapfrog_steps parameter.
    If `store_parameters_in_results` argument to the initializer was set to
    `True`, this only returns the value of the `num_leapfrog_steps` placed in
    the kernel results by the `bootstrap_results` method. The actual
    `num_leapfrog_steps` in that situation is governed by the
    `previous_kernel_results` argument to `one_step` method.
    Returns:
      num_leapfrog_steps: An integer `Tensor`.
    """
    return self._parameters['num_leapfrog_steps']

  @property
  def state_gradients_are_stopped(self):
    return self._parameters['state_gradients_are_stopped']

  @property
  def seed(self):
    return self._parameters['seed']

  @property
  def name(self):
    return self._parameters['name']

  @property
  def parameters(self):
    """Return `dict` of ``__init__`` arguments and their values."""
    return self._parameters

  @property
  def is_calibrated(self):
    return False

  @property
  def _store_parameters_in_results(self):
    return self._parameters['store_parameters_in_results']

  def expand(self, current):
    """Expands tensors to that they are of rank 2

    Arguments:
        * current: tensor to expand
    Returns:
        * expanded: expanded tensor

    """
    currentShape=tf.pad(
            tf.shape(current),
            paddings=[[tf.where(tf.rank(current) > 1, 0, 1), 0]],
            constant_values=1)
    expanded=tf.reshape(current, currentShape)
    return(expanded)

  #@tf.function
  def get_hessian(self, argv):
          n = None
          gradeint=None
          _gradients=None
          #G = None
          tensorList = []
          for x in range(len(argv)):
              tensorList.append(argv[x])
          with tf.GradientTape(persistent=True, watch_accessed_variables=False) as SecondTape:
              for x in tensorList:
                  SecondTape.watch(x)
              with tf.GradientTape(persistent=False, watch_accessed_variables=False) as FirstTape:
                  FirstTape.watch(tensorList)
                  argv = tensorList
                  theta = []
                  index = 0
                  for info in self.restoreShapes:
                      theta.append(tf.reshape(argv[index:index + info[1]], info[0]))
                      index += info[1]
                  potential = self.target_log_prob_fn(*theta)
              _gradients = FirstTape.gradient(potential, tensorList, unconnected_gradients=tf.UnconnectedGradients.NONE)
              SecondTape.watch(_gradients)
              #gradient=tf.convert_to_tensor(_gradients)
              gradient=tf.stack(_gradients)
              #n = array_ops.size(tensorList)
              n=len(tensorList)
          """
          loop_vars = [
                  array_ops.constant(0, tf.int32),
                  tensor_array_ops.TensorArray(tf.float32, n)
          ]

          gradientIter=iter(_gradients)
          _, hessian = control_flow_ops.while_loop(
              lambda j, _: j < n,
              lambda j, result: (j + 1,
                         result.write(j, SecondTape.gradient(next(gradientIter), tensorList, unconnected_gradients=tf.UnconnectedGradients.ZERO))),
              loop_vars
          )
          """

          hessian = SecondTape.jacobian(gradient, tensorList, unconnected_gradients=tf.UnconnectedGradients.ZERO)
          #_shape = array_ops.shape(tensorList)
          #_reshaped_hessian = array_ops.reshape(hessian.stack(),
          #                                         array_ops.concat((_shape, _shape), 0))
          #hessians = _reshaped_hessian
          #G = -hessians
          #print(hessian)
          G=-tf.convert_to_tensor(hessian)
          #print(G)

          """
          loop_vars = [
                array_ops.constant(0, tf.int32),
                tensor_array_ops.TensorArray(tf.float32, n)
            ]
          #print(gradient)
          #print(n)
          for x in range(n):
                #print(gradient[x])
          #print(gradient[0])
          #print(gradient[1])
          _, hessian = control_flow_ops.while_loop(
                lambda j, _: j < n,
                lambda j, result: (j + 1,
                                   result.write(j, SecondTape.gradient(gradient[j], tensorList, unconnected_gradients=tf.UnconnectedGradients.ZERO))),
                loop_vars
          )

          _shape = array_ops.shape(tensorList)
          _reshaped_hessian = array_ops.reshape(hessian.stack(),
                                                   array_ops.concat((_shape, _shape), 0))
          hessians = _reshaped_hessian
          G = -hessians
          """

          # Compute first-order derivatives and iterate for each x in xs.
          """
          hessians = []
          _gradients = gradients(ys, xs, **kwargs)
          for gradient, x in zip(_gradients, xs):
            # change shape to one-dimension without graph branching
            gradient = array_ops.reshape(gradient, [-1])

            # Declare an iterator and tensor array loop variables for the gradients.

            # Iterate over all elements of the gradient and compute second order
            # derivatives.
          """

          return(G, potential)

  #@tf.function
  def hamiltonian_energy_fn(self, *argv):
      #print()
      G,potential=self.get_hessian(argv)
      #print("potential", potential)
      #print()
      #print(G)
      #smallG=tf.minimum(G, 1e6)
      #smallG=tf.maximum(smallG, -1e6)

      #condG=tf.reduce_any(tf.math.not_equal(G, smallG))
      #G = smallG
      s, u, v = tf.linalg.svd(G)
      safeInvS=1/(tf.where(tf.reduce_any(tf.math.equal(s,0)), tf.ones_like(s),s))
      safeLog=tf.math.log(tf.where(tf.reduce_any(tf.math.less_equal(s, 0)), tf.ones_like(s), s))
      condA=tf.reduce_any(s==0)
      #condB=tf.reduce_any(tf.math.not_equal(G, smallG))
      #logDet = tf.where(tf.reduce_any([condA, condB, condG]),tf.constant(0.0),tf.reduce_sum(safeLog))
      logDet = tf.where(tf.reduce_any([condA]),tf.constant(0.0),tf.reduce_sum(safeLog))

      #invG=tf.where(tf.reduce_any([condA, condB, condG]), tf.linalg.diag(tf.ones_like(s)),
      #               tf.matmul(v ,tf.matmul(tf.linalg.diag(safeInvS), tf.transpose(u))))

      invG=tf.where(tf.reduce_any([condA]), tf.linalg.diag(tf.ones_like(s)),
                     tf.matmul(v ,tf.matmul(tf.linalg.diag(safeInvS), tf.transpose(u))))
      #print(-potential)
      #print(G)
      #print(invG)
      #print(logDet)
      #print()
      return(-potential, G, invG, logDet)

  #@tf.function
  def run_integrator(self,step_sizes, num_leapfrog_steps,current_momentum_parts, current_state_parts):
      integrator = GeneralLeapfrogIntegrator(
          self.hamiltonian_energy_fn, step_sizes, num_leapfrog_steps)
      [
          next_state_parts,
          initial_kinetic,
          final_kinetic,
          final_target_log_prob
      ] = integrator(current_momentum_parts,
                     current_state_parts)
      print("ik", initial_kinetic)
      print("fk", final_kinetic)
      return(next_state_parts,
          initial_kinetic,
          final_kinetic,
          final_target_log_prob)

  @mcmc_util.set_doc(RiemannManifoldHamiltonianMonteCarlo.one_step.__doc__)
  def one_step(self, current_state, previous_kernel_results):
    with tf.compat.v2.name_scope(
        mcmc_util.make_name(self.name, 'hmc', 'one_step')):
      if self._store_parameters_in_results:
        step_size = previous_kernel_results.step_size
        num_leapfrog_steps = previous_kernel_results.num_leapfrog_steps
      else:
        step_size = self.step_size
        num_leapfrog_steps = self.num_leapfrog_steps
      [
          current_state_parts,
          step_sizes,
          current_target_log_prob,
      ] = _prepare_args(
          self.target_log_prob_fn,
          current_state,
          step_size,
          previous_kernel_results.target_log_prob,
          maybe_expand=True,
          state_gradients_are_stopped=self.state_gradients_are_stopped)

      self.restoreShapes = []
      for x in current_state_parts:
          n = 1
          shape = x.shape
          for m in shape:
              n *= m
          self.restoreShapes.append([shape, n])
      current_state_parts = [tf.reshape(part, [-1]) for part in current_state_parts]
      current_state_parts = tf.concat(current_state_parts, -1)
      temp=[]
      #print(current_state_parts)
      for x in range(current_state_parts.shape[0]):
            temp.append(current_state_parts[x])
      current_state_parts=temp
      #print(current_state_parts)

      current_momentum_parts = []

      for x in current_state_parts:
          current_momentum_parts.append(tf.random.normal(
                shape=tf.shape(input=x),
                dtype=self._momentum_dtype or x.dtype.base_dtype,
                seed=self._seed_stream()))

      next_state_parts, initial_kinetic, final_kinetic, final_target_log_prob =self.run_integrator(step_sizes, num_leapfrog_steps,current_momentum_parts, current_state_parts)

      if self.state_gradients_are_stopped:
        next_state_parts = [tf.stop_gradient(x) for x in next_state_parts]

      def maybe_flatten(x):
        return x if mcmc_util.is_list_like(current_state) else x[0]

      independent_chain_ndims = distribution_util.prefer_static_rank(
          current_target_log_prob)

      next_state_parts = maybe_flatten(next_state_parts)

      new_kernel_results = previous_kernel_results._replace(
          log_acceptance_correction=_compute_log_acceptance_correction(
              initial_kinetic, final_kinetic,
              independent_chain_ndims),
          target_log_prob=final_target_log_prob
      )
      argv = next_state_parts#[0]
      next_state_parts = []
      index = 0
      #print(self.restoreShapes)
      for info in self.restoreShapes:
          next_state_parts.append(tf.reshape(argv[index:index + info[1]], info[0]))
          index += info[1]

      return next_state_parts, new_kernel_results

  @mcmc_util.set_doc(RiemannManifoldHamiltonianMonteCarlo.bootstrap_results.__doc__)
  def bootstrap_results(self, init_state):
      with tf.compat.v2.name_scope(
              mcmc_util.make_name(self.name, 'hmc', 'bootstrap_results')):
          if not mcmc_util.is_list_like(init_state):
              init_state = [init_state]
          if self.state_gradients_are_stopped:
              init_state = [tf.stop_gradient(x) for x in init_state]
          else:
              init_state = [tf.convert_to_tensor(value=x) for x in init_state]
          [
              init_target_log_prob,
              init_grads_target_log_prob,
          ] = mcmc_util.maybe_call_fn_and_grads(self.target_log_prob_fn, init_state)
          if self._store_parameters_in_results:
              return UncalibratedRiemannManifoldHamiltonianMonteCarloKernelResults(
                  log_acceptance_correction=tf.zeros_like(init_target_log_prob),
                  target_log_prob=init_target_log_prob,
                  step_size=tf.nest.map_structure(
                      lambda x: tf.convert_to_tensor(  # pylint: disable=g-long-lambda
                          value=x,
                          dtype=init_target_log_prob.dtype,
                          name='step_size'),
                      self.step_size),
                  num_leapfrog_steps=tf.convert_to_tensor(
                      value=self.num_leapfrog_steps,
                      dtype=tf.int32,
                      name='num_leapfrog_steps'))
          else:
              return UncalibratedRiemannManifoldHamiltonianMonteCarloKernelResults(
                  log_acceptance_correction=tf.zeros_like(init_target_log_prob),
                  target_log_prob=init_target_log_prob,
                  step_size=[],
                  num_leapfrog_steps=[]
              )

def _compute_log_acceptance_correction(current_kinetic,
                                       proposed_kinetic,
                                       independent_chain_ndims,
                                       name=None):
  """Helper to `kernel` which computes the log acceptance-correction.
  A sufficient but not necessary condition for the existence of a stationary
  distribution, `p(x)`, is "detailed balance", i.e.:
  ```none
  p(x'|x) p(x) = p(x|x') p(x')
  ```
  In the Metropolis-Hastings algorithm, a state is proposed according to
  `g(x'|x)` and accepted according to `a(x'|x)`, hence
  `p(x'|x) = g(x'|x) a(x'|x)`.
  Inserting this into the detailed balance equation implies:
  ```none
      g(x'|x) a(x'|x) p(x) = g(x|x') a(x|x') p(x')
  ==> a(x'|x) / a(x|x') = p(x') / p(x) [g(x|x') / g(x'|x)]    (*)
  ```
  One definition of `a(x'|x)` which satisfies (*) is:
  ```none
  a(x'|x) = min(1, p(x') / p(x) [g(x|x') / g(x'|x)])
  ```
  (To see that this satisfies (*), notice that under this definition only at
  most one `a(x'|x)` and `a(x|x') can be other than one.)
  We call the bracketed term the "acceptance correction".
  In the case of UncalibratedHMC, the log acceptance-correction is not the log
  proposal-ratio. UncalibratedHMC augments the state-space with momentum, z.
  Assuming a standard Gaussian distribution for momentums, the chain eventually
  converges to:
  ```none
  p([x, z]) propto= target_prob(x) exp(-0.5 z**2)
  ```
  Relating this back to Metropolis-Hastings parlance, for HMC we have:
  ```none
  p([x, z]) propto= target_prob(x) exp(-0.5 z**2)
  g([x, z] | [x', z']) = g([x', z'] | [x, z])
  ```
  In other words, the MH bracketed term is `1`. However, because we desire to
  use a general MH framework, we can place the momentum probability ratio inside
  the metropolis-correction factor thus getting an acceptance probability:
  ```none
                       target_prob(x')
  accept_prob(x'|x) = -----------------  [exp(-0.5 z**2) / exp(-0.5 z'**2)]
                       target_prob(x)
  ```
  (Note: we actually need to handle the kinetic energy change at each leapfrog
  step, but this is the idea.)
  Args:
    current_momentums: `Tensor` representing the value(s) of the current
      momentum(s) of the state (parts).
    proposed_momentums: `Tensor` representing the value(s) of the proposed
      momentum(s) of the state (parts).
    independent_chain_ndims: Scalar `int` `Tensor` representing the number of
      leftmost `Tensor` dimensions which index independent chains.
    name: Python `str` name prefixed to Ops created by this function.
      Default value: `None` (i.e., 'compute_log_acceptance_correction').
  Returns:
    log_acceptance_correction: `Tensor` representing the `log`
      acceptance-correction.  (See docstring for mathematical definition.)
  """
  with tf.compat.v2.name_scope(name or 'compute_log_acceptance_correction'):
      """
      #current_momentums=tf.reshape(current_momentums,proposed_momentums.shape)
      log_current_kinetic, log_proposed_kinetic = [], []
      for current_momentum, proposed_momentum in zip(
          current_momentums, proposed_momentums):
        axis = tf.range(independent_chain_ndims, tf.rank(current_momentum))
        log_current_kinetic.append(_log_sum_sq(current_momentum, axis))
        log_proposed_kinetic.append(_log_sum_sq(proposed_momentum, axis))
      current_kinetic = 0.5 * tf.exp(
          tf.reduce_logsumexp(
              input_tensor=tf.stack(log_current_kinetic, axis=-1), axis=-1))
      proposed_kinetic = 0.5 * tf.exp(
          tf.reduce_logsumexp(
              input_tensor=tf.stack(log_proposed_kinetic, axis=-1), axis=-1))
      """
      return mcmc_util.safe_sum([current_kinetic, -proposed_kinetic])#*0.0

def _prepare_args(target_log_prob_fn,
                  state,
                  step_size,
                  target_log_prob=None,
                  maybe_expand=False,
                  state_gradients_are_stopped=False):
  """Helper which processes input args to meet list-like assumptions."""
  state_parts = list(state) if mcmc_util.is_list_like(state) else [state]
  state_parts = [
      tf.convert_to_tensor(value=s, name='current_state') for s in state_parts
  ]
  if state_gradients_are_stopped:
    state_parts = [tf.stop_gradient(x) for x in state_parts]
  target_log_prob, _ = mcmc_util.maybe_call_fn_and_grads(
      target_log_prob_fn,
      state_parts,
      target_log_prob,
      None)
  step_sizes = (list(step_size) if mcmc_util.is_list_like(step_size)
                else [step_size])
  step_sizes = [
      tf.convert_to_tensor(
          value=s, name='step_size', dtype=target_log_prob.dtype)
      for s in step_sizes
  ]
  if len(step_sizes) == 1:
    step_sizes *= len(state_parts)
  if len(state_parts) != len(step_sizes):
    raise ValueError('There should be exactly one `step_size` or it should '
                     'have same length as `current_state`.')
  def maybe_flatten(x):
    return x if maybe_expand or mcmc_util.is_list_like(state) else x[0]
  return [
      maybe_flatten(state_parts),
      maybe_flatten(step_sizes),
      target_log_prob
  ]

def _log_sum_sq(x, axis=None):
  """Computes log(sum(x**2))."""
  return tf.reduce_logsumexp(
      input_tensor=2. * tf.math.log(tf.abs(x)), axis=axis)


"""
This class computes activity cliffs on all compounds
Derek van Tilborg, Eindhoven University of Technology, March 2022
"""

import numpy as np
from progress.bar import ShadyBar
from rdkit import Chem
from rdkit.Chem import AllChem
from rdkit.Chem import DataStructs
from Levenshtein import distance as levenshtein

from rdkit.Chem.Scaffolds.MurckoScaffold import MakeScaffoldGeneric as GraphFramework
from rdkit.Chem.Scaffolds.MurckoScaffold import GetScaffoldForMol

class Cliffs:
    """Find activity cliffs for a list of molecules and their bioactivity
        We use a three-part activity cliff definition consisting of Tanimoto similarity, generic Murcko scaffold
        similarity, and Levenshtein similarity. Compound pairs that have a similarity >= 0.7 and a fold-change
        in bioactivity of >= 10x, are considered a activity cliff. Because activity cliffs can be defined from
        different angles, we use a soft-consensus to determine if a compound is an activity cliff compound"""

    def __init__(self, smiles: list, activity: list, test_smiles: list = None, test_activity: list = None,
                 in_log10: bool = True, fold_threshold: int = 10, similarity_threshold: float = 0.9):
        """
            smiles: (lst) List of SMILES strings
            activity: (lst) List of bioactivities (standard in log10 [nM])
            test_smiles: (lst) List of SMILES strings for compounds in the test set (optional)
            test_activity: (lst) List of bioactivities (standard in log10 [nM]) for compounds in the test set (optional)
            in_log10: (bool) is the bioactivity provided in log10? If True (default) it will be deconverted to determine
            the fold change.
            fold_threshold: (int) how much must the fold change be before being an activity cliff? (default=10)
            similarity_threshold: (float) how similar must two compounds be before being an activity cliff (default=0.7)
        """
        self.in_log10 = in_log10
        self.smiles = smiles
        self.test_smiles = test_smiles
        if test_smiles is None:
            self.test_smiles = []
        self.all_smiles = self.smiles + self.test_smiles

        self.activity = activity
        self.test_activity = test_activity
        if test_activity is None:
            self.test_activity = []
        self.all_activity = self.activity + self.test_activity

        self.train_idx = list(range(len(smiles)))
        self.test_idx = list(range(len(smiles), len(self.all_smiles)))

        # Initiate a shitload of empty variables
        self.fc = None
        self.tanimoto_sim = None
        self.levenshtein_sim = None
        self.scaffold_sim = None
        self.tanimoto = None
        self.levenshtein = None
        self.scaffold = None
        self.soft_consensus = None
        self.tanimoto_tr = None
        self.tanimoto_tst = None
        self.scaffold_tr = None
        self.scaffold_tst = None
        self.levenshtein_tr = None
        self.levenshtein_tst = None
        self.soft_consensus_tr = None
        self.soft_consensus_tst = None
        self.db_fp = None
        self.db_scaf = None
        self.fold_threshold = fold_threshold
        self.similarity_threshold = similarity_threshold
        self.cliff_mols_tanimoto = None
        self.cliff_mols_tanimoto_tr = None
        self.cliff_mols_tanimoto_tst = None
        self.cliff_mols_scaffold = None
        self.cliff_mols_scaffold_tr = None
        self.cliff_mols_scaffold_tst = None
        self.cliff_mols_levenshtein = None
        self.cliff_mols_levenshtein_tr = None
        self.cliff_mols_levenshtein_tst = None
        self.cliff_mols_soft_consensus = None
        self.cliff_mols_soft_consensus_tr = None
        self.cliff_mols_soft_consensus_tst = None
        self.stats = None

        # Find cliffs, find which compounds are activity cliff compounds and subset the results if there is test data
        self.find_cliffs(fold_threshold, similarity_threshold)
        self.find_cliff_compounds()
        self.train_test_cliffs(self.train_idx, self.test_idx)
        self.get_stats()

    def find_fc(self, a: float, b: float):
        """Get the fold change of to bioactivities (deconvert from log10 if needed)"""

        if self.in_log10:
            a, b = 10**a, 10**b
        return max([a, b]) / min([a, b])

    def get_fc(self, waitbar: bool = True):
        """ Calculates the pairwise fold difference in compound activity given a list of activities"""

        act_len = len(self.all_activity)
        m = np.zeros([act_len, act_len])
        if waitbar:
            bar = ShadyBar('Calculating pairwise fold change                        ', max=act_len, check_tty=False)
        # Calculate upper triangle of matrix
        for i in range(act_len):
            if waitbar:
                bar.next()
            for j in range(i, act_len):
                m[i, j] = self.find_fc(self.all_activity[i], self.all_activity[j])

        if waitbar:
            bar.finish()

        # Fill in the lower triangle without having to loop (saves ~50% of time)
        m = m + m.T - np.diag(np.diag(m))
        # Fill the diagonal with 0's
        np.fill_diagonal(m, 0)

        self.fc = m

    def get_levenshtein_matrix(self, waitbar: bool = True, normalize: bool = True):
        """ Calculates a matrix of levenshtein similarity scores for a list of SMILES string"""

        smiles = self.all_smiles
        smi_len = len(smiles)
        if waitbar:
            bar = ShadyBar('Calculating pairwise Levenshtein similarity (normalized)', max=smi_len, check_tty=False)
        m = np.zeros([smi_len, smi_len])
        # Calculate upper triangle of matrix
        for i in range(smi_len):
            if waitbar:
                bar.next()
            for j in range(i, smi_len):
                if normalize:
                    m[i, j] = levenshtein(smiles[i], smiles[j]) / max(len(smiles[i]), len(smiles[j]))
                else:
                    m[i, j] = levenshtein(smiles[i], smiles[j])
        if waitbar:
            bar.finish()
        # Fill in the lower triangle without having to loop (saves ~50% of time)
        m = m + m.T - np.diag(np.diag(m))
        # Get from a distance to a similarity
        m = 1 - m

        # Fill the diagonal with 0's
        np.fill_diagonal(m, 0)

        self.levenshtein_sim = m

    def get_tanimoto_matrix(self, waitbar: bool = True, radius: int = 2, nBits: int = 1024):
        """ Calculates a matrix of Tanimoto similarity scores for a list of SMILES string"""

        # Make a fingerprint database
        self.db_fp = {}
        for smi in self.all_smiles:
            m = Chem.MolFromSmiles(smi)
            fp = AllChem.GetMorganFingerprintAsBitVect(m, radius=radius, nBits=nBits)
            self.db_fp[smi] = fp

        smi_len = len(self.all_smiles)
        if waitbar:
            bar = ShadyBar('Calculating pairwise Tanimoto similarity                ', max=smi_len, check_tty=False)
        m = np.zeros([smi_len, smi_len])
        # Calculate upper triangle of matrix
        for i in range(smi_len):
            if waitbar:
                bar.next()
            for j in range(i, smi_len):
                m[i, j] = DataStructs.TanimotoSimilarity(self.db_fp[self.all_smiles[i]],
                                                         self.db_fp[self.all_smiles[j]])
        if waitbar:
            bar.finish()
        # Fill in the lower triangle without having to loop (saves ~50% of time)
        m = m + m.T - np.diag(np.diag(m))
        # Fill the diagonal with 0's
        np.fill_diagonal(m, 0)

        self.tanimoto_sim = m

    def get_scaffold_matrix(self, waitbar: bool = True, radius: int = 2, nBits: int = 1024):
        """ Calculates a matrix of Tanimoto similarity scores for a list of SMILES string """

        # Make scaffold database
        self.db_scaf = {}
        for smi in self.all_smiles:
            m = Chem.MolFromSmiles(smi)
            try:
                skeleton = GraphFramework(m)
            except Exception:  # In the very rare case this doesnt work, use a normal scaffold
                print(f"Could not create a generic scaffold of {smi}, used a normal scaffold instead")
                skeleton = GetScaffoldForMol(m)
            skeleton_fp = AllChem.GetMorganFingerprintAsBitVect(skeleton, radius=radius, nBits=nBits)
            self.db_scaf[smi] = skeleton_fp

        smi_len = len(self.all_smiles)
        if waitbar:
            bar = ShadyBar('Calculating pairwise generic scaffold similarity        ', max=smi_len, check_tty=False)
        m = np.zeros([smi_len, smi_len])
        # Calculate upper triangle of matrix
        for i in range(smi_len):
            if waitbar:
                bar.next()
            for j in range(i, smi_len):
                m[i, j] = DataStructs.TanimotoSimilarity(self.db_scaf[self.all_smiles[i]],
                                                         self.db_scaf[self.all_smiles[j]])
        if waitbar:
            bar.finish()
        # Fill in the lower triangle without having to loop (saves ~50% of time)
        m = m + m.T - np.diag(np.diag(m))
        # Fill the diagonal with 0's
        np.fill_diagonal(m, 0)

        self.scaffold_sim = m

    def find_cliffs(self, fold_threshold: int = 10, similarity_threshold: float = 0.90):
        """Find Tanimoto, scaffold, and Levenshtein activity cliffs """

        # Calculate fold change and similarity
        self.get_fc()
        self.get_tanimoto_matrix()
        self.get_scaffold_matrix()
        self.get_levenshtein_matrix()

        # Determine which compound pairs are cliffs
        self.tanimoto = np.logical_and(self.fc > fold_threshold,
                                       self.tanimoto_sim > similarity_threshold).astype(int)

        self.levenshtein = np.logical_and(self.fc > fold_threshold,
                                          self.levenshtein_sim > similarity_threshold).astype(int)

        self.scaffold = np.logical_and(self.fc > fold_threshold,
                                       self.scaffold_sim > similarity_threshold).astype(int)

        # If a compound pair is a cliff in at least 1 method for soft consensus
        self.soft_consensus = self.tanimoto + self.scaffold + self.levenshtein
        self.soft_consensus[self.soft_consensus > 0] = 1

    def find_cliff_compounds(self):
        """ Find activity cliff compounds (having at least 1 cliff with any compound) for the different cliff types"""
        self.cliff_mols_tanimoto = [s for i, s in enumerate(self.all_smiles) if sum(self.tanimoto[i]) > 0]
        self.cliff_mols_scaffold = [s for i, s in enumerate(self.all_smiles) if sum(self.scaffold[i]) > 0]
        self.cliff_mols_levenshtein = [s for i, s in enumerate(self.all_smiles) if sum(self.levenshtein[i]) > 0]
        self.cliff_mols_soft_consensus = [s for i, s in enumerate(self.all_smiles) if sum(self.soft_consensus[i]) > 0]

    def train_test_cliffs(self, train_idx: list, test_idx: list):
        """ Subset the activity cliff matrices with the train/test indices"""

        # Tanimoto
        self.tanimoto_tr = self.tanimoto[np.ix_(train_idx, train_idx)]
        self.tanimoto_tst = self.tanimoto[np.ix_(test_idx, test_idx)]
        self.cliff_mols_tanimoto_tr = [s for s in self.cliff_mols_tanimoto if s in self.smiles]
        self.cliff_mols_tanimoto_tst = [s for s in self.cliff_mols_tanimoto if s in self.test_smiles]

        # Scaffold
        self.scaffold_tr = self.scaffold[np.ix_(train_idx, train_idx)]
        self.scaffold_tst = self.scaffold[np.ix_(test_idx, test_idx)]
        self.cliff_mols_scaffold_tr = [s for s in self.cliff_mols_scaffold if s in self.smiles]
        self.cliff_mols_scaffold_tst = [s for s in self.cliff_mols_scaffold if s in self.test_smiles]

        # Levenshtein
        self.levenshtein_tr = self.levenshtein[np.ix_(train_idx, train_idx)]
        self.levenshtein_tst = self.levenshtein[np.ix_(test_idx, test_idx)]
        self.cliff_mols_levenshtein_tr = [s for s in self.cliff_mols_levenshtein if s in self.smiles]
        self.cliff_mols_levenshtein_tst = [s for s in self.cliff_mols_levenshtein if s in self.test_smiles]

        # consensus
        self.soft_consensus_tr = self.soft_consensus[np.ix_(train_idx, train_idx)]
        self.soft_consensus_tst = self.soft_consensus[np.ix_(test_idx, test_idx)]
        self.cliff_mols_soft_consensus_tr = [s for s in self.cliff_mols_soft_consensus if s in self.smiles]
        self.cliff_mols_soft_consensus_tst = [s for s in self.cliff_mols_soft_consensus if s in self.test_smiles]

    def get_stats(self):
        """ Calculate various stats on the found activity cliffs"""
        self.stats = {
            "Fold change threshold":
                self.fold_threshold,
            "Structural similarity threshold":
                self.similarity_threshold,
            "n_tanimoto_cliffs":
                0 if self.tanimoto is None else sum(np.triu(self.tanimoto).flatten()),
            "n_scaffold_cliffs":
                0 if self.scaffold is None else sum(np.triu(self.scaffold).flatten()),
            "n_levenstein_cliffs":
                0 if self.levenshtein is None else sum(np.triu(self.levenshtein).flatten()),
            "n_soft_consensus_cliffs":
                0 if self.soft_consensus is None else sum(np.triu(self.soft_consensus).flatten()),
            "n_tanimoto_cliffs_train":
                0 if self.tanimoto_tr is None else sum(np.triu(self.tanimoto_tr).flatten()),
            "n_scaffold_cliffs_train":
                0 if self.scaffold_tr is None else sum(np.triu(self.scaffold_tr).flatten()),
            "n_levenstein_cliffs_train":
                0 if self.levenshtein_tr is None else sum(np.triu(self.levenshtein_tr).flatten()),
            "n_soft_consensus_cliffs_train":
                0 if self.soft_consensus_tr is None else sum(np.triu(self.soft_consensus_tr).flatten()),
            "n_tanimoto_cliffs_test":
                0 if self.tanimoto_tst is None else sum(np.triu(self.tanimoto_tst).flatten()),
            "n_scaffold_cliffs_test":
                0 if self.scaffold_tst is None else sum(np.triu(self.scaffold_tst).flatten()),
            "n_levenstein_cliffs_test":
                0 if self.levenshtein_tst is None else sum(np.triu(self.levenshtein_tst).flatten()),
            "n_soft_consensus_cliffs_test":
                0 if self.soft_consensus_tst is None else sum(np.triu(self.soft_consensus_tst).flatten()),
            "n_compounds":
                len(self.all_smiles),
            "n_compounds_train":
                len(self.all_smiles)-len(self.test_smiles),
            "n_compounds_test":
                len(self.test_smiles),
            "n_tanimoto_cliff_compounds":
                0 if self.cliff_mols_tanimoto is None else len(self.cliff_mols_tanimoto),
            "n_scaffold_cliff_compounds":
                0 if self.cliff_mols_scaffold is None else len(self.cliff_mols_scaffold),
            "n_levenshtein_cliff_compounds":
                0 if self.cliff_mols_levenshtein is None else len(self.cliff_mols_levenshtein),
            "n_soft_consensus_cliff_compounds":
                0 if self.cliff_mols_soft_consensus is None else len(self.cliff_mols_soft_consensus),
            "n_tanimoto_cliff_compounds_train":
                0 if self.cliff_mols_tanimoto_tr is None else len(self.cliff_mols_tanimoto_tr),
            "n_scaffold_cliff_compounds_train":
                0 if self.cliff_mols_scaffold_tr is None else len(self.cliff_mols_scaffold_tr),
            "n_levenshtein_cliff_compounds_train":
                0 if self.cliff_mols_levenshtein_tr is None else len(self.cliff_mols_levenshtein_tr),
            "n_soft_consensus_cliff_compounds_train":
                0 if self.cliff_mols_soft_consensus_tr is None else len(self.cliff_mols_soft_consensus_tr),
            "n_tanimoto_cliff_compounds_test":
                0 if self.cliff_mols_tanimoto_tst is None else len(self.cliff_mols_tanimoto_tst),
            "n_scaffold_cliff_compounds_test":
                0 if self.cliff_mols_scaffold_tst is None else len(self.cliff_mols_scaffold_tst),
            "n_levenshtein_cliff_compounds_test":
                0 if self.cliff_mols_levenshtein_tst is None else len(self.cliff_mols_levenshtein_tst),
            "n_soft_consensus_cliff_compounds_test":
                0 if self.cliff_mols_soft_consensus_tst is None else len(self.cliff_mols_soft_consensus_tst),
        }

    def __repr__(self):
        self.get_stats()
        out = ''
        for i in list(self.stats.keys()):
            if len(self.test_smiles) == 0:  # if no test data is provided, don't show test stats (they will be 0)
                if not i.endswith('train') and not i.endswith('test'):
                    out += f'{i}: {self.stats[i]}\n'
            else:
                out += f'{i}: {self.stats[i]}\n'
        return out