//
//  loadDatasets.swift
//  ECG
//
//  Created by Dave Fernandes on 2019-03-05.
//  Copyright © 2019 MintLeaf Software Inc. All rights reserved.
//

import Python
import TensorFlow

struct Example: TensorGroup {
    var labels: Tensor<Int32>
    var series: Tensor<Float>
}

func loadTimeSeries(from path: String) -> Example {
    let pickle = Python.import("pickle")
    let file = Python.open(path, "rb")
    let pyDict = pickle.load(file, encoding: "bytes")
    let dict = Dictionary<String, PythonObject>(pyDict)

    guard let series = dict?["x"],
        let labels = dict?["y"] else {
        fatalError()
    }

    let labelsTensor = Tensor<Int64>(numpy: labels)!
    let seriesTensor = Tensor<Float64>(numpy: series)!
    return Example(labels: Tensor<Int32>(labelsTensor), series: Tensor<Float32>(seriesTensor))
}

func loadDatasets() -> (training: Dataset<Example>, test: Dataset<Example>) {
    let trainingDataset = Dataset<Example>(elements: loadTimeSeries(from: "train_set.pickle"))
    let testDataset = Dataset<Example>(elements: loadTimeSeries(from: "test_set.pickle"))
    return (training: trainingDataset, test: testDataset)
}