USING: assocs io.encodings.ascii io.files kernel literals math
math.matrices math.parser math.ranges prettyprint sequences
sequences.extras sets splitting ;
IN: aoc.2021.05

<<

CONSTANT: vents $[                        ! {
    "input.txt" ascii file-lines [        !     { { 2 1 } { 2 8 } }
        "->" without " " split harvest    !     { { 6 4 } { 2 0 } }
        [ "," split [ dec> ] map ] map    !     ...
    ] map                                 ! }
]

: straight? ( pair pair -- ? ) [ = ] 2count 1 = ;
: straight ( seq -- newseq ) [ straight? ] assoc-filter ;

>>

CONSTANT: straight-vents $[ vents straight ]

! A matrix of zeros large enough to hold every vent
: <ocean-floor> ( -- matrix )
    vents concat unzip [ supremum 1 + ] bi@ <zero-matrix> ;

! This word should really already exist lol
: matrix-change-nth ( ..a pair matrix quot: ( ..a elt -- ..b newelt ) -- ..b )
    '[ matrix-nth @ ] [ matrix-set-nth ] 2bi ; inline

! Increment a matrix at a coordinate
: matrix-inc ( pair matrix -- ) [ 1 + ] matrix-change-nth ;

! Turn { { 0 0 } { 0 2 } } into { { 0 0 } { 0 1 } { 0 2 } } for instance
: expand-straight-line ( seq -- newseq )
    unzip [ first2 [a,b] ] bi@ cartesian-product concat ;

! Turn { { 6 4 } { 2 0 } } into { { 6 4 } { 5 3 } { 4 2 } { 3 1 } { 2 0 } } for instance
: expand-diagonal-line ( seq -- newseq )
    flip first2 [ first2 [a,b] ] bi@ zip ;

: expand-line ( seq -- newseq )
    dup first2 straight?
    [ expand-straight-line ] [ expand-diagonal-line ] if ;

: expand-lines ( seq -- newseq ) [ expand-line ] map-concat ;

: <mapped-ocean-floor> ( seq -- matrix )
    <ocean-floor> tuck [ matrix-inc ] curry each ;

: .danger-count ( seq -- )
    expand-lines <mapped-ocean-floor> concat [ 1 > ] count . ;

: part1 ( -- ) straight-vents .danger-count ;
: part2 ( -- ) vents .danger-count ;
: main ( -- ) part1 part2 ;

MAIN: main