{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "f4c3820b",
   "metadata": {},
   "source": [
    "# Joint Prior"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1d7d819b",
   "metadata": {},
   "outputs": [],
   "source": [
    "import superstats as sup"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6edee4fa",
   "metadata": {},
   "source": [
    "Next, we create a `JointPrior` object to specify priors for each observation model parameter. Superstats distinguishes between the following parameter types:\n",
    "\n",
    "- **`local_param`** — time-varying and estimated (e.g., the trajectory produced by a `StochasticTransition`)\n",
    "- **`hyper_param`** — time-invariant and estimated (e.g., the hyperpriors governing a transition, such as `sigma` or `slope`)\n",
    "- **`deterministic_param`** — time-varying but not estimated (produced by a `DeterministicTransition`)\n",
    "- **`shared_param`** — time-invariant and estimated (specified via a `Prior`, shared across all time steps)\n",
    "- **`fixed_param`** — time-invariant and not estimated (a fixed `scalar`)\n",
    "\n",
    "Each parameter type reflects a different combination of two properties: whether the parameter changes over time (**time-varying** vs. **time-invariant**) and whether it is inferred from data (**estimated**) or fixed (**not estimated**). The `JointPrior` object collects these specifications for all parameters into a single, coherent prior over the full model.\n",
    "\n",
    "```python\n",
    "joint_prior = sup.JointPrior(\n",
    "    param_1: sup.transition.RandomWalk(...), # StochasticTransition\n",
    "    param_2: sup.transition.Linear(...),     # DeterministicTransition\n",
    "    param_3: sup.Prior(...),                 # Shared parameter\n",
    "    param_4: 0.5                             # Fixed parameter\n",
    ")\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6cff8be0",
   "metadata": {},
   "source": [
    "## Stochastic Transitions"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "52f38506",
   "metadata": {},
   "source": [
    "Superstats has implemented the following stochastic transition models:\n",
    "\n",
    "- `RandomWalk`: with hyperparameters `sigma` for the scale of Gaussian noise and `delta` for linear additive drift.\n",
    "- `AutoRegression`: an AR(1) process with hyperparameters `sigma` for Gaussian noise, `phi` for the autoregressive coefficient, and `delta` for additive drift.\n",
    "- `OrnsteinUhlenbeck`: a mean-reverting process with hyperparameters `sigma` for diffusion noise, `mu` for the long-run mean, and `theta` for the mean-reversion speed.\n",
    "- `LevyFlight`: a random walk with alpha-stable noise, with hyperparameters `sigma` for noise scale, `delta` for additive drift, and `alpha` for tail heaviness/stability.\n",
    "- `Jump`: a jump process with `p_jump` controlling the probability of jumping to a new proposal value at each step.\n",
    "- `GaussianProcess`: trajectories sampled from a Gaussian process using configurable RBF, linear, periodic, or composite kernels.\n",
    "- `Mixture`: a mixture of two or more stochastic transition models, with Dirichlet-distributed mixture weights (does not work with `GaussianProcess`).\n",
    "\n",
    "For each hyperparameter, we can either specify a prior and estimate it from data, or fix it to a scalar value and exclude it from inference. When nothing is specified, Superstats uses sensible defaults, see `superstats.defaults.transition_defaults`.\n",
    "\n",
    "The following additional arguments can be specified:\n",
    "\n",
    "- `bounds` : Lower and upper bounds for the parameter. Parameters are generated in the unconstraint space and then constrained to this range via a scaled sigmoid.\n",
    "- `initial_prior` : Prior or fixed scalar for the initial parameter value. Note that the initial value is in the unconstrained space and then transformed via scaled sigmoid.\n",
    "\n",
    "**Example.**\n",
    "\n",
    "\n",
    "```python\n",
    "prior = sup.JointPrior(\n",
    "    param_1 = sup.transition.RandomWalk(\n",
    "        bounds=(-6, 6),\n",
    "        inital_prior=sup.Prior(\"normal\", loc=0, scale=2),\n",
    "        sigma=sup.Prior(\"halfnormal\", scale=0.2)\n",
    "    ),\n",
    "    param_2 = ...\n",
    ")\n",
    "```\n",
    "\n",
    "Note that parameter names should match the names used as arguments in the obeservation model simulator."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "afe9a98c",
   "metadata": {},
   "source": [
    "### Mixture"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5dfea85d",
   "metadata": {},
   "source": [
    "The `Mixture` transition allows mixing the following stochastic transition models:\n",
    "\n",
    "- `RandomWalk`\n",
    "- `AutoRegression`\n",
    "- `OrnsteinUhlenbeck`\n",
    "- `LevyFlight`\n",
    "- `Jump`\n",
    "\n",
    "The most interesting and sensible mixture is between one of the first four transitions and a `Jump`, since this lets a parameter follow smooth, gradual dynamics most of the time while occasionally undergoing a sudden discrete jump.\n",
    "\n",
    "A few notes on using the `Mixture` transition:\n",
    "\n",
    "- `bounds` and `initial_prior` must be defined once at initialization of the `Mixture` itself, and must **not** be specified again within the individual transitions it contains.\n",
    "\n",
    "- When a `Jump` transition is included in the `Mixture`, its `p_jump` is automatically fixed to $1.0$, since the `mixture_weights` already govern the probability of a jump occurring at a given time step. It would not be sensible for the `Jump` component to be selected at a given step and then, due to `p_jump < 1.0`, have a chance of no jump actually occurring.\n",
    "\n",
    "- As noted earlier, parameter trajectories are generated in an unconstrained space and then transformed via a scaled sigmoid to ensure they remain within `bounds`. This means that setting `proposal_prior=sup.prior.Prior(dist=\"logistic\", loc=0, scale=1)` for the `Jump` component results in a uniform distribution over the bounds after transformation — i.e., when a jump occurs, the parameter can take on any value within its bounds with equal probability.\n",
    "\n",
    "**Example.**\n",
    "\n",
    "```python\n",
    "prior = sup.JointPrior(\n",
    "    param_1 = sup.transition.RandomWalk(\n",
    "        bounds=(-6, 6),\n",
    "        initial_prior=sup.Prior(\"normal\", loc=0, scale=2),\n",
    "        sigma=sup.Prior(\"halfnormal\", scale=0.2)\n",
    "    ),\n",
    "    param_2 = sup.transition.Mixture(\n",
    "        bounds=(0, 1),\n",
    "        initial_prior=sup.Prior(\"normal\", loc=0, scale=1),\n",
    "        transitions=[\n",
    "            sup.transition.OrnsteinUhlenbeck(\n",
    "                sigma=sup.Prior(\"halfnormal\", scale=0.1),\n",
    "                theta=sup.Prior(\"halfnormal\", scale=0.05)\n",
    "            ),\n",
    "            sup.transition.Jump(\n",
    "                proposal_prior=sup.Prior(\"logistic\", loc=0, scale=1)\n",
    "            ),\n",
    "        ],\n",
    "        mixture_weights=sup.Prior(\"dirichlet\", alpha=[9, 1])\n",
    "    ),\n",
    "    param_3 = ...\n",
    ")\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "45be877c",
   "metadata": {},
   "source": [
    "## Deterministic Transitions"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "31662d0b",
   "metadata": {},
   "source": [
    "Coming soon..."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "806a1bb2",
   "metadata": {},
   "source": [
    "## Priors"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e2b5ea92",
   "metadata": {},
   "source": [
    "For all hyperparameters, and whenever we want to estimate a time-invariant observation model parameter, we can specify a standard prior via the `Prior` class. Superstats implements the following distributions:\n",
    "\n",
    "- **`normal`** — Gaussian, parameterized by `loc` (mean) and `scale` (standard deviation).\n",
    "- **`halfnormal`** — folded Gaussian, non-negative, parameterized by `scale`. Commonly used for scale/noise parameters (e.g., `sigma`).\n",
    "- **`uniform`** — uniform over `[low, high]`.\n",
    "- **`beta`** — Beta distribution on `[0, 1]`, parameterized by shape parameters `a` and `b`.\n",
    "- **`logistic`** — logistic distribution, parameterized by `loc` and `scale`. Often used as a `proposal_prior` in `Jump` transitions, since it maps to a uniform distribution over `bounds` after the scaled sigmoid transform.\n",
    "- **`dirichlet`** — distribution over the simplex, parameterized by a concentration vector `alpha`. Used, for example, to specify `mixture_weights` in a `Mixture` transition.\n",
    "\n",
    "Every `Prior` additionally supports an optional linear transform of the drawn samples via `scale_factor` and `shift`, i.e., `scale_factor * samples + shift`.\n",
    "\n",
    "Feel free to open an issue if a prior you'd like to use isn't implemented yet."
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "superstats (3.13.14.final.0)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.13.14"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}