"""pytensor implementation of the Wiener First Passage Time Distribution.

This code is based on Sam Mathias's Pytensor/Theano implementation
of the WFPT distribution here:
https://gist.github.com/sammosummo/c1be633a74937efaca5215da776f194b.
"""

import jax.numpy as jnp
import numpy as np
import pymc as pm
import pytensor
import pytensor.tensor as pt
from jax import debug as jax_debug
from jax import lax, vmap
from jax.scipy.stats import norm
from numpy import inf
from pymc.distributions.dist_math import check_parameters

from ..distribution_utils.dist import make_distribution

LOGP_LB = pm.floatX(-66.1)

π = np.pi
τ = 2 * π
sqrt_τ = pt.sqrt(τ)
log_π = pt.log(π)
log_τ = pt.log(τ)
log_4 = pt.log(4)


def _max(a: np.ndarray, b: np.ndarray) -> np.ndarray:
    return pt.max(pt.stack([a, b]), axis=0)


def k_small(rt: np.ndarray, err: float) -> np.ndarray:
    """Determine number of terms needed for small-t expansion.

    Parameters
    ----------
    rt
        A 1D numpy array of flipped RTs. (0, inf).
    err
        Error bound.

    Returns
    -------
    np.ndarray
        A 1D at array of k_small.
    """
    sqrt_rt = pt.sqrt(rt)
    log_rt = pt.log(rt)
    rt_log_2_sqrt_τ_rt_times_2 = rt * (log_4 + log_τ + log_rt)

    ks = 2 + pt.sqrt(-err * rt_log_2_sqrt_τ_rt_times_2)
    ks = _max(ks, sqrt_rt + 1)

    condition = 2 * sqrt_τ * sqrt_rt * err < 1
    ks = pt.switch(condition, ks, 2)

    return ks


def k_large(rt: np.ndarray, err: float) -> np.ndarray:
    """Determine number of terms needed for large-t expansion.

    Parameters
    ----------
    rt
        An 1D numpy array of flipped RTs. (0, inf).
    err
        Error bound

    Returns
    -------
    np.ndarray
        A 1D at array of k_large.
    """
    log_rt = pt.log(rt)
    sqrt_rt = pt.sqrt(rt)
    log_err = pt.log(err)

    π_rt_err = π * rt * err
    π_sqrt_rt = π * sqrt_rt

    kl = pt.sqrt(-2 * (log_π + log_err + log_rt)) / π_sqrt_rt
    kl = _max(kl, 1.0 / pt.sqrt(π_sqrt_rt))
    kl = pt.switch(π_rt_err < 1, kl, 1.0 / π_sqrt_rt)

    return kl


def compare_k(rt: np.ndarray, err: float) -> np.ndarray:
    """Compute and compare k_small with k_large.

    Parameters
    ----------
    rt
        An 1D numpy of flipped RTs. (0, inf).
    err
        Error bound.

    Returns
    -------
    np.ndarray
        A 1D boolean at array of which implementation should be used.
    """
    ks = k_small(rt, err)
    kl = k_large(rt, err)

    return ks < kl


def get_ks(k_terms: int, fast: bool) -> np.ndarray:
    """Return an array of ks.

    Returns an array of ks given the number of terms needed to approximate the sum of
    the infinite series.

    Parameters
    ----------
    k_terms
        number of terms needed
    fast
        whether the function is used in the fast of slow expansion.

    Returns
    -------
    np.ndarray
        An array of ks.
    """
    ks = (
        pt.arange(-pt.floor((k_terms - 1) / 2), pt.ceil((k_terms - 1) / 2) + 1)
        if fast
        else pt.arange(1, k_terms + 1).reshape((-1, 1))
    )

    return ks.astype(pytensor.config.floatX)


