#lang racket
(require "lib/common.rkt"
         data/queue
         threading)

(struct posn (x y) #:transparent)

(define (explode str)
  (map string (string->list str)))

(define (parse lines)
  (~>> lines
       (map explode)
       (map (λ~>> (map string->number)
                  (apply vector)))
       (apply vector))) 

(define (board-dim board)
  (cons (vector-length (vector-ref board 0))
        (vector-length board)))
(define (board-ref board x y)
  (vector-ref (vector-ref board y) x))

(define (adjacents board x y)
  (match-define (cons xmax ymax) (board-dim board))

  (define x-1 (- x 1))
  (define x+1 (+ x 1))
  (define y-1 (- y 1))
  (define y+1 (+ y 1))

  (for/list ([px (in-list (list x-1 x x+1 x))]
             [py (in-list (list y y-1 y y+1))]
             #:when (not (or (< px 0) (< py 0) (>= px xmax) (>= py ymax))))
    (posn px py)))

(define (board-< board x1 y1 x2 y2)
  (< (board-ref board x1 y1)
     (board-ref board x2 y2)))

(define (lowest-points board)
  (define (lowest? x y)
    (define adjs (adjacents board x y))
    (for/and ([adj (in-list adjs)])
      (board-< board x y (posn-x adj) (posn-y adj))))

  (match-define (cons m n) (board-dim board))
  (for*/list ([i (in-range m)]
              [j (in-range n)]
              #:when (lowest? i j))
    (posn i j)))

(define (day9a board)
  (define heights (map (λ (p) (board-ref board (posn-x p) (posn-y p)))
                       (lowest-points board)))
  (for/sum ([h (in-list heights)])
    (add1 h)))

(define (day9b board)
  (define lowest (lowest-points board))

  (define (bound? p2)
    (not (= (board-ref board (posn-x p2) (posn-y p2)) 9)))

  (define (run-bfs i j)
    (define seen (mutable-set))
    (define parents (make-hash))
    (define Q (make-queue))

    (set-add! seen (posn i j))
    (hash-set! parents (posn i j) (posn i j))
    (enqueue! Q (posn i j))

    (define (bfs/queue)
      (cond [(queue-empty? Q) #f]
            [else
             (define cur (dequeue! Q))
             (for ([adj (in-list (adjacents board (posn-x cur) (posn-y cur)))]
                   #:when (and (not (set-member? seen adj))
                               (bound? adj)))
               (hash-set! parents adj cur)
               (enqueue! Q adj)
               (set-add! seen adj))
             (bfs/queue)]))

    (bfs/queue)
    parents)

  (define sizes (map (λ (p) (hash-count (run-bfs (posn-x p) (posn-y p))))
                     lowest))

  (for/product ([s (in-list (take (sort sizes >) 3))])
    s))

(module+ main
  (call-with-input-file "data/day9.txt"
    (λ (prt)
      (define board (parse (port->lines prt)))
      (answer 9 1 (day9a board))
      (answer 9 2 (day9b board)))))

(module+ test
  (require rackunit)

  (define test-input
    (vector (vector 2 1 9 9 9 4 3 2 1 0)
            (vector 3 9 8 7 8 9 4 9 2 1)
            (vector 9 8 5 6 7 8 9 8 9 2)
            (vector 8 7 6 7 8 9 6 7 8 9)
            (vector 9 8 9 9 9 6 5 6 7 8)))

  (check-equal? (day9a test-input) 15)
  (check-equal? (day9b test-input) 1134))