{-# Language QuasiQuotes, NumericUnderscores, BlockArguments, RecordWildCards #-}
{-|
Module      : Main
Description : Day 20 solution
Copyright   : (c) Eric Mertens, 2022
License     : ISC
Maintainer  : emertens@gmail.com

<https://adventofcode.com/2022/day/20>

>>> :{
:main + "1
2
-3
3
-2
0
4
"
:}
3
1623178306

-}
module Main where

import Control.Monad (replicateM_)
import Data.Array.IO (IOUArray, newListArray, readArray, writeArray)
import Data.Array.Unboxed (UArray, (!), assocs, listArray)
import Data.Foldable (for_)
import Data.List (elemIndex)

import Advent (format, timesM)

-- |
-- >>> :main
-- 1591
-- 14579387544492
main :: IO ()
IO ()
main = do
    input <- [format|2022 20 (%d%n)*|]
    print =<< solve 1 input
    print =<< solve 10 (map (811_589_153 *) input)

solve :: Int -> [Int] -> IO Int
solve :: Int -> [Int] -> IO Int
solve Int
iterations [Int]
xs =
 do let n :: Int
n = [Int] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Int]
xs
        inputArray :: UArray Int Int
inputArray = (Int, Int) -> [Int] -> UArray Int Int
forall (a :: * -> * -> *) e i.
(IArray a e, Ix i) =>
(i, i) -> [e] -> a i e
listArray (Int
0, Int
nInt -> Int -> Int
forall a. Num a => a -> a -> a
-Int
1) [Int]
xs :: UArray Int Int
    ring <- Int -> IO Ring
newRing Int
n

    replicateM_ iterations $
      for_ (assocs inputArray) \(Int
i,Int
v) ->
       do let d :: Int
d = Int
v Int -> Int -> Int
forall a. Integral a => a -> a -> a
`mod` (Int
nInt -> Int -> Int
forall a. Num a => a -> a -> a
-Int
1)
              d' :: Int
d' = if Int
d Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Int
nInt -> Int -> Int
forall a. Integral a => a -> a -> a
`div`Int
2 then Int
dInt -> Int -> Int
forall a. Num a => a -> a -> a
-(Int
nInt -> Int -> Int
forall a. Num a => a -> a -> a
-Int
1) else Int
d
          a <- Int -> Ring -> IO Int
removeRing Int
i Ring
ring
          a' <- walk d' a ring
          insertBeforeRing i a' ring

    i0 <- case elemIndex 0 xs of
            Maybe Int
Nothing -> String -> IO Int
forall a. String -> IO a
forall (m :: * -> *) a. MonadFail m => String -> m a
fail String
"bad input"
            Just Int
i  -> Int -> IO Int
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Int
i
    i1 <- walk 1_000 i0 ring
    i2 <- walk 1_000 i1 ring
    i3 <- walk 1_000 i2 ring
    pure $! sum [inputArray!i1, inputArray!i2, inputArray!i3]

data Ring = Ring {
    Ring -> IOUArray Int Int
fwdLinks :: !(IOUArray Int Int), -- ^ forward links
    Ring -> IOUArray Int Int
bwdLinks :: !(IOUArray Int Int)  -- ^ backward links
}

-- | Build a new circular ring of given size
newRing :: Int {- ^ size -} -> IO Ring
newRing :: Int -> IO Ring
newRing Int
n =
  IOUArray Int Int -> IOUArray Int Int -> Ring
Ring
    (IOUArray Int Int -> IOUArray Int Int -> Ring)
-> IO (IOUArray Int Int) -> IO (IOUArray Int Int -> Ring)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (Int, Int) -> [Int] -> IO (IOUArray Int Int)
forall (a :: * -> * -> *) e (m :: * -> *) i.
(MArray a e m, Ix i) =>
(i, i) -> [e] -> m (a i e)
newListArray (Int
0,Int
nInt -> Int -> Int
forall a. Num a => a -> a -> a
-Int
1) ([Int
1..Int
nInt -> Int -> Int
forall a. Num a => a -> a -> a
-Int
1][Int] -> [Int] -> [Int]
forall a. [a] -> [a] -> [a]
++[Int
0])
    IO (IOUArray Int Int -> Ring) -> IO (IOUArray Int Int) -> IO Ring
