// This file is part of Eigen, a lightweight C++ template library
// for linear algebra.
//
// Copyright (C) 2014 Benoit Steiner <benoit.steiner.goog@gmail.com>
//
// This Source Code Form is subject to the terms of the Mozilla
// Public License v. 2.0. If a copy of the MPL was not distributed
// with this file, You can obtain one at http://mozilla.org/MPL/2.0/.
#ifndef EIGEN_CXX11_NEURAL_NETWORKS_SOFTMAX_H
#define EIGEN_CXX11_NEURAL_NETWORKS_SOFTMAX_H

namespace Eigen {

/** SoftMax
  * \ingroup CXX11_NeuralNetworks_Module
  *
  * \brief Applies a softmax
  *
  * The input parameter is expected to be a col-major tensor with a rank of 2 (depth and other).
  *
  * The result can be assigned to a tensor of rank and dimensions equal to that of the input. The result will be laid out in col-major order.
  *
*/

namespace {
struct SoftmaxOp {
  SoftmaxOp(const float beta) : beta_(beta) { }

  template <typename Input>
  typename Input::Dimensions dimensions(const Input& input) const {
    return input.dimensions();
  }

  template <typename Input, typename Output, typename Device>
  void eval(const Input& input, Output& output, const Device& device) const
  {
#if !defined(EIGEN_HAS_INDEX_LIST)
    // nvcc doesn't support cxx11
    Eigen::array<typename internal::traits<Input>::Index, 1> depth_dim;
    depth_dim[0] = 0;
    Eigen::array<typename internal::traits<Input>::Index, 2> bcast;
    bcast[0] = dimensions(input)[0];
    bcast[1] = 1;
    DSizes<typename internal::traits<Input>::Index, 2> dims2d;
    dims2d[0] = 1;
    dims2d[1] = dimensions(input)[1];
#else
    // Take advantage of cxx11 to give the compiler information it can use to
    // optimize the code.
    Eigen::IndexList<Eigen::type2index<0>> depth_dim;
    Eigen::IndexList<int, Eigen::type2index<1>> bcast;
    bcast.set(0, dimensions(input)[0]);
    Eigen::IndexList<Eigen::type2index<1>, typename internal::traits<Input>::Index> dims2d;
    dims2d.set(1, dimensions(input)[1]);
#endif

    output.device(device) = ((input - input.maximum(depth_dim).eval().reshape(dims2d).broadcast(bcast)) * beta_).exp();
    output.device(device) = output / (output.sum(depth_dim).eval().reshape(dims2d).broadcast(bcast));
  }

 private:
  const float beta_;
};
}


template <typename Input>
EIGEN_ALWAYS_INLINE
static const TensorCustomUnaryOp<const SoftmaxOp, const Input>
SoftMax(const Input& input, const float beta)
{
  EIGEN_STATIC_ASSERT(internal::traits<Input>::Layout == ColMajor, YOU_MADE_A_PROGRAMMING_MISTAKE);
  EIGEN_STATIC_ASSERT(internal::traits<Input>::NumDimensions == 2, YOU_MADE_A_PROGRAMMING_MISTAKE);

  const SoftmaxOp op(beta);
  return input.customOp(op);
}


} // end namespace Eigen

#endif // EIGEN_CXX11_NEURAL_NETWORKS_SOFTMAX_H