{-# Language QuasiQuotes, ImportQualifiedPost, 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\n2\n-3\n3\n-2\n0\n4\n"
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
    [Int]
input <- [format|2022 20 (%d%n)*|]
    Int -> IO ()
forall a. Show a => a -> IO ()
print (Int -> IO ()) -> IO Int -> IO ()
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< Int -> [Int] -> IO Int
solve Int
1 [Int]
input
    Int -> IO ()
forall a. Show a => a -> IO ()
print (Int -> IO ()) -> IO Int -> IO ()
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< Int -> [Int] -> IO Int
solve Int
10 ((Int -> Int) -> [Int] -> [Int]
forall a b. (a -> b) -> [a] -> [b]
map (Int
811_589_153 Int -> Int -> Int
forall a. Num a => a -> a -> a
*) [Int]
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
ring <- Int -> IO Ring
newRing Int
n

    Int -> IO () -> IO ()
forall (m :: * -> *) a. Applicative m => Int -> m a -> m ()
replicateM_ Int
iterations (IO () -> IO ()) -> IO () -> IO ()
forall a b. (a -> b) -> a -> b
$
      [(Int, Int)] -> ((Int, Int) -> IO ()) -> IO ()
forall (t :: * -> *) (f :: * -> *) a b.
(Foldable t, Applicative f) =>
t a -> (a -> f b) -> f ()
for_ (UArray Int Int -> [(Int, Int)]
forall (a :: * -> * -> *) e i.
(IArray a e, Ix i) =>
a i e -> [(i, e)]
assocs UArray Int Int
inputArray) \(Int
i,Int
v) ->
       do Int
a <- Int -> Ring -> IO Int
removeRing Int
i Ring
ring
          Int
a' <- Int -> Int -> Ring -> IO Int
walk Int
v Int
a Ring
ring
          Int -> Int -> Ring -> IO ()
insertBeforeRing Int
i Int
a' Ring
ring

    Int
i0 <- case Int -> [Int] -> Maybe Int
forall a. Eq a => a -> [a] -> Maybe Int
elemIndex Int
0 [Int]
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
    Int
i1 <- Int -> Int -> Ring -> IO Int
walk Int
1_000 Int
i0 Ring
ring
    Int
i2 <- Int -> Int -> Ring -> IO Int
walk Int
1_000 Int
i1 Ring
ring
    Int
i3 <- Int -> Int -> Ring -> IO Int
walk Int
1_000 Int
i2 Ring
ring
    Int -> IO Int
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Int -> IO Int) -> Int -> IO Int
forall a b. (a -> b) -> a -> b
$! [Int] -> Int
forall a. Num a => [a] -> a
forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
sum [UArray Int Int
inputArrayUArray Int Int -> Int -> Int
forall (a :: * -> * -> *) e i.
(IArray a e, Ix i) =>
a i e -> i -> e
!Int
i1, UArray Int Int
inputArrayUArray Int Int -> Int -> Int
forall (a :: * -> * -> *) e i.
(IArray a e, Ix i) =>
a i e -> i -> e
!Int
i2, UArray Int Int
inputArrayUArray Int Int -> Int -> Int
forall (a :: * -> * -> *) e i.
(IArray a e, Ix i) =>
a i e -> i -> e
!Int
i3]

data Ring = Ring {
    Ring -> Int
ringSize :: !Int,
    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 =
  Int -> IOUArray Int Int -> IOUArray Int Int -> Ring
Ring Int
n
    (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{Int
IOUArray Int Int
ringSize :: Ring -> Int
fwdLinks :: Ring -> IOUArray Int Int
bwdLinks :: Ring -> IOUArray Int Int
ringSize :: Int
fwdLinks :: IOUArray Int Int
bwdLinks :: IOUArray Int Int
..} =
 do Int
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
    Int
next <- 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
    IOUArray Int Int -> Int -> Int -> IO ()
forall (a :: * -> * -> *) e (m :: * -> *) i.
(MArray a e m, Ix i) =>
a i e -> i -> e -> m ()
writeArray IOUArray Int Int
fwdLinks Int
prev Int
next
    IOUArray Int Int -> Int -> Int -> IO ()
forall (a :: * -> * -> *) e (m :: * -> *) i.
(MArray a e m, Ix i) =>
a i e -> i -> e -> m ()
writeArray IOUArray Int Int
bwdLinks Int
next Int
prev
    Int -> IO Int
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Int
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{Int
IOUArray Int Int
ringSize :: Ring -> Int
fwdLinks :: Ring -> IOUArray Int Int
bwdLinks :: Ring -> IOUArray Int Int
ringSize :: Int
fwdLinks :: IOUArray Int Int
bwdLinks :: IOUArray Int Int
..} =
 do Int
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
    IOUArray Int Int -> Int -> Int -> IO ()
forall (a :: * -> * -> *) e (m :: * -> *) i.
(MArray a e m, Ix i) =>
a i e -> i -> e -> m ()
writeArray IOUArray Int Int
fwdLinks Int
node Int
next
    IOUArray Int Int -> Int -> Int -> IO ()
forall (a :: * -> * -> *) e (m :: * -> *) i.
(MArray a e m, Ix i) =>
a i e -> i -> e -> m ()
writeArray IOUArray Int Int
bwdLinks Int
node Int
prev
    IOUArray Int Int -> Int -> Int -> IO ()
forall (a :: * -> * -> *) e (m :: * -> *) i.
(MArray a e m, Ix i) =>
a i e -> i -> e -> m ()
writeArray IOUArray Int Int
fwdLinks Int
prev Int
node
    IOUArray Int Int -> Int -> Int -> IO ()
forall (a :: * -> * -> *) e (m :: * -> *) i.
(MArray a e m, Ix i) =>
a i e -> i -> e -> m ()
writeArray IOUArray Int Int
bwdLinks Int
next Int
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
start Ring{Int
IOUArray Int Int
ringSize :: Ring -> Int
fwdLinks :: Ring -> IOUArray Int Int
bwdLinks :: Ring -> IOUArray Int Int
ringSize :: Int
fwdLinks :: IOUArray Int Int
bwdLinks :: IOUArray Int Int
..}
  | Int
n' Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Int
len Int -> Int -> Int
forall a. Integral a => a -> a -> a
`div` Int
2 = Int -> (Int -> IO Int) -> Int -> IO Int
forall (m :: * -> *) a. Monad m => Int -> (a -> m a) -> a -> m a
timesM (Int
lenInt -> Int -> Int
forall a. Num a => a -> a -> a
-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
start
  | 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
start
  where
    len :: Int
len = Int
ringSize Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1
    n' :: Int
n' = Int
n Int -> Int -> Int
forall a. Integral a => a -> a -> a
`mod` Int
len