// Copyright 2019 Google Inc. All rights reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at:
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

import TensorFlowLiteCMetal

/// A delegate that uses the `Metal` framework for performing TensorFlow Lite graph operations with
/// GPU acceleration.
///
/// - Important: This is an experimental interface that is subject to change.
public final class MetalDelegate: Delegate {
  /// The configuration options for the `MetalDelegate`.
  public let options: Options

  // Conformance to the `Delegate` protocol.
  public private(set) var cDelegate: CDelegate

  /// Creates a new instance configured with the given `options`.
  ///
  /// - Parameters:
  ///   - options: Configurations for the delegate. The default is a new instance of
  ///       `MetalDelegate.Options` with the default configuration values.
  public init(options: Options = Options()) {
    self.options = options
    var delegateOptions = TFLGpuDelegateOptions()
    delegateOptions.allow_precision_loss = options.isPrecisionLossAllowed
    delegateOptions.wait_type = options.waitType.cWaitType
    delegateOptions.enable_quantization = options.isQuantizationEnabled
    cDelegate = TFLGpuDelegateCreate(&delegateOptions)
  }

  deinit {
    TFLGpuDelegateDelete(cDelegate)
  }
}

extension MetalDelegate {
  /// Options for configuring the `MetalDelegate`.
  public struct Options: Equatable, Hashable {
    /// Indicates whether the GPU delegate allows precision loss, such as allowing `Float16`
    /// precision for a `Float32` computation. The default is `false`.
    public var isPrecisionLossAllowed = false

    @available(
      *, deprecated, message: "Deprecated since TensorFlow Lite 2.4",
      renamed: "isPrecisionLossAllowed"
    )
    public var allowsPrecisionLoss: Bool {
      get { return isPrecisionLossAllowed }
      set(value) { isPrecisionLossAllowed = value }
    }

    /// A type indicating how the current thread should wait for work on the GPU to complete. The
    /// default is `passive`.
    public var waitType: ThreadWaitType = .passive

    /// Indicates whether the GPU delegate allows execution of an 8-bit quantized model. The default
    /// is `true`.
    public var isQuantizationEnabled = true

    /// Creates a new instance with the default values.
    public init() {}
  }
}

/// A type indicating how the current thread should wait for work scheduled on the GPU to complete.
public enum ThreadWaitType: Equatable, Hashable {
  /// The thread does not wait for the work to complete. Useful when the output of the work is used
  /// with the GPU pipeline.
  case none
  /// The thread waits until the work is complete.
  case passive
  /// The thread waits for the work to complete with minimal latency, which may require additional
  /// CPU resources.
  case active
  /// The thread waits for the work while trying to prevent the GPU from going into sleep mode.
  case aggressive

  /// The C `TFLGpuDelegateWaitType` for the current `ThreadWaitType`.
  var cWaitType: TFLGpuDelegateWaitType {
    switch self {
    case .none:
      return TFLGpuDelegateWaitTypeDoNotWait
    case .passive:
      return TFLGpuDelegateWaitTypePassive
    case .active:
      return TFLGpuDelegateWaitTypeActive
    case .aggressive:
      return TFLGpuDelegateWaitTypeAggressive
    }
  }
}