const char *help = "\
AdaBoost (c) Trebolloc & Co 2002\n\
\n\
This program will train an Adaboost of MLPs with tanh outputs for\n\
classification\n";

#include "FileDataSet.h"
#include "MseCriterion.h"
#include "Tanh.h"
#include "MseMeasurer.h"
#include "ClassMeasurer.h"
#include "TwoClassFormat.h"
#include "MultiClassFormat.h"
#include "OneHotClassFormat.h"
#include "StochasticGradient.h"
#include "GMTrainer.h"
#include "CmdLine.h"
#include "MLP.h"
#include "Boosting.h"
#include "WeightedSumMachine.h"

using namespace Torch;

int main(int argc, char **argv)
{
  char *train_file;
  char *test_file;
  int n_inputs;
  int n_targets;
  int n_hu;
  int n_boost;
  int max_load_train;
  int max_load_test;
  int seed_value;
  real accuracy;
  real learning_rate;
  real decay;
  int max_iter;
  char *dir_name;
  char *load_model;
  char *save_model;
  bool multi_class;
  bool one_hot;

  //=================== The command-line ==========================

  CmdLine cmd;

  cmd.info(help);

  cmd.addText("\nArguments:");
  cmd.addICmdArg("n_inputs", &n_inputs, "input dimension of the data");
  cmd.addICmdArg("n_targets", &n_targets, "output dimension of the data");
  cmd.addSCmdArg("train_file", &train_file, "the train file");
  cmd.addSCmdArg("test_file", &test_file, "the test file");

  cmd.addText("\nModel Options:");
  cmd.addICmdOption("-nhu", &n_hu, 25, "number of hidden units");
  cmd.addICmdOption("-nboost_max", &n_boost, 10, "maximum number of MLP in the boost");
  cmd.addBCmdOption("-one_hot", &one_hot, false, "class format = one_hot (default = two_class)");
  cmd.addBCmdOption("-multi_class", &multi_class, false, "class format = multi_class (default = two_class)");

  cmd.addText("\nLearning Options:");
  cmd.addICmdOption("-iter", &max_iter, 25, "max number of iterations");
  cmd.addRCmdOption("-lr", &learning_rate, 0.01, "learning rate");
  cmd.addRCmdOption("-e", &accuracy, 0.00001, "end accuracy");
  cmd.addRCmdOption("-lrd", &decay, 0, "learning rate decay");

  cmd.addText("\nMisc Options:");
  cmd.addICmdOption("-load_train", &max_load_train, -1, "max number of train examples to load");
  cmd.addICmdOption("-load_test", &max_load_test, -1, "max number of test examples to load");
  cmd.addICmdOption("-seed", &seed_value, -1, "initial seed for random generator");
  cmd.addSCmdOption("-dir", &dir_name, ".", "directory to save measures");
  cmd.addSCmdOption("-lm", &load_model, "", "start from given model file");
  cmd.addSCmdOption("-sm", &save_model, "", "save results into given model file");

  cmd.read(argc, argv);

  if (seed_value == -1)
    seed();
  else
    manual_seed((long)seed_value);

  //=================== DataSets ===================

  FileDataSet data(train_file, n_inputs, n_targets, false, max_load_train);
  data.setBOption("normalize inputs", true);
  data.init();

  FileDataSet test_data(test_file, n_inputs, n_targets, false, max_load_test);
  test_data.init();
  test_data.normalizeUsingDataSet(&data);

  // how is the class encoded in the datasets
  ClassFormat* class_format=NULL;
  if (one_hot)
    class_format = new OneHotClassFormat(&data);
  else if (multi_class)
    class_format = new MultiClassFormat(&data);
  else
    class_format = new TwoClassFormat(&data);

  //=================== The Model ==================
  // there will be one MLP created for each AdaBoost iteration. For each
  // of these MLP, we want a different MSE measurer and also a Class measurer.
  // we also need to train them, using a MSE criterion and a Stochastic 
  // gradient optimizer, given to a Trainer.

  MLP **mlp = new MLP *[n_boost];
  MseMeasurer **msemeasurer = new MseMeasurer *[n_boost];
  ClassMeasurer **classmeasurer = new ClassMeasurer *[n_boost];
  MseCriterion **mse = new MseCriterion *[n_boost];
  StochasticGradient **opt = new StochasticGradient *[n_boost];
  Trainer **trainer = new Trainer *[n_boost];
  List **trainers_measurers = new List *[n_boost];
  for (int i=0;i<n_boost;i++) {
    mlp[i] = new MLP(n_inputs,n_hu,n_targets);
    mlp[i]->setBOption("tanh outputs",true);
    mlp[i]->init();

    trainers_measurers[i] = NULL;

    char mse_name[100];
    sprintf(mse_name,"%s/the_mse%d_%d",dir_name,i,n_hu);
    msemeasurer[i] = new MseMeasurer(mlp[i]->outputs,&data, mse_name);
    msemeasurer[i]->init();
    addToList(&trainers_measurers[i],1,msemeasurer[i]);

    char class_name[100];
    sprintf(class_name,"%s/the_class_err%d_%d",dir_name,i,n_hu);
    classmeasurer[i] = new ClassMeasurer(mlp[i]->outputs,&data, class_format,class_name);
    classmeasurer[i]->init();
    addToList(&trainers_measurers[i],1,classmeasurer[i]);

    mse[i] = new MseCriterion(n_targets);
    mse[i]->init();

    opt[i] = new StochasticGradient();
    opt[i]->setIOption("max iter", max_iter);
    opt[i]->setROption("end accuracy", accuracy);
    opt[i]->setROption("learning rate", learning_rate);
    opt[i]->setROption("learning rate decay", decay);

    trainer[i] = (Trainer*)new GMTrainer(mlp[i], &data, mse[i], opt[i]);
  }

  WeightedSumMachine bmachine(trainer,n_boost,trainers_measurers);
  bmachine.init();

  // We also want to measure the performance of the AdaBoost itself

  List *measurers = NULL;

  char class_tr_name[100];
  sprintf(class_tr_name,"%s/the_boost_class_train%d",dir_name,n_hu);
  ClassMeasurer class_tr_m(bmachine.outputs,&data,class_format,class_tr_name);
  class_tr_m.init();
  addToList(&measurers,1,&class_tr_m);

  char class_te_name[100];
  sprintf(class_te_name,"%s/the_boost_class_test%d",dir_name,n_hu);
  ClassMeasurer class_te_m(bmachine.outputs,&test_data,class_format,class_te_name);
  class_te_m.init();
  addToList(&measurers,1,&class_te_m);

  Boosting boosting(&bmachine,&data,class_format);

  // =========== Training and/or Testing ===========

  if (strcmp(load_model,"")) {
    char load_model_name[100];
    sprintf(load_model_name,"%s/%s",dir_name,load_model);
    boosting.load(load_model_name);
    for (int i=0;i<n_boost;i++) {
      bmachine.n_trainers = i;
      boosting.n_trainers = i;
      boosting.test(measurers);
    }
  } else {
    boosting.train(measurers);
  }
  if (strcmp(save_model,"")) {
    char save_model_name[100];
    sprintf(save_model_name,"%s/%s",dir_name,save_model);
    boosting.save(save_model_name);
  }

  //=================== The End =====================

  for (int i=0;i<n_boost;i++) {
    delete mlp[i];
    delete mse[i];
    delete trainer[i];
    delete opt[i];
    freeList(&trainers_measurers[i]);
  }

  delete class_format;

  delete[] mlp;
  delete[] mse;
  delete[] trainer;
  delete[] opt;
  delete[] msemeasurer;
  delete[] classmeasurer;
  delete[] trainers_measurers;

  freeList(&measurers);

  return(0);
}