/*
 * Copyright 2018 Analytics Zoo Authors.
 *
 * 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.
 */

package com.intel.analytics.zoo.models.image.objectdetection.ssd

import com.intel.analytics.bigdl.dataset.{MiniBatch, Sample}
import com.intel.analytics.bigdl.tensor.Tensor
import com.intel.analytics.bigdl.tensor.TensorNumericMath.TensorNumeric
import com.intel.analytics.bigdl.transform.vision.image.ImageFeature
import com.intel.analytics.zoo.models.image.objectdetection.common.BboxUtil

/**
 * A batch of data feed into the model. The first size is batchsize
 *
 * @param input
 * @param target
 */
class SSDMiniBatch(val input: Tensor[Float], val target: Tensor[Float],
                   val imInfo: Tensor[Float] = null)
  extends MiniBatch[Float] {

  private val targetIndices = if (target != null) BboxUtil.getGroundTruthIndices(target) else null
  var imageFeatures: Array[ImageFeature] = _

  override def size(): Int = {
    input.size(1)
  }

  override def getInput(): Tensor[Float] = {
    input
  }

  override def getTarget(): Tensor[Float] = {
    target
  }

  override def slice(offset: Int, length: Int): MiniBatch[Float] = {
    val subInput = input.narrow(1, offset, length)
    val subTarget = if (target != null) {
      var i = 0
      val targetOffset = targetIndices(offset - 1)._1
      var targetLength = 0
      while (i < length) {
        targetLength += targetIndices(offset + i - 1)._2
        i += 1
      }
      target.narrow(1, targetOffset, targetLength)
    } else null

    SSDMiniBatch(subInput, subTarget)
  }

  override def set(samples: Seq[Sample[Float]])(implicit ev: TensorNumeric[Float])
  : SSDMiniBatch.this.type = {
    throw new NotImplementedError("do not use Sample here")
  }
}

object SSDMiniBatch {
  def apply(data: Tensor[Float], labels: Tensor[Float], imInfo: Tensor[Float] = null):
  SSDMiniBatch = new SSDMiniBatch(data, labels, imInfo)
}