def ftt01w_fast(tt: np.ndarray, w: float, k_terms: int) -> np.ndarray:
    """Perform fast computation of ftt01w.

    Density function for lower-bound first-passage times with drift rate set to 0 and
    upper bound set to 1, calculated using the fast-RT expansion.

    Parameters
    ----------
    tt
        Flipped, normalized RTs. (0, inf).
    w
        Normalized decision starting point. (0, 1).
    k_terms
        number of terms to use to approximate the PDF.

    Returns
    -------
    np.ndarray
        The approximated function f(tt|0, 1, w).
    """
    # Slightly changed the original code to mimic the paper and
    # ensure correctness
    k = get_ks(k_terms, fast=True)

    # A log-sum-exp trick is used here
    y = w + 2 * k.reshape((-1, 1))
    r = -pt.power(y, 2) / (2 * tt)
    c = pt.max(r, axis=0)
    p = pt.exp(c) * pt.sum(y * pt.exp(r - c), axis=0)
    # Normalize p
    p = p / pt.sqrt(2 * π * pt.power(tt, 3))

    return p


def ftt01w_slow(tt: np.ndarray, w: float, k_terms: int) -> np.ndarray:
    """Perform slow computation of ftt01w.

    Density function for lower-bound first-passage times with drift rate set to 0 and
    upper bound set to 1, calculated using the slow-RT expansion.

    Parameters
    ----------
    tt
        Flipped, normalized RTs. (0, inf).
    w
        Normalized decision starting point. (0, 1).
    k_terms
        number of terms to use to approximate the PDF.

    Returns
    -------
    np.ndarray
        The approximated function f(tt|0, 1, w).
    """
    k = get_ks(k_terms, fast=False)
    y = k * pt.sin(k * π * w)
    r = -pt.power(k, 2) * pt.power(π, 2) * tt / 2
    p = pt.sum(y * pt.exp(r), axis=0) * π

    return p


def ftt01w(
    rt: np.ndarray,
    a: float,
    w: float,
    err: float = 1e-7,
    k_terms: int = 10,
) -> np.ndarray:
    """Compute the approximate density of f(tt|0,1,w).

    Parameters
    ----------
    rt
        Flipped Response Rates. (0, inf).
    a
        Value of decision upper bound. (0, inf).
    w
        Normalized decision starting point. (0, 1).
    err
        Error bound.
    k_terms
        number of terms to use to approximate the PDF.

    Returns
    -------
    np.ndarray
        The Approximated density of f(tt|0,1,w).
    """
    tt = rt / a**2.0

    lambda_rt = compare_k(tt, err)

    p_fast = ftt01w_fast(tt, w, k_terms)
    p_slow = ftt01w_slow(tt, w, k_terms)

    p = pt.switch(lambda_rt, p_fast, p_slow)

    return p


def logp_ddm(
    data: np.ndarray,
    v: float,
    a: float,
    z: float,
    t: float,
    err: float = 1e-15,
    k_terms: int = 20,
    epsilon: float = 1e-15,
) -> np.ndarray:
    """Compute analytical likelihood for the DDM model with `sv`.

    Computes the log-likelihood of the drift diffusion model f(t|v,a,z) using
    the method and implementation of Navarro & Fuss, 2009.

    Parameters
    ----------
    data
        data: 2-column numpy array of (response time, response)
    v
        Mean drift rate. (-inf, inf).
    a
        Value of decision upper bound. (0, inf).
    z
        Normalized decision starting point. (0, 1).
    t
        Non-decision time [0, inf).
    err
        Error bound.
    k_terms
        number of terms to use to approximate the PDF.
    epsilon
        A small positive number to prevent division by zero or
        taking the log of zero.

    Returns
    -------
    np.ndarray
        The analytical likelihoods for DDM.
    """
    data = pt.reshape(data, (-1, 2)).astype(pytensor.config.floatX)
    rt = pt.abs(data[:, 0])
    response = data[:, 1]
    flip = response > 0
    a = a * 2.0
    v_flipped = pt.switch(flip, -v, v)  # transform v if x is upper-bound response
    z_flipped = pt.switch(flip, 1 - z, z)  # transform z if x is upper-bound response
    rt = rt - t
    negative_rt = rt < 0
    rt = pt.switch(negative_rt, epsilon, rt)

    p = pt.maximum(ftt01w(rt, a, z_flipped, err, k_terms), pt.exp(LOGP_LB))

    logp = pt.where(
        rt <= epsilon,
        LOGP_LB,
        pt.log(p)
        - v_flipped * a * z_flipped
        - (v_flipped**2 * rt / 2.0)
        - 2.0 * pt.log(a),
    )

    checked_logp = check_parameters(logp, a >= 0, msg="a >= 0")
    checked_logp = check_parameters(checked_logp, z >= 0, msg="z >= 0")
    checked_logp = check_parameters(checked_logp, z <= 1, msg="z <= 1")
    return checked_logp


