-//===-- Dataset.swift -----------------------------------------*- swift -*-===//
+// Copyright 2018 The TensorFlow Authors. All Rights Reserved.
 //
-// This source file is part of the Swift.org open source project
+// 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
 //
-// Copyright (c) 2014 - 2017 Apple Inc. and the Swift project authors
-// Licensed under Apache License v2.0 with Runtime Library Exception
+//     http://www.apache.org/licenses/LICENSE-2.0
 //
-// See https://swift.org/LICENSE.txt for license information
-// See https://swift.org/CONTRIBUTORS.txt for the list of Swift project authors
-//
-//===----------------------------------------------------------------------===//
-//
-// The dataset API.
-//
-//===----------------------------------------------------------------------===//
+// 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.
 
 /// The default graph seed.
 ///

// model TF's "global seed".
@usableFromInline @inline(__always)
func _tensorSeeds(_ seed: Tensor<Int64>) -> (Tensor<Int64>, Tensor<Int64>) {
  return (Tensor(_defaultGraphSeed), seed)
}

//===----------------------------------------------------------------------===//

/// element tensors.
@_fixed_layout
public struct Dataset<Element : TensorGroup> {
    public let _handle: VariantHandle

    @inlinable
    public init(_handle: VariantHandle) {
        self._handle = _handle
    }
}

public extension Dataset {
    @inlinable
    init(randomSeed: Int64) {
        let (seed1, seed2) = _tensorSeeds(Tensor(randomSeed))
        self.init(_handle: Raw.experimentalRandomDataset(
            seed: seed1,
            seed2: seed2,
            outputTypes: Element._typeList,
            outputShapes: Element._unknownShapeList))
    }
}

public extension Dataset {
    /// Creates a dataset from a batch of elements as a tensor.
    @inlinable
    init(elements: Element) {
        self.init(_handle: Raw.tensorSliceDataset(
            components: [elements],
            outputShapes: Element._unknownShapeList))
    }
}

extension Dataset : Sequence {
    public typealias Iterator = DatasetIterator<Element>

    /// Returns an iterator over the elements of this dataset.
    @inlinable
    public func makeIterator() -> DatasetIterator<Element> {
        let resource = Raw.anonymousIterator(
            outputTypes: Element._typeList,
            outputShapes: Element._unknownShapeList)
        Raw.makeIterator(dataset: _handle, iterator: resource)
        return DatasetIterator(_handle: resource)
    }
}

public extension Dataset {
    // Note that this Dataset API implementation uses an experimental tracing
    // feature, which is not robust and does not have great diagnostics yet.
    @inlinable
    func map<ResultElement : TensorGroup>(
        _ transform: (Element) -> ResultElement
    ) -> Dataset<ResultElement> {
        return Dataset<ResultElement>(_handle: Raw.mapDataset(
            inputDataset: _handle,
            otherArguments: Tensor<Int32>(0),
            f: transform,
            outputTypes: ResultElement._typeList,
            outputShapes: ResultElement._unknownShapeList,
            useInterOpParallelism: true,
            preserveCardinality: false))
    }

    @inlinable
    func map<ResultElement : TensorGroup>(
        parallelCallCount: Int,
        _ transform: (Element) -> ResultElement
    ) -> Dataset<ResultElement> {
        return Dataset<ResultElement>(_handle: Raw.parallelMapDataset(
            inputDataset: _handle,
            otherArguments: Tensor<Int32>(0),
            numParallelCalls: Tensor<Int32>(Int32(parallelCallCount)),
            f: transform,
            outputTypes: ResultElement._typeList,
            outputShapes: ResultElement._unknownShapeList,
            useInterOpParallelism: true,
            sloppy: false,
            preserveCardinality: false))
    }

    @inlinable
    func filter(
        _ isIncluded: (Element) -> Tensor<Bool>
    ) -> Dataset {
        return Dataset(_handle: Raw.filterDataset(
            inputDataset: _handle,
            otherArguments: Tensor<Int32>(0),
            predicate: isIncluded,
            outputTypes: Element._typeList,
            outputShapes: Element._unknownShapeList))
    }
}

public extension Dataset {
    @inlinable
    func shuffled(
        sampleCount: Int, randomSeed: Int64
    ) -> Dataset {
        let (seed1, seed2) = _tensorSeeds(Tensor(randomSeed))
        return Dataset(_handle: Raw.shuffleDataset(
            inputDataset: _handle,
            bufferSize: Tensor(Int64(sampleCount)),
            seed: seed1,
            seed2: seed2,
            outputTypes: Element._typeList,
            outputShapes: Element._unknownShapeList))
    }

    @inlinable
    func batched(_ batchSize: Int) -> Dataset {
        return Dataset(_handle: Raw.batchDataset(
            inputDataset: _handle,
            batchSize: Tensor(Int64(batchSize)),
            outputTypes: Element._typeList,
            outputShapes: Element._unknownShapeList))
    }
}

/// The type that allows iteration over a dataset's elements.
@_fixed_layout
public struct DatasetIterator<Element : TensorGroup> {
    @usableFromInline let _handle: ResourceHandle

    @usableFromInline
    internal init(_handle: ResourceHandle) {
        self._handle = _handle
    }
}

extension DatasetIterator : IteratorProtocol {
    /// Advances to the next element and returns it, or `nil` if no next element
    /// exists.
    @inlinable
    public mutating func next() -> Element? {
        let optional = Raw.iteratorGetNextAsOptional(
            iterator: _handle,
            outputTypes: Element._typeList,
            outputShapes: Element._unknownShapeList)
        guard Raw.optionalHasValue(optional: optional).scalarized() else {
            return nil
        }
        return Raw.optionalGetValue(
            optional: optional,
            outputShapes: Element._unknownShapeList)
    }
}

/// A 2-tuple-like struct that conforms to TensorGroup that represents a tuple
/// of 2 types conforming to TensorGroup.
@_fixed_layout
public struct Zip2TensorGroup<T : TensorGroup, U : TensorGroup> : TensorGroup {
    public var first: T
    public var second: U

    public init(_ first: T, _ second: U) {
        self.first = first
        self.second = second
    }
}

@inlinable
public func zip<T : TensorGroup, U : TensorGroup>(
    _ dataset1: Dataset<T>, _ dataset2: Dataset<U>
) -> Dataset<Zip2TensorGroup<T, U>> {
    let handle = Raw.zipDataset(
        inputDatasets: [dataset1._handle, dataset2._handle],
        outputTypes: Zip2TensorGroup<T, U>._typeList,
        outputShapes: Zip2TensorGroup<T, U>._unknownShapeList)
    return Dataset(_handle: handle)
}
