{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# A Discrete Switching System\n",
    "\n",
    "A la MHS, but all discrete.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from gtsam import DiscreteBayesNet, DiscreteKeys, DiscreteFactorGraph, Ordering\n",
    "from gtsam.symbol_shorthand import S\n",
    "from gtsam.symbol_shorthand import M\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def P(*args):\n",
    "    \"\"\" Create a DiscreteKeys instances from a variable number of DiscreteKey pairs.\"\"\"\n",
    "    # TODO: We can make life easier by providing variable argument functions in C++ itself.\n",
    "    dks = DiscreteKeys()\n",
    "    for key in args:\n",
    "        dks.push_back(key)\n",
    "    return dks\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import graphviz\n",
    "\n",
    "\n",
    "class show(graphviz.Source):\n",
    "    \"\"\" Display an object with a dot method as a graph.\"\"\"\n",
    "\n",
    "    def __init__(self, obj):\n",
    "        \"\"\"Construct from object with 'dot' method.\"\"\"\n",
    "        # This small class takes an object, calls its dot function, and uses the\n",
    "        # resulting string to initialize a graphviz.Source instance. This in turn\n",
    "        # has a _repr_mimebundle_ method, which then renders it in the notebook.\n",
    "        super().__init__(obj.dot())\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "nrStates = 3\n",
    "K = 5\n",
    "\n",
    "bayesNet = DiscreteBayesNet()\n",
    "for k in range(1, K):\n",
    "    key = S(k), nrStates\n",
    "    key_plus = S(k+1), nrStates\n",
    "    mode = M(k), 2\n",
    "    bayesNet.add(key_plus, P(mode, key), \"9/1/0 1/8/1 0/1/9  1/9/0 0/1/9 9/0/1\")\n",
    "\n",
    "bayesNet"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "show(bayesNet)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create a factor graph out of the Bayes net.\n",
    "factorGraph = DiscreteFactorGraph(bayesNet)\n",
    "show(factorGraph)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create a BayesTree out of the factor graph.\n",
    "ordering = Ordering()\n",
    "# First eliminate \"continuous\" states in time order\n",
    "for k in range(1, K+1):\n",
    "    ordering.push_back(S(k))\n",
    "for k in range(1, K):\n",
    "    ordering.push_back(M(k))\n",
    "print(ordering)\n",
    "bayesTree = factorGraph.eliminateMultifrontal(ordering)\n",
    "bayesTree"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "show(bayesTree)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": []
  }
 ],
 "metadata": {
  "interpreter": {
   "hash": "31f2aee4e71d21fbe5cf8b01ff0e069b9275f58929596ceb00d14d90e3e16cd6"
  },
  "kernelspec": {
   "display_name": "Python 3.8.9 64-bit",
   "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.9.7"
  },
  "orig_nbformat": 4
 },
 "nbformat": 4,
 "nbformat_minor": 2
}