(ns advent-of-code-2021.day12
  (:require  [clojure.java.io :as io]
             [clojure.string :as str]))


(let [edges (->> "day12/input.txt"
                 io/resource
                 slurp
                 str/split-lines
                 (mapcat (fn [line] ((juxt identity #(into [] (reverse %)))  
                                     (mapv keyword (str/split line #"-")))))
                 ((fn [edges] (reduce (fn [edge-map [from to]]
                                        (update edge-map from (fn [tos]
                                                                (if tos
                                                                  (conj tos to)
                                                                  #{to})))) 
                                      {}
                                      edges))))

      small-cave? (fn [node] (= (name node)  (clojure.string/lower-case (name node))))

      part-1-visit-state (fn visit-state 
                           [visited]
                           {:can-visit? (fn [node]
                                          (or (not (small-cave? node))
                                              (not (visited node))))
                            :with-visit (fn [node]
                                          (cond
                                            (not (small-cave? node)) (visit-state visited)
                                            (not (visited node)) (visit-state (conj visited node))))})

      part-1-start-path ['(:start) (part-1-visit-state #{:start})]

      part-2-visit-state (fn visit-state 
                           [visited visited-twice]
                           {:can-visit? (fn [node]
                                          (or (not (small-cave? node))
                                              (not (visited node))
                                              (and (not (#{:start} node)) 
                                                   (not visited-twice))))
                            :with-visit (fn [node]
                                          (cond
                                            (not (small-cave? node)) (visit-state visited visited-twice)
                                            (not (visited node)) (visit-state (conj visited node) visited-twice)
                                            (not visited-twice) (visit-state visited true)))})


      part-2-start-path ['(:start) (part-2-visit-state #{:start} false)]

      find-next-paths-for-one (fn [unfinished-path]
                                (let [[current-nodes visit-state] unfinished-path
                                      last-node (first current-nodes)
                                      next-nodes (edges last-node)
                                      visitable-nodes (filter (:can-visit? visit-state) next-nodes)]                        
                                  (mapv (fn [next-node]
                                               [(conj current-nodes next-node)                                               
                                                ((:with-visit visit-state) next-node)])
                                       visitable-nodes)))

      find-next-paths-for-multi (fn [unfinished-paths]
                                  (into []  (mapcat find-next-paths-for-one unfinished-paths)))     

      find-next-all-paths (fn [{finished-paths :finished unfinished-paths :unfinished}]
                            (let [next-paths (find-next-paths-for-multi unfinished-paths)
                                  next-finished-and-unfinished (group-by (fn [[current-nodes _]]
                                                                           (if (= :end (first current-nodes))
                                                                             :finished
                                                                             :unfinished))
                                                                         next-paths)]
                              (if finished-paths 
                                (update next-finished-and-unfinished 
                                        :finished                                      
                                        #(if % 
                                           (into finished-paths %)
                                           finished-paths))
                                next-finished-and-unfinished)))

      search-paths (fn [start-path]
                     (iterate find-next-all-paths {:unfinished  [start-path]}))]
  (map (fn [start-path]
         (->> (search-paths start-path)
              (drop-while #(seq (:unfinished %)))
              first
              :finished
              count))  [part-1-start-path part-2-start-path]))