{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Bean machine\n",
    "\n",
    "A bean machine is a triangular maze of pins through which balls are fed. At each node, the ball can decide to go left or right. We model this system using a triangular grid of nodes, where each node is a qubit. Each node takes a boolean value, which indicates whether the ball should pass on the left or right, each with 50% probability."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%matplotlib inline"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import pylab as pl\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "import networkx as nx\n",
    "from networkx.drawing.nx_agraph import write_dot, graphviz_layout"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Start with unit triangle"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "G = nx.DiGraph()\n",
    "G.add_edges_from([(\"q0\", \"q1\"), (\"q0\", \"q2\")])\n",
    "pl.figure(figsize=(2,2))\n",
    "pos=graphviz_layout(G, prog='dot')\n",
    "nx.draw(G, pos, with_labels=True, arrows=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from qiskit import BasicAer\n",
    "from qiskit import ClassicalRegister, QuantumRegister, QuantumCircuit\n",
    "from qiskit import execute\n",
    "from qiskit.tools.visualization import plot_state_city"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# set up Quantum Register and Classical Register for 3 qubits\n",
    "q = QuantumRegister(3, name=\"q\")\n",
    "c = ClassicalRegister(3)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "qc = QuantumCircuit(q, c)\n",
    "q0, q1, q2 = q\n",
    "\n",
    "qc.h(q0)\n",
    "qc.cx(q0, q1)  # Left branch\n",
    "qc.x(q0)  # Flip q0 to check for right branch\n",
    "qc.cx(q0, q2)  # Right branch\n",
    "\n",
    "qc.barrier()\n",
    "\n",
    "# flip q0 back\n",
    "qc.x(q0)\n",
    "\n",
    "# measure all qubits\n",
    "qc.barrier()\n",
    "qc.measure(q, c)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "qc.draw(output='mpl')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "job = execute(qc, backend=BasicAer.get_backend('qasm_simulator'), shots=10000)\n",
    "result = job.result()\n",
    "result.get_counts(qc)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "job = execute(qc, backend=BasicAer.get_backend('statevector_simulator'))\n",
    "result = job.result()\n",
    "state = result.get_statevector()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "state"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plot_state_city(state, title='Left or right')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "sum([s != 0 for s in state])  # Number of possible paths"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Scale it up one level"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "G = nx.DiGraph()\n",
    "G.add_edges_from([(\"q0\", \"q1\"), (\"q0\", \"q2\"), (\"q1\", \"q3\"), (\"q1\", \"q4\"), (\"q2\", \"q4\"), (\"q2\", \"q5\")])\n",
    "pl.figure(figsize=(4,3))\n",
    "pos=graphviz_layout(G, prog='dot')\n",
    "nx.draw(G, pos, with_labels=True, arrows=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# set up Quantum Register and Classical Register for 3 qubits\n",
    "q = QuantumRegister(6, name=\"q\")\n",
    "c = ClassicalRegister(6)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "qc = QuantumCircuit(q)\n",
    "q0, q1, q2, q3, q4, q5 = q\n",
    "\n",
    "# First layer. q0 is the control qubit.\n",
    "qc.h(q0)\n",
    "qc.cx(q0, q1)\n",
    "qc.x(q0)\n",
    "qc.cx(q0, q2)\n",
    "\n",
    "qc.barrier()\n",
    "\n",
    "# Second layer: right branch. q0 and q2 are control qubits.\n",
    "qc.h(q2)\n",
    "qc.ccx(q0, q2, q4)\n",
    "qc.x(q2)\n",
    "qc.ccx(q0, q2, q5)\n",
    "\n",
    "qc.barrier()\n",
    "\n",
    "# Second layer: left branch. q0 and q1 are control qubits.\n",
    "qc.x(q0)   # Flip q0 to use as control\n",
    "qc.h(q1)\n",
    "qc.ccx(q0, q1, q3)\n",
    "qc.x(q1)\n",
    "qc.ccx(q0, q1, q4)\n",
    "\n",
    "qc.barrier()\n",
    "\n",
    "# Flip things back\n",
    "qc.x(q1)\n",
    "qc.x(q2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "qc.draw(output='mpl')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Measure all qubits\n",
    "qc.add_register(c)  # Adding classical register\n",
    "qc.barrier()\n",
    "qc.measure(q, c)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "job = execute(qc, backend=BasicAer.get_backend('statevector_simulator'))\n",
    "result = job.result()\n",
    "state = result.get_statevector()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "state"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "sum([s != 0 for s in state])  # Number of possible paths"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "job = execute(qc, backend=BasicAer.get_backend('qasm_simulator'), shots=10000)\n",
    "result = job.result()\n",
    "counts = result.get_counts(qc)\n",
    "counts"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Get mapping from bitstrings to nodes\n",
    "nodes = [f\"{q.name}{n}\" for n in range(q.size)]\n",
    "node_indices = [nodes.index(x) for x in G.nodes() if G.out_degree(x)==0]\n",
    "bitstrings_to_nodes = {c: [nodes[i] for i in node_indices if list(c)[::-1][i] is \"1\"][0] for c in counts}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "node_indices"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "bitstrings_to_nodes"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "node_counts = pd.DataFrame(bitstrings_to_nodes.items(), columns=[\"bitstring\", \"node\"]\n",
    ").merge(pd.DataFrame(counts.items(), columns=[\"bitstring\", \"count\"]))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "node_counts"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "count_per_node = node_counts.groupby(\"node\").sum()\n",
    "count_per_node"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "count_per_node.plot(kind=\"bar\")"
   ]
  },
  {
   "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.7.3"
  },
  "pycharm": {
   "stem_cell": {
    "cell_type": "raw",
    "source": [],
    "metadata": {
     "collapsed": false
    }
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}