(library (lib)
  (export solve-part1 solve-part2)
  (import (chezscheme))

  ; Split list at first delimiter into `prefix' and `suffix'
  ; Return value is a pair like `((p r e f i x) s u f f i x)'
  (define (list-split-left delimiter? xs)
    (let loop ((prefix '())
               (suffix xs))
      (if (null? suffix)
        (cons prefix suffix)
        (let ((x (car suffix)))
          (if (delimiter? x)
              ; Found first delimiter, so return immediately
              (cons prefix (cdr suffix))
              ; `x' isn't a delimiter, so append it to `prefix'
              (loop (append prefix (list x)) (cdr suffix)))))))

  ; Split list at given delimiter into list of sublists
  (define (list-split delimiter? xs)
    (let loop ((pieces '())
               (rest xs))
      (if (null? rest)
          ; Fix order and remove all empty sublists from output
          ; (which are caused by consecutive delimiters in `xs')
          (reverse (remp null? pieces))
          ; Extract next piece from `rest' and prepend it to `pieces'
          (let ((next (list-split-left delimiter? rest)))
            (loop (cons (car next) pieces) (cdr next))))))

  ; Split string at whitespace into list of words
  (define (string-split str)
    (map list->string (list-split char-whitespace? (string->list str))))

  ; Remove the first character from `str'
  (define (string-strip-first str)
    (substring str 1 (string-length str)))

  ; Remove the last character from `str'
  (define (string-strip-last str)
    (substring str 0 (- (string-length str) 1)))

  ; Strip the parenthesis on each side and parse the number
  (define (parse-weight str)
    (string->number (string-strip-first (string-strip-last str))))

  ; Parse the current node's optional children into a list of names
  (define (parse-children words)
    (let loop ((input words)
               (output '()))
      (cond
        ; If there are no children, or we've exhausted the input list
        ((= (length input) 0)
          output)
        ; If this is is the last child, we can read its name directly
        ((= (length input) 1)
          (loop (cdr input) (cons (car input) output)))
        ; The input contains multiple words, possibly including "->"
        (else
          (if (string=? "->" (car input))
            ; Ignore the first "->"
            (loop (cdr input) output)
            ; This isn't the last child, so must strip comma from name
            (loop
              (cdr input)
              (cons (string-strip-last (car input)) output)))))))

  ; Record type representing a single tree node. Implicitly defines
  ; `make-tree-node', `tree-node-children', and `tree-node-weight'.
  (define-record-type tree-node
    (fields
      (immutable weight)
      (immutable children)))

  ; Parse each line of the input into a `(name . tree-node)' pair
  (define (parse lines)
    (let loop ((input lines)
               (output '()))
      (if (null? input)
          output
          (let ((words (string-split (car input))))
            (let ((name (car words))
                  (weight (parse-weight (cadr words)))
                  (children (parse-children (cddr words))))
              (loop
                (cdr input)
                (cons
                  (cons name (make-tree-node weight children))
                  output)))))))

  ; Given a node name, find its parent according to `tree'
  (define (find-parent name tree)
    (let loop ((candidates tree))
      (if (null? candidates)
          #f
          (if (find
                (lambda (child-name) (string=? name child-name))
                (tree-node-children (cdar candidates)))
              (caar candidates)
              (loop (cdr candidates))))))

  ; Find the root node name of the tree stored in `tree'
  (define (find-root tree)
    (let loop ((name (caar tree)))
      (let ((parent-name (find-parent name tree)))
        (if (not parent-name)
            name
            (loop parent-name)))))

  ; Return the `(name . weight)' pair for the `name'
  ; that has the lowest `weight' in `child-weights'.
  (define (find-lightest child-weights)
    (define target-weight (apply min (map cdr child-weights)))
    (find (lambda (cw) (= (cdr cw) target-weight)) child-weights))

  ; Return the `(name . weight)' pair for the `name'
  ; that has the highest `weight' in `child-weights'.
  (define (find-heaviest child-weights)
    (define target-weight (apply max (map cdr child-weights)))
    (find (lambda (cw) (= (cdr cw) target-weight)) child-weights))

  ; Count how many times `needle' occurs in `haystack'
  (define (count needle haystack)
    (length (filter (lambda (h) (equal? h needle)) haystack)))

  ; For `child-weights', i.e. a list of `(name . total-weight)' pairs
  ; for several nodes with a common parent, we find the "bad" child
  ; whose `total-weight' differs from the others. If there are only
  ; two children, this can't be decided; then `known-error' tells us
  ; whether the "bad" child is lighter or heavier than a "good" one.
  ; `known-error' is the bad-good weight difference, calculated by
  ; the caller at a time when there were more than two children.
  (define (odd-one-out child-weights known-error)
    (let ((lightest (find-lightest child-weights))
          (heaviest (find-heaviest child-weights)))
      (if (equal? lightest heaviest)
          ; If all children have the same weight, the parent must be bad
          #f
          (if known-error
              ; If we already know that the bad node is too light,
              ; always pick the lightest child, and vice versa.
              (if (< known-error 0)
                  lightest
                  heaviest)
              ; If we don't know yet whether the bad node is too light or heavy,
              ; pray that `(length child-weights)' > 2 and pick the rarest weight.
              (if (< (count (cdr lightest) (map cdr child-weights))
                     (count (cdr heaviest) (map cdr child-weights)))
                  lightest
                  heaviest)))))

  ; Add up the weight of `name' and all of its descendants
  (define (subtree-weight name tree)
    (define node (cdr (assoc name tree)))
    (let loop ((children (tree-node-children node))
               (total    (tree-node-weight   node)))
      (if (null? children)
          total
          (loop
            (cdr children)
            (+ total (subtree-weight (car children) tree))))))

  ; Given `name', create a list of `(child-name . total-weight)' pairs
  (define (filter-child-weights name tree)
    (map
      (lambda (child-name)
        (cons child-name (subtree-weight child-name tree)))
      (tree-node-children (cdr (assoc name tree)))))

  (define (solve-part1 lines)
    (define tree (parse lines))
    (find-root tree))

  (define (solve-part2 lines)
    (define tree (parse lines))
    (define root (find-root tree))
    ; Traverse the tree until we find the "bad" node with the wrong weight
    (let loop ((subtree-root root)
               (weight-error #f))
      ; Figure out which branch of `subtree-root' contains the bad node
      (let* ((child-weights (filter-child-weights subtree-root tree))
             (wrong-subtree (odd-one-out child-weights weight-error))
             (right-subtree (car (remove wrong-subtree child-weights))))
        (if (not wrong-subtree)
            ; If we can't find the bad branch, `subtree-root' itself is bad
            (- (tree-node-weight (cdr (assoc subtree-root tree))) weight-error)
            ; Recurse into the bad branch (N.B. `weight-error' doesn't change)
            (loop
              (car wrong-subtree)
              (- (cdr wrong-subtree) (cdr right-subtree)))))))

)