diff --git a/prototype.ipynb b/prototype.ipynb
index bc03684..985cf8f 100644
--- a/prototype.ipynb
+++ b/prototype.ipynb
@@ -78,9 +78,9 @@
    "outputs": [],
    "source": [
     "from random import shuffle\n",
-    "from scipy.stats import zipf, poisson\n",
+    "from scipy.stats import zipf, poisson, uniform\n",
     "from numpy.random import choice\n",
-    "from numpy import reshape, count_nonzero\n",
+    "from numpy import reshape, count_nonzero, exp\n",
     "from dataclasses import dataclass, field\n",
     "from typing import List, Tuple"
    ]
@@ -355,91 +355,14 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "class OptPickWave:\n",
+    "def flatten(l):\n",
     "    \"\"\"\n",
-    "    This is slightly cleverer version of the above, for optimisation, but not so nice to read, probably.\n",
-    "    \n",
-    "    Also includes other bits and pieces to make local optimisation algorithms work easier.\n",
+    "    TODO: do this better\n",
     "    \"\"\"\n",
-    "    \n",
-    "    def __init__(self, products):\n",
-    "        # each element of this list is a container and contains a list of orders\n",
-    "        self.container_orders = []\n",
-    "        self._cost = 0\n",
-    "        self._cost_stale = False\n",
-    "        self.products = products\n",
-    "    \n",
-    "    def add_order_to_new_container(self, order):\n",
-    "        self._cost_stale = True\n",
-    "        self.container_orders.append([order])\n",
-    "\n",
-    "    def add_order(self, order, container):\n",
-    "        self._cost_stale = True\n",
-    "        if container < 0 or container >= len(self.container_orders):\n",
-    "            raise Exception(\"Invalid container\")\n",
-    "        self.container_orders[container].append(order)\n",
-    "\n",
-    "    def remove_order(self, order):\n",
-    "        self._cost_stale = True\n",
-    "        self.container_orders = [[other_order for other_order in orders if other_order != order] for orders in self.container_orders]\n",
-    "        self.container_orders = [orders for orders in self.container_orders if len(orders) > 0]\n",
-    "    \n",
-    "    def move_order(self, order, container):\n",
-    "        current_container = -1\n",
-    "        for container, orders in enumerate(self.container_orders):\n",
-    "            if order in orders:\n",
-    "                current_container = container\n",
-    "                break\n",
-    "\n",
-    "        if current_container == -1:\n",
-    "            raise Exception(\"Order not found\")\n",
-    "        \n",
-    "        if current_container != container:\n",
-    "            self._cost_stale = True\n",
-    "            remove_order(order)\n",
-    "            self.container_orders[container].append(order)\n",
-    "    \n",
-    "    def compute_picking_cost(self, product_to_pick_face):\n",
-    "        if self._cost_stale:\n",
-    "            # number of containers\n",
-    "            self.n_containers = len(self.container_orders)\n",
-    "            # a dictionary that for each product contains how many times that's picked into each container\n",
-    "            self.product_picks = dict(zip(self.products, [[0] * self.n_containers for n in self.products]))\n",
-    "            # a list for each container, which tells which products are in that container\n",
-    "            self.products_in_containers = [[] for n in range(self.n_containers)]\n",
-    "            # list of packages, as (container, order) tuples\n",
-    "            self.packages = []\n",
-    "\n",
-    "            for container, orders in enumerate(self.container_orders):\n",
-    "                for order in orders:\n",
-    "                    for product in order.products():\n",
-    "                        self.product_picks[product][container] += 1\n",
-    "                    self.products_in_containers[container].extend(order.products())\n",
-    "                    self.packages.append((container, order))\n",
-    "\n",
-    "            # single pick face\n",
-    "            n_fixed = 1\n",
-    "\n",
-    "            # compute number of up or down steps, note they come in pairs\n",
-    "            n_move = 2 * max([product_to_pick_face[product].location()[1] for product, picks in self.product_picks.items() if count_nonzero(picks) > 0])\n",
-    "\n",
-    "            # move across once\n",
-    "            n_across = 1\n",
-    "\n",
-    "            # this is basically the same but needs number of containers\n",
-    "            n_picking = [[self.n_containers] + pick for product, pick in self.product_picks.items() if count_nonzero(pick) > 0]\n",
-    "\n",
-    "            # computer sorting operations\n",
-    "            q = self.n_containers\n",
-    "            r = len(self.packages)\n",
-    "            n_sorting_1 = [q, r]\n",
-    "            for container in range(self.n_containers):\n",
-    "                n_sorting_1 += [len(package[1].products()) if package[0] == container else  0 for package in self.packages]\n",
-    "\n",
-    "            self._cost = PickingCost(n_fixed, n_move, n_across, n_picking, [n_sorting_1])\n",
-    "            self._cost_stale = False\n",
-    "        \n",
-    "        return self._cost"
+    "    out = []\n",
+    "    for x in l:\n",
+    "        out += x\n",
+    "    return out"
    ]
   },
   {
@@ -448,15 +371,198 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "def generate_singleton_waves_opt(batch, products):\n",
-    "    pick_waves = []\n",
+    "def compute_picking_cost(self, container_orders, products, product_to_pick_face):\n",
+    "    # number of containers\n",
+    "    n_containers = len(container_orders)\n",
+    "    # a dictionary that for each product contains how many times that's picked into each container\n",
+    "    product_picks = dict(zip(products, [[0] * n_containers for n in products]))\n",
+    "    # a list for each container, which tells which products are in that container\n",
+    "    products_in_containers = [[] for n in range(n_containers)]\n",
+    "    # list of packages, as (container, order) tuples\n",
+    "    packages = []\n",
+    "\n",
+    "    for container, orders in enumerate(container_orders):\n",
+    "        for order in orders:\n",
+    "            for product in order.products():\n",
+    "                product_picks[product][container] += 1\n",
+    "            products_in_containers[container].extend(order.products())\n",
+    "            packages.append((container, order))\n",
+    "\n",
+    "    # single pick face\n",
+    "    n_fixed = 1\n",
+    "\n",
+    "    # compute number of up or down steps, note they come in pairs\n",
+    "    n_move = 2 * max([product_to_pick_face[product].location()[1] for product, picks in product_picks.items() if count_nonzero(picks) > 0])\n",
+    "\n",
+    "    # move across once\n",
+    "    n_across = 1\n",
+    "\n",
+    "    # this is basically the same but needs number of containers\n",
+    "    n_picking = [[n_containers] + pick for product, pick in product_picks.items() if count_nonzero(pick) > 0]\n",
+    "\n",
+    "    # computer sorting operations\n",
+    "    q = n_containers\n",
+    "    r = len(packages)\n",
+    "    n_sorting_1 = [q, r]\n",
+    "    for container in range(n_containers):\n",
+    "        n_sorting_1 += [len(package[1].products()) if package[0] == container else  0 for package in packages]\n",
+    "\n",
+    "    return PickingCost(n_fixed, n_move, n_across, n_picking, [n_sorting_1])"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "def chunk_list(lst, n):\n",
+    "    \"\"\"\n",
+    "    TODO: lifted wholesale from SO. check license\n",
+    "    \"\"\"\n",
+    "    for i in range(0, len(lst), n):\n",
+    "        yield lst[i:i + n]"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "class SimAnn:\n",
+    "    \"\"\"\n",
+    "    Simulated annealing optimiser for pick waves\n",
+    "    \"\"\"\n",
     "    \n",
-    "    for order in batch:\n",
-    "        pick_wave = OptPickWave(products)\n",
-    "        pick_wave.add_order_to_new_container(order)\n",
-    "        pick_waves.append(pick_wave)\n",
+    "    def __init__(self, batch, products, product_to_pick_face, max_containers, max_orders_per_wave):\n",
+    "        self.batch = batch.copy()\n",
+    "        self.products = products\n",
+    "        self.product_to_pick_face = product_to_pick_face\n",
+    "        self.max_containers = max_containers\n",
+    "        self.max_orders_per_wave = max_orders_per_wave\n",
+    "        \n",
+    "        self.order_to_wave_and_container = dict()\n",
+    "        self._cost = 0\n",
+    "        self._cost_stale = True\n",
+    "        \n",
+    "        self.steps = 0\n",
+    "        self.max_steps = 100\n",
+    "        \n",
+    "        self._heuristic_init()\n",
     "    \n",
-    "    return pick_waves"
+    "    def _heuristic_init(self):\n",
+    "        # place them in as many containers as we can (seems to be roughly the\n",
+    "        # best we can do with the basic strategies)\n",
+    "        # let's say we divide them up into sub-batches of max_orders_per_wave\n",
+    "        chunks = list(chunk_list(self.batch, self.max_orders_per_wave))\n",
+    "        \n",
+    "        wave = 0\n",
+    "        \n",
+    "        for chunk in chunks:\n",
+    "            for container, order in enumerate(chunk):\n",
+    "                self.order_to_wave_and_container[order] = (wave, container % self.max_containers)\n",
+    "            wave += 1\n",
+    "    \n",
+    "    def _accept_step(self, last_obj, new_obj):\n",
+    "        temp = (1 - self.steps / self.max_steps) * 10\n",
+    "        if new_obj < last_obj:\n",
+    "            pass #print(\"better obj\", exp((new_obj - last_obj) / temp))\n",
+    "        else:\n",
+    "            pass #print(\"worse obj \", exp((new_obj - last_obj) / temp))\n",
+    "        out = exp((new_obj - last_obj) / temp) < uniform.rvs()\n",
+    "        if out:\n",
+    "            print('moving...')\n",
+    "        return out\n",
+    "    \n",
+    "    def step(self):\n",
+    "        \"\"\"\n",
+    "        Perform one step of SA\n",
+    "        \"\"\"\n",
+    "        # pick an order and move it elsewhere\n",
+    "        order = choice(self.batch)\n",
+    "        wave = choice(self.waves)\n",
+    "        container = choice(range(self.max_containers))\n",
+    "\n",
+    "        # chance that we are doing a no-op is negligible\n",
+    "        last_wave = -1\n",
+    "        for wave in self.waves:\n",
+    "            if wave.contains_order(order):\n",
+    "                last_wave = wave\n",
+    "        \n",
+    "        assert(last_wave != -1)\n",
+    "        \n",
+    "        last_container = last_wave.get_order_container(order)\n",
+    "        \n",
+    "        last_obj = self.objective()\n",
+    "        \n",
+    "        print(order, wave, container, order == self.waves[-1])\n",
+    "        \n",
+    "        swap = False\n",
+    "        \n",
+    "        #print('doing')\n",
+    "        \n",
+    "        if last_wave.number_orders() == self.max_orders_per_wave or uniform.rvs() > 0.5:\n",
+    "            swap = True\n",
+    "            swap_order = choice(wave.all_orders())\n",
+    "            swap_old_container = wave.get_order_container(swap_order)\n",
+    "            swap_new_container = choice(range(self.max_containers))\n",
+    "            if wave == last_wave:\n",
+    "                wave.move_order(swap_order, swap_new_container)\n",
+    "            else:\n",
+    "                wave.remove_order(swap_order)\n",
+    "                last_wave.add_order(swap_order, swap_new_container)\n",
+    "        \n",
+    "        if wave == last_wave:\n",
+    "            wave.move_order(order, container)\n",
+    "        else:\n",
+    "            last_wave.remove_order(order)\n",
+    "            wave.add_order(order, container)\n",
+    "\n",
+    "        self._cost_stale = True\n",
+    "        \n",
+    "        new_obj = self.objective()\n",
+    "        \n",
+    "        if not self._accept_step(last_obj, new_obj):\n",
+    "            #print('undoing')\n",
+    "            if wave == last_wave:\n",
+    "                wave.move_order(order, last_container)\n",
+    "            else:\n",
+    "                wave.remove_order(order)\n",
+    "                last_wave.add_order(order, last_container)\n",
+    "            if swap:\n",
+    "                if wave == last_wave:\n",
+    "                    wave.move_order(swap_order, swap_old_container)\n",
+    "                else:\n",
+    "                    last_wave.remove_order(swap_order)\n",
+    "                    wave.add_order(swap_order, swap_old_container)\n",
+    "\n",
+    "                \n",
+    "            self._cost_stale = True\n",
+    "\n",
+    "        return self.objective()\n",
+    "            \n",
+    "    def objective(self):\n",
+    "        \"\"\"\n",
+    "        Return objective value\n",
+    "        \"\"\"\n",
+    "        if self._cost_stale:\n",
+    "            waves = []\n",
+    "            total_waves = max([wave_and_container[0] for order, wave_and_container in self.order_to_wave_and_container.items()])\n",
+    "            print(total_waves)\n",
+    "            print(self.order_to_wave_and_container.items())\n",
+    "            print([wave_and_container[0] for order, wave_and_container in self.order_to_wave_and_container.items()])\n",
+    "            for wave in range(total_waves):\n",
+    "                waves.append([list() for n in range(self.max_containers)])\n",
+    "            for order, wave_and_container in self.order_to_wave_and_container.items():\n",
+    "                wave, container = wave_and_container\n",
+    "                print(wave, container)\n",
+    "                waves[wave][container].append(order)\n",
+    "            \n",
+    "            self._cost = sum([total_time(compute_picking_cost(wave, self.products, self.product_to_pick_face)) for wave in waves]) / len(self.batch)\n",
+    "            self._cost_stale = False\n",
+    "            \n",
+    "        return self._cost"
    ]
   },
   {
@@ -465,7 +571,7 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "def generate_singleton_waves(batch, products):\n",
+    "def generate_singleton_waves(batch, products, max_containers, max_orders_per_wave):\n",
     "    \"\"\"\n",
     "    This function generates pick waves where each order is just picked by itself\n",
     "    \"\"\"\n",
@@ -485,18 +591,20 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "def generate_multi_order_waves(batch, products):\n",
+    "def generate_multi_order_waves(batch, products, max_containers, max_orders_per_wave):\n",
     "    \"\"\"\n",
     "    Pick each order into its own container\n",
     "    \"\"\"\n",
     "    pick_waves = []\n",
     "    \n",
-    "    pick_wave = PickWave(products, len(batch))\n",
-    "\n",
-    "    for container, order in enumerate(batch):\n",
-    "        pick_wave.add_order(order, container)\n",
-    "    \n",
-    "    pick_waves.append(pick_wave)\n",
+    "    chunks = list(chunk_list(batch, max_orders_per_wave))\n",
+    "    for chunk in chunks:\n",
+    "        pick_wave = PickWave(products, min(len(chunk), max_containers))\n",
+    "        \n",
+    "        for container, order in enumerate(chunk):\n",
+    "            pick_wave.add_order(order, container % max_containers)\n",
+    "        \n",
+    "        pick_waves.append(pick_wave)\n",
     "    \n",
     "    return pick_waves"
    ]
