! Copyright (C) 2022 Your name.
! See http://factorcode.org/license.txt for BSD license.
USING: aoc math.polynomials prettyprint hashtables math kernel assocs peg.ebnf sequences multiline math.parser strings arrays ;
IN: aoc.2022.21

EBNF: parse [=[
  n = "-"? [0-9]+ => [[ concat dec> 1array ]]
  var = [a-zA-Z]+ => [[ >string ]]
  op = ("+" => [[ [ p+ ] ]] |
        "-" => [[ [ p- ] ]] |
        "*" => [[ [ p* ] ]] |
        "/" => [[ [ p/mod drop ] ]])
  operation = var " "~ op " "~ var => [[ first3 swap 3array ]]
  kv = var ": "~ ( n | operation )
]=]

: test-input ( -- seq )
  P" work/aoc/2022/21/example.txt" load-file
  [ parse ] map
  >hashtable ;

: input ( -- seq )
  P" work/aoc/2022/21/input.txt" load-file
  [ parse ] map
  >hashtable ;

:: lookup-var ( var hm -- n )
  var hm
  [ [ dup length 2 > ! brittle but it works
      [ first3 [ [ hm lookup-var ] bi@ ] dip call( x x -- x ) ]
      when
    ] ?change-at ]
  [ at ] 2bi
;

: test-part1 ( -- n )
  "root" test-input lookup-var
;

: part1 ( -- n )
  "root" input lookup-var
;

: input2 ( -- seq )
  input
  [ "humn" { 0 1 } -rot swap set-at ] keep ;

: part2 ( -- n )
  ! The two in root
  "pnhm" input2 lookup-var
  "zvcm" input2 lookup-var
  p-
  first2 -1 * /
;

: day21 ( -- )
  part1 . part2 .  ;