 }
 
 func saveData(filename:String, images:Tensor<I>) throws {
-//    let image_filename = "/Users/ars3432/Downloads/train-images-idx3-ubyte"
-//    let images = loadData(filename: input)
     let data:Data = serialize(tensor: images)!
     try data.write(to: URL(fileURLWithPath: filename))
+    print("wrote \(data.count) bytes to \"\(filename)\"")
 }
 
-
-func parseOptions(args:[String], options:Dictionary<String, String>) {
-    for arg in args {
-        arg.self
-        if (options[arg] != nil) {
-            
+extension String {
+    subscript (r: CountableRange<Int>) -> String {
+        get {
+            let startIndex =  self.index(self.startIndex, offsetBy: r.lowerBound)
+            let endIndex = self.index(startIndex, offsetBy: r.upperBound - r.lowerBound)
+            return self[startIndex..<endIndex]
         }
     }
 }
 
-let data = try Data(contentsOf: URL(fileURLWithPath: "mnist.bin"))
-let images:Tensor<I> = deserialize(data: data)!
+func parseOptions(args:[String]) -> [String:String] {
+    var options = [String:String]()
 
-for i in 0..<10 {
-    var imageTensor = images[i, all, all]
-    imageTensor = imageTensor.reshape(Extent(imageTensor.shape[1], imageTensor.shape[2]))
+    var i = 1
+    while i < args.count-1 {
+        let option = String(args[i].characters.dropFirst())
+        let value = args[i+1]
+        options[option] = value
+        i += 2
+    }
     
-    // NB: Bug here .. reshape causes the view to get reset
-    print(imageTensor)
-    let image = toImage(imageTensor)
-    let data = image.tiffRepresentation!
-    try data.write(to: URL(fileURLWithPath: "image_\(i).tiff"))
+    return options
 }
 
-//if CommandLine.arguments[0] == "generated"
+let options = parseOptions(args: CommandLine.arguments)
+
+if options["mode"] == "display" {
+    let data = try Data(contentsOf: URL(fileURLWithPath: "mnist.bin"))
+    let images:Tensor<I> = deserialize(data: data)!
+
+    for i in 0..<10 {
+        var imageTensor = images[i, all, all]
+        imageTensor = imageTensor.reshape(Extent(imageTensor.shape[1], imageTensor.shape[2]))
+        
+        let image = toImage(imageTensor)
+        let data = image.tiffRepresentation!
+        try data.write(to: URL(fileURLWithPath: "image_\(i).tiff"))
+    }
+} else {
+    print("loading training images")
+    let trainImages = readImageData(url: URL(string: "http://yann.lecun.com/exdb/mnist/train-images-idx3-ubyte.gz")!)!
+    
+    print("loading training labels")
+    let trainLabels = readLabelData(url: URL(string: "http://yann.lecun.com/exdb/mnist/train-labels-idx1-ubyte.gz")!)!
+    
+    print("loading testing images")
+    let testImages = readImageData(url: URL(string: "http://yann.lecun.com/exdb/mnist/t10k-images-idx3-ubyte.gz")!)!
+    
+    print("loading testing labels")
+    let testLabels = readLabelData(url: URL(string: "http://yann.lecun.com/exdb/mnist/t10k-labels-idx1-ubyte.gz")!)!
+    
+    let data:Data? = serialize(tensors:[
+        "train_images": trainImages,
+        "train_labels": trainLabels,
+        "test_images":  testImages,
+        "test_labels":  testLabels])
+    
+    print("saving to \"\(options["out"])\"")
+    try data!.write(to: URL(fileURLWithPath: options["out"]!))
+}
+
+

    }
}

func loadData(filename:String) -> Tensor<I>? {
    let data = NSData(contentsOfFile: filename)
    
    let header = [UInt32](repeating: 0, count: 4)
    data?.getBytes(UnsafeMutableRawPointer(mutating: header), length: 16)
    
    let magic = Int32(bigEndian: Int32(header[0]))
    let num_images = Int(Int32(bigEndian: Int32(header[1])))
    let num_rows = Int(Int32(bigEndian: Int32(header[2])))
    let num_cols = Int(Int32(bigEndian: Int32(header[3])))
    
    if magic != 2051 { return nil }
    
    let num_bytes = num_rows*num_cols*num_images
    let buffer = [UInt8](repeating: 0, count: Int(num_bytes))
    
    let ptr = UnsafeMutableRawPointer(mutating: buffer)
    data?.getBytes(ptr, range: NSRange(location: 16, length: num_bytes))
    
    let storage = I(array: buffer)
    return Tensor<I>(storage: storage, shape: Extent(num_images, num_rows, num_cols))
}

// TODO: get rid of index .. just do a subview of the tensor

import Foundation
import AppKit
import Tensor
import Gzip

typealias D = NativeStorage<Float>
typealias I = NativeStorage<UInt8>
