USING: combinators combinators.short-circuit fry
io.encodings.ascii io.files kernel literals locals math
math.matrices math.vectors prettyprint sequences ;
IN: aoc.2020.11

CONSTANT: input $[ "input.txt" ascii file-lines ]

CONSTANT: adjacent {
    { -1 -1 } { 0 -1 } { 1 -1 }
    { -1  0 }          { 1  0 }
    { -1  1 } { 0  1 } { 1  1 }
}

: matrix-map-loc ( ... matrix quot: ( ... elt loc -- ... newelt ) -- ... newmatrix )
    [ dup dimension <coordinate-matrix> ] dip [ 2map ] curry 2map ; inline

: in-bounds? ( loc matrix -- ? )
    dimension [ { [ drop neg? not ] [ < ] } 2&& ] 2all? ;

: adj-locs ( loc -- seq ) adjacent [ v+ ] with map ;

: ib-locs ( loc matrix -- seq )
    [ adj-locs ] dip [ in-bounds? ] curry filter ;

: neighbors ( loc matrix -- seq )
    [ ib-locs ] [ matrix-nths ] bi ;

: occupied ( seq -- n ) [ CHAR: # = ] count ;

:: flip ( ch count n -- newch )
    {
        { [ count zero? ch CHAR: L = and ] [ CHAR: # ] }
        { [ count n >= ch CHAR: # = and ] [ CHAR: L ] }
        [ ch ]
    } cond ;

: step ( matrix -- newmatrix )
    dup [ neighbors occupied 4 flip ] curry matrix-map-loc ;

: simulate ( ... matrix quot: ( ... matrix -- ... newmatrix ) -- ... newmatrix )
    '[ dup @ tuck = not ] loop ; inline

: part1 ( -- ) input [ step ] simulate concat occupied . ;

: collided? ( loc matrix -- ? )
    { [ in-bounds? not ] [ matrix-nth CHAR: . = not ] } 2|| ;

: ray ( dir loc matrix -- collision-loc )
    [ 2dup collided? ]
    [ [ [ v+ ] keepd swap ] dip ] do until drop nip ;

: rays ( loc matrix -- seq )
    [ adjacent ] 2dip [ ray ] 2curry map ;

: ib-rays ( loc matrix -- seq )
    [ rays ] keep [ in-bounds? ] curry filter ;

: neighbors2 ( loc matrix -- seq )
    [ ib-rays ] [ matrix-nths ] bi ;

: step2 ( matrix -- newmatrix )
    dup [ neighbors2 occupied 5 flip ] curry matrix-map-loc ;

: part2 ( -- ) input [ step2 ] simulate concat occupied . ;