(asdf:load-system :adventofcode2020)
(in-package #:adventofcode2020)
(named-readtables:in-readtable :adventofcode2020)

(defun joltage-distribution (ratings)
  (let* ((extended (append '(0) ratings (mapcar λ(+ 3 _) (last ratings)))))
    (loop for a in extended
          for b in (cdr extended)
          counting (= 3 (- b a)) into threes
          counting (= 1 (- b a)) into ones
          finally (return (list ones threes)))))

(defun partition-ratings (ratings)
  (let ((groups nil) (g nil))
    (loop for a in ratings
          for b in (cdr (append ratings (mapcar λ(+ 3 _) (last ratings))))
          do (push a g)
          when (= 3 (- b a))
          do (progn
               (push (nreverse g) groups)
               (setf g nil)))
    (nreverse groups)))

; turns out you don't need the fancy memoized one
; (defparameter *count-deletions-table*
;   (alist-hash-table '((1 . 0) (2 . 0) (3 . 1) (4 . 3))))
; (defun count-deletions (len)
;   (let ((val (gethash len *count-deletions-table*)))
;     (if val val
;         (let ((new-val (* 2 (count-deletions (1- len)))))
;           (setf (gethash len *count-deletions-table*) new-val)
;           new-val))))
(defun count-deletions (len)
  (case len
    (1 0) (2 0) (3 1) (4 3)
    (t (* 2 (count-deletions (1- len))))))

(defun count-adapter-arrangements (ratings)
  (reduce #'* (mapcar (compose #'1+ #'count-deletions #'length)
                      (partition-ratings (cons 0 ratings)))))

(day 10 input
  (let ((ratings (sort (int-list-from input) #'<)))
    (part1 (apply #'* (joltage-distribution ratings)))
    (part2 (count-adapter-arrangements ratings))))

(def-suite day10)
(in-suite day10)

(defvar *simple-ratings*
  (sort '(16 10 15 5 1 11 7 19 6 12 4) #'<))

(defvar *more-ratings*
  (sort '(28 33 18 42 31 14 46 20 48 47
          24 23 49 45 19 38 39 11 1 32
          25 35 8 17 7 9 4 2 34 10 3) #'<))

(test joltage-distributions
  (is (equal '(7 5) (joltage-distribution *simple-ratings*)))
  (is (equal '(22 10) (joltage-distribution *more-ratings*))))

(test partition-ratings
  (is (equal
        '((1) (4 5 6 7) (10 11 12) (15 16) (19))
        (partition-ratings *simple-ratings*)))
  (is (equal
        '((1 2 3 4) (7 8 9 10 11) (14) (17 18 19 20) (23 24 25) 
                    (28) (31 32 33 34 35) (38 39) (42) (45 46 47 48 49))
        (partition-ratings *more-ratings*))))

(test count-deletions
  (is (equal
        '(0 3 1 0 0)
        (mapcar (compose #'count-deletions #'length)
                '((0 1) (4 5 6 7) (10 11 12) (15 16) (19)))))
  (is (equal
        '(6 6 0 3 1 0 6 0 0 6)
        (mapcar (compose #'count-deletions #'length)
                '((0 1 2 3 4) (7 8 9 10 11) (14) (17 18 19 20) (23 24 25) 
                    (28) (31 32 33 34 35) (38 39) (42) (45 46 47 48 49))))))

(test count-adapter-arrangements
  (is (equal 8 (count-adapter-arrangements *simple-ratings*)))
  (is (equal 19208 (count-adapter-arrangements *more-ratings*))))

(run! 'day10)