{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "{-# LANGUAGE NegativeLiterals #-}\n",
    "{-# LANGUAGE FlexibleContexts #-}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "import qualified Data.Vector as V\n",
    "import Data.Vector ((//), (!))\n",
    "import Data.List (unfoldr)\n",
    "import qualified Data.Map.Strict as M"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "type Memory = V.Vector Int\n",
    "type Redist = (Int, Int, Memory)\n",
    "type History = M.Map Memory Int"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[0,2,7,0]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "vec0 = V.fromList [0, 2, 7, 0]\n",
    "vec0"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "redistStart :: Memory -> Redist\n",
    "redistStart vec0 = (current, toDistribute, startVec)\n",
    "    where origin = V.maxIndex vec0\n",
    "          toDistribute = vec0!origin\n",
    "          current = (origin + 1) `mod` (length vec0)\n",
    "          startVec = vec0 // [(origin, 0)]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(3,7,[0,2,0,0])"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "redistStart $ V.fromList [0, 2, 7, 0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "redistR :: Redist -> Memory\n",
    "redistR (_, 0, vec) = vec\n",
    "redistR (i, n, vec) = redistR (i', n', vec')\n",
    "    where n' = n - 1\n",
    "          i' = (i + 1) `mod` (length vec)\n",
    "          vec' = vec // [(i, vec!i + 1)]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "redist = redistR . redistStart"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[2,4,1,2]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "redist $ V.fromList [0, 2, 7, 0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[3,1,2,3]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "redist $ redist $  V.fromList [0, 2, 7, 0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "redistSeq = unfoldr redistU\n",
    "    where redistU vec = Just (redist vec, redist vec)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[(0,[2,4,1,2]),(1,[3,1,2,3]),(2,[0,2,3,4]),(3,[1,3,4,1]),(4,[2,4,1,2]),(5,[3,1,2,3]),(6,[0,2,3,4]),(7,[1,3,4,1])]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "take 8 $ zip [0..] $ redistSeq $ V.fromList [0, 2, 7, 0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [],
   "source": [
    "findRepeat :: History -> [(Int, Memory)] -> Int\n",
    "findRepeat h ((n, x) : nxs) = if x `M.member` h \n",
    "                              then n + 1\n",
    "                              else findRepeat (M.insert x n h) nxs\n",
    "    where n0 = (M.!) h x"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "metadata": {},
   "outputs": [],
   "source": [
    "findRepeatB :: History -> [(Int, Memory)] -> Int\n",
    "findRepeatB h ((n, x) : nxs) = if x `M.member` h \n",
    "                              then n - n0\n",
    "                              else findRepeatB (M.insert x n h) nxs\n",
    "    where n0 = (M.!) h x"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "5"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "findRepeat M.empty $ zip [0..] $ redistSeq $ V.fromList [0, 2, 7, 0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "metadata": {},
   "outputs": [],
   "source": [
    "part1 :: [Int] -> Int\n",
    "-- part1 memlist = findRepeat M.empty $ zip [0..] $ redistSeq $ V.fromList memlist\n",
    "part1 = (findRepeat M.empty) . (zip [0..]) . redistSeq . V.fromList"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "metadata": {},
   "outputs": [],
   "source": [
    "part2 :: [Int] -> Int\n",
    "-- part1 memlist = findRepeat M.empty $ zip [0..] $ redistSeq $ V.fromList memlist\n",
    "part2 = (findRepeatB M.empty) . (zip [0..]) . redistSeq . V.fromList"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "metadata": {},
   "outputs": [],
   "source": [
    "main :: IO ()\n",
    "main = do \n",
    "        text <- readFile \"../../data/advent06.txt\"\n",
    "--         text0 <- readFile \"../../data/advent06.txt\"\n",
    "--         let text = \"0 2 7 0\"\n",
    "        let memory = map read $ words text\n",
    "        print $ part1 memory\n",
    "        print $ part2 memory"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "7864\n",
       "1695"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "main"
   ]
  },
  {
   "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
}