def logp_ddm_sdv(
    data: np.ndarray,
    v: float,
    a: float,
    z: float,
    t: float,
    sv: float,
    err: float = 1e-15,
    k_terms: int = 20,
    epsilon: float = 1e-15,
) -> np.ndarray:
    """Compute the log-likelihood of the drift diffusion model f(t|v,a,z).

    Using the method and implementation of Navarro & Fuss, 2009.

    Parameters
    ----------
    data
        2-column numpy array of (response time, response)
    v
        Mean drift rate. (-inf, inf).
    a
        Value of decision upper bound. (0, inf).
    z
        Normalized decision starting point. (0, 1).
    t
        Non-decision time [0, inf).
    sv
        Standard deviation of the drift rate [0, inf).
    err
        Error bound.
    k_terms
        number of terms to use to approximate the PDF.
    epsilon
        A small positive number to prevent division by zero or taking the log of zero.

    Returns
    -------
    np.ndarray
        The log likelihood of the drift diffusion model with the standard deviation
        of sv.
    """
    if sv == 0:
        return logp_ddm(data, v, a, z, t, err, k_terms, epsilon)

    data = pt.reshape(data, (-1, 2)).astype(pytensor.config.floatX)
    rt = pt.abs(data[:, 0])
    response = data[:, 1]
    flip = response > 0
    a = a * 2.0
    v_flipped = pt.switch(flip, -v, v)  # transform v if x is upper-bound response
    z_flipped = pt.switch(flip, 1 - z, z)  # transform z if x is upper-bound response
    rt = rt - t
    negative_rt = rt < 0
    rt = pt.switch(negative_rt, epsilon, rt)

    p = pt.maximum(ftt01w(rt, a, z_flipped, err, k_terms), pt.exp(LOGP_LB))

    logp = pt.switch(
        rt <= epsilon,
        LOGP_LB,
        pt.log(p)
        + (
            (a * z_flipped * sv) ** 2
            - 2 * a * v_flipped * z_flipped
            - (v_flipped**2) * rt
        )
        / (2 * (sv**2) * rt + 2)
        - 0.5 * pt.log(sv**2 * rt + 1)
        - 2 * pt.log(a),
    )

    checked_logp = check_parameters(logp, a >= 0, msg="a >= 0")
    checked_logp = check_parameters(checked_logp, z >= 0, msg="z >= 0")
    checked_logp = check_parameters(checked_logp, z <= 1, msg="z <= 1")
    checked_logp = check_parameters(checked_logp, sv > 0, msg="sv > 0")
    return checked_logp


ddm_bounds = {
    "v": (-inf, inf),
    "a": (0.0, inf),
    "z": (0.0, 1.0),
    "t": (0.0, inf),
}
ddm_sdv_bounds = ddm_bounds | {"sv": (0.0, inf)}

ddm_params = ["v", "a", "z", "t"]
ddm_sdv_params = ddm_params + ["sv"]

DDM: type[pm.Distribution] = make_distribution(
    rv="ddm",
    loglik=logp_ddm,
    list_params=ddm_params,
    bounds=ddm_bounds,
)

DDM_SDV: type[pm.Distribution] = make_distribution(
    rv="ddm_sdv",
    loglik=logp_ddm_sdv,
    list_params=ddm_sdv_params,
    bounds=ddm_sdv_bounds,
)


LOGP_MIN = jnp.exp(LOGP_LB)
LOGP_MAX = jnp.exp(-LOGP_LB)


