//============================================================================
//
// This file is part of the Shapes Generative Model Library (SGML) project.
//
// Copyright (c) 2010-2014 - Evangelos Kalogerakis, Siddhartha Chaudhuri (authors of the code) / Stanford University, UMass-Amherst
//
// SGML is free software: you can redistribute it and/or modify
// it under the terms of the GNU General Public License as published by
// the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.

// SGML is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU General Public License for more details.

// You should have received a copy of the GNU General Public License
// along with SGML.  If not, see <http://www.gnu.org/licenses/>.
//
//============================================================================

#ifndef __S12_SIGMOIDCPD_HPP
#define __S12_SIGMOIDCPD_HPP

namespace SNL {
namespace S12 {

class SigmoidCPD
{
  private:
    // X is input
    // Y is training output
    // R is the output probabilities

    // N is the number of samples
    // M is the number of parents (input dimensions)
    // P is the number of dimensions (after making the parents discrete), INCLUDES BIAS
    // K is the number of classes of Y

    SNL_MATRIX(double)  oldweights; // (K-1) x P

  public:
    size_t  K;
    SNL_MATRIX(double)  weights;    // (K-1) x P
    SNL_MATRIX(double) R;           // N X K

    SigmoidCPD()
    {
      K = 0;
    };
    SigmoidCPD(size_t _K)
    {
      K = _K;
    };

    SNL_VECTOR( SNL_VECTOR(bool) )  transformInputToBooleanInputPlusBias(const SNL_VECTOR( SNL_VECTOR(size_t) ) & _X, const SNL_VECTOR( SNL_VECTOR(size_t) ) & values_x)
    {
      size_t N = _X.size();
      size_t M = _X[0].size();

      if (values_x.size() != M)
      {
        std::cout << "transformInputToBooleanInputPlusBias: values_x size does not agree with number of parents X" << std::endl;
        exit(-1);
      }

      size_t P = 0;

      for (size_t m = 0; m < M; m++)
        P += values_x[m].size() - 1;

      P++; // for bias term

      SNL_VECTOR( SNL_VECTOR(bool) ) X(N);

      for (size_t n = 0; n < N; n++)
      {
        size_t p = 0;

        X[n].resize(P, false);

        for (size_t m = 0; m < M; m++)
        {
          for (size_t xk = 1; xk < values_x[m].size(); xk++, p++)
          {
            X[ n ][ p ] = _X[n][m] == values_x[m][xk];
          }
        }

        X[n][ p ] = true ; // bias term
      }

      return X;
    }


    void evaluate(const SNL_VECTOR( SNL_VECTOR(bool) ) & X)
    {
      size_t N = X.size();
      R.resize(N, K);

      for (size_t n = 0; n < N; n++)
      {
        double sumrn = 0.0;

        for (size_t k = 1; k < K; k++)
        {
          R[n][k] = min(exp(sum(mult1(weights.getRow(k - 1), X[n]))), SUPERMAXEPSILON);
          sumrn += R[n][k];
        }

        for (size_t k = 1; k < K; k++)
        {
          R[n][k] /= (1.0 + sumrn);
        }

        R[n][0] = 1.0 / (1.0 + sumrn);
      }
    }


    SNL_VECTOR( double ) evaluateDuringLearning(const SNL_VECTOR( SNL_VECTOR(bool) ) & X)
    {
      size_t N = X.size();
      evaluate(X);
      SNL_VECTOR( double ) RR(N * (K - 1));

      size_t kn = 0;

      for (size_t k = 1; k < K; k++)
      {
        for (size_t n = 0; n < N; n++, kn++)
        {
          RR[ kn ] = R[ n ][ k ];
        }
      }

      return RR;
    }