@@ -507,18 +615,17 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "def generate_batch_order_waves(batch, products):\n",
-    "    \"\"\"\n",
-    "    Pick each order into its own container\n",
-    "    \"\"\"\n",
+    "def generate_batch_order_waves(batch, products, max_containers, max_orders_per_wave):\n",
     "    pick_waves = []\n",
     "    \n",
-    "    pick_wave = PickWave(products, 1)\n",
+    "    chunks = list(chunk_list(batch, max_containers))\n",
+    "    for chunk in chunks:\n",
+    "        pick_wave = PickWave(products, 1)\n",
     "\n",
-    "    for order in batch:\n",
-    "        pick_wave.add_order(order, 0)\n",
-    "    \n",
-    "    pick_waves.append(pick_wave)\n",
+    "        for order in chunk:\n",
+    "            pick_wave.add_order(order, 0)\n",
+    "\n",
+    "        pick_waves.append(pick_wave)\n",
     "    \n",
     "    return pick_waves"
    ]
@@ -529,18 +636,18 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "def generate_multi_batch_order_waves(batch, containers, products):\n",
-    "    \"\"\"\n",
-    "    Pick each order into its own container\n",
-    "    \"\"\"\n",
+    "def generate_multi_batch_order_waves(batch, products, max_containers, max_orders_per_wave):\n",
     "    pick_waves = []\n",
     "    \n",