# Racing Diffusion Model (RDM)
def _jax_rdm_tpdf(t, b, v, A, eps_t: float = 1e-8):
    """RDM first-passage-time density for a single accumulator."""
    t_eff = jnp.maximum(t, eps_t)
    sqrt_t = jnp.sqrt(t_eff)

    alpha = (b - A - t_eff * v) / sqrt_t
    beta = (b - t_eff * v) / sqrt_t

    return (
        1.0
        / A
        * (
            -v * norm.cdf(alpha)
            + (1.0 / sqrt_t) * norm.pdf(alpha)
            + v * norm.cdf(beta)
            - (1.0 / sqrt_t) * norm.pdf(beta)
        )
    )


def _jax_rdm_tcdf(t, b, v, A, eps_t: float = 1e-8, eps_v: float = 1e-6):
    """RDM first-passage-time CDF for a single accumulator."""
    t_eff = jnp.maximum(t, eps_t)
    sqrt_t = jnp.sqrt(t_eff)
    v_eff = jnp.maximum(jnp.abs(v), eps_v) * jnp.sign(jnp.where(v == 0.0, 1.0, v))

    alpha_1 = jnp.sqrt(2.0) * (v_eff * t_eff - b) / jnp.sqrt(2.0 * t_eff)
    alpha_2 = jnp.sqrt(2.0) * (v_eff * t_eff - b + A) / jnp.sqrt(2.0 * t_eff)

    beta_1 = jnp.sqrt(2.0) * (-v_eff * t_eff - b) / jnp.sqrt(2.0 * t_eff)
    beta_2 = jnp.sqrt(2.0) * (-v_eff * t_eff - b + A) / jnp.sqrt(2.0 * t_eff)

    coeff = 1.0 / (2.0 * v_eff * A)
    F = (
        coeff * (norm.cdf(alpha_2) - norm.cdf(alpha_1))
        + sqrt_t / A * (alpha_2 * norm.cdf(alpha_2) - alpha_1 * norm.cdf(alpha_1))
        - coeff
        * (
            jnp.exp(2.0 * v_eff * (b - A)) * norm.cdf(beta_2)
            - jnp.exp(2.0 * v_eff * b) * norm.cdf(beta_1)
        )
        + sqrt_t / A * (norm.pdf(alpha_2) - norm.pdf(alpha_1))
    )

    return F


def _jax_rdm3_ll(t, ch, A, b, v0, v1, v2, eps_t: float = 1e-8):
    """Log-likelihood for 3-choice RDM at JAX level."""
    v_vector = jnp.stack(jnp.broadcast_arrays(v0, v1, v2))
    all_pdfs = vmap(lambda v: _jax_rdm_tpdf(t, b, v, A, eps_t=eps_t))(v_vector)
    all_cdfs = vmap(lambda v: _jax_rdm_tcdf(t, b, v, A, eps_t=eps_t))(v_vector)

    idx = jnp.arange(ch.shape[0])
    pdf_winner = all_pdfs[ch, idx]

    mask = jnp.arange(len(v_vector))[:, None] != ch
    survivors_losers = 1.0 - all_cdfs

    likelihood = pdf_winner * jnp.prod(jnp.where(mask, survivors_losers, 1.0), axis=0)

    return jnp.log(jnp.clip(likelihood, LOGP_MIN, LOGP_MAX))


def logp_rdm3(
    data: np.ndarray,
    A: float,
    b: float,
    v0: float,
    v1: float,
    v2: float,
    t: float,
    epsilon: float = 1e-10,
) -> np.ndarray:
    """Compute the log-likelihood of the RDM model with 3 drift rates."""
    data_reshaped = jnp.reshape(data, (-1, 2)).astype(pytensor.config.floatX)
    rt = jnp.abs(data_reshaped[:, 0])
    rt = rt - t
    response = data_reshaped[:, 1]
    response_int = response.astype(jnp.int_)

    is_negative_rt = rt <= 0.0
    rt_safe = jnp.where(is_negative_rt, jnp.asarray(epsilon, dtype=rt.dtype), rt)

    logp = _jax_rdm3_ll(rt_safe, response_int, A, b, v0, v1, v2, eps_t=epsilon)
    logp = jnp.where(is_negative_rt, jnp.asarray(LOGP_LB, dtype=logp.dtype), logp)
    logp = logp.squeeze()
    logp = jax_check_parameters(logp, b > A, msg="b > A")
    return logp


