text: #pragma once
  
  #include "Matrix.hpp"
  #include "Optimizer.hpp"
  
  class SoftMax{
  public:
    SoftMax(){};
    SoftMax(const int inputDim, const int classNum):
      weight(MatD::Zero(inputDim, classNum)), bias(VecD::Zero(classNum))
    {}
    SoftMax(const int inputDim, const int classNum, const int exception_, const Real gamma_, const Real mPlus_, const Real mMinus_):
      weight(MatD::Zero(inputDim, classNum)), bias(VecD::Zero(classNum)), exception(exception_), gamma(gamma_), mPlus(mPlus_), mMinus(mMinus_)
    {}
  
    class Grad;
  
    MatD weight; VecD bias;
  
    void calcDist(const VecD& input, VecD& output);
    Real calcLoss(const VecD& output, const int label);
    Real calcLoss(const VecD& output, const VecD& goldOutput);
    void backward(const VecD& input, const VecD& output, const int label, VecD& deltaFeature, SoftMax::Grad& grad);
    void backward(const VecD& input, const VecD& output, const VecD& goldOutput, VecD& deltaFeature, SoftMax::Grad& grad);
    void backwardAttention(const VecD& input, const VecD& output, const VecD& deltaOut, VecD& deltaFeature, SoftMax::Grad& grad);
    void sgd(const SoftMax::Grad& grad, const Real learningRate);
    void save(std::ofstream& ofs);
    void load(std::ifstream& ifs);
  
    void operator += (const SoftMax& softmax);
    void operator /= (const Real val);
  
    //for ranking loss
    int exception;
    Real gamma, mPlus, mMinus;
    void calcScore(const VecD& input, VecD& output);
    Real calcRankingLoss(const VecD& output, const int label);
    void backwardRankingLoss(const VecD& input, const VecD output, const int label, VecD& deltaFeature, SoftMax::Grad& grad);
  };
  
  class SoftMax::Grad{
  public:
    Grad(): gradHist(0){}
    Grad(const SoftMax& softmax):
      gradHist(0)
    {
      this->weight = MatD::Zero(softmax.weight.rows(), softmax.weight.cols());
      this->bias = VecD::Zero(softmax.bias.rows());
    }
  
    SoftMax::Grad* gradHist;
    
    MatD weight; VecD bias;
  
    void init(){
      this->weight.setZero();
      this->bias.setZero();
    }
  
    Real norm(){
      return this->weight.squaredNorm()+this->bias.squaredNorm();
    }
  
    void l2reg(const Real lambda, const SoftMax& s){
      this->weight += lambda*s.weight;
    }
  
    void l2reg(const Real lambda, const SoftMax& s, const SoftMax& target){
      this->weight += lambda*(s.weight-target.weight);
      this->bias += lambda*(s.bias-target.bias);
    }
  
    void sgd(const Real learningRate, SoftMax& softmax){
      Optimizer::sgd(this->weight, learningRate, softmax.weight);
      Optimizer::sgd(this->bias, learningRate, softmax.bias);
    }
  
    void adagrad(const Real learningRate, SoftMax& softmax, const Real initVal = 1.0){
      if (this->gradHist == 0){
        this->gradHist = new SoftMax::Grad(softmax);
        this->gradHist->weight.fill(initVal);
        this->gradHist->bias.fill(initVal);
      }
  
      Optimizer::adagrad(this->weight, learningRate, this->gradHist->weight, softmax.weight);
      Optimizer::adagrad(this->bias, learningRate, this->gradHist->bias, softmax.bias);
    }
  
    void momentum(const Real learningRate, const Real m, SoftMax& softmax){
      if (this->gradHist == 0){
        this->gradHist = new SoftMax::Grad(softmax);
        this->gradHist->weight.fill(0.0);
        this->gradHist->bias.fill(0.0);
      }
  
      Optimizer::momentum(this->weight, learningRate, m, this->gradHist->weight, softmax.weight);
      Optimizer::momentum(this->bias, learningRate, m, this->gradHist->bias, softmax.bias);
    }
  
    void operator += (const SoftMax::Grad& grad){
      this->weight += grad.weight;
      this->bias += grad.bias;
    }
  
    void operator /= (const Real val){
      this->weight /= val;
      this->bias /= val;
    }
  };
  