// Beispielhafter Aufruf: BUILD_x86_64/progs/testSemanticSegmentation3D -config <CONFIGFILE>

/**
* @file testSemanticSegmentation3D.cpp
* @brief test semantic segmentation routines for 3d images and 2d images
* @author Erik Rodner, Björn Fröhlich, Sven Sickert
* @date 03/20/2008
*/

#ifdef NICE_USELIB_OPENMP
#include <omp.h>
#endif

#include "core/basics/Config.h"
#include "core/basics/StringTools.h"
#include "core/image/Morph.h"
#include "core/image/MultiChannelImage3DT.h"

#include "vislearning/baselib/ICETools.h"

#include "semseg/semseg/SemSegContextTree3D.h"
#include "semseg/semseg/SemSegTools.h"

#include <fstream>
#include <vector>

#undef DEBUG

using namespace OBJREC;

using namespace NICE;

using namespace std;

void startClassification (SemanticSegmentation *semseg,
                          std::vector< NICE::Matrix > & M_vec,
                          const Config & conf,
                          const LabeledSet* testFiles,
                          const ClassNames & classNames,
                          const set<int> & forbidden_classes,
                          std::map<int,int> & classMapping,
                          const unsigned short cvRuns)
{
    bool write_results = conf.gB ( "debug", "write_results", false );
    bool writeProbMaps = conf.gB ( "debug", "write_prob_maps", false );
    if (cvRuns > 1)
        write_results = false;

    bool run_3Dseg = conf.gB( "SSContextTree", "run_3dseg", false);
    string output_type = conf.gS ( "debug", "output_type", "ppm" );

    vector< int > zsizeVec;
    semseg->getDepthVector ( testFiles, zsizeVec, run_3Dseg );
    int depthCount = 0, idx = 0;
    vector< string > filelist;
    NICE::MultiChannelImageT<int> segresult;
    NICE::MultiChannelImageT<int> gt;

    for (LabeledSet::const_iterator it = testFiles->begin(); it != testFiles->end(); it++)
    {
        for (std::vector<ImageInfo *>::const_iterator jt = it->second.begin();
             jt != it->second.end(); jt++)
        {
            ImageInfo & info = *(*jt);
            std::string file = info.img();
            filelist.push_back ( file );
            depthCount++;

            NICE::ImageT<int> lm;
            NICE::ImageT<int> lm_gt;
            if ( info.hasLocalizationInfo() )
            {
                const LocalizationResult *l_gt = info.localization();

                lm.resize ( l_gt->xsize, l_gt->ysize );
                lm.set ( 0 );

                lm_gt.resize ( l_gt->xsize, l_gt->ysize );
                lm_gt.set ( 0 );

                l_gt->calcLabeledImage ( lm, classNames.getBackgroundClass() );
#ifdef DEBUG
                cout << "testSemanticSegmentation3D: Generating Labeled NICE::Image (Ground-Truth)" << endl;
#endif
                l_gt->calcLabeledImage ( lm_gt, classNames.getBackgroundClass() );
            }
            segresult.addChannel ( lm );
            gt.addChannel ( lm_gt );

            int depthBoundary = 0;
            if ( run_3Dseg )
            {
                depthBoundary = zsizeVec[idx];
            }

            if ( depthCount < depthBoundary ) continue;

            NICE::MultiChannelImage3DT<double> probabilities;

            semseg->classify ( filelist, segresult, probabilities );

            // save to file
            for ( int z = 0; z < segresult.channels(); z++ )
            {
                NICE::ColorImage orig ( filelist[z] );
                NICE::ColorImage rgb;
                NICE::ColorImage rgb_gt;

                for ( int y = 0 ; y < orig.height(); y++ )
                    for ( int x = 0 ; x < orig.width(); x++ )
                    {
                        lm.setPixel ( x, y, segresult.get ( x, y, ( uint ) z ) );
                        if ( run_3Dseg )
                            lm_gt.setPixel ( x, y, gt.get ( x, y, ( uint ) z ) );
                    }

                // confusion matrix
                NICE::Matrix M ( classMapping.size(), classMapping.size() );
                M.set ( 0 );
                SemSegTools::updateConfusionMatrix (
                            lm, lm_gt, M, forbidden_classes, classMapping );
                M_vec.push_back ( M );

                classNames.labelToRGB ( lm, rgb );
                classNames.labelToRGB ( lm_gt, rgb_gt );

                if ( write_results )
                {
                    std::string fname = StringTools::baseName ( filelist[z], false );
                    std::string outStr;
                    SemSegTools::saveResultsToImageFile ( &conf, "debug",
                                orig, rgb_gt, rgb, fname, outStr );

                    // write Probability maps
                    if (writeProbMaps)
                    {
                        NICE::ColorImage prob_map( probabilities.width(), probabilities.height() );
                        prob_map.set(0,0,0);
                        int iNumChannels = probabilities.channels();
                        for ( int idxProbMap = 0; idxProbMap < iNumChannels; idxProbMap++)
                        {
                            for ( int y = 0 ; y < probabilities.height(); y++ )
                            {
                                for ( int x = 0 ; x < probabilities.width(); x++ )
                                {
                                    double probVal = probabilities.get( x, y, z, idxProbMap ) * 255.0;
                                    for ( int c = 0 ; c < 3 ; c++ )
                                        prob_map.setPixel( x, y, c, round(probVal) );
                                }
                            }
                            std::stringstream ssFileProbMap;
                            ssFileProbMap << outStr << "_probs." << "c-" << classNames.code( idxProbMap ) << "." << output_type;
                            //classNames
                            prob_map.write ( ssFileProbMap.str() );
                        }
                    }
                }
            }

            // prepare for new 3d image
            filelist.clear();
            segresult.reInit(0,0,0);
            gt.reInit(0,0,0);
            depthCount = 0;
            idx++;
        }
    }

    segresult.freeData();
}

