/**
 * @file SemSegObliqueTree.h
 * @brief Semantic Segmentation using Oblique Trees
 * @author Sven Sickert
 * @date 10/17/2014

*/

#include <iostream>

#include "SemSegObliqueTree.h"
#include "SemSegTools.h"

#include "vislearning/classifier/fpclassifier/randomforest/FPCRandomForests.h"
#include "vislearning/features/fpfeatures/ConvolutionFeature.h"
#include "vislearning/baselib/cc.h"

using namespace OBJREC;

using namespace std;
using namespace NICE;

//###################### CONSTRUCTORS #########################//

SemSegObliqueTree::SemSegObliqueTree () : SemanticSegmentation ()
{
    conf = NULL;

    saveLoadData = false;
    fileLocation = "classifier.data";
    colorMode = 0;

    fpc = new FPCRandomForests ();
}

SemSegObliqueTree::SemSegObliqueTree (
        const Config *conf,
        const ClassNames *classNames )
    : SemanticSegmentation( conf, classNames )
{
    initFromConfig( conf );
}

//###################### DESTRUCTORS ##########################//

SemSegObliqueTree::~SemSegObliqueTree ()
{
    if ( fpc != NULL )
        delete fpc;
}

//#################### MEMBER FUNCTIONS #######################//

void SemSegObliqueTree::preprocessChannels (
        CachedExample *ce,
        bool isColor ) const
{
    NICE::MultiChannelImage3DT<int> * img = NULL;
    NICE::MultiChannelImage3DT<double> * imgD = NULL;

    imgD = & ce->getDChannel( CachedExample::D_EOH );
    assert( imgD->channels() == 0 );

    if ( isColor )
    {
        img = & ce->getIChannel( CachedExample::I_COLOR );

        imgD->reInit ( img->width(), img->height(), img->depth(), 3 );

        for ( int z = 0; z < img->depth(); z++ )
            for ( int y = 0; y < img->height(); y++ )
                for ( int x = 0; x < img->width(); x++ )
                {
                    double r = (double)img->get( x, y, z, 0);
                    double g = (double)img->get( x, y, z, 1);
                    double b = (double)img->get( x, y, z, 2);

                    if ( colorMode == 1 )
                    {
                        double h,s,v;
                        ColorConversion::ccRGBtoHSV(r, g, b, &h, &s, &v);
                        imgD->set( x, y, z, h, 0);
                        imgD->set( x, y, z, s, 1);
                        imgD->set( x, y, z, v, 2);
                    }
                    else if ( colorMode == 2 )
                    {
                        double cX, cY, cZ, cL, ca, cb;
                        r /= 255.0;
                        g /= 255.0;
                        b /= 255.0;
                        ColorConversion::ccRGBtoXYZ( r, g, b, &cX, &cY, &cZ, 0 );
                        ColorConversion::ccXYZtoCIE_Lab( cX, cY, cZ, &cL, &ca, &cb, 0 );
                        imgD->set( x, y, z, cL, 0);
                        imgD->set( x, y, z, ca, 1);
                        imgD->set( x, y, z, cb, 2);
                    }
                    else
                    {
                        imgD->set( x, y, z, r/255.0, 0 );
                        imgD->set( x, y, z, g/255.0, 1 );
                        imgD->set( x, y, z, b/255.0, 2 );
                    }
                }

        // remove integer channels
        img->freeData();
    }
    else
    {
        img = & ce->getIChannel( CachedExample::I_GRAYVALUES );

        // gray values to range [0,1]
        imgD->reInit ( img->width(), img->height(), img->depth(), 1 );

        for ( int z = 0; z < img->depth(); z++ )
            for ( int y = 0; y < img->height(); y++ )
                for ( int x = 0; x < img->width(); x++ )
                {
                    double g = (double)img->get( x, y, z, 0) / 255.0;
                    imgD->set( x, y, z, g, 0);
                }

        // remove integer channel
        img->freeData();
    }

    img = NULL;
    imgD = NULL;
}


void SemSegObliqueTree::initFromConfig( const Config *_conf,
                                        const string &s_confSection )
{
    conf = _conf;
    saveLoadData = conf->gB ( s_confSection, "save_load_data", false );
    colorMode = conf->gI ( s_confSection, "color_mode", 0 );
    fileLocation = conf->gS ( s_confSection, "datafile", "classifier.data" );

    fpc = new FPCRandomForests ( _conf, "FPCRandomForests" );
    fpc->setMaxClassNo( classNames->getMaxClassno() );
}

