#lang curly-fn racket

(require "scripts/aoc.rkt")

;; solution for day 7

(define (part1 G wanted)
  (define-vertex-property G has-wanted #:init #f)
  (do-dfs G
    ;; only visit when the to vertex is not already marked with has-wanted
    #:visit?: (if (and $from (has-wanted $to))
                (begin (has-wanted-set! $from #t) #f)
                #t)
    ;; propagate has-wanted if the child has it
    #:epilogue: (when (and $from (or (has-wanted $to) (string=? wanted $to)))
                  (has-wanted-set! $from #t)))
  ;; count the results
  (for/sum ([(_ v) (in-hash (has-wanted->hash))] #:when v) 1))

(define (part2 G wanted)
  ;; initially, each vertex just counts the weight of itself, which is 1
  (define-vertex-property G all-weight #:init 1)

  ;; accumulates the result in the vertex property tracking all weights
  ;; adds the total weight of $to multiplied by the edge weight of $from -> $to to the total weight
  ;; of $from
  ;; this will be called for each neighbor of $from, and $from's total weight starts at 1
  (define (increment $from $to)
    (when $from
      (all-weight-set! $from
                       (+ (all-weight $from)
                          (* (edge-weight G $from $to) (all-weight $to))))))

  (do-dfs G
    #:order (lambda (_) (list wanted))
    ;; always check already-visited nodes
    #:process-unvisited?: #t
    ;; add already-visited nodes to the total
    #:process-unvisited: (increment $from $to)
    ;; add nodes we finished visiting to the total
    #:epilogue: (increment $from $to))
  ;; subtract 1 because we counted the shiny gold bag itself, when we only need the weight of
  ;; what's inside it
  (sub1 (all-weight wanted)))

(module+ main
  (define input-lines (file->lines "inputs/7"))
  (define edges
    (apply append
           (for/list ([line (in-list input-lines)])
             (match line
               [(pregexp #px"^([a-z ]+) bags contain ([^\\.]+).$" (list _ src-color contains))
                 (if (string=? "no other bags" contains)
                   '()
                   (for/list ([item (in-list (string-split contains ", "))])
                     (match item
                       [(pregexp #px"^([0-9]+) ([a-z ]+) bag[s]?"
                                 (list _ (app string->number number) dst-color))
                        (list number src-color dst-color)]
                       [x (error "not shonks" x)])))]
               [x (error "not shonks" x)]))))
  (define input (weighted-graph/directed edges))
  (define wanted "shiny gold")

  ;; part 1
  (answer 7 1 (part1 input wanted))

  ;; part 2
  (answer 7 2 (part2 input wanted))

  (displayln "meow"))