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

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

>>> :{
:main +
"467..114..
...*......
..35..633.
......#...
617*......
.....+.58.
..592.....
......755.
...$.*....
.664.598.."
:}
4361
467835

-}
module Main (main) where

import Advent (getInputArray, ordNub, arrIx)
import Advent.Coord (Coord, left, neighbors, right)
import Data.Array.Unboxed (UArray, assocs)
import Data.Char (isDigit)
import Data.Map qualified as Map
import Data.Maybe (fromMaybe)

-- | Parse the input schematic and print answers to both parts.
--
-- >>> :main
-- 527144
-- 81463996
main :: IO ()
IO ()
main =
 do input <- Int -> Int -> IO (UArray Coord Char)
getInputArray Int
2023 Int
3
    let numbers = UArray Coord Char -> [(Int, [(Coord, Char)])]
extractNumbers UArray Coord Char
input
    print (sum [partNo | (partNo, _:_) <- numbers])
    print (sum [a * b | [a, b] <- gearNumbers numbers])

-- | Extract the numbers from the diagram and the parts adjacent to them.
extractNumbers :: UArray Coord Char -> [(Int, [(Coord, Char)])]
extractNumbers :: UArray Coord Char -> [(Int, [(Coord, Char)])]
extractNumbers UArray Coord Char
input =
  [ (String -> Int
forall a. Read a => String -> a
read String
digits, [Coord] -> [(Coord, Char)]
forall {t :: * -> *}. Foldable t => t Coord -> [(Coord, Char)]
partsNear [Coord]
cs)
  | (Coord
c, Char
digit) <- UArray Coord Char -> [(Coord, Char)]
forall (a :: * -> * -> *) e i.
(IArray a e, Ix i) =>
a i e -> [(i, e)]
assocs UArray Coord Char
input
  , Char -> Bool
isDigit Char
digit, Bool -> Bool
not (Char -> Bool
isDigit (Coord -> Char
lkp (Coord -> Coord
left Coord
c))) -- left-boundary of number
  , let ([Coord]
cs, String
digits) = [(Coord, Char)] -> ([Coord], String)
forall a b. [(a, b)] -> ([a], [b])
unzip (Coord -> [(Coord, Char)]
numbersAfter Coord
c)
  ]
  where
    lkp :: Coord -> Char
lkp = Char -> Maybe Char -> Char
forall a. a -> Maybe a -> a
fromMaybe Char
'.' (Maybe Char -> Char) -> (Coord -> Maybe Char) -> Coord -> Char
forall b c a. (b -> c) -> (a -> b) -> a -> c
. UArray Coord Char -> Coord -> Maybe Char
forall (a :: * -> * -> *) e i (f :: * -> *).
(IArray a e, Ix i, Alternative f) =>
a i e -> i -> f e
arrIx UArray Coord Char
input
    numbersAfter :: Coord -> [(Coord, Char)]
numbersAfter Coord
start =
      [ (Coord
c, Char
digit)
      | Coord
c <- (Coord -> Coord) -> Coord -> [Coord]
forall a. (a -> a) -> a -> [a]
iterate Coord -> Coord
right Coord
start
      , let digit :: Char
digit = Coord -> Char
lkp Coord
c
      , then (a -> Bool) -> [a] -> [a]
((Coord, Char) -> Bool) -> [(Coord, Char)] -> [(Coord, Char)]
forall a. (a -> Bool) -> [a] -> [a]
takeWhile by Char -> Bool
isDigit Char
digit
      ]
    partsNear :: t Coord -> [(Coord, Char)]
partsNear t Coord
cs =
      [ (Coord
c, Char
sym)
      | Coord
c <- [Coord] -> [Coord]
forall a. Ord a => [a] -> [a]
ordNub ((Coord -> [Coord]) -> t Coord -> [Coord]
forall (t :: * -> *) a b. Foldable t => (a -> [b]) -> t a -> [b]
concatMap Coord -> [Coord]
neighbors t Coord
cs)
      , let sym :: Char
sym = Coord -> Char
lkp Coord
c
      , Char -> Bool
isPart Char
sym
      ]

-- | Make lists of the numbers next to each gear in the schematic
gearNumbers :: [(Int, [(Coord, Char)])] -> [[Int]]
gearNumbers :: [(Int, [(Coord, Char)])] -> [[Int]]
gearNumbers [(Int, [(Coord, Char)])]
numbers =
  Map Coord [Int] -> [[Int]]
forall k a. Map k a -> [a]
Map.elems (([Int] -> [Int] -> [Int]) -> [(Coord, [Int])] -> Map Coord [Int]
forall k a. Ord k => (a -> a -> a) -> [(k, a)] -> Map k a
Map.fromListWith [Int] -> [Int] -> [Int]
forall a. [a] -> [a] -> [a]
(++)
    [(Coord
part, [Int
partNo]) | (Int
partNo, [(Coord, Char)]
parts) <- [(Int, [(Coord, Char)])]
numbers, (Coord
part, Char
'*') <- [(Coord, Char)]
parts])

-- | Things that aren't digits or periods.
isPart :: Char -> Bool
isPart :: Char -> Bool
isPart Char
x = Bool -> Bool
not (Char -> Bool
isDigit Char
x) Bool -> Bool -> Bool
&& Char
x Char -> Char -> Bool
forall a. Eq a => a -> a -> Bool
/= Char
'.'