(load "../common.scm")

; returns a vector of the same length of input string
; where all values are 0 except for when the corresponding
; entry in the string is '#' which will have a 1 instead.
(define (string->bin-vector str)
  (vector-unfold
    (lambda (i)
      (if (char=? (string-ref str i) #\#) 1 0))
    (string-length str)))

; returns a binary matrix loaded from the file
; all the lines are expected to be the same lengths
(define (file->matrix file)
  (let loop [(data '())]
    (if (eof-object? (peek-char file))
      (matrix-from-data (reverse-list->vector data))
      (loop (cons (string->bin-vector (get-line file)) data)))))

; palette enhancing algorihtm: 
; for each pixel in input image
; gather all pixels around it to form a 3x3 matrix
; turn that matrix into a binary number, row by row
; look up new pixel value in input palette
(define (enhance image palette steps)
  (let* [(new-width (+ (matrix-width image) 4))
	 (new-height (+ (matrix-height image) 4))
	 (new-image (filled-matrix new-width new-height 0))
	 (top-left (matrix-get image 0 0 0))
	 ; if palette 0 is 1, outside becomes all 1 each odd iteration
	 (outside (if (= (vector-ref palette 0) 1) (mod steps 2) 0))]
    (update-matrix!
      new-image
      (lambda (x y v)
	(vector-ref palette (+ (* (matrix-get image (- x 3) (- y 3) outside) 256)
			       (* (matrix-get image (- x 2) (- y 3) outside) 128)
			       (* (matrix-get image (- x 1) (- y 3) outside) 64)
			       (* (matrix-get image (- x 3) (- y 2) outside) 32)
			       (* (matrix-get image (- x 2) (- y 2) outside) 16)
			       (* (matrix-get image (- x 1) (- y 2) outside) 8)
			       (* (matrix-get image (- x 3) (- y 1) outside) 4)
			       (* (matrix-get image (- x 2) (- y 1) outside) 2)
			       (* (matrix-get image (- x 1) (- y 1) outside) 1)))))
    (if (<= steps 1)
      new-image
      (enhance new-image palette (- steps 1)))))

; for debugging purposes
; output a PGM 1bpp image file from a given binary matrix
(define (matrix->image-file matrix filename)
  (when (file-exists? filename)
    (delete-file filename))
  (call-with-output-file
    filename
    (lambda (file)
      (fprintf file "P1~%~a ~a~%" (matrix-width matrix) (matrix-height matrix))
      (vector-for-each
	(lambda (y line)
	  (vector-for-each 
	    (lambda (x pixel) (fprintf file "~a " (- 1 pixel))) line)
	  (fprintf file "~%"))
	(matrix-data matrix)))))

(call-with-input-file
  "input"
  (lambda (file)
    (let* [(enhancer (string->bin-vector (get-line file)))
	  (image (begin
		   (get-line file) ; skip empty line
		   (file->matrix file)))]
      (let [(enhanced-2 (enhance image enhancer 2))]
	(printf "part 1:~% Lit pixels after enhancing twice: ~a~%"
		(matrix-sum enhanced-2))
	(let [(enhanced-50 (enhance enhanced-2 enhancer 48))]
	  (printf "part 2:~% Lit pixels after enhancing fifty times: ~a~%"
		  (matrix-sum enhanced-50)))))))