# -*- coding: utf-8 -*-

__all__ = ["TTVOrbit", "compute_expected_transit_times"]

import aesara_theano_fallback.tensor as tt
import numpy as np

from ..utils import as_tensor_variable
from .keplerian import KeplerianOrbit


def compute_expected_transit_times(min_time, max_time, period, t0):
    """Compute the expected transit times within a dataset

    Args:
        min_time (float): The start time of the dataset
        max_time (float): The end time of the dataset
        period (array): The periods of the planets
        t0 (array): The reference transit times for the planets

    Returns:
        A list of arrays of expected transit times for each planet

    """
    periods = np.atleast_1d(period)
    t0s = np.atleast_1d(t0)
    transit_times = []
    for period, t0 in zip(periods, t0s):
        min_ind = np.floor((min_time - t0) / period)
        max_ind = np.ceil((max_time - t0) / period)
        times = t0 + period * np.arange(min_ind, max_ind, 1)
        times = times[(min_time <= times) & (times <= max_time)]
        transit_times.append(times)
    return transit_times


class TTVOrbit(KeplerianOrbit):
    """A generalization of a Keplerian orbit with transit timing variations

    Only one of the arguments ``ttvs`` or ``transit_times`` can be given and
    the other will be computed from the one that was provided.

    In practice the way this works is that the time axis is shifted to account
    for the TTVs before solving for a standard Keplerian orbit. To identify
    which times corrorspond to which transits, this will find the closest
    labelled transit to the timestamp and adjust the timestamp accordingly.
    This means that things will go *very badly* if the time between
    neighboring transits is larger than ``2*period``.

    Args:
        ttvs: A list (with on entry for each planet) of "O-C" vectors for each
            transit of each planet in units of days. "O-C" means the
            difference between the observed transit time and the transit time
            expected for a regular periodic orbit.
        transit_times: A list (with on entry for each planet) of transit times
            for each transit of each planet in units of days. These times will
            be used to compute the implied (least squares) ``ttv_period`` and
            ``t0``. It is possible to supply a separate ``period`` parameter
            that will set the shape of the transits, but care should be taken
            to make sure that ``period`` and ``ttv_period`` don't diverge
            because things will break if the time between neighboring transits
            is larger than ``2*period``.
        transit_inds: A list of integer value tensors giving the transit
            number for each transit in ``transit_times'' or ``ttvs``. This is
            useful when not all transits are observed. This should always be
            zero indexed.
        delta_log_period: If using the ``transit_times`` argument, this
            parameter specifies the difference (in natural log) between the
            leqast squares period and the effective period of the transit.

    """

    def __init__(self, *args, **kwargs):
        ttvs = kwargs.pop("ttvs", None)
        transit_times = kwargs.pop("transit_times", None)
        transit_inds = kwargs.pop("transit_inds", None)
        if ttvs is None and transit_times is None:
            raise ValueError(
                "one of 'ttvs' or 'transit_times' must be " "defined"
            )
        if ttvs is not None:
            self.ttvs = [as_tensor_variable(ttv, ndim=1) for ttv in ttvs]
            if transit_inds is None:
                self.transit_inds = [
                    tt.arange(ttv.shape[0]) for ttv in self.ttvs
                ]
            else:
                self.transit_inds = [
                    tt.cast(as_tensor_variable(inds, ndim=1), "int64")
                    for inds in transit_inds
                ]

        else:
            # If transit times are given, compute the least squares period and
            # t0 based on these times.
            self.transit_times = []
            self.ttvs = []
            self.transit_inds = []
            period = []
            t0 = []
            for i, times in enumerate(transit_times):
                times = as_tensor_variable(times, ndim=1)
                if transit_inds is None:
                    inds = tt.arange(times.shape[0])
                else:
                    inds = tt.cast(
                        as_tensor_variable(transit_inds[i]), "int64"
                    )
                self.transit_inds.append(inds)

                # A convoluted version of linear regression; don't ask
                N = times.shape[0]
                sumx = tt.sum(inds)
                sumx2 = tt.sum(inds**2)
                sumy = tt.sum(times)
                sumxy = tt.sum(inds * times)
                denom = N * sumx2 - sumx**2
                slope = (N * sumxy - sumx * sumy) / denom
                intercept = (sumx2 * sumy - sumx * sumxy) / denom
                expect = intercept + inds * slope

                period.append(slope)
                t0.append(intercept)
                self.ttvs.append(times - expect)
                self.transit_times.append(times)

            kwargs["t0"] = tt.stack(t0)

            # We'll have two different periods: one that is the mean difference
            # between transit times and one that is a parameter that sets the
            # transit shape. If a "period" parameter is not given, these will
            # be the same. Users will probably want to put a prior relating the
            # two periods if they use separate values.
            self.ttv_period = tt.stack(period)
            if "period" not in kwargs:
                if "delta_log_period" in kwargs:
                    kwargs["period"] = tt.exp(
                        tt.log(self.ttv_period)
                        + kwargs.pop("delta_log_period")
                    )
                else:
                    kwargs["period"] = self.ttv_period

        super(TTVOrbit, self).__init__(*args, **kwargs)

        if ttvs is not None:
            self.ttv_period = self.period
            self.transit_times = [
                self.t0[i] + self.period[i] * self.transit_inds[i] + ttv
                for i, ttv in enumerate(self.ttvs)
            ]

        # Compute the full set of transit times
        self.all_transit_times = []
        for i, inds in enumerate(self.transit_inds):
            expect = self.t0[i] + self.period[i] * tt.arange(inds.max() + 1)
            self.all_transit_times.append(
                tt.set_subtensor(expect[inds], self.transit_times[i])
            )

        # Set up a histogram for identifying the transit offsets
        self._bin_edges = [
            tt.concatenate(
                (
                    [tts[0] - 0.5 * self.ttv_period[i]],
                    0.5 * (tts[1:] + tts[:-1]),
                    [tts[-1] + 0.5 * self.ttv_period[i]],
                )
            )
            for i, tts in enumerate(self.all_transit_times)
        ]
        self._bin_values = [
            tt.concatenate(([tts[0]], tts, [tts[-1]]))
            for i, tts in enumerate(self.all_transit_times)
        ]

    def _get_model_dt(self, t):
        vals = []
        for i in range(len(self.ttvs)):
            inds = tt.extra_ops.searchsorted(self._bin_edges[i], t)
            vals.append(self._bin_values[i][inds])
        return tt.stack(vals, -1)

    def _warp_times(self, t, _pad=True):
        # This is the key function that takes the TTVs into account by
        # stretching the time axis around each transit
        if _pad:
            return tt.shape_padright(t) - self._get_model_dt(t)
        return t - self._get_model_dt(t)