    // Y is N x 1, X is N x M
    void learn(const SNL_VECTOR(size_t) & _Y, const SNL_VECTOR(size_t) & values_y, const SNL_VECTOR( SNL_VECTOR(size_t) ) & _X, const SNL_VECTOR( SNL_VECTOR(size_t) ) & values_x)
    {
      size_t N = _Y.size();
      size_t M = _X[0].size();
      size_t K = values_y.size();

      if (_X.size() != N)
      {
        std::cout << "Sigmoid CPD: X and Y do not have the same number of training samples" << std::endl;
        exit(-1);
      }

      if (values_x.size() != M)
      {
        std::cout << "Sigmoid CPD: values_x size does not agree with number of parents X" << std::endl;
        exit(-1);
      }

      // prepare Y matrix
      size_t kn = 0;
      SNL_VECTOR( bool ) Y(N * (K - 1));

      for (size_t k = 1; k < K; k++)
      {
        for (size_t n = 0; n < N; n++, kn++)
          Y[kn] = _Y[n] == values_y[k];
      }

      size_t P = 0;

      for (size_t m = 0; m < M; m++)
        P += values_x[m].size() - 1;

      P++; // for bias term

      // prepare X matrix
      SNL_VECTOR( SNL_VECTOR(bool) ) XX = transformInputToBooleanInputPlusBias(_X, values_x);
      SNL_MATRIX( bool ) X(N * (K - 1), P * (K - 1), false);
      kn = 0;
      size_t pk = 0;

      for (size_t k = 1; k < K; k++)
      {
        for (size_t n = 0; n < N; n++, kn++)
        {
          size_t p = 0;

          for (size_t m = 0; m < M; m++)
          {
            for (size_t xk = 1; xk < values_x[m].size(); xk++, p++)
            {
              X[ kn ][ p + pk ] = _X[n][m] == values_x[m][xk];
            }
          }

          X[kn][ p + pk ] = true ; // bias term
        }

        pk += P;
      }


      // W diagonal helper matrix
      SNL_MATRIX(double) W(N * (K - 1), N * (K - 1), 0.0);


      weights.resize(K - 1, P);
      oldweights.resize(K - 1, P);
      weights.fill(0);
      oldweights.fill(1);

      bool restart = true;
      int ITER = 0;
      int ITER_LIMIT = 50;
      double step_size = 1.0;

      while (sumabs(subtract(weights, oldweights)) > EPSILON && ITER++ <= ITER_LIMIT)
      {
        if (restart)
        {
          weights.fill(0);
          oldweights.fill(1);

          restart = false;
        }

        SNL_VECTOR(double) RR = evaluateDuringLearning(XX);

        if (M == 0 || K <= 1)
          break;

        for (size_t k1 = 1; k1 < K; k1++)
        {
          for (size_t k2 = 1; k2 < K; k2++)
          {
            for (size_t n = 0; n < N; n++)
            {
              if (k1 == k2)
                W[(k1 - 1)*N + n ][(k2 - 1)*N + n ] = R[n][k1] * (1.0 - R[n][k1]);
              else
                W[(k1 - 1)*N + n ][(k2 - 1)*N + n ] = -R[n][k1] * R[n][k2];
            }
          }
        }

        bool problem = false;
        SNL_MATRIX( double ) ixwx = invert(regularize(mult1(mult2(transpose(X), W), X), EPSILON), &problem);
        SNL_VECTOR(double) weights_vec = mult(mult1(ixwx, transpose(X)) , subtract2(Y, RR));

        oldweights = weights;
        pk = 0;

        for (size_t k = 1; k < K; k++)
        {
          for (size_t p = 0; p < P; p++, pk++)
          {
            if (ITER == 1)
              weights[k - 1][p] += weights_vec[pk];
            else
              weights[k - 1][p] += step_size * weights_vec[pk];

            if (abs(weights[k - 1][p]) > 100)
            {
              restart = true;
              ITER = 0;
              ITER_LIMIT += 100;
              step_size /= 1.25;
              break;
            }
          }

          if (restart)
            break;
        }


      } // end of loop of IRLS


      //std::cout << ITER << std::endl;
      //std::cout << R << std::endl;
      //std::cout << weights << std::endl;
      //system("pause");
    }

};



inline bool sigmoidTest()
{
  std::cout << "Testing Sigmoid CPD functionality...";
  SNL_VECTOR( SNL_VECTOR(size_t) ) X;
  SNL_VECTOR( SNL_VECTOR(size_t) ) X_values;
  SNL_VECTOR(size_t) x(3);
  x[0] = 2;
  x[1] = 1;
  x[2] = 1;
  X.push_back(x);
  x[0] = 2;
  x[1] = 1;
  x[2] = 2;
  X.push_back(x);
  x[0] = 2;
  x[1] = 1;
  x[2] = 1;
  X.push_back(x);
  x[0] = 2;
  x[1] = 1;
  x[2] = 2;
  X.push_back(x);
  x[0] = 2;
  x[1] = 1;
  x[2] = 1;
  X.push_back(x);
  x[0] = 2;
  x[1] = 1;
  x[2] = 2;
  X.push_back(x);
  x[0] = 4;
  x[1] = 1;
  x[2] = 0;
  X.push_back(x);
  x[0] = 4;
  x[1] = 1;
  x[2] = 1;
  X.push_back(x);
  x[0] = 4;
  x[1] = 1;
  x[2] = 2;
  X.push_back(x);
  x[0] = 4;
  x[1] = 1;
  x[2] = 1;
  X.push_back(x);
  x[0] = 4;
  x[1] = 1;
  x[2] = 2;
  X.push_back(x);
  x[0] = 0;
  x[1] = 0;
  x[2] = 1;
  X.push_back(x);
  x[0] = 0;
  x[1] = 0;
  x[2] = 1;
  X.push_back(x);
  x[0] = 0;
  x[1] = 0;
  x[2] = 1;
  X.push_back(x);

  //SNL_VECTOR( size_t ) T = X[0];
  //T[0] = 4;
  //T[1] = 2;
  //std::cout << multmatrixwithvec( X, T ) << std::endl;
  //system("pause");


  //SNL_MATRIX(size_t) T = X;
  //T[0][1] = 5;
  //T[1][2] = 2;
  //T[0][2] = 3;
  //std::cout << mult( X, transpose(T) ) << std::endl;
  //system("pause");

  SNL_VECTOR(size_t) x1(3);
  x1[0] = 0;
  x1[1] = 2;
  x1[2] = 4;
  X_values.push_back(x1);
  SNL_VECTOR(size_t) x2(2);
  x2[0] = 0;
  x2[1] = 1;
  X_values.push_back(x2);
  SNL_VECTOR(size_t) x3(3);
  x3[0] = 0;
  x3[1] = 1;
  x3[2] = 2;
  X_values.push_back(x3);

  SNL_VECTOR( size_t ) Y(14);
  Y[0] = 4;
  Y[1] = 2;
  Y[2] = 4;
  Y[3] = 2;
  Y[4] = 4;
  Y[5] = 2;
  Y[6] = 1;
  Y[7] = 3;
  Y[8] = 1;
  Y[9] = 3;
  Y[10] = 1;
  Y[11] = 4;
  Y[12] = 4;
  Y[13] = 4;
  //Y[0] = 1;
  //Y[1] = 0;
  //Y[2] = 1;
  //Y[3] = 0;
  //Y[4] = 1;
  //Y[5] = 0;
  //Y[6] = 0;
  //Y[7] = 1;
  //Y[8] = 0;
  //Y[9] = 1;
  //Y[10] = 0;
  //Y[11] = 1;
  //Y[12] = 1;
  //Y[13] = 1;
  SNL_VECTOR( size_t ) Y_values(4);
  Y_values[0] = 1;
  Y_values[1] = 2;
  Y_values[2] = 3;
  Y_values[3] = 4;

  SigmoidCPD scpd(4);
  scpd.learn(Y, Y_values, X, X_values);

  if (equals(add(maximumId_perrow(scpd.R), size_t(1)), Y))
  {
    std::cout << "Pass!" << std::endl;
    return true;
  }
  else
  {
    std::cout << "Fail." << std::endl;
    return false;
  }
};

} // namespace S12
} // namespace SNL

#endif
