(import (srfi :43)) ; vector extensions
(load "../common.scm")

; constructor of low point objects
; from their position and altitude
(define (make-low-point x y altitude)
  (list x y altitude))
(define (low-point-x lp) (car lp))
(define (low-point-y lp) (cadr lp))
(define (low-point-altitude lp) (caddr lp))

; on a given heightmap, gives the list of points that are
; strictly lower than their left/right/up/down neighbors
(define (find-low-points heightmap)
  (let [(low-points '())]
    (vector-for-each
      (lambda (y line)
	(vector-for-each
	  (lambda (x altitude)
	    (let [(left (matrix-get heightmap (- x 1) y))
		  (right (matrix-get heightmap (+ x 1) y))
		  (up (matrix-get heightmap x (- y 1)))
		  (down (matrix-get heightmap x (+ y 1)))]
	      (when (and (< altitude left)
			 (< altitude right)
			 (< altitude up)
			 (< altitude down))
		(set! low-points (cons (make-low-point x y altitude) low-points)))))
	  line))
      (matrix-data heightmap))
    low-points))

; compute risk level from low point list
(define (risk-level low-points)
  (apply + (map 1+ (map low-point-altitude low-points))))

; starting with from a given low-point
; call closure with and all points of the corresponding basin
(define (visit-basin heightmap low-point closure)
  (let* [(width (matrix-width heightmap))
	 (height (matrix-height heightmap))
	 (width-1 (- width 1))
	 (height-1 (- height 1))
	 (visited (make-vector (* width height) 0))]
    (define (offset x y)
      (+ (* y width) x))
    (let flood-fill [(x (low-point-x low-point))
		     (y (low-point-y low-point))]
      (let [(ofs (offset x y))]
	(when (= (vector-ref visited ofs) 0)
	  (vector-set! visited ofs 1)
	  (when (< (matrix-get heightmap x y) 9)
	    (closure x y)
	    (when (> x 0)        (flood-fill (- x 1) y))
	    (when (< x width-1)  (flood-fill (+ x 1) y))
	    (when (> y 0)        (flood-fill x (- y 1)))
	    (when (< y height-1) (flood-fill x (+ y 1)))))))))

; compute size of basin starting at given low point
(define (basin-size heightmap low-point)
  (let [(size 0)]
    (visit-basin heightmap low-point
		 (lambda (x y)
		   (set! size (+ 1 size))))
    size))

; returns list of basin sizes corresponding to given low points
(define (compute-basin-sizes heightmap low-points)
  (map
    (lambda (lp)
      (basin-size heightmap lp))
    low-points))

; write a PGM file for given heightmap and low points
; basin limits are grey
; basins are red-ish (brighter the higher the point)
; low points are bright green
;
; for debugging purposes
(define (dump-heightmap heightmap low-points filename)
  (when (file-exists? filename)
    (delete-file filename))
  (call-with-output-file
    filename
    (lambda (file)
      ; write header to file: "P3" followed by width, height and max value of channel
      (fprintf file "P3~%~a ~a~%255~%"
	       (matrix-width heightmap)
	       (matrix-height heightmap))

      (let* [(width (matrix-width heightmap))
	     (width*3 (* width 3))
	     (height (matrix-height heightmap))
	     (data (make-vector (* width*3 width)))]

	; from (x,y) returns of offset of pixel's red channel in data vector
	; green and blue follow as ofs+1 and ofs+2
	(define (offset x y)
	  (+ (* width*3 y)
	     (* 3 x)))

	; show heightmap as gray levels
	(let y-loop [(y 0)]
	  (let x-loop [(x 0)]
	    (let [(ofs (offset x y))
		  (gray-level (* (matrix-get heightmap x y) 20))]
	      ; red channel
	      (vector-set! data ofs gray-level)
	      ; green channel
	      (vector-set! data (+ 1 ofs) gray-level)
	      ; blue channel
	      (vector-set! data (+ 2 ofs) gray-level))
	    (when (< x (- width 1))
	      (x-loop (+ x 1))))
	  (when (< y (- height 1))
	    (y-loop (+ y 1))))

	(for-each
	  (lambda (lp)
	    ; show basins as redish areas
	    (visit-basin heightmap lp
			 (lambda (x y)
			   (let [(ofs (offset x y))]
			     (vector-set! data ofs 255))))
	    ; show low points as bright green
	    (let [(ofs (offset (low-point-x lp) (low-point-y lp)))]
	      (vector-set! data ofs 0)
	      (vector-set! data (+ 1 ofs) 255)
	      (vector-set! data (+ 2 ofs) 0)))
	  low-points)

	; write pixels to file
	(vector-for-each
	  (lambda (index bits)
	    (fprintf file "~a " bits)
	    (when (= (mod index width*3) 0)
	      fprintf file "~%"))
	  data)))))

(call-with-input-file
  "input"
  (lambda (file)
    (let* [(heightmap (load-matrix file))
	   (low-points (find-low-points heightmap))]
      (printf "part 1:~% Risk level: ~a~%"
	      (risk-level low-points))
      (let [(basin-sizes (sort > (compute-basin-sizes heightmap low-points)))]
	(dump-heightmap heightmap low-points "heightmap.pgm")
	(printf "part 2:~% Three largest basin sizes multiplied: ~a~%"
		(* (car basin-sizes)
		   (cadr basin-sizes)
		   (caddr basin-sizes)))))))