USING: assocs combinators.short-circuit
compiler.tree.propagation.call-effect eval io.encodings.ascii
io.files kernel listener literals make math prettyprint qw
sequences sets splitting.extras ;
IN: aoc.2020.18

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

CONSTANT: precedence
{
    { "(" 0 }
    { ")" 0 }
    { "+" 2 }
    { "*" 1 }
}

: tokenize ( str -- tokens ) "( )" split* { " " } without ;

: %s ( seq -- ) % " " % ;

: op? ( token -- ? ) qw{ ( ) * + } member? ;

: (valid-stack?) ( stack -- ? )
    { [ last "(" = ] [ last2 [ precedence at ] bi@ < ] } 1|| ;

: valid-stack? ( stack -- ? )
    dup length 2 < [ drop t ] [ (valid-stack?) ] if ;

: slurp ( stack -- stack' )
    dup pop drop [ dup pop dup "(" = ] [ %s ] until drop ;

DEFER: ?op%
: collapse ( stack -- stack' )
    dup [ pop ] [ pop ] bi %s suffix! ?op% ;

: pop-op ( stack -- stack' )
    dup last ")" = [ slurp ] [ collapse ] if ;

: ?op% ( stack -- stack/stack' )
    dup valid-stack? [ pop-op ] unless ;

: handle-token ( stack token -- stack )
    dup op? [ suffix! ?op% ] [ %s ] if ;

: infix>postfix ( str -- newstr )
    V{ } clone swap tokenize [ [ handle-token ] each ] "" make
    swap reverse " " join append ;

: main ( -- )    ! main works for both part 1 and part 2; just change +'s precedence level to 1 for part 1.
    input
    [ infix>postfix [ ( -- n ) (eval) ] with-interactive-vocabs ]
    map-sum . ;