# Pastebin yKmWeiQu
#include <iostream>
#include <mlpack/core.hpp>
#include <mlpack/core/util/cli.hpp>


#include <mlpack/core/optimizers/rmsprop/rmsprop.hpp>
#include <mlpack/methods/ann/init_rules/gaussian_init.hpp>
#include <mlpack/methods/ann/rbm/binary_rbm_policy.hpp>
#include <mlpack/methods/ann/rbm/spike_slab_rbm_policy.hpp>
#include <mlpack/methods/ann/rbm.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>
#include <ctime>

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

PROGRAM_INFO("RBM", "MNIST");
PARAM_STRING_IN("input-dataset", "path/to/MNIST.", "i",
    "mnist.arm");
PARAM_STRING_IN("input-parameters", "path/to/output.", "w",
    "parameters.arm");

PARAM_STRING_IN("output-parameters", "path/to/output.", "p",
    "parameters.arm");

PARAM_STRING_IN("output-dataset", "path/to/output.", "o",
    "output.txt");

PARAM_FLAG("train", "train or not", "t");

PARAM_FLAG("sample", "sample or not", "s");

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

  arma::mat trainData;
  std::size_t hiddenLayerSize = 500;
  std::size_t poolSize = 1;
  double radius = 0;
  double tempRadius = 0;
  int numEpochs = 200;
  bool persistent = false;
  bool useMonitoringCost = false;
  bool shuffle = true;
  double learningRate = 0.02;
  int batchSize = 20;
  int numSteps = 3;
  int mSteps = 1;
  int numSamplingSteps = 1;
  GaussianInitialization gaussian(0, 1e-4);
  bool bTrain = true;
  bool bSample = true;
  std::string input_dataset, output_parameters, input_parameters, output_dataset;

  bTrain = CLI::HasParam("train");

  bSample = CLI::HasParam("sample");

  if (!bTrain && !bSample)
  {
    Log::Fatal << "Invalid parameters" << std::endl;
    return 0;
  }

  if (bTrain && (!CLI::HasParam("input-dataset") || !CLI::HasParam("output-parameters")))
  {
    Log::Fatal << "Invalid parameters" << std::endl;
    return 0;
  }

  if (bSample && !bTrain && (!CLI::HasParam("input-dataset") || !CLI::HasParam("input-parameters") || !CLI::HasParam("output-dataset")))
  {
    Log::Fatal << "Invalid parameters" << std::endl;
    return 0;
  }
  
  if (bSample && bTrain && (!CLI::HasParam("input-dataset") || !CLI::HasParam("output-dataset")))
  {
    Log::Fatal << "Invalid parameters" << std::endl;
    return 0;
  }

  math::RandomSeed(std::time(NULL));

  Log::Info << std::boolalpha << "hiddenLayerSize = " << hiddenLayerSize << " numEpochs = "
            << numEpochs << " persistent = " << persistent << " shuffle = " << shuffle
            << " learningRate = " << learningRate << " batchSize = " << batchSize
            << " numSteps = " << numSteps << " poolSize = " << poolSize << std::endl;

  input_dataset = CLI::GetParam<std::string>("input-dataset");
  Log::Info << "Input dataset = " << input_dataset << std::endl;
  if (bTrain)
  {
    output_parameters = CLI::GetParam<std::string>("output-parameters");
    Log::Info << "Output parameters = '" << output_parameters << "'" << std::endl;
  }
  else
  {
    input_parameters = CLI::GetParam<std::string>("input-parameters");
    Log::Info << "Input parameters = '" << input_parameters << "'" << std::endl;
  }
  if (bSample)
  {
    output_dataset = CLI::GetParam<std::string>("output-dataset");
    Log::Info << "Output dataset = '" << output_dataset << "'" << std::endl;
  }
  Log::Info << "Loading dataset..." << std::endl;
  trainData.load(input_dataset);

  trainData = trainData.submat(0, 0, trainData.n_rows - 1, 100 - 1);

