#include "MlpLearnAlgorithm.hpp"

#include "Dataset.hpp"
#include "Population.hpp"
#include "MlpFactory.hpp"
#include "MlpIndividual.hpp"
#include "MlpObjectives.hpp"
#include "Mlp.hpp"
#include "Random.hpp"

using namespace jymlp;


#include <boost/property_tree/ptree.hpp>
using boost::property_tree::ptree;

#define BOOST_FILESYSTEM_NO_DEPRECATED
#include <boost/filesystem.hpp>

#include <iostream>
#include <fstream>
#include <string>
#include <memory>
#include <random>
using namespace std;

void
MlpLearnAlgorithm::impl_setup(const ptree & all){
    // FIXME: TODO
}

void
MlpLearnAlgorithm::setup(const ptree & all,
                         const string & outputdirname,
                         const string & datadirname){
    Algorithm::setup(all,outputdirname);
    _hack_copy_of_ptree = all;
    _datadirname = datadirname;
}

void
MlpLearnAlgorithm::run(){
    /* TODO: Most of these first things should go to impl_setup! */
    const ptree & all = _hack_copy_of_ptree;
    const ptree & sub = all.get_child("MLP_learn");

    string datafname = _datadirname + sub.get<string>("datafile");

    string task = sub.get<string>("task");
    size_t popsz = sub.get<int>("evolution.pop_size");
    int maxiter = sub.get<int>("evolution.max_iter");
    string eoper = sub.get<string>("evolution.evol_oper");

    int bpiter = sub.get<int>("improvement.imp_steps");
    double bpstep = sub.get<double>("improvement.imp_inilen");
    string improve_method = sub.get<string>("improvement.imp_obj");

    shared_ptr<Dataset> ds(new Dataset(datafname,-1.0,+1.0)); // Note binvec encoding!

    istringstream iss("init_size: " + sub.get<string>("init.init_arch"));
    MlpFactory mf(iss);

    mt19937 mt;

    MlpObjectives meobj(sub.get<string>("evolution.evol_obj"),ds->getNClasses());

    for (size_t i=0;i<meobj.nobj();++i){
        cout << "f" << (i+1) << ": " << meobj.name(i) << endl;
    }
    // FIXME: Sanity checks!! E.g., must have equal number of output layers and classes etc.
    

    Population parent_pop(meobj.nobj());      // Evolving population
    Population history(meobj.nobj());  // Complete trace of the whole run
    //Population archive;  // Archive of truly non-dominated solutions


    // Mock-up for different improvements:
    vector<double> cweights(ds->getNClasses());
    for (auto & d: cweights) d=1.0/ds->getNClasses(); // default to equal class weights.

    unique_ptr<MlpBackpropUnaryOperator> nextop;

    // Create, write and read some:
    vector<size_t> unparented(1); unparented[0] = 0; // Create unparented.
    for (size_t i=0;i<popsz;++i){
        MlpIndividual *ni = new MlpIndividual(meobj,
                                              unique_ptr<Mlp>(mf.create(&mt)),
                                              ds);
        ni->setId(nextID());
        if (improve_method == "mse") {
            nextop = unique_ptr<MlpBackpropUnaryOperator>
                (new MlpBackpropUnaryOperator(bpiter,bpstep,cweights,false,false,false));
        } else if (improve_method == "mse(direction=randomclass)"){
            randomize1(cweights, &mt);
            nextop = unique_ptr<MlpBackpropUnaryOperator>
                (new MlpBackpropUnaryOperator(bpiter,bpstep,cweights,false,false,false));
        } else if (improve_method == "mse(direction=rerandomclass)"){
            randomize1(cweights, &mt);
            nextop = unique_ptr<MlpBackpropUnaryOperator>
                (new MlpBackpropUnaryOperator(bpiter,bpstep,cweights,true,false,false));
        } else if (improve_method == "ssw(amount=rerandom);mse(direction=rerandomclass)"){
            randomize1(cweights, &mt);
            nextop = unique_ptr<MlpBackpropUnaryOperator>
                (new MlpBackpropUnaryOperator(bpiter,bpstep,cweights,true,true,false));
        } else if (improve_method == "mee/mse"){
            randomize1(cweights, &mt);
            nextop = unique_ptr<MlpBackpropUnaryOperator>
                (new MlpBackpropUnaryOperator(bpiter,bpstep,cweights,true,true,true));
        } else {
            cerr << "Unknown improvement method: " << improve_method << endl;
            exit(1); // FIXME: Implement exception facilities.
        }
        ni->addUnaryOperator(move(nextop));
        markAsParents(ni->getId(),unparented);
        parent_pop.adopt(unique_ptr<MlpIndividual>(ni));
    }

    parent_pop.evaluate();
    // Start by NSGA-II start step:
    parent_pop.assignRankAndCrowdingDistance();

    ofstream ofs(_outputdirname + "pop0.dat", ios::trunc);
    parent_pop.toStream(ofs);
    ofs.close();

    // Trace also the 0 pop:
    for (size_t ii = 0; ii<popsz; ++ii){
        traceMeasures(parent_pop.getIndividualID(ii),parent_pop.getTrace(ii));
    }

    // Mock-up run, to be made flexible:
    history.deepCopyFrom(parent_pop);


    for (int i=1; i <= maxiter; ++i) {
        cerr << "Performing iteration " << i << "/" << maxiter << endl;

        // TODO: selection and collaboration here.

        // Life-time learning / unary operations:
        Population child_pop(meobj.nobj());
        child_pop.deepCopyFrom(parent_pop);
        child_pop.improve(&mt);
        child_pop.evaluate();

        // Mockup update of IDs, and tracing parents:
        vector<size_t> mockpar(1); // as of yet, only one parent possible. TODO: Crossovers etc.
        // FIXME: Parent information can only be updated by the combination process!
        for (size_t ii = 0; ii<popsz; ++ii){
            mockpar[0] = child_pop.getIndividualID(ii);
            child_pop.setIndividualID(ii, nextID());
            markAsParents(child_pop.getIndividualID(ii),mockpar);
        }

        history.deepCopyFrom(child_pop);


        if (eoper=="none"){
            parent_pop.clear();
            parent_pop.adoptFrom(child_pop);
            // FIXME: The following doesn't seem to work for single-objective stuff(what gives?):
            // parent_pop.fillNondominatedSortFrom(child_pop);
        } else if (eoper=="ns"){
            // Prepare next population; this will assign ranks:
            Population mixed_pop(meobj.nobj());
            mixed_pop.adoptFromTwo(parent_pop, child_pop);
            //parent_pop.clear();
            parent_pop.fillNondominatedSortFrom(mixed_pop);
        }

        // Trace measures:
        for (size_t ii = 0; ii<popsz; ++ii){
            traceMeasures(parent_pop.getIndividualID(ii),parent_pop.getTrace(ii));
        }
    }

    // Write things for post-run exploration.. TODO: Real-time logging and tracing!
    ofstream info(_outputdirname + "info.txt", ios::trunc);
    info << "Format: 2nd line is data file name." << endl;
    info << sub.get<string>("datafile");

    ofstream ofs3(_outputdirname + "pop_final.dat", ios::trunc);
    parent_pop.toStream(ofs3);

    // TODO: History is unnecessary, if separate populations are stored?
    // (The auxiliary non-dominated archive is a completely different story...)
    ofstream ofs2(_outputdirname + "pop_history.dat", ios::trunc);
    history.toStream(ofs2);

    // FIXME: Should update this during the run, not only at the end:
    /* And.. it segfaults.. hmm.. should debug this properly
    Population nondom_all(meobj.nobj());
    nondom_all.fillNondominatedSortFrom(history);
    ofstream ofsa(_outputdirname + "pop_nondom.dat", ios::trunc);
    nondom_all.toStream(ofsa);
    */

    // Parent graph:
    ofstream ofs4(_outputdirname + "parents.dat", ios::trunc);
    writeParentHistory(ofs4); ofs4.close();

    // Traced measures:
    ofstream ofs5(_outputdirname + "traced.dat", ios::trunc);
    writeTracedMeasures(ofs5);
}