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

import TensorFlow
import Python
import Foundation

let batchSize: Int = 200
let maxEpochs: Int = 4

let (trainDataset, testDataset) = loadDatasets()
let testBatches = testDataset.batched(1000)

var model = ECGModel()
let optimizer = Adam(for: model, learningRate: 0.001, decay: 0)

// Training loop
for epoch in 1...maxEpochs {
    print("Epoch \(epoch), training...")
    
    var trainingLossSum: Float = 0
    var trainingBatchCount = 0
    let trainingShuffled = trainDataset.shuffled(sampleCount: 500000, randomSeed: Int64(epoch))
    let t0 = Date()
    
    // Loop over mini-batches in training set
    for batch in trainingShuffled.batched(batchSize) {
        let gradients = gradient(at: model) {
            (model: ECGModel) -> Tensor<Float> in
            
            let thisLoss = loss(model: model, examples: batch)
            trainingLossSum += thisLoss.scalarized()
            trainingBatchCount += 1
            return thisLoss
        }
        optimizer.update(&model.allDifferentiableVariables, along: gradients)
    }
    
    let t1 = Date()
    print("  training loss: \(trainingLossSum / Float(trainingBatchCount))  step: \(trainingBatchCount * epoch) (\(t1.timeIntervalSince(t0)) sec)")
    
    var testLossSum: Float = 0
    var testBatchCount = 0
    
    // Loop over test set
    for batch in testBatches {
        testLossSum += loss(model: model, examples: batch).scalarized()
        testBatchCount += 1
    }
    print("  test loss: \(testLossSum / Float(testBatchCount))")
}

// Print metrics and confusion matrix
var yActual = [Int32]()
var yPredicted = [Int32]()

for batch in testBatches {
    let labelValues = batch.labels.scalars
    let predictedValues = model.predictedClasses(for: batch.series).scalars
    yActual.append(contentsOf: labelValues)
    yPredicted.append(contentsOf: predictedValues)
}

let skm = Python.import("sklearn.metrics")
let report = skm.classification_report(yActual, yPredicted)
print(report)
let confusionMatrix = skm.confusion_matrix(yActual, yPredicted)
print(confusionMatrix)