def jax_check_parameters(logp, condition, msg="Condition failed", print_msg=False):
    """Check parameters for validity in JAX.

    Equivalent to pymc.check_parameters with can_be_replaced_by_ninf=True.

    Note: This uses `jax.debug.print`, which is safe to use under JAX
    transformations but may produce repeated messages if called inside vmapped
    or jitted code. Returning -inf effectively rejects the sample.
    """
    violations = jnp.logical_not(condition)

    if print_msg:

        def _print_message(_):
            jax_debug.print(
                "jax_check_parameters: {msg}, {n_violations} violations",
                msg=msg,
                n_violations=jnp.sum(violations),
            )

        lax.cond(jnp.any(violations), _print_message, lambda _: None, operand=None)

    return jnp.where(condition, logp, -jnp.inf)


rdm3_params = ["A", "b", "v0", "v1", "v2", "t"]

rdm3_bounds = {
    "A": (0.0, inf),
    "b": (0.1, inf),
    "v0": (0.001, inf),
    "v1": (0.001, inf),
    "v2": (0.001, inf),
    "t": (0.0, inf),
}

RDM3: type[pm.Distribution] = make_distribution(
    rv="racing_diffusion_3",
    loglik=logp_rdm3,
    list_params=rdm3_params,
    bounds=rdm3_bounds,
)


# LBA
def _pt_normpdf(t):
    return (1 / pt.sqrt(2 * pt.pi)) * pt.exp(-(t**2) / 2)


def _pt_normcdf(t):
    return (1 / 2) * (1 + pt.erf(t / pt.sqrt(2)))


def _pt_tpdf(t, A, b, v, s):
    g = (b - A - t * v) / (t * s)
    h = (b - t * v) / (t * s)
    f = (
        -v * _pt_normcdf(g)
        + s * _pt_normpdf(g)
        + v * _pt_normcdf(h)
        - s * _pt_normpdf(h)
    ) / A
    return f


def _pt_tcdf(t, A, b, v, s):
    e1 = ((b - A - t * v) / A) * _pt_normcdf((b - A - t * v) / (t * s))
    e2 = ((b - t * v) / A) * _pt_normcdf((b - t * v) / (t * s))
    e3 = ((t * s) / A) * _pt_normpdf((b - A - t * v) / (t * s))
    e4 = ((t * s) / A) * _pt_normpdf((b - t * v) / (t * s))
    F = 1 + e1 - e2 + e3 - e4
    return F


def _pt_lba3_ll(t, ch, A, b, v0, v1, v2):
    s = 0.1
    __min = pt.exp(LOGP_LB)
    __max = pt.exp(-LOGP_LB)
    k = len([0, 1, 2])
    like = pt.zeros((*t.shape, k))
    running_idx = pt.arange(t.shape[0])

    like_1 = (
        _pt_tpdf(t, A, b, v0, s)
        * (1 - _pt_tcdf(t, A, b, v1, s))
        * (1 - _pt_tcdf(t, A, b, v2, s))
    )
    like_2 = (
        (1 - _pt_tcdf(t, A, b, v0, s))
        * _pt_tpdf(t, A, b, v1, s)
        * (1 - _pt_tcdf(t, A, b, v2, s))
    )
    like_3 = (
        (1 - _pt_tcdf(t, A, b, v0, s))
        * (1 - _pt_tcdf(t, A, b, v1, s))
        * _pt_tpdf(t, A, b, v2, s)
    )

    like = pt.stack([like_1, like_2, like_3], axis=-1)

    # One should RETURN this because otherwise it will be pruned from graph
    # like_printed = pytensor.printing.Print('like')(like)

    prob_neg = _pt_normcdf(-v0 / s) * _pt_normcdf(-v1 / s) * _pt_normcdf(-v2 / s)
    return pt.log(pt.clip(like[running_idx, ch] / (1 - prob_neg), __min, __max))