forall a b. IO (a -> b) -> IO a -> IO b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> (Int, Int) -> [Int] -> IO (IOUArray Int Int)
forall (a :: * -> * -> *) e (m :: * -> *) i.
(MArray a e m, Ix i) =>
(i, i) -> [e] -> m (a i e)
newListArray (Int
0,Int
nInt -> Int -> Int
forall a. Num a => a -> a -> a
-Int
1) ((Int
nInt -> Int -> Int
forall a. Num a => a -> a -> a
-Int
1)Int -> [Int] -> [Int]
forall a. a -> [a] -> [a]
:[Int
0..Int
nInt -> Int -> Int
forall a. Num a => a -> a -> a
-Int
2])

-- | Remove a node from the ring and return the index of the node before/after it.
removeRing ::
  Int {- ^ node to remove from ring -} ->
  Ring {- ^ ring -} ->
  IO Int {- ^ index after removed node -}
removeRing :: Int -> Ring -> IO Int
removeRing Int
i Ring{IOUArray Int Int
fwdLinks :: Ring -> IOUArray Int Int
bwdLinks :: Ring -> IOUArray Int Int
fwdLinks :: IOUArray Int Int
bwdLinks :: IOUArray Int Int
..} =
 do prev <- IOUArray Int Int -> Int -> IO Int
forall (a :: * -> * -> *) e (m :: * -> *) i.
(MArray a e m, Ix i) =>
a i e -> i -> m e
readArray IOUArray Int Int
bwdLinks Int
i
    next <- readArray fwdLinks i
    writeArray fwdLinks prev next
    writeArray bwdLinks next prev
    pure next

insertBeforeRing ::
  Int {- ^ node to insert -} ->
  Int {- ^ node to insert before -} ->
  Ring {- ^ ring -} ->
  IO ()
insertBeforeRing :: Int -> Int -> Ring -> IO ()
insertBeforeRing Int
node Int
next Ring{IOUArray Int Int
fwdLinks :: Ring -> IOUArray Int Int
bwdLinks :: Ring -> IOUArray Int Int
fwdLinks :: IOUArray Int Int
bwdLinks :: IOUArray Int Int
..} =
 do prev <- IOUArray Int Int -> Int -> IO Int
forall (a :: * -> * -> *) e (m :: * -> *) i.
(MArray a e m, Ix i) =>
a i e -> i -> m e
readArray IOUArray Int Int
bwdLinks Int
next
    writeArray fwdLinks node next
    writeArray bwdLinks node prev
    writeArray fwdLinks prev node
    writeArray bwdLinks next node

walk ::
  Int {- ^ step count (positive for forward, negative for backward) -} ->
  Int {- ^ starting index -} ->
  Ring {- ^ ring -} ->
  IO Int {- ^ ending index -}
walk :: Int -> Int -> Ring -> IO Int
walk Int
n Int
i Ring{IOUArray Int Int
fwdLinks :: Ring -> IOUArray Int Int
bwdLinks :: Ring -> IOUArray Int Int
fwdLinks :: IOUArray Int Int
bwdLinks :: IOUArray Int Int
..}
  | Int
n Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
0     = Int -> (Int -> IO Int) -> Int -> IO Int
forall (m :: * -> *) a. Monad m => Int -> (a -> m a) -> a -> m a
timesM (-Int
n) (IOUArray Int Int -> Int -> IO Int
forall (a :: * -> * -> *) e (m :: * -> *) i.
(MArray a e m, Ix i) =>
a i e -> i -> m e
readArray IOUArray Int Int
bwdLinks) Int
i
  | Bool
otherwise = Int -> (Int -> IO Int) -> Int -> IO Int
forall (m :: * -> *) a. Monad m => Int -> (a -> m a) -> a -> m a
timesM   Int
n  (IOUArray Int Int -> Int -> IO Int
forall (a :: * -> * -> *) e (m :: * -> *) i.
(MArray a e m, Ix i) =>
a i e -> i -> m e
readArray IOUArray Int Int
fwdLinks) Int
i