{-# 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 Control.Monad (unless)
import Data.List (findIndex)
import Data.Map qualified as Map
import Data.Maybe (fromJust)

data D = DL | DR

stageTH

-- |
--
-- >>> :main
-- 20777
-- 13289612809129
main :: IO ()
IO ()
main =
 do (steps, nodes) <- [format|2023 8 @D*%n%n(%s = %(%s, %s%)%n)*|]
    let 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
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 -> [String]
mkPath String
"AAA"
    let 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']

    unless (all (isTrivial part2 (length steps)) paths2) (fail "input not trivial")

    print (findIndex' part1 path1)
    print (foldl1 lcm (map (findIndex' part2) 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