(ns advent.pathfind
  (:require [play-clj.core :refer :all])
  (:import (java.lang Math)))

(def scale 2)

(def cost-comparator (fn [[a] [b]]
                       (let [a (long a)
                             b (long b)]
                         (unchecked-subtract a b))))

(defn from-scale [[x y]]
  [(* scale x) (* scale y)])

(defn to-scale [[x y]]
  [(quot x scale) (quot y scale)])

(defn printmap [my-map & [skip]]
  (let [skip (or skip 1)]
    (doseq [row (take-nth skip my-map)]
      (println (take-nth skip (map {1 \space 0 "W" "X" "X" "." "."} row))))))

(defn random-map []
  (-> (vec (take 640 (repeatedly (fn [] (vec (take 480 (repeatedly (fn [] (rand-nth [1 1 1 1 1 0])))))))))
      (update-in [1 1] (constantly 1))
      (update-in [639 479] (constantly 1))))

(defn neighbors [^long x ^long y my-map]
  (let [left-x (unchecked-subtract x 1)
        right-x (unchecked-add x 1)
        top-y (unchecked-subtract y 1)
        below-y (unchecked-add y 1)
        candidates [[left-x top-y]    [x top-y]  [right-x top-y]
                    [left-x y]                   [right-x y]
                    [left-x below-y] [x below-y] [right-x below-y]]
        height (count (first my-map))
        width (count my-map)]
    (remove (fn [[x y]]
              (= 0 (nth (nth my-map x) y)))
            (filter (fn [[x y]]
                      (let [x (long x)
                            y (long y)]
                        (and (< -1 x width)
                             (< -1 y height))))
                    candidates))))

(defn resolve-path [came-from play-loc target-loc]
  (let [came-from (into {} came-from)]
    (if (nil? (came-from target-loc))
      nil
      (loop [path []
             current-node target-loc]
        (if (or (= current-node play-loc)
                (nil? current-node))
          (reverse (conj path (from-scale current-node)))
          (recur
           (conj path (from-scale current-node))
           (came-from current-node)))))))

(def d2 (- (Math/sqrt 2) 2))

(defn heuristic ^long [^long goal-x ^long goal-y ^long current-x ^long current-y]
  (let [dist-x (if (< goal-x current-x)
                 (unchecked-subtract current-x goal-x)
                 (unchecked-subtract goal-x current-x ))
        dist-y (if (< goal-y current-y)
                 (unchecked-subtract current-y goal-y)
                 (unchecked-subtract goal-y current-y))
        min-dist (double (min dist-x dist-y))]
    (unchecked-add (unchecked-add dist-x dist-y)
                   (long (unchecked-multiply (double d2) min-dist)))))

(defn visit-all [my-map play-loc target-loc]
  (let [play-loc (to-scale play-loc)
        [tx ty :as target-loc] (to-scale target-loc)
        tx (long tx)
        ty (long ty)]
    (if (= 0 (nth (nth my-map tx) ty))
      nil
      (let [cost-so-far ^java.util.HashMap (java.util.HashMap. {play-loc 0})
            came-from ^java.util.HashMap (java.util.HashMap.)
            fronteir ^java.util.PriorityQueue (java.util.PriorityQueue. (/ (* 320 240) scale) cost-comparator)]
        (.offer fronteir [0 play-loc])
        (loop [current-loc (.poll fronteir)]
          (if (or (nil? current-loc)
                  (= (nth current-loc 1) target-loc))
            (resolve-path came-from play-loc target-loc)
            (let [[_ [cx cy :as current-loc]] current-loc
                  cx (long cx)
                  cy (long cy)]
              (doseq [[nx ny :as neighbor] (neighbors cx cy my-map)
                      :let [nx (long nx)
                            ny (long ny)]]
                (let [cost-for-neighbor (.get cost-so-far neighbor)
                      new-cost (unchecked-add (long (.get cost-so-far current-loc))
                                              (long (nth (nth my-map nx) ny)))]
                  (when (or (nil? cost-for-neighbor)
                            (< new-cost cost-for-neighbor))
                    (.put came-from neighbor current-loc)
                    (.put cost-so-far neighbor new-cost)
                    (.offer fronteir [(unchecked-add new-cost (heuristic tx ty nx ny))
                                      neighbor]))))
              (recur (.poll fronteir)))))))))

(defn print-resolved [path my-map]
  (doseq [row (reduce (fn [acc path]
                        (if path
                          (update-in acc (reverse path) (constantly "X"))
                          acc))
                      my-map
                      path)]
    (println (map {1 \space 5 "." 0 "W" "X" "X"} row)))
  nil)

(defn test-pathfind []
  (let [my-map (random-map)
        path (visit-all my-map [1 1] [639 479 ])]
    (println "Test")
    (print-resolved path my-map)
    (println path)))

(defn map-from-resource [filename]
  (let [pm (pixmap filename)
        black (color 0 0 0 255)
        painful (color 255 0 0 255)
        result (transient [])
        scale (long scale)
        height (long (pixmap! pm :get-height))]
    (doseq [^long x (range (/ (pixmap! pm :get-width) scale))
            ^long y (range (/ height scale)) 
            :let [
                  current-color (color (pixmap! pm :get-pixel (unchecked-multiply scale x) (unchecked-subtract height (unchecked-multiply scale y))))]]
      (conj! result (cond
                     (color! current-color :equals black) 0
                     (color! current-color :equals painful) 2
                     :else 1)))
    (partition (/ (pixmap! pm :get-height) scale) (persistent! result) )))