! Copyright (C) 2022 Your name.
! See http://factorcode.org/license.txt for BSD license.
USING: peg.ebnf multiline strings math math.parser accessors kernel prettyprint
  sequences vectors arrays assocs combinators io.files io.encodings.utf8 math.order ;

IN: aoc2022-7

TUPLE: ntree val children parent ;
C: <ntree> ntree

! Tree operations
: add-assoc-tree ( child parent -- parent )
  [ children>> push ] [ [ second ] dip >>parent parent>> ]  2bi
;

: make-node ( x -- ntree )
  V{ } clone f <ntree> ;

: assoc-node ( x -- assoc-ntree )
  [ second ] [ make-node ] bi 2array ;

: goto-dir ( tree ac -- tree/f )
  swap children>> at ;

: goto-root ( tree -- tree )
  dup parent>> [ parent>> goto-root ] [  ] if ; inline recursive

! parser
SYMBOLS: cd ls dir ;

EBNF: cmd-parser [=[
  filename= [a-zA-Z./]+ => [[ >string ]]
  n= [0-9]+ => [[ dec> ]]
  sizename = n " "~ filename
  dir= "dir" => [[ drop dir ]] " "~ filename
  pairs= ((sizename | dir) "\n"*~)*
  cd= "cd" => [[ drop cd ]] " "~ filename
  ls= "ls" => [[ drop ls ]] "\n"~ pairs
  cmds= ("$ "~ (cd | ls) "\n"*~)*
]=]

: input ( -- ast )
  P" work/aoc2022-7/input.txt" utf8 file-contents cmd-parser ;

! handle instructions

! Make a dummy node with "/" as child so that we can
! cd into "/". Other option cd "/" -> goto-root and no dummy root
: root-node ( -- tree )
  V{ dir "/" } assoc-node 0 make-node add-assoc-tree ;

: ls-op ( seq(n,file) parent -- parent )
  [ assoc-node swap add-assoc-tree ] reduce ;

: cd-op ( tree ac -- tree )
  dup ".." = [ drop parent>> ] [ goto-dir ] if
  ;

: interpret ( tree pair(inst,ac) -- tree )
  first2 swap
  {
    { cd [ cd-op ] }
    { ls [ swap ls-op ] }
  } case ;

: build-file-tree ( -- tree )
  input
  root-node
  [ interpret ] reduce
  goto-root ;

: file? ( tree -- ? )
  children>> empty? ;

: dir? ( tree -- ? )
  file? not ;

: annotate-size-directories ( tree -- n )
  [ val>> ]
  [
    dup dir?
    [ children>> [ second annotate-size-directories ] map-sum ]
    [ val>> first ]
    if
  ]
  bi
  [ 0 rot set-nth ] keep
;

: part1-answer ( tree -- n )
  [ children>> [ second ] map
    [ dir? ] filter
    [ part1-answer ] map-sum
  ]
  [ val>> first dup 100000 < [ ] [ drop 0 ] if ]
  bi
  +
;

: build-annotated-tree ( -- tree n )
  build-file-tree
  "/" goto-dir
  dup annotate-size-directories ;

: part1 ( -- tree )
  build-annotated-tree
  drop part1-answer ;

! best candidate for removal
: best-candidate ( tree n -- n )
  dup
  '[ children>> [ second ] map
    [ dir? ] filter
    [ val>> first _ > ] filter
    [ _ best-candidate ] map
  ]
  [
    val>> first
  ]
  bi
  [ min ] reduce
; inline recursive

: part2 ( -- n )
  build-annotated-tree
  70000000 swap -
  30000000 swap -
  best-candidate ;

: day7 ( -- n n ) part1 part2 ;