{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import z3"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We don't need to solve for all of the trajectory intersections at once; picking just a handful will do.  \n",
    "\n",
    "There are initially six unknowns (the three components each of $\\overrightarrow{x_{rock}}$ and $\\overrightarrow{v_{rock}}$), and each additional equation adds three equations of the form $x_{i\\cdot rock} + v_{i \\cdot rock} t_n = x_{i\\cdot hail\\cdot n} + v_{i \\cdot hail\\cdot n} t_n$, which also contains one new unknown, the time of intersection $t_n$.\n",
    "\n",
    "| Hail | Equations | Unknowns |\n",
    "| ---- | ----------- | ----------- |\n",
    "| 1 | 3 | 7 |\n",
    "| 2 | 6 | 8 |\n",
    "| 3 | 9 | 9 |\n",
    "\n",
    "We just need 3 hail trajectories.  The first three in the set are enough."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[[156689809620606, 243565579389165, 455137247320393, -26, 48, -140],\n",
       " [106355761063908, 459832650718033, 351953299411025, 73, -206, -52],\n",
       " [271915251832336, 487490927073225, 398003502953444, 31, -414, -304]]"
      ]
     },
     "execution_count": 2,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "hail = [[int(i) for i in l.replace('@',',').split(',')]\n",
    "                for l in open('input')]\n",
    "\n",
    "hail_sample = hail[0:3]\n",
    "\n",
    "hail_position = [h[0:3] for h in hail_sample]\n",
    "hail_velocity = [h[3:] for h in hail_sample]\n",
    "\n",
    "hail_sample"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "start_position = z3.RealVector('p', 3)\n",
    "start_velocity = z3.RealVector('v', 3)\n",
    "hit_time = z3.RealVector('t', 3)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "For each of the hailstone trajectories in the sample, the thrown rock's coordinates will be equal to the hailstone's coordinates at a given time $t$: $x_{i\\cdot rock} + v_{i \\cdot rock} t = x_{i\\cdot hail} + v_{i \\cdot hail} t$ for $i \\in \\{x,\\,y,\\,z\\}$."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "solver = z3.Solver()\n",
    "solver.add(*[start_position[i] + start_velocity[i] * t == hp[i] + hv[i] * t\n",
    "             for t, hp, hv in zip(hit_time, hail_position, hail_velocity) for i in range(3)])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Is the solver satisfied?  That is, do none of the constraints conflict with each other?"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<b>sat</b>"
      ],
      "text/plain": [
       "sat"
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "solver.check()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The solver model returns the state that satisfies all the constraints:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "[t<sub>2</sub> = 474506740599,\n",
       " t<sub>0</sub> = 931974028142,\n",
       " v<sub>0</sub> = -337,\n",
       " t<sub>1</sub> = 829702369046,\n",
       " v<sub>1</sub> = -6,\n",
       " v<sub>2</sub> = 155,\n",
       " p<sub>0</sub> = 446533732372768,\n",
       " p<sub>2</sub> = 180204909018503,\n",
       " p<sub>1</sub> = 293892176908833]"
      ],
      "text/plain": [
       "[t__2 = 474506740599,\n",
       " t__0 = 931974028142,\n",
       " v__0 = -337,\n",
       " t__1 = 829702369046,\n",
       " v__1 = -6,\n",
       " v__2 = 155,\n",
       " p__0 = 446533732372768,\n",
       " p__2 = 180204909018503,\n",
       " p__1 = 293892176908833]"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "solver.model()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "and all that's left to do is extract the results from the state and calculate the answer."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "920630818300104"
      ],
      "text/plain": [
       "920630818300104"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "solver.model().eval(sum(start_position))"
   ]
  }
 ],
 "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.12.0"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}