USING: arrays assocs assocs.extras debugger grouping
io.encodings.ascii io.files kernel literals math math.matrices
math.statistics math.vectors peg.parsers peg.search prettyprint
ranges sequences sequences.extras sequences.product vectors ;
IN: aoc2022.day14

CONSTANT: input
    $[ "vocab:aoc2022/day14/input.txt" ascii file-lines ]

: parse ( str -- seq )
    integer-parser search 2 <groups> ;

: expand ( loc loc -- seq )
    [ [a..b] ] 2map <product-sequence> ;

: rocks ( seq -- newseq )
    2 <clumps> [ first2 expand ] map-concat ;

MEMO: cavern ( -- seq )
    input [ parse rocks ] map-concat ;

MEMO: rows ( -- n )
    cavern values supremum 1 + ;

MEMO: bounds ( -- left right )
    cavern keys minmax ;

: cols ( -- n )
    bounds swap - 1 + ;

MEMO: entry ( -- loc )
    0 500 bounds drop - 2array ;

: rock-coords ( -- seq )
    cavern [ bounds drop - ] map-keys assoc-invert ;

: cave-area ( -- seq )
    rows cols CHAR: . <matrix> [ >vector ] map >vector ;

MEMO: cave ( -- seq )
    cave-area CHAR: # rock-coords pick matrix-set-nths ;

MEMO: new-cave ( -- seq )
    cave clone [ clone ] map ;

! Sand first tries to move straight down, then left-down,
! then right-down.
CONSTANT: move-order
    { { 1 0 } { 1 -1 } { 1 1 } }

: sand ( -- loc )
    V{ } ;

: spawn ( -- )
    entry 0 sand copy ;

: clear? ( loc -- ? )
    cave matrix-nth CHAR: . = ;

: reify ( -- )
    CHAR: o sand cave matrix-set-nth spawn ;

: next-move ( -- loc )
    move-order [ sand v+ clear? ] find nip ;

: fall ( -- )
    sand next-move dup [ [ + ] 2map! drop ] [ 2drop reify ] if ;

: oob? ( loc -- ? )
    dup [ neg? ] any? swap cave dimension v>= vany? or ;

: sim ( -- )
    spawn [ next-move sand v+ oob? entry clear? not or ]
    [ fall ] until ;

: part1 ( -- )
    [ sim ] try cave concat [ CHAR: o = ] count . ;

: void ( n -- seq )
    [ CHAR: . ] V{ } replicate-as ;

: deepen ( -- )
    cols void dup clone cave push cave push ;

: extend-sides ( n -- )
    dup 1 entry rot '[ _ + ] change-nth
    cave swap void
    [ [ append ] curry map ]
    [ [ prepend ] curry map ] bi
    [ cave nth 0 swap copy ] each-index ;

: build-floor ( -- )
    cave last [ drop CHAR: # ] map! drop ;

: prep-part-2 ( -- )
    deepen 200 extend-sides build-floor ;

: part2 ( -- )
    prep-part-2 spawn part1 ;

: day14 ( -- )
    part1 part2 ;

MAIN: day14