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

; convert wire's letter to corresponding bit of 7-segment display
(define (signal->bit c)
  (case c
    [#\a #b0000001]
    [#\b #b0000010]
    [#\c #b0000100]
    [#\d #b0001000]
    [#\e #b0010000]
    [#\f #b0100000]
    [#\g #b1000000]))

; read a set of wires and return corresponding bits
(define (parse-wires file)
  (let loop [(n 0)]
    (let [(c (get-char file))]
      (if (and (not (eof-object? c))
	       (char-alphabetic? c))
	[loop (bitwise-ior (signal->bit c) n)]
	n))))

; reads a set of 10 signals and 4 digits
; returned as a two lists of bitfields
(define (parse-line file)
  (let loop [(signals '()) (digits '()) (n 0)]
    (cond
      [(< n 10)
       (loop (cons (parse-wires file) signals) digits (+ n 1))]
      [(= n 10)
       (get-char file) ; consume |
       (get-char file) ; consume space
       (loop signals digits (+ n 1))]
      [(< n 15)
       (loop signals (cons (parse-wires file) digits) (+ n 1))]
      [else
	(list (reverse signals) (reverse digits))])))

; read all the sets of signals/digits
(define (parse-file file)
  (let loop [(lines '())]
    (if (eof-object? (peek-char file))
      lines
      (loop (cons (parse-line file) lines)))))

; count digits that are easily recognized as 1, 4, 7 or 8
; since they are the only combinations of 2, 4, 3 and 7 wires
; it's as easy as counting the number of bits set in the bitfield
(define (count-easy-digits lines)
  (let [(count 0)]
    (for-each
      (lambda (line)
	(for-each
	  (lambda (digit)
	    (let [(bit-count (bitwise-bit-count digit))]
	      (when (or
		      (= bit-count 2)  ; that's a 1
		      (= bit-count 4)  ; that's a 4
		      (= bit-count 3)  ; that's a 7
		      (= bit-count 7)) ; that's an 8
		(set! count (+ count 1)))))
	  (cadr line)))
      lines)
    count))

; returns a vector of 10 encodings for numbers 0 to 9 from the 10 signals received
(define (build-decoder signals)
  (let [(decoder (make-vector 10))
	(encode-2-3-5 '())
	(encode-0-6-9 '())]
    (for-each 
      (lambda (signal)
	(case (bitwise-bit-count signal)
	  (2 (vector-set! decoder 1 signal))
	  (3 (vector-set! decoder 7 signal))
	  (4 (vector-set! decoder 4 signal))
	  (5 (set! encode-2-3-5 (cons signal encode-2-3-5)))
	  (6 (set! encode-0-6-9 (cons signal encode-0-6-9)))
	  (7 (vector-set! decoder 8 signal))))
      signals)
    ; between the 3 possible encodings of 3,
    ; only the one with both the bits of 1 set can be correct
    (let [(encode-1 (vector-ref decoder 1))]
      (vector-set! decoder 3
		   (fold-left
		     (lambda (encoding signal)
		       (if (= (bitwise-and signal encode-1) encode-1)
			 signal
			 encoding))
		     0 encode-2-3-5)))
    ; between the 3 possible encodings of 9,
    ; only the one with all the bits of 4 set can be correct
    (let [(encode-4 (vector-ref decoder 4))]
      (vector-set! decoder 9
		   (fold-left
		     (lambda (encoding signal)
		       (if (= (bitwise-and signal encode-4) encode-4)
			 signal
			 encoding))
		     0 encode-0-6-9)))
    ; (xor (encode 4) (encode 1) yields the bits corresponding to
    ; the top left and middle segments
    (let [(tl+m (bitwise-xor (vector-ref decoder 1)
			     (vector-ref decoder 4)))]
      ; between the 3 possible encodings of 5,
      ; only the one with both of those bits set can be correct
      (vector-set! decoder 5
		   (fold-left
		     (lambda (encoding signal)
		       (if (= (bitwise-and signal tl+m) tl+m)
			 signal
			 encoding))
		     0 encode-2-3-5))
      ; between the 3 possible encodings of 6,
      ; only the one that is not nine and has top-left and mid segment set can be correct
      (let [(encode-9 (vector-ref decoder 9))]
	(vector-set! decoder 6
		     (fold-left
		       (lambda (encoding signal)
			 (if (and (= (bitwise-and signal tl+m) tl+m)
				  (not (= signal encode-9)))
			   signal
			   encoding))
		       0 encode-0-6-9))))
    ; encoding for two is the last possible encoding
    (let [(encode-3 (vector-ref decoder 3))
	  (encode-5 (vector-ref decoder 5))]
      (vector-set! decoder 2
		   (fold-left
		     (lambda (encoding signal)
		       (if (not (or (= signal encode-3)
				    (= signal encode-5)))
			 signal
			 encoding))
		     0 encode-2-3-5)))
    ; encoding for zero is the last possible encoding
    (let [(encode-6 (vector-ref decoder 6))
	  (encode-9 (vector-ref decoder 9))]
      (vector-set! decoder 0
		   (fold-left
		     (lambda (encoding signal)
		       (if (not (or (= signal encode-6)
				    (= signal encode-9)))
			 signal
			 encoding))
		     0 encode-0-6-9)))
    decoder)) 

; with given decoder vector, decode input digit
(define (decode decoder input)
  (vector-index (lambda (encoding) (= encoding input))
		decoder))

; returns the number read on the 4x 7-segment display
; by decoding the encoding of digits with the signals
; and applying the decoded bitfields to the digits received
(define (decode-digits line)
  (let [(signals (car line))
	(digits (cadr line))]
    (let [(decoder (build-decoder signals))]
      (fold-left
	(lambda (result digit)
	  (+ 
	    (* 10 result)
	    (decode decoder digit)))
	0
	digits))))

(call-with-input-file
  "input"
  (lambda (file)
    (let [(lines (parse-file file))]
      (printf "part 1:~% Number of easy digits (1, 4, 7 or 8): ~a~%"
	      (count-easy-digits lines))
      (printf "part 2:~% Sum of all decoded digits: ~a~%"
	      (apply + (map decode-digits lines))))))