{-# Language QuasiQuotes, TemplateHaskell, ImportQualifiedPost, LambdaCase #-}
{-|
Module      : Main
Description : Day 8 solution
Copyright   : (c) Eric Mertens, 2023
License     : ISC
Maintainer  : emertens@gmail.com

<https://adventofcode.com/2023/day/8>

-}
module Main (main) where

import Advent (format, stageTH)
import Data.Map qualified as Map
import Data.Maybe (fromJust)
import Data.List (findIndex)
import Control.Monad (unless)

data D = DL | DR

stageTH

-- |
--
-- >>> :main
-- 20777
-- 13289612809129
main :: IO ()
IO ()
main =
 do ([D]
steps, [(String, String, String)]
nodes) <- [format|2023 8 @D*%n%n(%s = %(%s, %s%)%n)*|]
    let nodes' :: Map String (D -> String)
nodes' = [(String, D -> String)] -> Map String (D -> String)
forall k a. Ord k => [(k, a)] -> Map k a
Map.fromList [(String
k, \case D
DL -> String
a; D
DR -> String
b) | (String
k,String
a,String
b) <- [(String, String, String)]
nodes]
    let mkPath :: String -> [String]
mkPath String
start = (String -> D -> String) -> String -> [D] -> [String]
forall b a. (b -> a -> b) -> b -> [a] -> [b]
scanl (Map String (D -> String)
nodes' Map String (D -> String) -> String -> D -> String
forall k a. Ord k => Map k a -> k -> a
Map.!) String
start ([D] -> [D]
forall a. HasCallStack => [a] -> [a]
cycle [D]
steps)

    let path1 :: [String]
path1  = String -> [String]
mkPath String
"AAA"
    let paths2 :: [[String]]
paths2 = [String -> [String]
mkPath String
start | (String
start, String
_, String
_) <- [(String, String, String)]
nodes, String -> Char
forall a. HasCallStack => [a] -> a
last String
start Char -> Char -> Bool
forall a. Eq a => a -> a -> Bool
== Char
'A']

    Bool -> IO () -> IO ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
unless (([String] -> Bool) -> [[String]] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
all ((String -> Bool) -> Int -> [String] -> Bool
forall a. Eq a => (a -> Bool) -> Int -> [a] -> Bool
isTrivial String -> Bool
part2 ([D] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [D]
steps)) [[String]]
paths2) (String -> IO ()
forall a. String -> IO a
forall (m :: * -> *) a. MonadFail m => String -> m a
fail String
"input not trivial")

    Int -> IO ()
forall a. Show a => a -> IO ()
print ((String -> Bool) -> [String] -> Int
forall a. (a -> Bool) -> [a] -> Int
findIndex' String -> Bool
part1 [String]
path1)
    Int -> IO ()
forall a. Show a => a -> IO ()
print ((Int -> Int -> Int) -> [Int] -> Int
forall a. (a -> a -> a) -> [a] -> a
forall (t :: * -> *) a. Foldable t => (a -> a -> a) -> t a -> a
foldl1 Int -> Int -> Int
forall a. Integral a => a -> a -> a
lcm (([String] -> Int) -> [[String]] -> [Int]
forall a b. (a -> b) -> [a] -> [b]
map ((String -> Bool) -> [String] -> Int
forall a. (a -> Bool) -> [a] -> Int
findIndex' String -> Bool
part2) [[String]]
paths2))

part1, part2 :: String -> Bool
part1 :: String -> Bool
part1 String
x = String
"ZZZ" String -> String -> Bool
forall a. Eq a => a -> a -> Bool
== String
x
part2 :: String -> Bool
part2 String
x = Char
'Z' Char -> Char -> Bool
forall a. Eq a => a -> a -> Bool
== String -> Char
forall a. HasCallStack => [a] -> a
last String
x

-- Verifies that we actually got one of the trivial input files.
-- * The goal must be reached after a number of cycles that is a multiple of the steps
-- * The next goal must be the same as the previous and must be reachable in the
--   same number of steps
--
-- This guarantees that the path must actually cycle infinitely and that
-- there is exactly one goal state in the cycle.
isTrivial :: Eq a => (a -> Bool) -> Int -> [a] -> Bool
isTrivial :: forall a. Eq a => (a -> Bool) -> Int -> [a] -> Bool
isTrivial a -> Bool
p Int
n [a]
xs =
  case [ (Int
i,a
x) | (Int
i, a
x) <- [Int] -> [a] -> [(Int, a)]
forall a b. [a] -> [b] -> [(a, b)]
zip [Int
0..] [a]
xs, a -> Bool
p a
x ] of
    (Int
i1,a
g1) : (Int
i2,a
g2) : [(Int, a)]
_ -> Int
i1 Int -> Int -> Int
forall a. Integral a => a -> a -> a
`rem` Int
n Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
0 Bool -> Bool -> Bool
&& Int
2 Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
i1 Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
i2 Bool -> Bool -> Bool
&& a
g1 a -> a -> Bool
forall a. Eq a => a -> a -> Bool
== a
g2
    [(Int, a)]
_ -> Bool
False

findIndex' :: (a -> Bool) -> [a] -> Int
findIndex' :: forall a. (a -> Bool) -> [a] -> Int
findIndex' a -> Bool
p = Maybe Int -> Int
forall a. HasCallStack => Maybe a -> a
fromJust (Maybe Int -> Int) -> ([a] -> Maybe Int) -> [a] -> Int
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (a -> Bool) -> [a] -> Maybe Int
forall a. (a -> Bool) -> [a] -> Maybe Int
findIndex a -> Bool
p