{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Think Bayes\n",
    "\n",
    "This notebook presents example code and exercise solutions for Think Bayes.\n",
    "\n",
    "Copyright 2018 Allen B. Downey\n",
    "\n",
    "MIT License: https://opensource.org/licenses/MIT"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Configure Jupyter so figures appear in the notebook\n",
    "%matplotlib inline\n",
    "\n",
    "# Configure Jupyter to display the assigned value after an assignment\n",
    "%config InteractiveShell.ast_node_interactivity='last_expr_or_assign'\n",
    "\n",
    "# import classes from thinkbayes2\n",
    "from thinkbayes2 import Hist, Pmf, Suite, MakeMixture"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 38,
   "metadata": {},
   "outputs": [],
   "source": [
    "def add(pmf1, pmf2):\n",
    "    res = Pmf()\n",
    "    for v1, p1 in pmf1.Items():\n",
    "        for v2, p2 in pmf2.Items():\n",
    "            res[v1, v2] = p1 * p2\n",
    "    return res"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "from sympy import symbols\n",
    "p_citizen, p_cv, p_ncv, p_error = symbols('p_citizen, p_cv, p_ncv, p_error')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "def make_binary(p, name1, name2):\n",
    "    return Pmf({name1: p, name2: 1-p})"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 67,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Pmf({'citizen': p_citizen, 'non-citizen': -p_citizen + 1})"
      ]
     },
     "execution_count": 67,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "citizen_status = ['citizen', 'non-citizen']\n",
    "pmf_citizen = make_binary(p_citizen, *citizen_status)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 68,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Pmf({'error': p_error, 'no-error': -p_error + 1})"
      ]
     },
     "execution_count": 68,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "error_status = ['error', 'no-error']\n",
    "pmf_error = make_binary(p_error, *error_status)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 69,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "('citizen', 'error') p_citizen*p_error\n",
      "('citizen', 'no-error') p_citizen*(-p_error + 1)\n",
      "('non-citizen', 'error') p_error*(-p_citizen + 1)\n",
      "('non-citizen', 'no-error') (-p_citizen + 1)*(-p_error + 1)\n"
     ]
    }
   ],
   "source": [
    "pmf_citizen_report = add(pmf_citizen, pmf_error)\n",
    "pmf_citizen_report.Print()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 70,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Pmf({'vote': p_cv, 'no-vote': -p_cv + 1})"
      ]
     },
     "execution_count": 70,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "vote_status = ['vote', 'no-vote']\n",
    "pmf_cv = make_binary(p_cv, *vote_status)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 71,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "('no-vote', 'error') p_error*(-p_cv + 1)\n",
      "('no-vote', 'no-error') (-p_cv + 1)*(-p_error + 1)\n",
      "('vote', 'error') p_cv*p_error\n",
      "('vote', 'no-error') p_cv*(-p_error + 1)\n"
     ]
    }
   ],
   "source": [
    "pmf_cv_report = add(pmf_cv, pmf_error)\n",
    "pmf_cv_report.Print()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 72,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Pmf({'vote': p_ncv, 'no-vote': -p_ncv + 1})"
      ]
     },
     "execution_count": 72,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "pmf_ncv = make_binary(p_ncv, *vote_status)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 73,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "('no-vote', 'error') p_error*(-p_ncv + 1)\n",
      "('no-vote', 'no-error') (-p_error + 1)*(-p_ncv + 1)\n",
      "('vote', 'error') p_error*p_ncv\n",
      "('vote', 'no-error') p_ncv*(-p_error + 1)\n"
     ]
    }
   ],
   "source": [
    "pmf_ncv_report = add(pmf_ncv, pmf_error)\n",
    "pmf_ncv_report.Print()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 66,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "(('citizen', 'error'), ('no-vote', 'error')) p_citizen*p_error**2*(-p_cv + 1)\n",
      "(('citizen', 'error'), ('no-vote', 'no-error')) p_citizen*p_error*(-p_cv + 1)*(-p_error + 1)\n",
      "(('citizen', 'error'), ('vote', 'error')) p_citizen*p_cv*p_error**2\n",
      "(('citizen', 'error'), ('vote', 'no-error')) p_citizen*p_cv*p_error*(-p_error + 1)\n",
      "(('citizen', 'no-error'), ('no-vote', 'error')) p_citizen*p_error*(-p_cv + 1)*(-p_error + 1)\n",
      "(('citizen', 'no-error'), ('no-vote', 'no-error')) p_citizen*(-p_cv + 1)*(-p_error + 1)**2\n",
      "(('citizen', 'no-error'), ('vote', 'error')) p_citizen*p_cv*p_error*(-p_error + 1)\n",
      "(('citizen', 'no-error'), ('vote', 'no-error')) p_citizen*p_cv*(-p_error + 1)**2\n",
      "(('non-citizen', 'error'), ('no-vote', 'error')) p_error**2*(-p_citizen + 1)*(-p_ncv + 1)\n",
      "(('non-citizen', 'error'), ('no-vote', 'no-error')) p_error*(-p_citizen + 1)*(-p_error + 1)*(-p_ncv + 1)\n",
      "(('non-citizen', 'error'), ('vote', 'error')) p_error**2*p_ncv*(-p_citizen + 1)\n",
      "(('non-citizen', 'error'), ('vote', 'no-error')) p_error*p_ncv*(-p_citizen + 1)*(-p_error + 1)\n",
      "(('non-citizen', 'no-error'), ('no-vote', 'error')) p_error*(-p_citizen + 1)*(-p_error + 1)*(-p_ncv + 1)\n",
      "(('non-citizen', 'no-error'), ('no-vote', 'no-error')) (-p_citizen + 1)*(-p_error + 1)**2*(-p_ncv + 1)\n",
      "(('non-citizen', 'no-error'), ('vote', 'error')) p_error*p_ncv*(-p_citizen + 1)*(-p_error + 1)\n",
      "(('non-citizen', 'no-error'), ('vote', 'no-error')) p_ncv*(-p_citizen + 1)*(-p_error + 1)**2\n"
     ]
    }
   ],
   "source": [
    "mix = Pmf()\n",
    "\n",
    "for val1, p1 in pmf_citizen_report.Items():\n",
    "    c, e = val1\n",
    "    pmf = pmf_cv_report if c=='citizen' else pmf_ncv_report\n",
    "    for val2, p2 in pmf.Items():\n",
    "        mix[val1, val2] = p1 * p2\n",
    "        \n",
    "mix.Print()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 74,
   "metadata": {},
   "outputs": [],
   "source": [
    "def report(state, alternatives):\n",
    "    val, error = state\n",
    "    if error != 'error':\n",
    "        return val\n",
    "    alt1, alt2 = alternatives\n",
    "    return alt1 if val==alt2 else alt2"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 75,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'non-citizen'"
      ]
     },
     "execution_count": 75,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "report(('citizen', 'error'), citizen_status)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 76,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'citizen'"
      ]
     },
     "execution_count": 76,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "report(('citizen', 'no-error'), citizen_status)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 81,
   "metadata": {},
   "outputs": [],
   "source": [
    "pmf_report = Pmf()\n",
    "\n",
    "for (cstate, vstate), p in mix.Items():\n",
    "    creport = report(cstate, citizen_status)\n",
    "    vreport = report(vstate, vote_status)\n",
    "    pmf_report[creport, vreport] += p"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 82,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "('citizen', 'no-vote') p_citizen*p_cv*p_error*(-p_error + 1) + p_citizen*(-p_cv + 1)*(-p_error + 1)**2 + p_error**2*p_ncv*(-p_citizen + 1) + p_error*(-p_citizen + 1)*(-p_error + 1)*(-p_ncv + 1)\n",
      "('citizen', 'vote') p_citizen*p_cv*(-p_error + 1)**2 + p_citizen*p_error*(-p_cv + 1)*(-p_error + 1) + p_error**2*(-p_citizen + 1)*(-p_ncv + 1) + p_error*p_ncv*(-p_citizen + 1)*(-p_error + 1)\n",
      "('non-citizen', 'no-vote') p_citizen*p_cv*p_error**2 + p_citizen*p_error*(-p_cv + 1)*(-p_error + 1) + p_error*p_ncv*(-p_citizen + 1)*(-p_error + 1) + (-p_citizen + 1)*(-p_error + 1)**2*(-p_ncv + 1)\n",
      "('non-citizen', 'vote') p_citizen*p_cv*p_error*(-p_error + 1) + p_citizen*p_error**2*(-p_cv + 1) + p_error*(-p_citizen + 1)*(-p_error + 1)*(-p_ncv + 1) + p_ncv*(-p_citizen + 1)*(-p_error + 1)**2\n"
     ]
    }
   ],
   "source": [
    "pmf_report.Print()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "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.6.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}