(ns day20.part1
  (:use clojure.test)
  (:require
   [clojure.pprint :refer [pprint]]
   [day20.os-tree :as ost]))

(defn mix [input]
  (letfn [(rotate-nth [index n dist]
            "Move the item at position n by dist within index, wrapping if necessary"
            (let [input-size (count input)
                  item (ost/select index n)
                  prev-idx (mod n (dec input-size))
                  insertion-idx (mod (+ prev-idx dist) (dec input-size))]
              (-> index
                  (ost/delete n)
                  (ost/insert insertion-idx item))))]

    (let [initial-arrangement
          (ost/construct
           (vec (range (count input))))

          final-arrangement
          (reduce-kv
           (fn [index i n]
             (let [target (ost/search index i)]
               (rotate-nth index target n)))
           initial-arrangement
           input)]

      (map
       input
       (ost/in-order
        final-arrangement)))))

(defn -main [file]
  (let [input (->> file
                   slurp
                   clojure.string/trim
                   clojure.string/split-lines
                   (mapv read-string))]
    (->> input
         mix
         repeat
         (apply concat)
         (drop-while (complement zero?))
         (take-nth 1000)
         rest
         (take 3)
         (apply +))))

(ns day20.os-tree
  "Implements a binary tree that represents an ordered sequence
   and allows for efficient access, insertion and deletion by index.
   Each node is represented by a vector with the following format:

     [value index left-child right-child]"

  (:use clojure.test))

(defn construct [items]
  "Create a balanced tree from a vector of items"
  (if (zero? (count items))
    nil
    (let [midpoint (quot (count items) 2)
          left-items (subvec items 0 midpoint)
          right-items (subvec items (inc midpoint) (count items))]
      (vector
       (items midpoint)
       (count left-items)
       (construct left-items)
       (construct right-items)))))

(defn insert [[node-val node-idx left right :as node] idx new-val]
  "Insert new-val at position idx within the tree"
  (cond
    (nil? node)
    [new-val 0 nil nil]

    (= node-idx idx)
    (let [updated-right
          (insert
           right
           (- idx node-idx)
           node-val)]

      [new-val node-idx left updated-right])

    (< node-idx idx)
    (let [updated-right
          (insert
           right
           (dec (- idx node-idx))
           new-val)]

      [node-val node-idx left updated-right])

    :else
    [node-val
     (inc node-idx)
     (insert left idx new-val)
     right]))

(defn select [[v node-idx left right] idx]
  "Return the value located at idx"
  (cond
    (= node-idx idx) v

    (< node-idx idx)
    (select
     right
     (dec (- idx node-idx)))

    :else
    (select left idx)))

(defn search [[v node-idx left right :as node] needle]
  "Return the index of needle"
  (when-not (nil? node)
    (or
     (when (= v needle) node-idx)
     (search left needle)
     (when-let [right-result (search right needle)]
       (+ right-result node-idx 1)))))

(defn delete [[node-val node-idx left right :as node] idx]
  "Remove the value located at idx from the tree"
  (letfn [(replace-self []
            "Replace the node that is being deleted with
             a descendant (or nothing)."
            (cond
              (and (nil? left) (nil? right)) nil

              (nil? right) left

              :else
              (let [replacement (select right 0)]
                [replacement
                 node-idx
                 left
                 (delete right 0)])))]

    (cond
      (= node-idx idx) (replace-self)

      (< node-idx idx)
      [node-val
       node-idx
       left
       (delete right (dec (- idx node-idx)))]

      :else
      [node-val
       (dec node-idx)
       (delete left idx)
       right])))

(defn depth [[node-val node-idx left right :as node]]
  "Determine the maximum depth of the tree"
  (if (nil? node)
    0
    (max
     (inc (depth left))
     (inc (depth right)))))

(defn in-order [[node-val _ left right :as node]]
  "Return the contents of the tree in index order"
  (if (nil? node)
    nil
    (concat
     (in-order left)
     (cons node-val (in-order right)))))