//  for (std::size_t c = 0; c < trainData.n_cols; c++)
//    trainData.col(c) /= 255.0;

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

  radius *= 2;

  Log::Info << "Radius: " << radius << std::endl;
  Log::Info << "Dataset dimensions: (" << trainData.n_rows << ", " << trainData.n_cols << ")" << std::endl;
  Log::Info << "Dataset maximum = " << trainData.max() << std::endl;

  
  arma::mat slabPenalty(poolSize, hiddenLayerSize);
  slabPenalty.fill(10);

  SpikeSlabRBMPolicy ss_rbm(trainData.n_rows, hiddenLayerSize, poolSize, slabPenalty, radius);
  RBM<GaussianInitialization, SpikeSlabRBMPolicy> modelssRBM(trainData, gaussian, ss_rbm,
      numSteps, mSteps, useMonitoringCost, persistent);

  modelssRBM.Reset();

  if (bTrain)
  {
    std::size_t numIterations = trainData.n_cols * numEpochs;
    numIterations /= batchSize;
numIterations = 1;
    MiniBatchSGD msgd(batchSize, learningRate, numIterations, 0, shuffle);
    modelssRBM.Policy().VisiblePenalty().fill(3);
    modelssRBM.Policy().SpikeBias().zeros();

    // test the reset function

    Log::Info << "Training..." << std::endl;
    modelssRBM.Train(trainData, msgd);

    Log::Info << "Saving parameters to " << output_parameters << "..." << std::endl;
    modelssRBM.Parameters().save(output_parameters);
    Log::Info << "Parameters saved!" << std::endl;
  }
  else
  {
    arma::mat parameters;

    Log::Info << "Loading parameters..." << std::endl;
    parameters.load(input_parameters);

    modelssRBM.Parameters() = parameters;
  }

  Log::Info << "Lambda.Min = "  << modelssRBM.Policy().VisiblePenalty().min() << " Lambda.Max = " << modelssRBM.Policy().VisiblePenalty().max() << std::endl;

  if (bSample)
  {
    Log::Info << "Sampling..." << std::endl;
  //  arma::mat input  = trainData.col(24342);
    int numSamples = 10;
    int numChains = 20;
    arma::mat mainOutput(28 * numSamples, 28 * numChains);
    arma::mat tmp, output;

    for (int c = 0; c < numChains; c++)
    {
      arma::mat input  = trainData.col(math::RandInt(0, trainData.n_cols));
      arma::mat copy = input;
      copy.reshape(28, 28);
      copy = copy.t();
      mainOutput.submat(0, c * 28, 27, c * 28 + 27) = copy;
  //    arma::mat window(mainOutput.memptr() + c * 28 * numSamples, 28 * 28, 1, false, false);
      Log::Info << "Sampling chain " << c << "..." << std::endl;

      for (int s = 1; s < numSamples; s++)
      {
//        Log::Info << "Sampling sample " << s << "..." << std::endl;
        for (int i = 0; i < numSamplingSteps; i++)
        {
          modelssRBM.SampleHidden(std::move(input), std::move(tmp));
          modelssRBM.SampleVisible(std::move(tmp), std::move(input));
        }
        modelssRBM.Policy().VisibleMean(std::move(tmp), std::move(output));
        arma::mat copy = output;
        copy.reshape(28, 28);
        copy = copy.t();

        copy -= copy.min();

        copy /= copy.max() - copy.min();
/*
        for (size_t e = 0; e < copy.n_elem; e++)
        {
          if (copy[e] < 0)
            copy[e] = 0;
          if (copy[e] > 1)
            copy[e] = 1;
        }
*/
        mainOutput.submat(s * 28, c * 28, s * 28 + 27, c * 28 + 27) = copy;
      }
    }
  //  model.HiddenLayer().Forward(std::move(tmp), std::move(output));
  //  model.Gibbs(std::move(trainData.col(10)), std::move(output));

    Log::Info << "Saving output to " << output_dataset << "..." << std::endl;
  //  output = trainData.col(10).t();
  //  output.reshape(28, 28);
    mainOutput.save(output_dataset, arma::raw_ascii);
    Log::Info << "Output saved!" << std::endl;

  }

  return 0;
}