#include <mlpack/core.hpp>

#include <mlpack/core/optimizers/rmsprop/rmsprop.hpp>
#include <mlpack/methods/ann/init_rules/gaussian_init.hpp>
#include <mlpack/methods/ann/rbm/spike_slab_rbm_policy.hpp>
#include <mlpack/methods/ann/rbm.hpp>
#include <mlpack/methods/ann/rbm/binary_rbm_policy.hpp>
#include <mlpack/methods/softmax_regression/softmax_regression.hpp>
#include <mlpack/core/optimizers/minibatch_sgd/minibatch_sgd.hpp>
#include <mlpack/core/optimizers/sgd/sgd.hpp>
#include <mlpack/core/optimizers/lbfgs/lbfgs.hpp>

using namespace mlpack;
using namespace mlpack::ann;
using namespace mlpack::optimization;
using namespace mlpack::regression;

PROGRAM_INFO("ssRBM", "cifar");


int main(int argc, char** argv)
{
  CLI::ParseCommandLine(argc, argv);

  size_t hiddenLayerSize = 256;
  size_t numPatches = 50;
  size_t batchSize = 20;
  double stepSize = 0.001;
  size_t numEpoches = 10;
  double tolerance = 1e-5;
  bool shuffle = true;
  size_t poolSize = 4;
  double slabPenalty = 10.5;
  double lambdaBias = 10;
  double radius = 0;
  double tempRadius = 0;

  /*
   * Training Phase
   * We have 40000 patches for 400 training images and 192 feature vector.
   * 
   */

  std::string trainDataFile = "./random_train.txt";
  std::string testDataFile = "./convfeat_test.txt";
  std::string LabelsFile = "./labels.txt";

  arma::mat trainData, testData, trainLabelsTemp, testLabelsTemp;
  
  trainData.load(trainDataFile);
  testData.load(testDataFile);
  trainLabelsTemp.load(LabelsFile);
  // since we are training on 1000 images
  testLabelsTemp = trainLabelsTemp.rows(900, 999);
  trainLabelsTemp = trainLabelsTemp.rows(0, 899);

  std::cout << testLabelsTemp(0) << std::endl;
  std::cout << arma::size(trainData) << std::endl;
  std::cout << arma::size(testData) << std::endl;

  trainData = trainData.cols(0, 10);
  testData = testData.cols(0, 10);


  trainData = trainData.t();
  testData = testData.t();
  
  GaussianInitialization gaussian(0, 0.1);

  arma::Row<size_t> trainLabels = arma::zeros<arma::Row<size_t>>(1,
      trainLabelsTemp.n_rows);
  arma::Row<size_t> testLabels = arma::zeros<arma::Row<size_t>>(1,
      testLabelsTemp.n_rows);

  for (size_t i = 0; i < trainLabelsTemp.n_rows; ++i)
  {
    trainLabels(i) = arma::as_scalar(trainLabelsTemp.row(i));
  }

  for (size_t i = 0; i < testLabelsTemp.n_rows; ++i)
  {
    testLabels(i) = arma::as_scalar(testLabelsTemp.row(i));
  }


  // Calculate radius
  for (size_t i = 0; i < trainData.n_cols; i++)
  {
    tempRadius = arma::norm(trainData.col(i));
    if (radius < tempRadius)
      radius = tempRadius;
  }
  radius *= 1.5;

  
  size_t iter =  (numEpoches * trainData.n_cols) / batchSize;
  std::cout << "iter#" << iter << std::endl;  

  SpikeSlabRBMPolicy<> rbmPolicy(trainData.n_rows, hiddenLayerSize, poolSize,
      slabPenalty, radius);

  RBM<GaussianInitialization, SpikeSlabRBMPolicy<>> modelssRBM(trainData,
      gaussian, rbmPolicy, 1, 1, true, false);

  MiniBatchSGD msgd(batchSize, stepSize, iter, tolerance, shuffle);
  modelssRBM.Reset();
  modelssRBM.Policy().VisiblePenalty().fill(lambdaBias);
  modelssRBM.Policy().SpikeBias().fill(-5);
  modelssRBM.Train(trainData, msgd);
  
  // hiddenLayer size * number of patches in one image
  // 256 * 100 * 1000
  arma::cube hiddenMeanTest(hiddenLayerSize , numPatches, testData.n_cols / numPatches);
  arma::cube hiddenMeanTrain(hiddenLayerSize , numPatches, trainData.n_cols / numPatches);
  arma::mat ssRBMFeaturesTest(hiddenLayerSize * numPatches, testData.n_cols / numPatches);
  arma::mat ssRBMFeaturesTrain(hiddenLayerSize * numPatches, trainData.n_cols / numPatches);
  // slice size 256 * 100 * 1000
  for (size_t i = 0; i < hiddenMeanTest.n_slices; i++)
    for (size_t k = 0; k < numPatches; k++)
    {
      arma::mat tmp;
      modelssRBM.SampleHidden(std::move(testData.col(i * numPatches + k)),
          std::move(tmp));
      hiddenMeanTest.slice(i).col(k) = tmp.rows(0, hiddenLayerSize - 1);
    }

  for (size_t i = 0; i < hiddenMeanTest.n_slices; ++i)
  {
    ssRBMFeaturesTest.col(i) = arma::vectorise(hiddenMeanTest.slice(i), 1).t();
  }

  for (size_t i = 0; i < hiddenMeanTrain.n_cols; i++)
    for (size_t k = 0; k < numPatches; k++)
    {
      
      arma::mat tmp;
      modelssRBM.SampleHidden(std::move(trainData.col(i * numPatches + k)), 
          std::move(tmp));
      hiddenMeanTrain.slice(i).col(k) = tmp.rows(0, hiddenLayerSize - 1);;
      
    }
  for (size_t i = 0; i < hiddenMeanTrain.n_slices; ++i)
  {
    ssRBMFeaturesTrain.col(i) = arma::vectorise(hiddenMeanTrain.slice(i), 1).t();
  }

  std::cout << "ssRBMFeaturesTrain = " << arma::size(ssRBMFeaturesTrain) << std::endl;
  std::cout << "ssRBMFeaturesTest = " << arma::size(ssRBMFeaturesTest) << std::endl;
  

  const size_t numClasses = 10; // Number of classes.
  const size_t numBasis = 5; // Parameter required for L-BFGS algorithm.
  const size_t numIterations = 1000; // Maximum number of iterations.


  arma::Row<size_t> predictions1, predictions2;
  L_BFGS optimizer1(numBasis, numIterations);
  SoftmaxRegression regressor1(ssRBMFeaturesTrain, trainLabels, numClasses,
  0.001, false, optimizer1);
  double classificationAccurayssRBM = regressor1.ComputeAccuracy(ssRBMFeaturesTest, testLabels);
  arma::Row<size_t> tmpLabel;
  regressor1.Predict(ssRBMFeaturesTest, tmpLabel);
  std::cout << "tmpLabel Size = " << arma::size(tmpLabel) << std::endl;
  tmpLabel.print();
  std::cout << "Real lables" << std::endl;
  testLabels.print();
  std::cout << "RBM Accuracy" << classificationAccurayssRBM << std::endl;
}