import Foundation

import Glibc

dlopen("/home/ubuntu/swift/usr/lib/swift/linux/libPath.so", RTLD_NOW)

import Path

print(Path.home)

dlopen("/home/ubuntu/swift/usr/lib/swift/linux/libJust.so", RTLD_NOW)

import Just

func shell(_ launchPath: String, _ arguments: [String]) -> String?
{
    let task = Process()
    task.launchPath = launchPath
    task.arguments = arguments

    let pipe = Pipe()
    task.standardOutput = pipe
    task.launch()

    let data = pipe.fileHandleForReading.readDataToEndOfFile()
    let output = String(data: data, encoding: String.Encoding.utf8)

    return output
}

var from = "http://yann.lecun.com/exdb/mnist"

func get_file(_ from:String, _ fn:String) {
    let cts = Just.get("\(from)/\(fn).gz").content!
    try! cts.write(to: URL.init(fileURLWithPath: "\(fn).gz"))
    shell("/bin/gunzip", ["-fq", (Path.cwd/"\(fn).gz").string])
}

protocol ConvertableFromByte {
    init(_ d:UInt8)
}

extension Float : ConvertableFromByte{}
extension Int32 : ConvertableFromByte{}

import TensorFlow

func get_data<T:ConvertableFromByte & TensorFlowScalar>(_ fn:String, _ skip:Int) -> Tensor<T> {
    let data = try! Data.init(contentsOf: URL.init(fileURLWithPath: fn)).dropFirst(skip)
    return Tensor(data.map(T.init))
}

var trn_fn = "train-images-idx3-ubyte"
var trn_lab_fn = "train-labels-idx1-ubyte"
var val_fn = "t10k-images-idx3-ubyte"
var val_lab_fn = "t10k-labels-idx1-ubyte"

get_file(from, trn_fn)
get_file(from, trn_lab_fn)

get_file(from, val_fn)
get_file(from, val_lab_fn)

var trn_arr:Tensor<Float> = get_data(trn_fn, 16)/255.0
var trn_lab:Tensor<Int32> = get_data(trn_lab_fn, 8)

var val_arr:Tensor<Float> = get_data(val_fn, 16)/255.0
var val_lab:Tensor<Int32> = get_data(val_lab_fn, 8)

print(val_lab[100..<300])

val_lab.scalarCount