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

"""This module has an EXPERIMENTAL and UNTESTED implementation of analytic
marginalizion over eccentricity as a function of transit duration (since the
mapping is not 1-to-1). The implementation as written probably DOES NOT WORK,
but I've (DFM) left it here for my own future reference. Feel free to ping me
if you're interested in what's going on here!"""

__all__ = ["duration_to_eccentricity"]

from itertools import product

import aesara_theano_fallback.tensor as tt
import numpy as np
import pymc3 as pm

from .keplerian import KeplerianOrbit, _get_consistent_inputs


def duration_to_eccentricity(
    func, duration, ror, **kwargs
):  # pragma: no cover
    num_planets = kwargs.pop("num_planets", 1)
    orbit_type = kwargs.pop("orbit_type", KeplerianOrbit)
    name = kwargs.get("name", "dur_ecc")

    inputs = _get_consistent_inputs(
        kwargs.get("a", None),
        kwargs.get("period", None),
        kwargs.get("rho_star", None),
        kwargs.get("r_star", None),
        kwargs.get("m_star", None),
        kwargs.get("rho_star_units", None),
        kwargs.get("m_planet", 0.0),
        kwargs.get("m_planet_units", None),
    )
    a, period, rho_star, r_star, m_star, m_planet = inputs
    b = kwargs.get("b", 0.0)
    s = tt.sin(kwargs["omega"])
    umax_inv = tt.switch(tt.lt(s, 0), tt.sqrt(1 - s ** 2), 1.0)

    const = (
        period * tt.shape_padright(r_star) * tt.sqrt((1 + ror) ** 2 - b ** 2)
    )
    const /= np.pi * a

    u = duration / const

    e1 = -s * u ** 2 / ((s * u) ** 2 + 1)
    e2 = tt.sqrt((s ** 2 - 1) * u ** 2 + 1) / ((s * u) ** 2 + 1)

    models = []
    logjacs = []
    logprobs = []
    for args in product(*(zip("np", (-1, 1)) for _ in range(num_planets))):
        labels, signs = zip(*args)

        # Compute the eccentricity branch
        ecc = tt.stack([e1[i] + signs[i] * e2[i] for i in range(num_planets)])

        # Work out the Jacobian
        valid_ecc = tt.and_(tt.lt(ecc, 1.0), tt.ge(ecc, 0.0))
        logjac = tt.switch(
            tt.all(valid_ecc),
            tt.sum(
                0.5 * tt.log(1 - ecc ** 2)
                + 2 * tt.log(s * ecc + 1)
                - tt.log(tt.abs_(s + ecc))
                - tt.log(const)
            ),
            -np.inf,
        )
        ecc = tt.switch(valid_ecc, ecc, tt.zeros_like(ecc))

        # Create a sub-model to capture this component
        with pm.Model(name="dur_ecc_" + "_".join(labels)) as model:
            pm.Deterministic("ecc", ecc)
            orbit = orbit_type(ecc=ecc, **kwargs)
            logprob = tt.sum(func(orbit))

        models.append(model)
        logjacs.append(logjac)
        logprobs.append(logprob)

    # Compute the marginalized likelihood
    logjacs = tt.stack(logjacs)
    logprobs = tt.stack(logprobs)
    logprob = tt.switch(
        tt.gt(1.0 / u, umax_inv),
        tt.sum(pm.logsumexp(logprobs + logjacs)),
        -np.inf,
    )
    pm.Potential(name + "_logp", logprob)
    pm.Deterministic(name + "_logjacs", logjacs)
    pm.Deterministic(name + "_logprobs", logprobs)

    # Loop over models and compute the marginalized values for all the
    # parameters and deterministics
    norm = tt.sum(pm.logsumexp(logjacs))
    logw = tt.switch(
        tt.gt(1.0 / u, umax_inv),
        logjacs - norm,
        -np.inf + tt.zeros_like(logjacs),
    )
    pm.Deterministic(name + "_logw", logw)
    for k in models[0].named_vars.keys():
        name = k[len(models[0].name) + 1 :]
        pm.Deterministic(
            name,
            sum(
                tt.exp(logw[i]) * model.named_vars[model.name + "_" + name]
                for i, model in enumerate(models)
            ),
        )
