{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Stochaskell, version 0.1.0\n",
       "Copyright (C) 2015-2020 David A Roberts\n",
       "This program comes with ABSOLUTELY NO WARRANTY.\n",
       "This is free software, and you are welcome to redistribute it\n",
       "under certain conditions; see the LICENSE for details.\n",
       "\n",
       "Using installation directory at \n",
       "  /home/jovyan/stochaskell"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "{-# LANGUAGE RebindableSyntax, DeriveGeneric, DeriveFunctor, DeriveTraversable, DerivingVia, FlexibleInstances #-}\n",
    "import Language.Stochaskell\n",
    "stochaskell"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "import Language.Stochaskell.Expression\n",
    "data Tree t\n",
    "  = Branch t t\n",
    "  | Leaf R\n",
    "  deriving (Show, Generic, Generic1, Functor, Foldable, Traversable)\n",
    "  deriving Show1 via (Generically1 Tree)\n",
    "instance                 ExprType    (Tree (FixE Tree))\n",
    "instance                 Constructor (Tree (FixE Tree))\n",
    "instance (ExprType t) => ExprType    (Tree (Expression t))\n",
    "instance (ExprType t) => Constructor (Tree (Expression t))\n",
    "type T = FixE' Tree"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "production :: R -> P (Expression (Tree R))\n",
    "production p = mixture'\n",
    "  [(p, do\n",
    "    return . fromConcrete $ Branch (p/2) (p/2))\n",
    "  ,(1-p, do\n",
    "    x <- uniform 0 1\n",
    "    return . fromConcrete $ Leaf x)\n",
    "  ]\n",
    "\n",
    "pcfg :: P T\n",
    "pcfg = unfoldP production 1"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Branch (Fix (Branch (Fix (Leaf 0.4898595150062972)) (Fix (Branch (Fix (Leaf 0.2523685666712354)) (Fix (Leaf 0.631889727678022)))))) (Fix (Branch (Fix (Leaf 0.1620737376264192)) (Fix (Branch (Fix (Leaf 0.46874567676577183)) (Fix (Leaf 0.2778033925412521))))))"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "tree <- simulate pcfg\n",
    "print tree"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "reduction :: Tree R -> P R\n",
    "reduction (Branch a b) = do\n",
    "  return (a + b)\n",
    "reduction (Leaf a) = do\n",
    "  b <- bernoulli (1/2)\n",
    "  return (if b then a else 0)\n",
    "\n",
    "treeSubsetSum :: T -> P R\n",
    "treeSubsetSum = foldP (fromCaseP' reduction)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "0.958605191772069"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "simulate (treeSubsetSum tree)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Haskell",
   "language": "haskell",
   "name": "haskell"
  },
  "language_info": {
   "codemirror_mode": "ihaskell",
   "file_extension": ".hs",
   "name": "haskell",
   "pygments_lexer": "Haskell",
   "version": "8.6.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}