 
 (mat/set-current-implementation :vectorz)
 
-;; Primary Protocols
+;; Neural Protocols
+
 (defprotocol NeuralLayer
   (forward [this input] "Pass data into the layer and return its output.")
   (backward [this input output-gradient] "Back propagate error gradients through the layer."))
 
-(defprotocol CostFn
-  (cost [this activation target])
-  (delta [this activation z target]))
+(defprotocol LossFn
+  (loss [this v target])
+  (delta [this v target]))
+
+;; Helpers
+(defn exp
+  [a]
+  (mat/emap #(Math/exp %) a))
+
+(defn exp!
+  [a]
+  (mat/emap! #(Math/exp %) a))
+
+(defn log
+  [a]
+  (mat/emap #(Math/log %) a))
+
+(defn log!
+  [a]
+  (mat/emap! #(Math/log %) a))
+
+(defn rand-vector
+  "Produce a vector with guassian random elements having mean of 0.0 and std of 1.0."
+  [n]
+  (let [rgen (Random.)]
+    (mat/array (repeatedly n #(.nextGaussian rgen)))))
+
+(defn rand-matrix
+  [m n]
+  (let [rgen (Random.)]
+    (mat/array (repeatedly m (fn [] (repeatedly n #(.nextGaussian rgen)))))))
+
+(defn weight-matrix
+  [m n]
+  (let [stdev (/ 1.0 n)]
+    (mat/emap #(* stdev %) (rand-matrix m n))))
+
+;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;
+;; Activation Functions
+;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;
+
+; NOTE: This form of computing a sigmoid over a matrix is faster than mapping
+; over each individual element.  Using the mutable versions of most functions
+; saves .
+(defn sigmoid
+  "y =  1 / (1 + e^(-z))
+  Produces an output between 0 and 1."
+  [z]
+  (mat/div! (mat/add! (exp! (mat/negate z)) 1.0)))
+
+(defn sigmoid!
+  "y =  1 / (1 + e^(-z))
+  Produces an output between 0 and 1."
+  [z]
+  (mat/div! (mat/add! (exp! (mat/negate! z)) 1.0)))
+
+(defn sigmoid'
+  [z]
+  (let [sz (sigmoid z)]
+    (mat/emul sz (mat/sub 1.0 sz))))
+
+(defrecord SigmoidActivation [output]
+  NeuralLayer
+  (forward [this input]
+    (mat/assign! output input)
+    (sigmoid! output))
+
+  ; z = dot(weights, input) + b
+  ; activation = sigmoid(z)
+  (backward [this input output-gradient]
+    (mat/emul output-gradient (sigmoid' input))))
+
+(defn sigmoid-activation
+  [shape]
+  (map->SigmoidActivation {:output (mat/zero-array shape)}))
+
+(defn softmax
+  "Used for multinomial classification (choose 1 of n classes), where the output
+  layer can be interpreted as class probabilities."
+  [z]
+  (mat/div z (mat/esum z)))
+
+(defrecord LinearLayer [weights biases weight-gradients bias-gradients outputs]
+  NeuralLayer
+  (forward [this input]
+    (mat/assign! outputs biases)
+    (mat/add! outputs (mat/mmul weights input)))
+
+  (backward [this input output-gradient]
+    (mat/fill! weight-gradients 0.0)
+    (mat/assign! bias-gradients output-gradient)
+
+    ; Compute the error gradients with respect to the input data by multiplying the
+    ; output errors backwards through the weights, and then compute the error
+    ; with respect to the weights by multiplying the input error times the
+    ; output error.
+    (let [input-grad (mat/mmul (mat/transpose weights) output-gradient)]
+      (mat/assign! weight-gradients (mat/outer-product output-gradient input-grad)))))
+
+(defn linear-layer
+  [n-inputs n-outputs]
+  (let [weights (weight-matrix n-outputs n-inputs)
+        biases (rand-matrix 1 n-outputs)]
+    (map->LinearLayer
+      {:weights weights
+       :biases biases
+       :weight-gradients (mat/zero-array (mat/shape weights))
+       :bias-gradients (mat/zero-array (mat/shape biases))
+       :outputs (mat/zero-array (mat/shape biases))})))
 
 (defrecord SequentialNetwork [layers config]
   NeuralLayer

;; 4) update weights
;;  * multiply gradient by the learning-rate to smooth out jitter in updates.
(defn backprop
  [{:keys [biases weights cost-fn n-layers] :as net} input expected-output]
  (let [bias-gradients (map #(mat/zero-array (mat/shape %)) biases)
        weight-gradients (map #(mat/zero-array (mat/shape %)) weights)
        [activations zs] (reduce

          ;(println (format "sample score: %5.2f" (float (/ sample-score sample-size))))
          (recur new-net (inc epoch)))))))

(defn network
  [layer-sizes & opts]
  {:n-layers (count layer-sizes)
   :biases (map rand-vector (next layer-sizes))
   :weights (map weight-matrix (rest layer-sizes) (drop-last layer-sizes))
   :loss-fn (get opts :loss-fn (QuadraticLoss.))})

(def trained* (atom nil))

(defn mnist-labels
