! Copyright (C) 2023 CapitalEx.
! See https://factorcode.org/license.txt for BSD license.
USING: arrays assocs combinators.short-circuit kernel make
math.parser math.vectors multiline peg.ebnf sequences
sequences.deep strings ;
IN: aoc.2023.3

: input ( -- string )
    "vocab:aoc/2023/3/input.txt" utf8 file-contents ;


CONSTANT: neighbors {
    { -1 -1 } { 0 -1 } {  1 -1 }
    { -1  0 }          {  1  0 }
    { -1  1 } { 0  1 } {  1  1 }
}

SYMBOL: +empty+
SYMBOL: id


TUPLE: grid-number value position ;
: <grid-number> ( value -- grid-number )
    f \ grid-number boa ;

EBNF: parse-row [=[
    Digit     = "0"|"1"|"2"|"3"|"4"|"5"|"6"|"7"|"8"|"9"
    Number    = Digit+ => [[ concat [ length ] [ dec> ] bi <grid-number> <array> ]]
    Emtpy     = "."    => [[ drop +empty+ ]]
    Symbol    = .      => [[ 1string ]]
    Row       = (Number|Emtpy|Symbol)* => [[ flatten ]] 
]=]

: collect-elements ( 2d-seq -- )
    [ '[ _ 2array 2dup swap dup grid-number? [ position<< ] [ 2drop ] if ,, ] each-index ] each-index ;

: hash-grid ( 2d-seq -- hash-grid )
    [ collect-elements ] LH{ } make ;

: parse-schematic ( string -- hash-grid )
    split-lines [ parse-row ] map hash-grid ;

: symbol? ( obj -- ? )
    { [ grid-number? not ] [ +empty+ = not ] } 1&& ;

: find-symbols ( hash-grid -- symbols )
    [ nip symbol? ] assoc-filter ;

: search-neighbors ( hash-grid position -- )
    neighbors [ v+ of , ] 2with each ;

: collect-neighbors ( hash-grid symbols -- seq )
    [ [ drop search-neighbors ] V{ } make ] withd V{ } assoc>map  ;

: find-part-numbers ( hash-grid symbols -- vector )
    collect-neighbors flatten members [ grid-number? ] [ value>> ] filter-map ;

: solve-part-one ( -- )
    input parse-schematic
    dup find-symbols 
    find-part-numbers 
    sum . ;

: find-asterisk ( assoc -- assoc )
     [ nip "*" = ] assoc-filter ;

: find-gears ( seq -- seq  )
    [ members [ grid-number? ] filter ] map 
    [ length 2 = ] filter ;

: sum-of-gear-ratios ( seq -- seq )
    [ [ first value>> ] [ second value>> ] bi * ] map-sum ;

: solve-part-two ( -- )
    input parse-schematic
    dup find-symbols 
    find-asterisk 
    collect-neighbors 
    find-gears 
    sum-of-gear-ratios . ;