#include <mlpack/core.hpp>
#include <mlpack/core/util/cli.hpp>

#include <mlpack/methods/ann/init_rules/gaussian_init.hpp>
#include <mlpack/methods/ann/gan.hpp>
#include <mlpack/methods/softmax_regression/softmax_regression.hpp>
#include <mlpack/core/optimizers/minibatch_sgd/minibatch_sgd.hpp>

using namespace mlpack;
using namespace mlpack::ann;
using namespace mlpack::math;
using namespace mlpack::optimization;
using namespace mlpack::regression;
using namespace std::placeholders;

PROGRAM_INFO("GAN", "DIGITS");
PARAM_STRING_IN("dataset", "path/to/dataset.", "i", "mnist.arm");
PARAM_STRING_IN("output-dataset", "path/to/output.", "o", "output.txt");

PARAM_INT_IN("dataset-max-cols", "dataset max cols.", "m",  -1);

PARAM_INT_IN("generator-hidden-layer-size", "hidden layer size.", "G",  128);
PARAM_INT_IN("discriminator-hidden-layer-size", "hidden layer size.", "D",  128);

PARAM_INT_IN("batch-size", "batch-size.", "b", 100);

PARAM_INT_IN("generator-update-step", "generator-update-step.", "g", 100);
PARAM_INT_IN("noise-dim", "noise dim.", "N", 25);
PARAM_INT_IN("num-samples", "num samples.", "n", 10);
PARAM_DOUBLE_IN("step-size", "learning rate.", "r", 0.01);
PARAM_INT_IN("num-epoches", "num-epoches.", "e", 20);
PARAM_DOUBLE_IN("tolerance", "tolerance.", "t", 1e-5);
PARAM_FLAG("shuffle", "shuffle or not", "s");