-    "    pick_wave = PickWave(products, containers)\n",
+    "    chunks = list(chunk_list(batch, max_orders_per_wave))\n",
     "\n",
-    "    for container, order in enumerate(batch):\n",
-    "        pick_wave.add_order(order, container % containers)\n",
-    "    \n",
-    "    pick_waves.append(pick_wave)\n",
+    "    for chunk in chunks:\n",
+    "        pick_wave = PickWave(products, max_containers)\n",
+    "\n",
+    "        for container, order in enumerate(chunk):\n",
+    "            pick_wave.add_order(order, container % max_containers)\n",
+    "\n",
+    "        pick_waves.append(pick_wave)\n",
     "    \n",
     "    return pick_waves"
    ]
@@ -555,7 +662,7 @@
     "no_pick_faces = 100\n",
     "\n",
     "# zipf's law shape parameter\n",
-    "zipf_shape = 1.07"
+    "zipf_shape = 1.09"
    ]
   },
   {
@@ -585,21 +692,10 @@
    "outputs": [],
    "source": [
     "# Generate a batch of orders\n",
-    "batch = [generate_new_order(poisson(25).rvs(), product_probabilities) for n in range(60)]"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "# Picking orders individually\n",
-    "singleton_waves = generate_singleton_waves_opt(batch, products)\n",
-    "total = PickingCost()\n",
-    "for wave in singleton_waves:\n",
-    "    total += wave.compute_picking_cost(product_to_pick_face)\n",
-    "total_time(total)"
+    "batch = [generate_new_order(poisson(25).rvs(), product_probabilities) for n in range(600)]\n",
+    "\n",
+    "max_containers = 5\n",
+    "max_orders_per_wave = 50"
    ]
   },
   {
@@ -611,11 +707,11 @@
    "outputs": [],
    "source": [
     "# Picking orders individually\n",
-    "singleton_waves = generate_singleton_waves(batch, products)\n",
+    "singleton_waves = generate_singleton_waves(batch, products, max_containers, max_orders_per_wave)\n",
     "total = PickingCost()\n",
     "for wave in singleton_waves:\n",
     "    total += wave.compute_picking_cost(product_to_pick_face)\n",
-    "total_time(total)"
+    "total_time(total) / len(batch)"
    ]
   },
   {
@@ -627,11 +723,12 @@
    "outputs": [],
    "source": [
     "# Picking all orders at the same time\n",
-    "multi_order_waves = generate_multi_order_waves(batch, products)\n",
+    "multi_order_waves = generate_multi_order_waves(batch, products, max_containers, max_orders_per_wave)\n",
+    "print(multi_order_waves[0].n_containers)\n",
     "total = PickingCost()\n",
     "for wave in multi_order_waves:\n",
     "    total += wave.compute_picking_cost(product_to_pick_face)\n",
-    "total_time(total)"
+    "total_time(total) / len(batch)"
    ]
   },
   {
@@ -641,11 +738,11 @@
    "outputs": [],
    "source": [
     "# Picking all orders at the same time but sort at end\n",
-    "multi_order_waves = generate_batch_order_waves(batch, products)\n",
+    "multi_order_waves = generate_batch_order_waves(batch, products, max_containers, max_orders_per_wave)\n",
     "total = PickingCost()\n",
     "for wave in multi_order_waves:\n",
     "    total += wave.compute_picking_cost(product_to_pick_face)\n",
-    "total_time(total)"
+    "total_time(total) / len(batch)"
    ]
   },
   {
@@ -655,19 +752,131 @@
    "outputs": [],
    "source": [
     "# Picking all orders at the same time into at most 5 containers but sort at end\n",
-    "multi_order_waves = generate_multi_batch_order_waves(batch, 5, products)\n",
+    "multi_order_waves = generate_multi_batch_order_waves(batch, products, max_containers, max_orders_per_wave)\n",
     "total = PickingCost()\n",
     "for wave in multi_order_waves:\n",
     "    total += wave.compute_picking_cost(product_to_pick_face)\n",
-    "total_time(total)"
+    "total_time(total) / len(batch)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {
+    "scrolled": false
+   },
+   "outputs": [],
+   "source": [
+    "sa = SimAnn(batch, products, product_to_pick_face, max_containers, max_orders_per_wave)\n",
+    "\n",
+    "sa.objective()"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "sa.objective()"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "len(sa.waves[0].container_orders)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "sa.step()"
    ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "sa.step()"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "sa.objective()"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {
+    "scrolled": true
+   },
+   "outputs": [],
+   "source": [
+    "[sa.step() for x in range(50)]"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "sa.waves"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "sa.waves[-1].container_orders"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "max_containers"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "sa.waves[0].container_orders"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": []
   }
  ],
  "metadata": {
   "kernelspec": {
-   "display_name": "Python 3",
+   "display_name": "pickoptimiser",
    "language": "python",
-   "name": "python3"
+   "name": "pickoptimiser"
   },
   "language_info": {
    "codemirror_mode": {