/**
 test semantic segmentation routines
*/
int main ( int argc, char **argv )
{
    std::set_terminate ( __gnu_cxx::__verbose_terminate_handler );

    Config conf ( argc, argv );

    ResourceStatistics rs;

    /*---------------CONFIGURATION---------------*/
    unsigned short crossValRuns = conf.gI ( "debug", "cross_val_runs", 1 );
    /*-------------------------------------------*/

#ifdef DEBUG
    cerr << "Writing Results to " << resultdir << endl;
#endif

    std::vector< NICE::Matrix > M_vec;

    MultiDataset md ( &conf );

    const ClassNames & classNames = md.getClassNames ( "train" );
    set<int> forbidden_classes;
    classNames.getSelection ( conf.gS ( "analysis", "forbidden_classes", "" ),
                              forbidden_classes );

    vector<bool> usedClasses ( classNames.numClasses(), true );
    for ( set<int>::const_iterator it = forbidden_classes.begin();
          it != forbidden_classes.end(); ++it)
    {
        usedClasses [ *it ] = false;
    }

    map<int,int> classMapping, classMappingInv;
    int j = 0;
    for ( int i = 0; i < usedClasses.size(); i++ )
        if (usedClasses[i])
        {
            classMapping[i] = j;
            classMappingInv[j] = i;
            j++;
        }


    // initialize semantic segmentation method
    SemanticSegmentation *semseg = NULL;

    // TRAINING AND TESTING
    if ( crossValRuns == 1 )
    {
        semseg = new SemSegContextTree3D ( &conf, &classNames );

        // STANDARD EVALUATION
        cout << "\nTRAINING" << endl;
        cout << "########\n" << endl;
        semseg->train( &md );

        cout << "\nCLASSIFICATION" << endl;
        cout << "##############\n" << endl;
        const LabeledSet *testFiles = md["test"];
        startClassification (semseg, M_vec, conf, testFiles, classNames,
                             forbidden_classes, classMapping, crossValRuns );

        delete semseg;
    }
    else
    {
        // CROSS-VALIDATION
        for (int cval = 1; cval <= crossValRuns; cval++)
        {
            semseg = new SemSegContextTree3D ( &conf, &classNames );

            stringstream ss;
            ss << cval;
            string cvaltrain = "train_cv" + ss.str();
            string cvaltest = "test_cv" + ss.str();

            cout << "\nTRAINING " << cval << endl;
            cout << "###########\n" << endl;
            const LabeledSet *trainFiles = md[cvaltrain];
            semseg->train( trainFiles );

            cout << "\nCLASSIFICATION " << cval << endl;
            cout << "#################\n" << endl;
            const LabeledSet *testFiles = md[cvaltest];
            startClassification (semseg, M_vec, conf, testFiles, classNames,
                                 forbidden_classes, classMapping, crossValRuns );

            delete semseg;
        }
    }

    // resource statistics
    SemSegTools::computeResourceStatistics ( rs );

    NICE::Matrix M ( classMapping.size(), classMapping.size() );
    M.set ( 0 );
    for ( int s = 0; s < ( int ) M_vec.size(); s++ )
    {
        NICE::Matrix M_tmp = M_vec[s];
        for ( int r = 0; r < ( int ) M_tmp.rows(); r++ )
            for ( int c = 0; c < ( int ) M_tmp.cols(); c++ )
                M ( r, c ) += M_tmp ( r, c );
    }
    // evaluation & analysis
    SemSegTools::computeClassificationStatistics(
                M, classNames, forbidden_classes, classMappingInv );

    return 0;
}