 ;; uses so each convolution kernel has a row of weights and a bias entry in the bias
 ;; vector.
 ;;
-;; This means the convolution step is:
-;; (m/matrix-multiply convolved-rows (m/transpose weights))
-;; The result is an interleaved matrix that looks like:
-;; [[cr1k1 cr1k2 cr1k3 cr1k4] (repeats for num convolved rows...)].
-;; Note that simply calling as-vector means our output is in the exact same format
-;; as our input where each convolution kernel output plane is interleaved.  It also means
-;; we can use our linear layer with no changes directly after a convolution layer.
-;;
 
 
 #?(:cljs (register-module cortex.impl.layers.Convolutional))

      (m/join (m/as-vector (:weight-gradient this)) (m/as-vector (:bias-gradient this)))))


(defn max-pooling-forward-default!
  "The forward pass writes to two things; output and output-indexes.  The indexes
record which item from the input row we actually looked at.  Returns
[output output-indexes].  Note that this does not change the number of channels
so it needs to remember the max per-channel."
  [input output output-indexes
   {:keys [^long height ^long width
           ^long padx ^long pady
           ^long k-width ^long k-height
           ^long stride-w ^long stride-h
           ^long num-channels] :as conv-layer-config}]
  ;;Each input row has k-width*num-channels*k-height in it.
  ;;Each output index gets num-channels written to it.
  (let [conv-sequence (convolution-operation-sequence conv-layer-config)
        height (long height)
        width (long width)
        padx (long padx)
        pady (long pady)
        k-width (long k-width)
        k-height (long k-height)
        stride-w (long stride-w)
        stride-h (long stride-h)
        num-channels (long num-channels)
        output-width (get-padded-strided-dimension width padx k-width stride-w)]
    (doseq [conv-item conv-sequence]
      (let [{:keys [^long input-x ^long input-y ^long conv-x ^long conv-y
                    ^long output-x ^long output-y]} conv-item
            input-x (long input-x)
            input-y (long input-y)
            conv-x (long conv-x)
            conv-y (long conv-y)
            output-x (long output-x)
            output-y (long output-y)
            output-offset (* num-channels (+ (* output-y output-width) output-x))]
        (loop [conv-x conv-x]
          (when (< conv-x k-width)
            (let [input-offset-x (+ input-x conv-x)
                  valid-input? (and (>= input-y 0)
                                    (< input-y height)
                                    (>= input-offset-x 0)
                                    (< input-offset-x width))
                  kernel-index (+ (* conv-y k-width) conv-x)
                  input-offset (* num-channels (+ (* input-y width) input-offset-x))]
              (loop [chan 0]
                (when (< chan num-channels)
                  (let [input-val (double (if valid-input?
                                            (m/mget input (+ input-offset chan))
                                            0.0))
                        output-offset (+ output-offset chan)
                        existing-value (double (m/mget output output-offset))]
                    (when (or (= kernel-index 0)
                              (> input-val existing-value))
                      (m/mset! output output-offset input-val)
                      (m/mset! output-indexes output-offset kernel-index)))
                  (recur (inc chan)))))
            (recur (inc conv-x))))))
    [output output-indexes]))


#?(:cljs(defn max-pooling-forward!
          [input output output-indexes conv-config]
          (max-pooling-forward-default! input output output-indexes conv-config))

   :clj(defn max-pooling-forward!
         [input output output-indexes conv-config]
         (let [input-ary (mp/as-double-array input)
               output-ary (mp/as-double-array output)
               output-indexes-ary (mp/as-double-array output-indexes)
               {:keys [width height k-width k-height padx pady stride-w stride-h
                       num-channels]} conv-config]
           (if (and input-ary output-ary output-indexes-ary)
             (ConvOps/maxPooling width height k-width k-height padx pady
                                 stride-w stride-h num-channels
                                 input-ary output-ary output-indexes-ary)
             (max-pooling-forward-default! input output output-indexes conv-config)))))


(defn max-pooling-backward!
  "Calculates the input gradient using the inverse of the convolution step
combined with the output gradient and the output indexes which tell you which kernel
index the output came from."
  [output-gradient output-indexes input-gradient
   {:keys [^long width ^long height
           ^long k-width ^long k-height
           ^long padx ^long pady
           ^long stride-w ^long stride-h
           ^long num-channels] :as conv-layer-config}]
  (m/mset! input-gradient 0.0)
  (let [output-width (get-padded-strided-dimension width padx k-width stride-w)
        output-height (get-padded-strided-dimension height pady k-height stride-h)
        num-pixels (* output-width output-height)]
    (loop [pixel 0]
      (when (< pixel num-pixels)
        (let [output-offset (* pixel num-channels)
              output-x (rem pixel output-width)
              output-y (quot pixel output-width)]
          (loop [chan 0]
            (when (< chan num-channels)
              (let [output-offset (+ output-offset chan)
                    kernel-index (long (m/mget output-indexes output-offset))
                    conv-x (rem kernel-index k-width)
                    conv-y (quot kernel-index k-width)
                    input-x (- (+ (* output-x stride-w) conv-x) padx)
                    input-y (- (+ (* output-y stride-h) conv-y) pady)]
                (when (and (>= input-y 0)
                           (< input-y height)
                           (>= input-x 0)
                           (< input-x width))
                  (let [^double output-gradient-value (m/mget output-gradient output-offset)
                        input-offset (+ (* num-channels (+ (* input-y width)
                                                           input-x))
                                        chan)
                        ^double input-gradient-value (m/mget input-gradient input-offset)]
                    (m/mset! input-gradient input-offset (+ input-gradient-value
                                                            output-gradient-value)))))
              (recur (inc chan)))))
        (recur (inc pixel))))))


;;Max pooling layer.  There are other pooling layer types (average,a sochiastic)
;;that may be implemented later but for now we only need max pooling.

                          conv-layer-config]
    cp/PModule
    (cp/calc [this input]
      (planar-convolution-forward! this input))

    (cp/output [m]
      (m/as-vector (:output m)))

    cp/PNeuralTraining
    (forward [this input]
      (cp/calc this input))

    (backward [this input output-gradient]
      (planar-convolution-backward! this input output-gradient))

    (input-gradient [this]
      (m/as-vector (:input-gradient this)))

    cp/PParameters
    (parameters [this]
