text: /*
   * Copyright 2016 The BigDL 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.bigdl.utils.tf
  
  import java.nio.ByteOrder
  
  import com.intel.analytics.bigdl.nn._
  import com.intel.analytics.bigdl.nn.abstractnn.{AbstractModule, DataFormat}
  import com.intel.analytics.bigdl.tensor.Tensor
  import com.intel.analytics.bigdl.utils.T
  import Tensorflow._
  import BigDLToTensorflow._
  import org.tensorflow.framework.{DataType, NodeDef}
  
  import scala.collection.mutable.ArrayBuffer
  
  /**
   * Wrapper of logic to convert module to tensorflow node definition
   */
  trait BigDLToTensorflow {
  
    /**
     * Convert the module to a tensorflow nodedef
     * @return Mapped nodedef list, the first is the output node
     */
    def toTFDef(module: AbstractModule[_, _, _], inputs: Seq[NodeDef],
                byteOrder: ByteOrder): Seq[NodeDef]
  }
  
  object BigDLToTensorflow {
    /**
     * This method is just for test purpose. Do not use the bigdl.saveNHWC for real use case
     * @return
     */
    private[tf] def processSaveDim(dim: Int): Int = {
      if (System.getProperty("bigdl.enableNHWC", "false").toBoolean) {
        if (dim == 2) return 4
        if (dim == 3) return 2
        if (dim == 4) return 3
        dim
      } else {
        dim
      }
    }
  
    /**
     * This method is just for test purpose. Do not use the bigdl.enableNHWC for real use case
     * @return
     */
    private[tf] def getDataFormat(): TensorflowDataFormat = {
      if (System.getProperty("bigdl.enableNHWC", "false").toBoolean) {
        TensorflowDataFormat.NHWC
      } else {
        TensorflowDataFormat.NCHW
      }
    }
  }
  
  object InputToTF extends BigDLToTensorflow {
    override def toTFDef(module: AbstractModule[_, _, _], inputs: Seq[NodeDef],
                         byteOrder: ByteOrder): Seq[NodeDef] = {
      require(inputs.length == 1, "Input only accept one input")
  
      Seq(identity(inputs(0), module.getName()))
    }
  }
  
  object ReLUToTF extends BigDLToTensorflow {
    override def toTFDef(module: AbstractModule[_, _, _], inputs: Seq[NodeDef],
                         byteOrder: ByteOrder): Seq[NodeDef] = {
      require(inputs.length == 1, "Relu only accept one input")
  
      Seq(relu(inputs(0), module.getName()))
    }
  }
  
  object LinearToTF extends BigDLToTensorflow {
    override def toTFDef(module: AbstractModule[_, _, _], inputs: Seq[NodeDef],
                         byteOrder: ByteOrder): Seq[NodeDef] = {
      require(inputs.length == 1, "Linear only accept one input")
      val linear = module.asInstanceOf[Linear[_]]
      val weight = const(linear.weight.t().contiguous(), linear.getName() + "/weight", byteOrder)
      val weightReader = identity(weight, linear.getName() + "/weightReader")
      val mm = matmul(inputs(0), weightReader, linear.getName() + "/matmul")
      val bias = const(linear.bias, linear.getName() + "/bias", byteOrder)
      val biasReader = identity(bias, linear.getName() + "/biasReader")
      val addNode = add(mm, biasReader, linear.getName() + "/add")
      Seq(addNode, biasReader, bias, mm, weightReader, weight)
    }
  }
  
  object SpatialConvolutionToTF extends BigDLToTensorflow {
    override def toTFDef(module: AbstractModule[_, _, _], inputs: Seq[NodeDef],
                         byteOrder: ByteOrder): Seq[NodeDef] = {
      import scala.language.existentials
      require(inputs.length == 1, "SpatialConvolution only accept one input")
      val spatialConv = module.asInstanceOf[SpatialConvolution[_]]
      if (spatialConv.nGroup == 1) {
        val (dataFormat, filterTensor) = if (spatialConv.format == DataFormat.NCHW) {
          (TensorflowDataFormat.NCHW,
            spatialConv.weight.select(1, 1)
              .transpose(2, 3).transpose(3, 4)
              .transpose(1, 2).transpose(2, 3)
              .transpose(3, 4).contiguous())
        } else {
          (TensorflowDataFormat.NHWC, spatialConv.weight.select(1, 1))
        }
  
        val filter = const(filterTensor, spatialConv.getName() + "/filter", byteOrder)
        val filterReader = identity(filter, spatialConv.getName() + "/filterReader")
        val conv = conv2D(inputs(0), filterReader, spatialConv.strideW, spatialConv.strideH,
          spatialConv.kernelW, spatialConv.kernelH, spatialConv.padW, spatialConv.padH,
          dataFormat, spatialConv.getName() + "/conv2D")
        if (spatialConv.bias != null) {
          val bias = const(spatialConv.bias, spatialConv.getName() + "/bias", byteOrder)
          val biasReader = identity(bias, spatialConv.getName() + "/biasReader")
          val add = biasAdd(conv, biasReader, dataFormat,
            spatialConv.getName() + "/biasAdd")
          Seq(add, biasReader, bias, conv, filterReader, filter)
        } else {
          Seq(conv, filterReader, filter)
        }
      } else {
        require(spatialConv.format == DataFormat.NCHW, "Only NCHW support conv group")
        val nodes = new ArrayBuffer[NodeDef]()
        val splitDim = const(Tensor.scalar[Int](1), spatialConv.getName() + "/split_dim",
          ByteOrder.LITTLE_ENDIAN)
        val splits = split(splitDim, inputs(0), spatialConv.nGroup, spatialConv.getName() + "/split")
        nodes.append(splitDim)
        nodes.appendAll(splits)
        val axis = const(Tensor.scalar[Int](1), spatialConv.getName() + "/concat/axis",
          ByteOrder.LITTLE_ENDIAN)
        nodes.append(axis)
        val outputs = (0 until spatialConv.nGroup).map(g => {
          val filterTensor = spatialConv.weight.select(1, g + 1)
            .transpose(2, 3).transpose(3, 4)
            .transpose(1, 2).transpose(2, 3)
            .transpose(3, 4).contiguous()
  
          val filter = const(filterTensor, spatialConv.getName() + s"/group$g/filter", byteOrder)
          val filterReader = identity(filter, spatialConv.getName() + s"/group$g/filterReader")
          val conv = conv2D(splits(g), filterReader, spatialConv.strideW, spatialConv.strideH,
            spatialConv.kernelW, spatialConv.kernelH, spatialConv.padW, spatialConv.padH,
            TensorflowDataFormat.NCHW, spatialConv.getName() + s"/group$g/conv2D")
          if (spatialConv.bias != null) {
            val bias = const(spatialConv.bias.narrow(1,
              g * spatialConv.nOutputPlane / spatialConv.nGroup + 1,
              spatialConv.nOutputPlane / spatialConv.nGroup),
              spatialConv.getName() + s"/group$g/bias", byteOrder)
            val biasReader = identity(bias, spatialConv.getName() + s"/group$g/biasReader")
            val add = biasAdd(conv, biasReader, TensorflowDataFormat.NCHW,
              spatialConv.getName() + s"/group$g/biasAdd")
            nodes.append(add, biasReader, bias, conv, filterReader, filter)
            add
          } else {
            nodes.append(conv, filterReader, filter)
            conv
          }
        }) ++ Seq(axis)
  
        val concatNode = concat(outputs, spatialConv.getName() + "/concat/output")
        Seq(concatNode) ++ nodes
      }
    }
  }
  
  object TemporalConvolutionToTF extends BigDLToTensorflow {
    override def toTFDef(module: AbstractModule[_, _, _], inputs: Seq[NodeDef],
                         byteOrder: ByteOrder): Seq[NodeDef] = {
      require(inputs.length == 1, "SpatialConvolution only accept one input")
      val spatialConv = module.asInstanceOf[TemporalConvolution[_]]
  
      val const1 = const(Tensor.scalar[Int](1), spatialConv.getName() + "/dim1", byteOrder)
      val expandDimsInput = expandDims(inputs.head, const1,
        spatialConv.getName() + "/expandDimsInput")
  
      val filterTensor = spatialConv.weight
        .view(spatialConv.outputFrameSize, spatialConv.kernelW, spatialConv.inputFrameSize)
        .transpose(2, 3).transpose(1, 3).contiguous()
  
      val filter = const(filterTensor, spatialConv.getName() + "/filter", byteOrder)
      val filterReader = identity(filter, spatialConv.getName() + "/filterReader")
  
      val const2 = const(Tensor.scalar[Int](0), spatialConv.getName() + "/dim2", byteOrder)
      val expandDimsWeight = expandDims(filterReader, const2,
        spatialConv.getName() + "/expandDimsWeight")
  
      val conv = conv2D(expandDimsInput, expandDimsWeight, spatialConv.strideW, 1,
        spatialConv.kernelW, 1, 0, 0,
        getDataFormat(), spatialConv.getName() + "/conv2D")
  
      val sq = squeeze(conv, Seq(1), spatialConv.getName() + "/squeeze")
  
      val bias = const(spatialConv.bias, spatialConv.getName() + "/bias", byteOrder)
      val biasReader = identity(bias, spatialConv.getName() + "/biasReader")
      val add = biasAdd(sq, biasReader, getDataFormat(),
        spatialConv.getName() + "/biasAdd")
      Seq(add, biasReader, bias, conv, filterReader, filter, sq,
        expandDimsInput, expandDimsWeight, const1, const2)
    }
  }
  
  object SqueezeToTF extends BigDLToTensorflow {
    override def toTFDef(module: AbstractModule[_, _, _], inputs: Seq[NodeDef],
                         byteOrder: ByteOrder): Seq[NodeDef] = {
      require(inputs.length == 1, "Squeeze only accept one input")
      val sq = module.asInstanceOf[Squeeze[_]]
      Seq(squeeze(inputs(0), sq.dims.map(processSaveDim(_) - 1), sq.getName()))
    }
  }
  
  object TanhToTF extends BigDLToTensorflow {
    override def toTFDef(module: AbstractModule[_, _, _], inputs: Seq[NodeDef],
                         byteOrder: ByteOrder): Seq[NodeDef] = {
      require(inputs.length == 1, "Tanh only accept one input")
      Seq(tanh(inputs(0), module.getName()))
    }
  }
  
  object ReshapeToTF extends BigDLToTensorflow {
    override def toTFDef(module: AbstractModule[_, _, _], inputs: Seq[NodeDef],
                         byteOrder: ByteOrder): Seq[NodeDef] = {
      require(inputs.length == 1, "Reshape only accept one input")
      val rh = module.asInstanceOf[Reshape[_]]
      val size = Tensor[Int](rh.size.length)
      var i = 0
      while(i < rh.size.length) {
        size.setValue(i + 1, rh.size(i))
        i += 1
      }
      val shape = const(size, rh.getName() + "/shape", byteOrder)
      val reshapeNode = reshape(inputs(0), shape, rh.getName())
      Seq(reshapeNode, shape)
    }
  }
  
  object ViewToTF extends BigDLToTensorflow {
    override def toTFDef(module: AbstractModule[_, _, _], inputs: Seq[NodeDef],
                         byteOrder: ByteOrder): Seq[NodeDef] = {
      require(inputs.length == 1, "Reshape only accept one input")
      val viewLayer = module.asInstanceOf[View[_]]
      val size = Tensor[Int](viewLayer.sizes.length + 1).setValue(1, -1)
      var i = 1
      while(i < viewLayer.sizes.length + 1) {
        size.setValue(i + 1, viewLayer.sizes(i - 1))
        i += 1
      }
  
      val shape = const(size, viewLayer.getName() + "/shape", byteOrder)
      val reshapeNode = reshape(inputs(0), shape, viewLayer.getName())
      Seq(reshapeNode, shape)
    }
  }
  
  object MaxpoolToTF extends BigDLToTensorflow {
    override def toTFDef(module: AbstractModule[_, _, _], inputs: Seq[NodeDef],
                         byteOrder: ByteOrder): Seq[NodeDef] = {
      require(inputs.length == 1, "Maxpool only accept one input")
      val layer = module.asInstanceOf[SpatialMaxPooling[_]]
      val dataFormat = if (layer.format == DataFormat.NHWC) {
        TensorflowDataFormat.NHWC
      } else {
        TensorflowDataFormat.NCHW
      }
      Seq(maxPool(inputs(0), layer.kW, layer.kH, layer.padW, layer.padH,
        layer.dW, layer.dH, dataFormat, layer.getName(), layer.ceilMode))
    }
  }
  
  object PaddingToTF extends BigDLToTensorflow {
    override def toTFDef(module: AbstractModule[_, _, _], inputs: Seq[NodeDef],
                         byteOrder: ByteOrder): Seq[NodeDef] = {
      require(inputs.length == 1, "Padding only accept one input")
      val layer = module.asInstanceOf[Padding[_]]
      require(layer.nIndex == 1, "only support padding nIndex == 1")
      require(layer.nInputDim > 0, "nInputDim must be explicit specified")
      val padding = Tensor[Int](layer.nInputDim, 2).zero()
      if (layer.pad < 0) {
        padding.setValue(layer.dim, 1, -layer.pad)
      }
      else {
        padding.setValue(layer.dim, 2, layer.pad)
      }
      val paddingsNode = const(padding, layer.getName() + "/padding", byteOrder)
      val padNode = pad(inputs(0), paddingsNode, layer.getName() + "/output")
      Seq(padNode, paddingsNode)
    }
  }
  
  object AvgpoolToTF extends BigDLToTensorflow {
    override def toTFDef(module: AbstractModule[_, _, _], inputs: Seq[NodeDef],
                         byteOrder: ByteOrder): Seq[NodeDef] = {
      require(inputs.length == 1, "Avgpool only accept one input")
      val layer = module.asInstanceOf[SpatialAveragePooling[_]]
      val dataFormat = if (layer.format == DataFormat.NHWC) {
        TensorflowDataFormat.NHWC
      } else {
        TensorflowDataFormat.NCHW
      }
      Seq(avgPool(inputs(0), layer.kW, layer.kH, layer.padW, layer.padH,
        layer.dW, layer.dH, dataFormat, layer.getName(), layer.ceilMode))
    }
  }
  
  object SigmoidToTF extends BigDLToTensorflow {
    override def toTFDef(module: AbstractModule[_, _, _], inputs: Seq[NodeDef],
                         byteOrder: ByteOrder): Seq[NodeDef] = {
      require(inputs.length == 1, "Sigmoid only accept one input")
      Seq(sigmoid(inputs(0), module.getName()))
    }
  }
  
  object DropoutToTF extends BigDLToTensorflow {
    override def toTFDef(module: AbstractModule[_, _, _], inputs: Seq[NodeDef],
                         byteOrder: ByteOrder): Seq[NodeDef] = {
      require(inputs.length == 1, "Dropout only accept one input")
      val layer = module.asInstanceOf[Dropout[_]]
      require(layer.isTraining() == false, "only support evaluating mode dropout")
      require(inputs.length == 1, "require only one tensor input")
      Seq(identity(inputs(0), layer.getName()))
    }
  }
  
  object ScaleToTF extends BigDLToTensorflow {
    override def toTFDef(module: AbstractModule[_, _, _], inputs: Seq[NodeDef],
      byteOrder: ByteOrder): Seq[NodeDef] = {
      val layer = module.asInstanceOf[Scale[_]]
      val weight = const(layer.cmul.weight, layer.getName() + "/mul/weight", ByteOrder.LITTLE_ENDIAN)
      val mulNode = multiply(weight, inputs(0), layer.getName() + "/mul/mul")
      val bias = const(layer.cadd.bias, layer.getName() + "/add/bias", ByteOrder.LITTLE_ENDIAN)
      val output = add(mulNode, bias, layer.getName() + "/add/add")
      Seq(output, bias, mulNode, weight)
    }
  }
  
  object CAddTableToTF extends BigDLToTensorflow {
    override def toTFDef(module: AbstractModule[_, _, _], inputs: Seq[NodeDef],
                         byteOrder: ByteOrder): Seq[NodeDef] = {
      Seq(addN(inputs, module.getName()))
    }
  }
  
  object CMultTableToTF extends BigDLToTensorflow {
    override def toTFDef(module: AbstractModule[_, _, _], inputs: Seq[NodeDef],
                         byteOrder: ByteOrder): Seq[NodeDef] = {
      require(inputs.length == 2, "Tensorflow only support two tensor multiply together")
  
      Seq(multiply(inputs(0), inputs(1), module.getName()))
    }
  }
  
  object JoinTableToTF extends BigDLToTensorflow {
    override def toTFDef(module: AbstractModule[_, _, _], inputs: Seq[NodeDef],
                         byteOrder: ByteOrder): Seq[NodeDef] = {
      val layer = module.asInstanceOf[JoinTable[_]]
      val axis = const(Tensor.scalar[Int](layer.dimension - 1), layer.getName() + "/axis", byteOrder)
      val updateInputs = new ArrayBuffer[NodeDef]()
      updateInputs ++= inputs.reverse
      updateInputs.append(axis)
      Seq(concat(updateInputs, layer.getName()), axis)
    }
  }
  
  object MeanToTF extends BigDLToTensorflow {
    override def toTFDef(module: AbstractModule[_, _, _], inputs: Seq[NodeDef],
                         byteOrder: ByteOrder): Seq[NodeDef] = {
      require(inputs.length == 1, "Mean only accept one input")
      val layer = module.asInstanceOf[Mean[_]]
      require(layer.squeeze == true, "Mean must squeeze input")
      val dimsTensor = Tensor[Int](layer.dimension)
      dimsTensor.setValue(1, layer.dimension - 1)
  
      val dims = const(dimsTensor, layer.getName() + "/dims", byteOrder)
      val mean = reduceMean(inputs(0), dims, false, layer.getName() + "/output")
      Seq(mean, dims)
    }
  }
  
  object SoftMaxToTF extends BigDLToTensorflow {
    override def toTFDef(module: AbstractModule[_, _, _], inputs: Seq[NodeDef],
                         byteOrder: ByteOrder): Seq[NodeDef] = {
      require(inputs.length == 1, "Softmax only accept one input")
      Seq(softmax(inputs(0), module.getName()))
    }
  }
  
  object LogSoftMaxToTF extends BigDLToTensorflow {
    override def toTFDef(module: AbstractModule[_, _, _], inputs: Seq[NodeDef],
                         byteOrder: ByteOrder): Seq[NodeDef] = {
      require(inputs.length == 1, "LogSoftmax only accept one input")
      Seq(logSoftmax(inputs(0), module.getName()))
    }
  }
  
  object BatchNorm2DToTF extends BigDLToTensorflow {
    override def toTFDef(module: AbstractModule[_, _, _], inputs: Seq[NodeDef],
                         byteOrder: ByteOrder): Seq[NodeDef] = {
      require(inputs.length == 1, "BatchNorm only accept one input")
      val layer = module.asInstanceOf[SpatialBatchNormalization[_]]
      require(!layer.isTraining(), "Only support evaluate mode batch norm")
      // reshape to nchw
      val size = Tensor[Int](layer.nDim)
      for (i <- 0 until layer.nDim) {
        size.setValue(i + 1, 1)
      }
  
      size(2) = layer.runningVar.size(1)
      if (layer.weight != null) {
        val shapeVar = const(size, layer.getName() + "/reshape_1/shape", byteOrder)
        val shapeMean = const(size, layer.getName() + "/reshape_2/shape", byteOrder)
        val shapeScale = const(size, layer.getName() + "/reshape_3/shape", byteOrder)
        val shapeOffset = const(size, layer.getName() + "/reshape_4/shape", byteOrder)
  
        val varNode = const(layer.runningVar, layer.getName() + "/var", byteOrder)
        val mean = const(layer.runningMean, layer.getName() + "/mean", byteOrder)
        val scale = const(layer.weight, layer.getName() + "/scale", byteOrder)
        val offset = const(layer.bias, layer.getName() + "/offset", byteOrder)
        val reshapeVar = reshape(varNode, shapeVar, s"${layer.getName()}/reshape_1")
        val reshapeMean = reshape(mean, shapeMean, s"${layer.getName()}/reshape_2")
        val reshapeScale = reshape(scale, shapeScale, s"${layer.getName()}/reshape_3")
        val reshapeOffset = reshape(offset, shapeOffset, s"${layer.getName()}/reshape_4")
        // construct graph
        val sqrtVar = rsqrt(reshapeVar, layer.getName() + "/sqrtvar")
        val mul0 = multiply(reshapeScale, sqrtVar, layer.getName() + "/mul0")
        val mul1 = multiply(inputs(0), mul0, layer.getName() + "/mul1")
        val mul2 = multiply(reshapeMean, mul0, layer.getName() + "/mul2")
        val sub = subtract(reshapeOffset, mul2, layer.getName() + "/sub")
        val output = add(mul1, sub, layer.getName() + "/output")
        Seq(output, sub, mul2, mul1, mul0, reshapeOffset, reshapeMean, reshapeScale,
          shapeOffset, shapeMean, shapeScale, offset, scale, mean,
          sqrtVar, reshapeVar, shapeVar, varNode)
      } else {
        val shapeVar = const(size, layer.getName() + "/reshape_1/shape", byteOrder)
        val shapeMean = const(size, layer.getName() + "/reshape_2/shape", byteOrder)
  
        val varNode = const(layer.runningVar, layer.getName() + "/var", byteOrder)
        val mean = const(layer.runningMean, layer.getName() + "/mean", byteOrder)
        val reshapeVar = reshape(varNode, shapeVar, s"${layer.getName()}/reshape_1")
        val reshapeMean = reshape(mean, shapeMean, s"${layer.getName()}/reshape_2")
        // construct graph
        val sqrtVar = rsqrt(reshapeVar, layer.getName() + "/sqrtvar")
        val mul1 = multiply(inputs(0), sqrtVar, layer.getName() + "/mul1")
        val mul2 = multiply(reshapeMean, sqrtVar, layer.getName() + "/mul2")
        val output = subtract(mul1, mul2, layer.getName() + "/output")
        Seq(output, mul2, mul1, reshapeMean, shapeMean, mean, sqrtVar, reshapeVar, shapeVar, varNode)
      }
    }
  }
  
  object LRNToTF extends BigDLToTensorflow {
    override def toTFDef(module: AbstractModule[_, _, _], inputs: Seq[NodeDef],
      byteOrder: ByteOrder): Seq[NodeDef] = {
      val layer = module.asInstanceOf[SpatialCrossMapLRN[_]]
      if (layer.format == DataFormat.NHWC) {
        Seq(lrn(inputs(0), (layer.size - 1) / 2, layer.k.toFloat, (layer.alpha / layer.size).toFloat,
          layer.beta.toFloat, module.getName()))
      } else {
        val perm1 = const(Tensor[Int](T(0, 2, 3, 1)), module.getName() + "/perm1",
          ByteOrder.LITTLE_ENDIAN)
        val transpose1 = transpose(inputs(0), perm1, module.getName() + "/transpose1")
        val lrnNode = lrn(transpose1, (layer.size - 1) / 2, layer.k.toFloat,
          (layer.alpha / layer.size).toFloat,
          layer.beta.toFloat, module.getName() + "/lrn")
        val perm2 = const(Tensor[Int](T(0, 3, 1, 2)), module.getName() + "/perm2",
          ByteOrder.LITTLE_ENDIAN)
        val output = transpose(lrnNode, perm2, module.getName() + "/transpose2")
        Seq(output, perm1, transpose1, lrnNode, perm2)
      }
    }
  }
  