(import (srfi :43)) ; vector extensions

; hash function for a list of two characters
(define (char-pair-hash cp)
  (+ (* (char->integer (car cp)) 79)
     (char->integer (cadr cp))))

; parses lines of "AB -> C" with A, B and C being characters
; returns hashtable where keys are ABs and values are Cs
(define (read-table file)
  (let loop [(table (make-hashtable char-pair-hash equal? 32))]
    (let [(line (get-line file))]
      (if (eof-object? line)
	table
	(let [(key (list (string-ref line 0) (string-ref line 1)))
	      (value (string-ref line 6))]
	  (hashtable-set! table key value)
	  (loop table))))))

; generate new polymer from inserting elements in between each pair according to given table
(define (insert polymer table)
  (let [(end (- (string-length polymer) 1))]
    (let loop [(i 0) (output "")]
      (if (< i end)
	(loop (+ i 1)
	      (let* [(pair (list (string-ref polymer i) (string-ref polymer (+ i 1))))
		     (element (hashtable-ref table pair #f))]
		(if element
		  (string-append output (string (car pair)) (string element))
		  (string-append output (string (car pair))))))
	(string-append output (substring polymer end (+ end 1)))))))

; parameter counts is a vector where each entry is a list (C N)
; where C is a key and N is a number.
; returns a vector where entry of ele has increased by count.
(define (add-ele-count counts ele count)
  (let [(index (vector-index (lambda (p) (equal? (car p) ele)) counts))]
    (if index
      (begin
	(vector-set! counts index (list ele (+ count (cadr (vector-ref counts index)))))
	counts)
      (vector-append counts (vector (list ele count))))))

; count amount of each character in input string.
; returns vector of lists where car is a character and cadr is the amount of it in input
; sorted from most common to least common
(define (count-chars str)
  (let [(count (vector))]
    (string-for-each (lambda (c) (set! count (add-ele-count count c 1))) str)
    count))

; print result to console
(define (display-subtracted-most-least-elements elements)
  (let* [(elements (vector-sort (lambda (l r) (> (cadr l) (cadr r))) elements))
	 (most-common (vector-ref elements 0))
	 (least-common (vector-ref elements (- (vector-length elements) 1)))]
    (printf "Most common ~a (~a) - least common ~a (~a) = ~a~%"
	    (car most-common) (cadr most-common)
	    (car least-common) (cadr least-common)
	    (- (cadr most-common) (cadr least-common)))))

; returns a vector counting the number of each consecutive characters in str
(define (count-pairs str)
  (let [(len-1 (- (string-length str) 1))]
    (let loop [(n 0) (pairs '#())]
      (if (< n len-1)
	(loop (+ n 1)
	      (let [(pair (list (string-ref str n) (string-ref str (+ n 1))))]
		(add-ele-count pairs pair 1)))
	pairs))))

; count the number of elements where given pairs-count is a vector
; with amounts of pairs of elements
; #(((a b) N) ((c d) N') ...)
(define (count-elements pairs-count)
  (vector-fold
    (lambda (index elements-count pair-count)
      (let [(pair (car pair-count))
	    (count (cadr pair-count))]
	(set! elements-count (add-ele-count elements-count (car pair) count))
	(add-ele-count elements-count (cadr pair) count)))
    '#() pairs-count))

; directly computes resulting element amounts of
; inserting elements in input polymer up to depth steps
; returns vector of lists where car is an element's character and cadr is the amount of it
(define (polymerize polymer table depth)
  (let step [(d 0)
	     (polymer-pairs (count-pairs polymer))
	     (candidate-pairs '#())]
    (if (< d depth)
      (let [(new-polymer-pairs '#()) (new-candidates '#())]
	(define (multiply index pair-count)
	  (let* [(pair (car pair-count))
		 (count (cadr pair-count))
		 (element (hashtable-ref table pair #f))]
	    (if element
	      (let [(new-pair (list (car pair) element))
		    (new-candidate (list element (cadr pair)))]
		(set! new-polymer-pairs (add-ele-count new-polymer-pairs new-pair count))
		(set! new-candidates (add-ele-count new-candidates new-candidate count)))
	      (add-ele-count pair count))))
	(vector-for-each multiply polymer-pairs)
	(vector-for-each multiply candidate-pairs)
	(step (+ d 1) new-polymer-pairs new-candidates))
      (let [(elements (count-elements polymer-pairs))
	    (last (string-ref polymer (- (string-length polymer) 1)))]
	(add-ele-count elements last 1)))))

(call-with-input-file
  "input"
  (lambda (file)
    (let* [(polymer (get-line file))
	   (table (begin
		    (get-line file) ; skip empty line
		    (read-table file)))]
      (let loop [(n 0) (polymer polymer)]
	(if (< n 10)
	  (loop (+ n 1) (insert polymer table))
	  (begin
	    (printf "part 1:~% After 10 polymerization steps: ~% ")
	    (display-subtracted-most-least-elements (count-chars polymer)))))
      (printf "part 2:~% After 40 polymerization steps: ~% ")
      (display-subtracted-most-least-elements (polymerize polymer table 40)))))