def _pt_lba2_ll(t, ch, A, b, v0, v1):
    s = 0.1
    __min = pt.exp(LOGP_LB)
    __max = pt.exp(-LOGP_LB)
    k = len([0, 1])
    like = pt.zeros((*t.shape, k))
    running_idx = pt.arange(t.shape[0])

    like_1 = _pt_tpdf(t, A, b, v0, s) * (1 - _pt_tcdf(t, A, b, v1, s))
    like_2 = (1 - _pt_tcdf(t, A, b, v0, s)) * _pt_tpdf(t, A, b, v1, s)

    like = pt.stack([like_1, like_2], axis=-1)

    # One should RETURN this because otherwise it will be pruned from graph
    # like_printed = pytensor.printing.Print('like')(like)

    prob_neg = _pt_normcdf(-v0 / s) * _pt_normcdf(-v1 / s)
    return pt.log(pt.clip(like[running_idx, ch] / (1 - prob_neg), __min, __max))


def logp_lba2(
    data: np.ndarray,
    A: float,
    b: float,
    v0: float,
    v1: float,
) -> np.ndarray:
    """Compute the log-likelihood of the LBA model with 2 drift rates."""
    data = pt.reshape(data, (-1, 2)).astype(pytensor.config.floatX)
    rt = pt.abs(data[:, 0])
    response = data[:, 1]
    response_int = pt.cast(response, "int32")
    logp = _pt_lba2_ll(rt, response_int, A, b, v0, v1).squeeze()
    checked_logp = check_parameters(logp, b > A, msg="b > A")
    return checked_logp


def logp_lba3(
    data: np.ndarray,
    A: float,
    b: float,
    v0: float,
    v1: float,
    v2: float,
) -> np.ndarray:
    """Compute the log-likelihood of the LBA model with 3 drift rates."""
    data = pt.reshape(data, (-1, 2)).astype(pytensor.config.floatX)
    rt = pt.abs(data[:, 0])
    response = data[:, 1]
    response_int = pt.cast(response, "int32")
    logp = _pt_lba3_ll(rt, response_int, A, b, v0, v1, v2).squeeze()
    checked_logp = check_parameters(logp, b > A, msg="b > A")
    return checked_logp


lba2_params = ["A", "b", "v0", "v1"]
lba3_params = ["A", "b", "v0", "v1", "v2"]

lba2_bounds = {
    "A": (0.0, inf),
    "b": (0.2, inf),
    "v0": (0.0, inf),
    "v1": (0.0, inf),
}

lba3_bounds = {
    "A": (0.0, inf),
    "b": (0.2, inf),
    "v0": (0.0, inf),
    "v1": (0.0, inf),
    "v2": (0.0, inf),
}

LBA2: type[pm.Distribution] = make_distribution(
    rv="lba2",
    loglik=logp_lba2,
    list_params=lba2_params,
    bounds=lba2_bounds,
)

LBA3: type[pm.Distribution] = make_distribution(
    rv="lba3",
    loglik=logp_lba3,
    list_params=lba3_params,
    bounds=lba3_bounds,
)


