/** @file
 *  @brief Test weighted training sets with neural networks.
 *
 *  $Id: nntest.cpp 2073 2005-05-17 07:56:53Z ling $
 */

#include <iostream>
#include <fstream>
#include "lemga/random.h"
#include "lemga/nnlayer.h"
#include "lemga/feedforwardnn.h"
#include <lemga/stump.h>
#include <lemga/adaboost.h>

using namespace lemga;

int main (unsigned int argc, char* argv[]) {
    if (argc < 7) {
        std::cerr << "Usage: " << argv[0] << " dat wgt tdat trials"
                  << " #_input #_hidden ... #_output [output]\n"
                  << "    or " << argv[0] << " dat wgt tdat trials"
                  << " #_input #_boosting [output]\n";
        return -1;
    }
    UINT trials = std::atoi(argv[4]);
    const UINT n_in = std::atoi(argv[5]);
    const UINT n_out = 1;
    const char* outf = "test.err";
    if (argv[argc-1][0] < '0' || argv[argc-1][0] >'9')
        outf = argv[--argc];

    set_seed(0);

    LearnModel *lm;
    if (argc > 7) {
        /* constructing neural nets with layers */
        FeedForwardNN nn;

        std::cout << "Neural network " << n_in;
        UINT l_in = n_in;
        for (UINT i = 6; i < argc; ++i) {
            UINT l_out = atoi(argv[i]);
            std::cout << 'x' << l_out;

            NNLayer l(l_in, l_out);
            l.set_weight_range(-0.2, 0.2);
            nn.add_top(l);

            l_in = l_out;
        }
        assert(l_in == n_out);
        std::cout << " constructed\n";
        nn.set_batch_mode();
        nn.set_train_method(nn.CONJUGATE_GRADIENT);
        nn.set_parameter(0.1, 1e-4, 1000);
        lm = nn.clone();
    } else {
        /* boosting */
        Stump st(n_in);
        UINT n_boost = atoi(argv[6]);
        AdaBoost ab(false);
        ab.set_max_models(n_boost);
        ab.set_base_model(st);
        ab.use_gradient_descent(false);
        std::cout << "AdaBoost with " << n_boost << " iterations\n";
        lm = ab.clone();
    }

    /* load training and test data */
    std::ifstream fd(argv[1]);
    if (!fd.is_open()) {
        std::cerr << "nntest: training data open error\n"; return -2;
    }
    pDataSet trd = load_data(fd, (1L<<30)-1, n_in, n_out);

    fd.clear(); fd.close(); fd.open(argv[2]);
    if (!fd.is_open()) {
        std::cerr << "nntest: training weights open error\n"; return -2;
    }
    pDataWgt trw;
    {
        DataWgt* pw = new DataWgt;
        for (;;) {
            REAL x;
            if (!(fd >> x)) break;
            pw->push_back(x);
        }
        trw = pw;
    }
    std::cout << trw->size() << " weights\n";

    fd.clear(); fd.close(); fd.open(argv[3]);
    if (!fd.is_open()) {
        std::cerr << "nntest: test data open error\n"; return -2;
    }
    pDataSet ted = load_data(fd, (1L<<30)-1, n_in, n_out);

    std::cout << trd->size() << " training samples and "
              << ted->size() << " test samples loaded\n";
    fd.close();

    std::ofstream fe(outf);
    while (trials--) {
        /* train the neural network */
        lm->initialize();
        lm->set_train_data(trd, trw);
        lm->train();

        /* test the network */
        UINT tre = 0, tee = 0;
        for (UINT i = 0; i < trd->size(); ++i)
            tre += ((*lm)(trd->x(i))[0] * trd->y(i)[0] <= 0);
        for (UINT i = 0; i < ted->size(); ++i)
            tee += ((*lm)(ted->x(i))[0] * ted->y(i)[0] <= 0);
        std::cout << "training error (no wgt): " << tre*100.0/trd->size()
                  << "%, test error: " << tee*100.0/ted->size() << "%\n";

        fe << tre << ' ' << tee << '\n';
    }

    return 0;
}