int main(int argc, char* argv[])
{
  CLI::ParseCommandLine(argc, argv);
  size_t generatorHiddenLayerSize = 128;
  size_t discriminatorHiddenLayerSize = 128;
  size_t generatorOutputSize;
  size_t discriminatorOutputSize = 1;
  size_t batchSize = 100;
  size_t noiseDim = 25;
  size_t generatorUpdateStep = 10;
  size_t numSamples = 10;
  double stepSize = 0.01;
  double beta1 = 0.9;
  double beta2 = 0.999;
  double eps = 1e-8;
  size_t numEpoches = 20;
  double tolerance = 1e-5;
  int datasetMaxCols = -1;
  bool shuffle = true;

  // Load the dataset
  arma::mat trainData;

  if (!CLI::HasParam("dataset"))
    Log::Fatal << "Input dataset is undefined!" << std::endl;

  if (!CLI::HasParam("output-dataset"))
    Log::Fatal << "Output dataset is undefined!" << std::endl;

  std::string dataset = CLI::GetParam<std::string>("dataset");
  std::string output_dataset = CLI::GetParam<std::string>("output-dataset");

  if (CLI::HasParam("dataset-max-cols"))
    datasetMaxCols = CLI::GetParam<int>("dataset-max-cols");

  if (CLI::HasParam("generator-hidden-layer-size"))
    generatorHiddenLayerSize = CLI::GetParam<int>("generator-hidden-layer-size");

  if (CLI::HasParam("discriminator-hidden-layer-size"))
    discriminatorHiddenLayerSize = CLI::GetParam<int>("discriminator-hidden-layer-size");

  if (CLI::HasParam("batch-size"))
    batchSize = CLI::GetParam<int>("batch-size");

  if (CLI::HasParam("generator-update-step"))
    generatorUpdateStep = CLI::GetParam<int>("generator-update-step");

  if (CLI::HasParam("noise-dim"))
    noiseDim = CLI::GetParam<int>("noise-dim");

  if (CLI::HasParam("num-samples"))
    numSamples = CLI::GetParam<int>("num-samples");

  if (CLI::HasParam("step-size"))
    stepSize = CLI::GetParam<double>("step-size");

  if (CLI::HasParam("num-epoches"))
    numEpoches = CLI::GetParam<int>("num-epoches");

  if (CLI::HasParam("tolerance"))
    tolerance = CLI::GetParam<double>("tolerance");

  shuffle = CLI::HasParam("shuffle");

  Log::Info << "dataset = '" << dataset << "'" << std::endl;
  Log::Info << "output_dataset = '" << output_dataset << "'" << std::endl;
  Log::Info << std::boolalpha
      << "generatorHiddenLayerSize = " << generatorHiddenLayerSize
      << " discriminatorHiddenLayerSize = " << discriminatorHiddenLayerSize
      << " batchSize = " << batchSize
      << " generatorUpdateStep = " << generatorUpdateStep
      << " noiseDim = " << noiseDim
      << " numSamples = " << numSamples
      << " stepSize = " << stepSize
      << " numEpoches = " << numEpoches
      << " tolerance = " << tolerance
      << " shuffle = " << shuffle << std::endl;

  Log::Info << "Loading dataset '" << dataset << "'..." << std::endl;
  trainData.load(dataset);

  if (datasetMaxCols > 0)
    trainData = trainData.cols(0, datasetMaxCols - 1);

  Log::Info << "Dataset loaded (" << trainData.n_rows << ", " << trainData.n_cols << ")" << std::endl;
  generatorOutputSize = trainData.n_rows;

  FFN<CrossEntropyError<>> discriminator;
  discriminator.Add<Linear<>>(trainData.n_rows, discriminatorHiddenLayerSize);
  discriminator.Add<LeakyReLU<>>(0.2);
  discriminator.Add<Linear<>>(discriminatorHiddenLayerSize, discriminatorHiddenLayerSize / 2);
  discriminator.Add<LeakyReLU<>>(0.2);
  discriminator.Add<Linear<>>(discriminatorHiddenLayerSize / 2, discriminatorHiddenLayerSize / 4);
  discriminator.Add<LeakyReLU<>>(0.2);
  discriminator.Add<Linear<>>(discriminatorHiddenLayerSize / 4, discriminatorOutputSize);
  discriminator.Add<SigmoidLayer<>>();

  // Generator network
  FFN<CrossEntropyError<>> generator;
  generator.Add<Linear<>>(noiseDim, generatorHiddenLayerSize / 4);
  generator.Add<LeakyReLU<>>(0.2);
  generator.Add<Linear<>>(generatorHiddenLayerSize / 4, generatorHiddenLayerSize / 2);
  generator.Add<LeakyReLU<>>(0.2);
  generator.Add<Linear<>>(generatorHiddenLayerSize / 2, generatorOutputSize);
  generator.Add<SigmoidLayer<>>();

  // Intialisation function
  GaussianInitialization gaussian(0, 0.1);
  // Optimizer

  size_t numIterations = trainData.n_cols * numEpoches;
  numIterations /= batchSize;

  AdamBatchSGD optimizer(batchSize, stepSize, numIterations, tolerance, shuffle);
  std::function<double ()> noiseFunction = [] () { return math::Random(0, 1); };
  // GAN model
  GAN<FFN<CrossEntropyError<>>, GaussianInitialization, std::function<double ()> > gan(trainData, generator, discriminator, gaussian, noiseFunction, noiseDim,
      batchSize, generatorUpdateStep, 30);

  Log::Info << "Training..." << std::endl;
  gan.Train(optimizer);

  // Generate samples
  Log::Info << "Sampling..." << std::endl;
  arma::mat noise(noiseDim, 1);
  size_t dim = std::sqrt(trainData.n_rows);
  arma::mat generatedData(2 * dim, dim * numSamples);
  for (size_t i = 0; i < numSamples; i++)
  {
    arma::mat samples;
    noise.imbue( [&]() { return math::Random(0, 1); } );
//    noise.randn();
    generator.Forward(noise, samples);

    samples.reshape(dim, dim);
    samples = samples.t();

    generatedData.submat(0, i * dim, dim - 1, i * dim + dim - 1) = samples;


    samples = trainData.col(math::RandInt(0, trainData.n_cols));
    samples.reshape(dim, dim);
    samples = samples.t();

    generatedData.submat(dim, i * dim, 2 * dim - 1, i * dim + dim - 1) = samples;
  }

  Log::Info << "Saving output to " << output_dataset << "..." << std::endl;
  generatedData.save(output_dataset, arma::raw_ascii);
  Log::Info << "Output saved!" << std::endl;

  return 0;
}