{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "{-# LANGUAGE NegativeLiterals #-}\n",
    "{-# LANGUAGE FlexibleContexts #-}\n",
    "{-# LANGUAGE OverloadedStrings #-}\n",
    "{-# LANGUAGE TypeFamilies #-}\n",
    "{-# LANGUAGE BangPatterns #-}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "import Prelude hiding (Left, Right)\n",
    "import Data.List\n",
    "import qualified Data.Set as S"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "type Point = (Int, Int)\n",
    "type Infection = S.Set Point\n",
    "\n",
    "data Direction = Up | Right | Down | Left deriving (Show, Eq, Enum)\n",
    "\n",
    "data World = World { infected :: Infection\n",
    "                   , position :: Point\n",
    "                   , direction :: Direction\n",
    "                   , infectionCount :: Int\n",
    "                   } deriving (Eq, Show)\n",
    "                   "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "text <- readFile \"../../data/advent22.txt\"\n",
    "grid = lines text"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "sampleGrid = lines \"..#\\n#..\\n...\\n\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "initialInfected g = S.fromList [(r, c) | r <- [0..(length g - 1)], c <- [0..((length . head) g - 1)],\n",
    "                                        g!!r!!c == '#']"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "fromList [(0,2),(1,0)]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "initialInfected sampleGrid"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "initialPosition g = (length g `div` 2, (length . head) g `div` 2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(1,1)"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "initialPosition sampleGrid"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "leftOf Up = Left\n",
    "leftOf x = pred x\n",
    "\n",
    "rightOf Left = Up\n",
    "rightOf x = succ x"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Down"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "leftOf Left"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [],
   "source": [
    "delta :: Direction -> Point\n",
    "delta Up = (-1, 0)\n",
    "delta Right = (0, 1)\n",
    "delta Down = (1, 0)\n",
    "delta Left = (0, -1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [],
   "source": [
    "(+:) (r, c) (dr, dc) = (r + dr, c + dc)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [],
   "source": [
    "initialWorld grid = World \n",
    "    { infected = initialInfected grid\n",
    "    , position = initialPosition grid\n",
    "    , direction = Up\n",
    "    , infectionCount = 0\n",
    "    }"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "World {infected = fromList [(0,2),(1,0)], position = (1,1), direction = Up, infectionCount = 0}"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "initialWorld sampleGrid"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [],
   "source": [
    "step world = World {infected = inf', position = pos', direction = dir', infectionCount = ic'}\n",
    "    where here = position world\n",
    "          infectedHere = here `S.member` infected world\n",
    "          dir' = if infectedHere then rightOf (direction world)\n",
    "                                 else leftOf (direction world)\n",
    "          inf' = if infectedHere then S.delete here $ infected world\n",
    "                                 else S.insert here $ infected world\n",
    "          ic'  = if infectedHere then infectionCount world\n",
    "                                 else infectionCount world + 1\n",
    "          pos' = here +: delta dir'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "World {infected = fromList [(0,2),(1,1)], position = (0,0), direction = Up, infectionCount = 1}"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "step $ step $ initialWorld sampleGrid"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [],
   "source": [
    "progress n = (!! n) . iterate step "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "5587"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "infectionCount $ progress 10000 $ initialWorld sampleGrid"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "5182"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "infectionCount $ progress 10000 $ initialWorld grid"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Haskell",
   "language": "haskell",
   "name": "haskell"
  },
  "language_info": {
   "codemirror_mode": "ihaskell",
   "file_extension": ".hs",
   "name": "haskell",
   "version": "8.0.2"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}