(ql:quickload '(fiveam uiop cl-ppcre fset lparallel arrows))
(setf lparallel:*kernel* (lparallel:make-kernel 8))

(defparameter *materials* '(:geode :obsidian :clay :ore))
(defun material (string)
  (cond
    ((string= string "ore") :ore)
    ((string= string "clay") :clay)
    ((string= string "obsidian") :obsidian)
    ((string= string "geode") :geode)))

(defun parse-robot-recipe (line)
  (cl-ppcre:register-groups-bind ((#'material robot) (#'parse-integer fq) (#'material fm) (#'parse-integer sq) (#'material sm))
      ("(\\w+) robot costs (\\d+) (\\w+)(?: and (\\d+) (\\w+))?" line)
    (let ((first-material (fset:with (fset:empty-bag) fm fq)))
      (fset:map (robot
                 (if sm (fset:with first-material sm sq)
                     first-material))))))


(defun can-build-robot-pm (requirements resources)
  (fset:subbag? requirements resources))

(defun consume (requirements resources)
  (fset:bag-difference resources requirements))

(defun material-collect (robots resources)
  (fset:bag-sum robots resources))

(defun resource-increase (current-state time-left)
  (destructuring-bind (robots resources) current-state
    (reduce (lambda (acc material)
              (fset:with acc material
                         (* time-left (fset:multiplicity robots material))))
            *materials*
            :initial-value resources)))

(defun bot-requirements (type recipes)
  (fset:lookup recipes type))

(defun build-robot (type robots resources recipes)
  (let ((req (bot-requirements type recipes)))
    (when (can-build-robot-pm req resources)
      (let ((new-robots (fset:with robots type))
            (resources-left (consume req resources)))
        (values new-robots (material-collect robots resources-left))))))

(defun max-cost (recipes)
  (fset:reduce
   (lambda (acc key value)
     (declare (ignore key))
     (fset:union value acc))
   recipes
   :initial-value (fset:empty-bag)))

(defun can-build (bots)
  (append
   (list :clay :ore)
   (when (plusp (fset:multiplicity bots :clay)) (list :obsidian))
   (when (plusp (fset:multiplicity bots :obsidian)) (list :geode))))

(defun propose-builds (bots max-costs)
  (loop for try in (can-build bots)
        when (or (eq :geode try)
                 (< (fset:multiplicity bots try) ;; every cycle I not yet produce
                    (fset:multiplicity max-costs try))) ;; enough for any build
          collect try))

(defun probe (current-state recipes max-cost)
  (destructuring-bind (robots resources) current-state
    (reduce
     (lambda (acc bot-type)
       (multiple-value-bind (new-robots new-resources)
           (build-robot bot-type robots resources recipes)
         (if new-robots (fset:with acc (list new-robots new-resources)) acc)))
     (propose-builds robots max-cost)
     :initial-value (fset:set (list robots (material-collect robots resources))))))

(defun advance (robot-resource-seq blueprints max-cost)
  (fset:reduce
   (lambda (new-states current-state)
     (fset:union new-states
            (probe current-state blueprints max-cost)))
   robot-resource-seq
   :initial-value (fset:empty-set)))

(defun materials> (m1 m2)
  (eq :greater
      (reduce (lambda (result material)
                (if (eq result :equal)
                    (let ((c1 (fset:multiplicity m1 material))
                          (c2 (fset:multiplicity m2 material)))
                      (cond
                        ((< c1 c2) :less)
                        ((= c1 c2) :equal)
                        ((> c1 c2) :greater)))
                    result))
              *materials*
              :initial-value :equal)))

(defun leading-configs (robot-resource-seq time-left truncate)
  ;; (format t "~a~%" (fset:size robot-resource-seq))
  (fset:subseq
   (fset:sort robot-resource-seq
              #'materials> :key (lambda (state) (resource-increase state time-left)))
   0 truncate))

(defun find-best (robot-resource-seq blueprints time-left &optional (state-truncation 1024) (max-cost (max-cost blueprints)))
  (if (zerop time-left)
      (arrows:-> robot-resource-seq
                 ;; (fset:subseq 0 3)
                 (fset:first)
                 (cadr)
                 (fset:multiplicity :geode))
      (find-best (leading-configs
                  (advance robot-resource-seq blueprints max-cost)
                  (1- time-left) state-truncation)
                 blueprints (1- time-left)
                 state-truncation
                 max-cost)))

(defun get-blueprints (filename)
  (mapcar
   (lambda (l)
     (let ((blueprint (uiop:split-string l :separator ":.")))
       (reduce #'fset:map-union
               (mapcar #'parse-robot-recipe (butlast (cdr blueprint))))))
   (uiop:read-file-lines filename)))

(defun solver (blueprints time-left)
  (lparallel:pmap 'list
                  (lambda (recipes)
                    (find-best (fset:seq (list (fset:bag :ore) (fset:empty-bag)))
                               recipes time-left 1024))
                  blueprints))

(defun part1 (blueprints time-left)
  (loop for geodes in (solver blueprints time-left)
        for idx from 1
        sum (* idx geodes)))

(defun part2 (blueprints time-left)
  (apply #'* (solver blueprints time-left)))

(fiveam:test solution
  ;; part1
  (fiveam:is (= 33 (part1 (get-blueprints "eg-in") 24)))
  (fiveam:is (= 1266 (part1 (get-blueprints "input") 24)))
  ;; part1
  (fiveam:is (= 3472 (part2 (get-blueprints "eg-in") 32)))
  (fiveam:is (= 5800 (part2 (subseq (get-blueprints "input") 0 3) 32))))