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

<https://adventofcode.com/2020/day/23>

-}
module Main (main) where

import Advent.Format (format)
import Control.Monad (unless, zipWithM_)
import Data.Array.IO (IOUArray, getBounds, newArray_, readArray, writeArray)
import Data.Char (digitToInt)
import Data.Foldable (for_)

-- | The array maps cup numbers (indexes) to the next cup
-- in the sequence (elements).
type Ring = IOUArray Int Int

-- | Construct a /ring of cups/ given an initial arrangement and filled to
-- the given size.
newRing ::
  Int   {- ^ ring size           -} ->
  [Int] {- ^ initial arrangement -} ->
  IO Ring
newRing :: Int -> [Int] -> IO Ring
newRing Int
n [Int]
order =
  do a <- (Int, Int) -> IO Ring
forall i. Ix i => (i, i) -> IO (IOUArray i Int)
forall (a :: * -> * -> *) e (m :: * -> *) i.
(MArray a e m, Ix i) =>
(i, i) -> m (a i e)
newArray_ (Int
1,Int
n)
     for_ [1..n-1] \Int
i -> Ring -> Int -> Int -> IO ()
forall (a :: * -> * -> *) e (m :: * -> *) i.
(MArray a e m, Ix i) =>
a i e -> i -> e -> m ()
writeArray Ring
a Int
i (Int
iInt -> Int -> Int
forall a. Num a => a -> a -> a
+Int
1)
     zipWithM_ (writeArray a) order (tail order)
     if n == length order
        then writeArray a (last order) (head order)
        else writeArray a n (head order) >> writeArray a (last order) (1+maximum order)
     pure a

readRing ::
  Ring ->
  Int {- ^ length       -} ->
  Int {- ^ starting cup -} ->
  IO [Int]
readRing :: Ring -> Int -> Int -> IO [Int]
readRing Ring
a Int
n Int
i
  | Int
n Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Int
0 = [Int] -> IO [Int]
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure []
  | Bool
otherwise =
    do j    <- Ring -> Int -> IO Int
forall (a :: * -> * -> *) e (m :: * -> *) i.
(MArray a e m, Ix i) =>
a i e -> i -> m e
readArray Ring
a Int
i
       rest <- readRing a (n-1) j
       pure (j : rest)

-- |
-- >>> :main
-- 47382659
-- 42271866720
main :: IO ()
IO ()
main =
  do inp <- (Char -> Int) -> [Char] -> [Int]
forall a b. (a -> b) -> [a] -> [b]
map Char -> Int
digitToInt ([Char] -> [Int]) -> IO [Char] -> IO [Int]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [format|2020 23 %s%n|]
     p1 inp
     p2 inp

play ::
  Ring ->
  Int {- ^ iterations -} ->
  Int {- ^ current cup -} ->
  IO ()
play :: Ring -> Int -> Int -> IO ()
play Ring
a Int
i Int
cur =
  Bool -> IO () -> IO ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
unless (Int
iInt -> Int -> Bool
forall a. Eq a => a -> a -> Bool
==Int
0)
  do -- extract a group of three cups
     g1 <- Ring -> Int -> IO Int
forall (a :: * -> * -> *) e (m :: * -> *) i.
(MArray a e m, Ix i) =>
a i e -> i -> m e
readArray Ring
a Int
cur
     g2 <- readArray a g1
     g3 <- readArray a g2

     -- find next cup and link current one to it
     nx <- readArray a g3
     writeArray a cur nx

     -- find the new destination label
     (_,hi) <- getBounds a
     let dec Int
1 = Int
hi
         dec Int
x = Int
x Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1
         dest = (Int -> Bool) -> (Int -> Int) -> Int -> Int
forall a. (a -> Bool) -> (a -> a) -> a -> a
until (\Int
x -> Int
x Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
/= Int
g1 Bool -> Bool -> Bool
&& Int
x Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
/= Int
g2 Bool -> Bool -> Bool
&& Int
x Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
/= Int
g3) Int -> Int
dec (Int -> Int
dec Int
cur)

     -- splice the group back in at dest
     writeArray a g3 =<< readArray a dest
     writeArray a dest g1

     play a (i-1) nx

p1 :: [Int] -> IO ()
p1 :: [Int] -> IO ()
p1 [Int]
inp =
  do ring <- Int -> [Int] -> IO Ring
newRing ([Int] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Int]
inp) [Int]
inp
     let sz = [Int] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Int]
inp
     play ring 100 (head inp)

     xs <- readRing ring (sz-1) 1
     putStrLn (concatMap show xs)

p2 :: [Int] -> IO ()
p2 :: [Int] -> IO ()
p2 [Int]
inp =
  do let sz :: Int
sz   =  Int
1_000_000
         iter :: Int
iter = Int
10_000_000

     ring <- Int -> [Int] -> IO Ring
newRing Int
sz [Int]
inp
     play ring iter (head inp)

     xs <- readRing ring 2 1
     print (product xs)