"""Type definitions for the HSSM package."""

from os import PathLike
from typing import Any, Callable, Literal, NotRequired, Optional, TypedDict, Union

import bambi as bmb
import numpy as np
from pymc import Distribution
from pytensor.graph.op import Op

LogLik = Union[str, PathLike, Callable, Op, type[Distribution]]
ParamSpec = Union[float, dict[str, Any], bmb.Prior, None]

SupportedModels = Literal[
    "ddm",
    "ddm_sdv",
    "full_ddm",
    "angle",
    "angle_extended",
    "levy",
    "ornstein",
    "weibull",
    "race_no_bias_angle_4",
    "ddm_seq2_no_bias",
    "gamma_drift",
    "gamma_drift_angle",
    "lba3",
    "lba4",
    "lba2",
    "racing_diffusion_3",
    "poisson_race",
    "softmax_inv_temperature_2",
    "softmax_inv_temperature_3",
    "softmax_inv_temperature_4",
]


LoglikKind = Literal["analytical", "approx_differentiable", "blackbox"]


class NDTEdgeShift(TypedDict):
    """Declare which parameter shifts the lower edge of the response-time support.

    The likelihood admits response times from ``t - scale * <param>`` upwards.
    """

    param: str
    scale: float


class LoglikConfig(TypedDict):
    """Type for the value of LoglikConfig."""

    loglik: LogLik
    backend: Optional[Literal["jax", "pytensor"]]
    default_priors: dict[str, ParamSpec]
    bounds: dict[str, tuple[float, float]]
    extra_fields: Optional[list[str]]
    ndt_edge_shift: NotRequired[NDTEdgeShift]
    # A JAX implementation of this same likelihood, f(data, *list_params) -> (N,),
    # for callers that need JAX (the RLSSM registry); plain HSSM uses `loglik`.
    jax_loglik: NotRequired[Callable[..., Any]]


LoglikConfigs = dict[LoglikKind, LoglikConfig]


class DefaultConfig(TypedDict):
    """Type for the value of DefaultConfig."""

    response: list[str]
    list_params: list[str]
    choices: list[int]
    description: Optional[str]
    likelihoods: LoglikConfigs


DefaultConfigs = dict[SupportedModels, DefaultConfig]

LogLikeFunc = Callable[..., np.ndarray]
LogLikeGrad = Callable[..., np.ndarray]