USING: combinators continuations fry io.encodings.ascii io.files
kernel math math.parser prettyprint sequences splitting ;
IN: aoc.2019.02

CONSTANT: instruction-offset 4
CONSTANT: halt-opcode 99
ERROR: unknown-opcode n ;

: opcode>quot ( n -- quot )
    {
        { 1 [ [ + ] ] }
        { 2 [ [ * ] ] }
        [ unknown-opcode ]
    } case ;

: intcode>seq ( str -- seq ) "," split [ string>number ] map ;

: meta-nth ( n seq -- elt ) [ nth ] keep nth ;
: under+ ( m obj n -- m+n obj ) [ + ] curry dip ;
: at-address ( m seq n -- o ) under+ meta-nth ;

: execute-instruction ( n seq -- seq' )
    [ instruction-offset * ] dip tuck {
        [ 1 at-address ]
        [ 2 at-address ]
        [ nth opcode>quot call( x x -- x ) ]        
        [ 3 under+ nth ]
        [ nip set-nth ]
    } 2cleave ;

: get-opcode ( m seq -- n ) [ instruction-offset * ] dip nth ;

: (run-intcode) ( seq -- seq' )
    0 swap [ 2dup get-opcode halt-opcode = ]
    [ dupd execute-instruction 1 under+ ] until nip ;

: set-nth* ( seq m n -- seq' ) rot [ set-nth ] keep ;
: set-noun ( seq n -- seq' ) 1 set-nth* ;
: set-verb ( seq n -- seq' ) 2 set-nth* ;

: run-intcode ( str noun verb -- seq )
    [ intcode>seq ] [ set-noun ] [ set-verb ] tri*
    (run-intcode) ;

: parse-input ( -- seq )
    "resource:work/aoc/2019/02/input.txt" ascii file-contents ;

: part1 ( -- ) parse-input 12 2 run-intcode first . ;

: part2 ( -- )
    [
        100 <iota> dup parse-input
        [ -rot [ run-intcode ] 2keep pick first 19690720 = 
        [ [ 100 * ] dip + . drop return ] [ 3drop ] if ]
        curry cartesian-each
    ] with-return ;