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

// Model is from: https://arxiv.org/pdf/1805.00794.pdf

import TensorFlow

public struct ConvUnit<Scalar: TensorFlowFloatingPoint> : Layer {
    var conv1: Conv1D<Scalar>
    var conv2: Conv1D<Scalar>
    var pool: MaxPool1D<Scalar>
    
    public init(kernelSize: Int, channels: Int) {
        conv1 = Conv1D<Scalar>(filterShape: (kernelSize, channels, channels), padding: .same, activation: relu)
        conv2 = Conv1D<Scalar>(filterShape: (kernelSize, channels, channels), padding: .same)
        pool = MaxPool1D<Scalar>(poolSize: kernelSize, stride: 2, padding: .valid)
    }
    
    @differentiable
    public func callAsFunction(_ input: Tensor<Scalar>) -> Tensor<Scalar> {
        var tmp = input.sequenced(through: conv1, conv2)
        tmp = pool(relu(tmp + input))
        return tmp
    }
}

public struct ConvModel : Layer {
    var conv1: Conv1D<Float>
    var convUnit = [ConvUnit<Float>]()
    var dense1: Dense<Float>
    var dense2: Dense<Float>
    
    @noDerivative let convUnitCount = 5
    
    public init() {
        conv1 = Conv1D<Float>(filterShape: (5, 1, 32), stride: 1, padding: .same)
        for _ in 0..<convUnitCount {
            convUnit.append(ConvUnit<Float>(kernelSize: 5, channels: 32))
        }
        dense1 = Dense<Float>(inputSize: 64, outputSize: 32, activation: relu)
        dense2 = Dense<Float>(inputSize: 32, outputSize: 5)
    }
    
    @differentiable
    public func callAsFunction(_ input: Tensor<Float>) -> Tensor<Float> {
        var tmp = conv1(input.expandingShape(at: 2))
        
        for i in 0..<convUnitCount {
            let unit = convUnit[i]
            tmp = unit(tmp)
        }
        
        tmp = tmp.reshaped(to: [-1, 64])
        tmp = tmp.sequenced(through: dense1, dense2)
        return tmp
    }
    
    public func predictedClasses(for input: Tensor<Float>) -> Tensor<Int32> {
        return model.inferring(from: input).argmax(squeezingAxis: 1)
    }
}

typealias ECGModel = ConvModel

@differentiable(wrt: model)
func loss(model: ECGModel, examples: Example) -> Tensor<Float> {
    let logits = model(examples.series)
    return softmaxCrossEntropy(logits: logits, labels: examples.labels)
}