/** training function */
void SemSegObliqueTree::train ( const MultiDataset *md )
{
    if ( saveLoadData && FileMgt::fileExists( fileLocation ) )
    {
        read( fileLocation, 1 );
    }
    else
    {
        Examples examples;

        // image storage
        vector<CachedExample *> imgexamples;

        // create pixel-wise training examples
        SemSegTools::collectTrainingExamples (
          conf,
          "FPCRandomForests",
          * ( ( *md ) ["train"] ),
          *classNames,
          examples,
          imgexamples,
          run3Dseg );

        assert ( examples.size() > 0 );

        FeaturePool fp;
        ConvolutionFeature cf ( conf );
        cf.explode( fp );

        for ( vector<CachedExample *>::iterator cei = imgexamples.begin();
              cei != imgexamples.end(); cei++ )
            preprocessChannels ( *cei, cf.isColorMode() );

        // start training using random forests
        fpc->train( fp, examples);

        // save trained classifier to file
        if (saveLoadData) save( fileLocation, 1 );

        // Cleaning up
        for ( vector<CachedExample *>::iterator i = imgexamples.begin();
              i != imgexamples.end();
              i++ )
            delete ( *i );

        fp.destroy();
    }
}

/** classification function */
void SemSegObliqueTree::semanticseg(
        CachedExample *ce,
        ImageT<int> &segresult,
        NICE::MultiChannelImageT<double> &probabilities )
{
    // for speed optimization
    FPCRandomForests *fpcrf = dynamic_cast<FPCRandomForests *> ( fpc );

    int xsize, ysize;
    ce->getImageSize ( xsize, ysize );
    probabilities.reInit ( xsize, ysize, classNames->getMaxClassno() + 1 );
    segresult.resize ( xsize, ysize );

    vector<DecisionTree *> forest = fpcrf->getForest ();
    DecisionNode *root = forest[0]->getRoot ();
    ConvolutionFeature* cf = dynamic_cast<ConvolutionFeature*> (root->f);

    preprocessChannels( ce, cf->isColorMode() ); //FIXME!!!

    Example pce ( ce, 0, 0 );
    for ( int y = 0 ; y < ysize ; y++ )
        for ( int x = 0 ; x < xsize ; x++ )
        {
            pce.x = x;
            pce.y = y;
            ClassificationResult r = fpcrf->classify ( pce );
            segresult.setPixel ( x, y, r.classno );
            for ( int i = 0 ; i < ( int ) probabilities.channels(); i++ )
                probabilities[i](x,y) = r.scores[i];
        }
}

/** classification function 3d */
void SemSegObliqueTree::semanticseg(
        OBJREC::CachedExample *ce,
        NICE::MultiChannelImageT<int> &segresult,
        NICE::MultiChannelImage3DT<double> &probabilities )
{
    // for speed optimization
    FPCRandomForests *fpcrf = dynamic_cast<FPCRandomForests *> ( fpc );

    int xsize, ysize, zsize;
    ce->getImageSize3 ( xsize, ysize, zsize );
    probabilities.reInit ( xsize, ysize, zsize, classNames->getMaxClassno() + 1 );
    segresult.reInit ( xsize, ysize, (unsigned int)zsize );

    vector<DecisionTree *> forest = fpcrf->getForest ();
    DecisionNode *root = forest[0]->getRoot ();
    ConvolutionFeature* cf = dynamic_cast<ConvolutionFeature*> (root->f);

    preprocessChannels( ce, cf->isColorMode() ); //FIXME!!!

    Example pce ( ce, 0, 0, 0 );
    for ( int z = 0 ; z < zsize; z++ )
        for ( int y = 0 ; y < ysize ; y++ )
            for ( int x = 0 ; x < xsize ; x++ )
            {
                pce.x = x;
                pce.y = y;
                pce.z = z;
                ClassificationResult r = fpcrf->classify ( pce );
                segresult.set ( x, y, r.classno, (unsigned int)z );
                for ( int i = 0 ; i < ( int ) probabilities.channels(); i++ )
                    probabilities.set ( x, y, z, r.scores[i], i );
            }
}

///////////////////// INTERFACE PERSISTENT /////////////////////
// interface specific methods for store and restore
///////////////////// INTERFACE PERSISTENT /////////////////////

void SemSegObliqueTree::restore( istream &is, int format )
{
    //dirty solution to circumvent the const-flag
    const_cast<ClassNames*>(this->classNames)->restore ( is, format );

    fpc->restore( is, format );
}

void SemSegObliqueTree::store ( ostream &os, int format ) const
{
    classNames->store( os, format );
    fpc->store( os, format );
}

void SemSegObliqueTree::clear ( )
{
    fpc->clear();
}