def logp_poisson_race(
    data: np.ndarray,
    r1: float,
    r2: float,
    k1: float,
    k2: float,
    t: float,
    epsilon: float = 1e-15,
) -> np.ndarray:
    """Compute analytical log-likelihoods for a 2-accumulator Poisson race.

    Each accumulator time follows a Gamma distribution with continuous
    shape parameter k and rate r. The per-trial likelihood decomposes
    into the density of the winning accumulator evaluated at the observed
    decision time and the survival function of the losing accumulator
    at the same time.

    Implemented as in https://link.springer.com/article/10.3758/BF03212980
    with two modifications:
      1. We allow continuous shape parameters (k1, k2) rather than just integers.
      2. We do not condition on the underlying stimulus condition.

    Parameters
    ----------
    data
        2-column tensor of (response time, response). Response > 0 indicates
        accumulator 2 (the second accumulator, parameters ``r2``/``k2``);
        otherwise accumulator 1 (the first accumulator, ``r1``/``k1``).
        This matches simulators that encode choices as ``-1`` for winner index 0
        and ``+1`` for winner index 1.
    r1, r2
        Rates (> 0) for the two accumulators.
    k1, k2
        Shape parameters (> 0) for the two accumulators.
    t
        Non-decision time [0, inf).
    epsilon
        A small positive number to prevent division by zero or taking ``log(0)``.

    Returns
    -------
    np.ndarray
        Per-trial log-likelihoods (shape: ``(n_trials,)``).

        Note that this function constructs a symbolic PyTensor graph; when used
        inside a PyMC/PyTensor model the returned object is a symbolic tensor,
        and evaluating/compiling the graph yields an ndarray.
    """
    epsilon = pm.floatX(epsilon)
    one = pm.floatX(1.0)
    data = pt.reshape(data, (-1, 2)).astype(pytensor.config.floatX)

    rt = pt.abs(data[:, 0])
    response = data[:, 1]
    flip = response > 0

    rt = rt - t
    negative_rt = rt <= 0
    rt_safe = pt.maximum(rt, epsilon)

    r_c = pt.switch(flip, r2, r1)
    r_l = pt.switch(flip, r1, r2)
    k_c = pt.switch(flip, k2, k1)
    k_l = pt.switch(flip, k1, k2)

    r_c_safe = pt.maximum(r_c, epsilon)

    log_pdf = (
        k_c * pt.log(r_c_safe)
        + (k_c - 1.0) * pt.log(rt_safe)
        - r_c * rt_safe
        - pt.gammaln(k_c)
    )

    survival = pt.gammaincc(k_l, r_l * rt_safe)
    survival = pt.clip(survival, epsilon, one)
    log_survival = pt.log(survival)

    logp = log_pdf + log_survival
    logp = pt.switch(negative_rt, LOGP_LB, logp)

    checked = check_parameters(logp, r1 > 0, msg="r1 > 0")
    checked = check_parameters(checked, r2 > 0, msg="r2 > 0")
    checked = check_parameters(checked, k1 > 0, msg="k1 > 0")
    checked = check_parameters(checked, k2 > 0, msg="k2 > 0")
    checked = check_parameters(checked, t >= 0, msg="t >= 0")
    return checked


# set bounds
poisson_race_bounds = {
    "r1": (np.finfo(float).eps, np.inf),
    "r2": (np.finfo(float).eps, np.inf),
    "k1": (np.finfo(float).eps, np.inf),
    "k2": (np.finfo(float).eps, np.inf),
    "t": (0.0, np.inf),
}
poisson_race_params = list(poisson_race_bounds)

# build distribution
POISSON_RACE = make_distribution(
    rv="poisson_race",
    loglik=logp_poisson_race,
    list_params=poisson_race_params,
    bounds=poisson_race_bounds,
)


def softmax_inv_temperature(data: np.ndarray, beta: np.ndarray, *logits):
    """Compute the log-likelihood of the Inverse Softmax Temperature Model.

    Parameters
    ----------
    data
        1D array of responses (choices).
    beta
        A scalar for the softmax temperature (0, inf).
    *logits
        Logits for each choice excluding logit0.

    Returns
    -------
    pt.TensorVariable
        The log likelihood of the Inverse Softmax Temperature Model.
    """
    choices = pt.where(data < 1, 0.0, data).astype("int32")
    zeros = pt.zeros_like(data, dtype=pytensor.config.floatX)

    logits_stacked = pt.stack(
        [
            zeros,  # logit0 is always 0
            *(logit + zeros for logit in logits),
        ],
    )

    logits_scaled = logits_stacked * beta
    choice_logits = logits_scaled[choices, pt.arange(data.shape[0])]
    log_prob_choices = choice_logits - pt.logsumexp(
        logits_scaled, axis=0, keepdims=False
    )

    return log_prob_choices
