 import Foundation
 import TensorFlow
 
-/// Returns the images tensor and labels tensor.
-public func readMnist(
-    imagesFile: String, labelsFile: String
-) -> (Tensor<Float>, Tensor<Int32>) {
+/// Reads MNIST images and labels from specified file paths.
+public func readMnist(imagesFile: String, labelsFile: String) -> (images: Tensor<Float>, labels: Tensor<Int32>) {
     print("Reading data.")
-    let imageData =
-        try! Data(contentsOf: URL(fileURLWithPath: imagesFile)).dropFirst(16)
-    let labelData =
-        try! Data(contentsOf: URL(fileURLWithPath: labelsFile)).dropFirst(8)
+    let imageData = try! Data(contentsOf: URL(fileURLWithPath: imagesFile)).dropFirst(16)
+    let labelData = try! Data(contentsOf: URL(fileURLWithPath: labelsFile)).dropFirst(8)
     let images = imageData.map { Float($0) }
     let labels = labelData.map { Int32($0) }
     let rowCount = Int32(labels.count)

    return (imagesTensor.toAccelerator(), labelsTensor.toAccelerator())
}

func main() {
    // Get script directory. This is necessary for MNIST.swift to work when
    // invoked from any directory.
    let currentDirectory =

    // of returning it from readMnist() crashes the compiler:
    // https://bugs.swift.org/browse/SR-7706
    // let batchSize = Tensor<Float>(Float(images.shape[0]))
    let batchSize = Float(images.shape[0])

    // Hyper-parameters.
    let learningRate: Float = 0.2
    var loss = Float.infinity

    // Training loop.
    print("Begin training for \(iterationCount) iterations.")

    for i in 0...iterationCount {
        // Forward pass.
        let z1 = images • parameters.w1 + parameters.b1
        let h1 = sigmoid(z1)
        let z2 = h1 • parameters.w2 + parameters.b2
        let predictions = sigmoid(z2)

        // Backward pass. This will soon be replaced by automatic
        // differentiation.
        let dz2 = (predictions - labels) / batchSize
        let dw2 = h1.transposed() • dz2
        let db2 = dz2.sum(squeezingAxes: 0)
        let dz1 = matmul(dz2, parameters.w2.transposed()) * h1 * (1 - h1)
        let dw1 = images.transposed() • dz1
        let db1 = dz1.sum(squeezingAxes: 0)
        let gradients = MNISTParameters(w1: dw1, w2: dw2, b1: db1, b2: db2)

        // Update parameters.
        parameters.update(withGradients: gradients) { param, grad in
          param -= grad * learningRate
        }

        // Update the sigmoid-based cross-entropy loss, where we treat the 10
        // class labels as independent. This is unnecessary for the MNIST case,
        // where we want to predict a single label. In that case we should
        // consider switching to a softmax-based cross-entropy loss.
        let part1 = -labels * log(predictions)
        let part2 = -(1 - labels) * log(1 - predictions)
        loss = (part1 + part2).sum() / batchSize
      
        print("Loss:", loss)
    }
}

var parameters = MNISTParameters()
// Start training.
train(&parameters, iterationCount: 20)
