(ns day5
  (:require [clojure.string :as str]))

(defn init-stacks [number-of-stacks]
  (repeat number-of-stacks []))

(defn put-in-stacks [stacks line]
  (map-indexed (fn [i cargo] (conj (nth stacks i) cargo)) line))

(defn fill-in-stacks [lines]
  (loop [stacks (init-stacks (count (first lines)))
         l lines]
    (if (empty? l)
      stacks
      (recur (put-in-stacks stacks (first l)) (rest l)))))


(defn extract-cargo [line]
  (map second line))

(defn parse-ints [v]
  (map #(Integer/parseInt %) v))

(defn parse-input []
  (let [in-str (slurp "resources/day5.in")
        lines (str/split-lines in-str)
        lines-with-stacks (reverse (drop-last (take-while #(not= % "") lines)))
        split-stacks (map #(partition 4 4 " " %) lines-with-stacks)
        clean-stacks (map extract-cargo split-stacks)
        stacks (fill-in-stacks clean-stacks)
        stacks (map #(take-while (fn [x] (not= x \space)) %) stacks)
        lines-with-transitions (rest (drop-while #(not= % "") lines))
        transitions (map #(str/replace % #"[(move)(from)(to)]" "") lines-with-transitions)
        transitions (map str/trim transitions)
        transitions (map #(str/split % #"  ") transitions)]
    {:stacks (map vec stacks), :transitions (map parse-ints transitions)}))

(defn replace-in-coll [coll n replacement]
  (concat (take n coll) (list replacement) (nthnext coll (inc n))))

(defn move-cargo [stacks from to]
  (let [from-idx (- from 1)
        to-idx (- to 1)
        from-stack (nth stacks from-idx)
        to-stack (nth stacks to-idx)
        cargo (last from-stack)
        new-from-stack (vec (drop-last from-stack))
        new-to-stack (vec (conj to-stack cargo))
        replaced-from (replace-in-coll stacks from-idx new-from-stack)
        new-stacks (replace-in-coll replaced-from to-idx new-to-stack)]
    new-stacks))

(defn batch-cargo [stack batch-size]
  (let [start-idx (- (count stack) batch-size)]
    (nthnext stack start-idx)))

(defn move-cargo-in-batches [stacks [batch-size from to]]
  (let [from-idx (- from 1)
        to-idx (- to 1)
        from-stack (nth stacks from-idx)
        to-stack (nth stacks to-idx)
        cargo (batch-cargo from-stack batch-size)
        new-from-stack (vec (drop-last batch-size from-stack))
        new-to-stack (vec (concat to-stack cargo))
        replaced-from (replace-in-coll stacks from-idx new-from-stack)
        new-stacks (replace-in-coll replaced-from to-idx new-to-stack)]
    new-stacks))

(defn expand-transition [[how-many from to]]
  (repeat how-many [from to]))

(defn expand-transitions [transitions]
  (->> transitions
       (map expand-transition)
       (flatten)
       (partition 2)))

(defn top-from-stacks [stacks]
  (apply str (map last stacks)))

(defn do-transitions [stacks transitions]
  (loop [s stacks
         t transitions]
    (let [transition (first t)
          from (first transition)
          to (second transition)]
      (if (empty? t)
        s
        (recur (move-cargo s from to) (rest t))))))

(defn do-batch-transitions [stacks transitions]
  (loop [s stacks
         t transitions]
    (let [transition (first t)]
      (if (empty? t)
        s
        (recur (move-cargo-in-batches s transition) (rest t))))))


(defn part1-solution []
  (let [input (parse-input)
        stacks (:stacks input)
        transitions (:transitions input)
        after-moving (->> (expand-transitions transitions)
                          (do-transitions stacks))]
    (top-from-stacks after-moving)))

(defn part2-solution []
  (let [input (parse-input)
        stacks (:stacks input)
        transitions (:transitions input)
        after-moving (do-batch-transitions stacks transitions)]
    (top-from-stacks after-moving)))


(comment
  (parse-input)
  (put-in-stacks (init-stacks 9) '(\B \W \N \P \D \V \G \L \T))
  (put-in-stacks '([\B] [\W] [\N] [\P] [\D] [\V] [\G] [\L] [\T]) '(\P \G \R \Z \Z \C \Z \G \P))
  (:stacks (parse-input))
  (move-cargo (:stacks (parse-input)) 1 2)
  (expand-transitions (take 1 (:transitions (parse-input))))
  (expand-transitions (:transitions (parse-input)))
  (top-from-stacks (:stacks (parse-input)))
  (top-from-stacks (do-transitions (:stacks (parse-input)) (expand-transitions (:transitions (parse-input)))))
  (part1-solution)
  (batch-cargo [1 2 3 4 5 6] 2)
  (batch-cargo [1 2 3 4 5 6] 3)
  (batch-cargo [1 2 3 4 5 6] 4)
  (batch-cargo [1 2 3 4 5 6] 1)
  (move-cargo-in-batches (:stacks (parse-input)) (first (:transitions (parse-input))))
  (part2-solution)
  )