(defpackage :aoc/2024/16
  (:use :cl :aoc :alexandria :trivia :lla :queues)
  (:export
   #:parse-input
   #:best-score
   #:best-path
   #:solve-part2
   #:sample-data
   #:sample-data2
   #:part1
   #:part2
   ))

(in-package :aoc/2024/16)

(defun parse-line (line)
  (coerce line 'list))


(defun parse-input (lines)
  (make-array
   (list (length lines) (length (first lines)))
   :initial-contents (mapcar #'parse-line lines)))

(defparameter input-text (test-input 2024 16))
(defparameter input-data (parse-input input-text))

(defparameter sample-text (aoc:split-lines "###############
#.......#....E#
#.#.###.#.###.#
#.....#.#...#.#
#.###.#####.#.#
#.#.#.......#.#
#.#.#####.###.#
#...........#.#
###.#.#####.#.#
#...#.....#.#.#
#.#.#.###.#.#.#
#.....#...#.#.#
#.###.#.#.#.#.#
#S..#.....#...#
###############"))

(defstruct pos
  (x 0 :type fixnum)
  (y 0 :type fixnum))

(defun pos-move (pos dir)
  (case dir
    (#\^ (make-pos :x (pos-x pos) :y (1- (pos-y pos))))
    (#\v (make-pos :x (pos-x pos) :y (1+ (pos-y pos))))
    (#\< (make-pos :x (1- (pos-x pos)) :y (pos-y pos)))
    (#\> (make-pos :x (1+ (pos-x pos)) :y (pos-y pos)))
    ))

(defun show-map (map &optional (dest t))
  (let ((w (array-dimension map 1))
	(h (array-dimension map 0)))
    (format
     dest "~{~A~}"
     (loop
       for row from 0 below h
       collect (format nil "~A~%"
		       (coerce
			(loop for col from 0 below w
			      collect (aref map row col))
			'string))))))

(defun get-pos (map pos)
  (aref map (pos-y pos) (pos-x pos)))

(defun set-pos (map pos tile)
  (setf (aref map (pos-y pos) (pos-x pos)) tile))

(defun find-tile (map char)
  (let ((w (array-dimension map 1))
	(h (array-dimension map 0)))
    (loop
      for row from 0 below h
      for col = (loop
		  for col from 0 below w
		  if (eq (aref map row col) char) return col)
      if col
	return (make-pos :x col :y row)
      )))

(defun swap-pos (map p np)
  "swaps content of map. returns np"
  (let ((o (get-pos map p)))
    (setf (aref map (pos-y p) (pos-x p)) (get-pos map np))
    (setf (aref map (pos-y np) (pos-x np)) o)
    np))


(defun wet-finger-distance (a b)
  "estimate the distance"
  (+
   (abs (- (pos-x a) (pos-x b)))
   (abs (- (pos-y a) (pos-y b)))
   (and
    (eq (pos-x a) (pos-x b))
    (eq (pos-y a) (pos-y b))
    1000)))

(defun rotate-left (dir)
  (case dir
    (#\^ #\<)
    (#\v #\>)
    (#\< #\v)
    (#\> #\^)
    ))
(defun rotate-right (dir)
  (case dir
    (#\^ #\>)
    (#\v #\<)
    (#\< #\^)
    (#\> #\v)
    ))

(defun valid-pos-p (map pos)
  (not (eq #\# (get-pos map pos))))

(defstruct reindeer pos dir)

(defstruct state path score)

(defun next-moves (map reindeer)
  "return conses of next reindeer and cost to get there"
  (let ((d (reindeer-dir reindeer))
	(p (reindeer-pos reindeer)))
    (remove-if-not
     #'(lambda (r) (valid-pos-p map (reindeer-pos (car r))))
     (list
      (cons (make-reindeer :pos (pos-move p d) :dir d) 1)
      (cons (make-reindeer :pos p :dir (rotate-left d)) 1000)
      (cons (make-reindeer :pos p :dir (rotate-right d)) 1000)))))


(defun distance-estimate-f (finish)
  "estimate the distance for a reindeer from given finish"
  (let ((a finish))
    (lambda (reindeer)
      (let ((b (reindeer-pos reindeer) ))
	(+
	 (abs (- (pos-x a) (pos-x b)))
	 (abs (- (pos-y a) (pos-y b)))
	 (if (or
	      (eq (pos-x a) (pos-x b))
	      (eq (pos-y a) (pos-y b)))
	     0
	     1000))))))

(defun compare-states-f (finish)
  (let ((dist-est (distance-estimate-f finish)))
    (lambda (s1 s2)
      (< (+ (state-score s1)
	    (funcall dist-est (car (state-path s1))))
	 (+ (state-score s2)
	    (funcall dist-est (car (state-path s2))))))))

(defun best-path (map)
  (let* ((start (find-tile map #\S))
	 (finish (find-tile map #\E))
	 (state-comparison (compare-states-f finish))
	 (best (make-hash-table :test #'equalp))
	 (todo (make-queue :priority-queue :compare state-comparison)))
    (qpush todo (make-state
		 :path (list (make-reindeer :pos start :dir #\>))
		 :score 0))    
    (loop
      for s = (qpop todo)
      until (or
	     (null s)
	     (equalp (reindeer-pos (first (state-path s))) finish))
      finally (return s)
      do (loop
	   for next in (next-moves map (first (state-path s)))
	   for score = (+ (state-score s) (cdr next))
	   if (> (gethash (car next) best 1000000000) score)
	     do (progn
		  (qpush todo
			 (make-state
			  :path (cons (car next) (state-path s))
			  :score score))
		  (setf (gethash (car next) best) score)))

      )))

(defun best-score (map)
  (state-score (best-path map)))

(defparameter sample-data 
  (parse-input sample-text))

(defun part1 (data)
  (format nil "~A" (best-score data)))


(defun best-paths (map)
  (let* ((start (find-tile map #\S))
	 (finish (find-tile map #\E))
	 (state-comparison (compare-states-f finish))
	 (best (make-hash-table :test #'equalp))
	 (best-score 100000000)
	 (pos-on-best-path (make-hash-table :test #'equalp))
	 (todo (make-queue :priority-queue :compare state-comparison)))
    (qpush todo (make-state
		 :path (list (make-reindeer :pos start :dir #\>))
		 :score 0))    
    (loop
      for s = (qpop todo)
      until (null s)
      for path = (state-path s)
      for score = (state-score s)
      finally (return (loop for k being the hash-keys in pos-on-best-path collect k))
      if (<= score best-score)
	do (if (equalp finish (reindeer-pos (first (state-path s))))
	       (progn
		 (if (< score best-score)
		     (progn
		       (setf best-score score)
		       (clrhash pos-on-best-path)))
		 (loop
		   for p in path
		   do (incf (gethash (reindeer-pos p) pos-on-best-path 0))))
	       (loop
		 for next in (next-moves map (first path))
		 for next-score = (+ score (or  (cdr next) 0))
		 if (>= (gethash (car next) best 1000000000) next-score)
		   do (progn
			(qpush todo
			       (make-state
				:path (cons (car next) (state-path s))
				:score next-score))
			(setf (gethash (car next) best) next-score))))

      )))

(defun show-paths (map paths)
  (let ((nmap (copy-array map)))
    (loop
      for p in paths
      do (set-pos nmap p #\O))
    (show-map nmap)))

(defun solve-part2 (map)
  (length (best-paths map)))

(defun part2 (data)
  (format nil "~A" (solve-part2 data) ))

(defun solve-day ()
  (format t "part1: ~A~%" (part1 input-data))
  (format t "part2: ~A~%" (part2 input-data)))

(defun submit ()
  (let ((p1 (part1 input-data))
        (p2 (part2 input-data)))
    (if p1 (submit-part1 2024 16 p1))
    (if p2 (submit-part2 2024 16 p2))))