{-# 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 Data.Map (Map)
import Data.Map qualified as Map
import Data.Char (isDigit)

import Advent (getInputMap, ordNub)
import Advent.Coord (Coord, left, neighbors, right)

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

-- | Extract the numbers from the diagram and the parts adjacent to them.
extractNumbers :: Map Coord Char -> [(Int, [(Coord, Char)])]
extractNumbers :: Map Coord Char -> [(Int, [(Coord, Char)])]
extractNumbers Map 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) <- Map Coord Char -> [(Coord, Char)]
forall k a. Map k a -> [(k, a)]
Map.assocs Map 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 Coord
i = Char -> Coord -> Map Coord Char -> Char
forall k a. Ord k => a -> k -> Map k a -> a
Map.findWithDefault Char
'.' Coord
i Map 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
'.'