/**
* @file SemSegTools.cpp
* @brief tools for semantic segmentation
* @author Erik Rodner, Sven Sickert
* @date 03/19/2009

*/
#include <iostream>

#include "SemSegTools.h"

using namespace OBJREC;

using namespace std;
using namespace NICE;

#undef DEBUG_LOCALIZATION
#undef DEBUG

void SemSegTools::segmentToOverlay (
        const NICE::Image *orig,
        const NICE::ColorImage & segment,
        NICE::ColorImage & result )
{
    int xsize = orig->width();
    int ysize = orig->height();

    result.resize( xsize, ysize );
    std::vector< NICE::MatrixT<double> > channelMat;

    double alpha = .3;

    for (int c = 0; c < 3; c++)
    {
        NICE::MatrixT<double> chan ( xsize, ysize );
        channelMat.push_back( chan );
    }

    for (int y = 0; y < ysize; y++)
        for (int x = 0; x < xsize; x++)
        {
            uchar val = orig->getPixelQuick(x,y);
            for (int c = 0; c < 3; c++)
                channelMat[c](x,y) = alpha*(double)val
                                   + (1.0-alpha)*(double)segment.getPixel( x, y, c );
        }

    for (int y = 0; y < ysize; y++)
        for (int x = 0; x < xsize; x++)
            for (int c = 0; c < 3; c++)
            {
                int val = channelMat[c](x,y);
                result.setPixel( x, y, c, (uchar)val);
            }

}

void SemSegTools::updateConfusionMatrix(
        const Image &img,
        const Image &gt,
        Matrix &M,
        const std::set<int> &forbiddenClasses )
{
    double subsamplex = gt.width() / ( double ) img.width();
    double subsampley = gt.height() / ( double ) img.height();

    for ( int y = 0 ; y < gt.height() ; y++ )
        for ( int x = 0 ; x < gt.width() ; x++ )
        {
            int xx = ( int ) ( x / subsamplex );
            int yy = ( int ) ( y / subsampley );

            if ( xx < 0 ) xx = 0;

            if ( yy < 0 ) yy = 0;

            if ( xx > img.width() - 1 ) xx = img.width() - 1;

            if ( yy > img.height() - 1 ) yy = img.height() - 1;

            int cimg = img.getPixel ( xx, yy );
            int gimg = gt.getPixel ( x, y );

            if ( forbiddenClasses.find ( gimg ) == forbiddenClasses.end() )
            {
                M ( gimg, cimg ) ++;
            }
        }
}

void SemSegTools::computeClassificationStatistics(
        Matrix &confMat,
        const ClassNames &classNames,
        const std::set<int> &forbiddenClasses )
{
    double overallTrue = 0.0;
    double sumAll  = 0.0;

    // print confusion matrix & get overall recognition rate
    std::cout << "Confusion Matrix:" <<  std::endl;
    for ( int r = 0; r < (int) confMat.rows(); r++ )
    {
        for ( int c = 0; c < (int) confMat.cols(); c++ )
        {
            if ( r == c )
                overallTrue += confMat( r, c );

            sumAll += confMat( r, c );
            std::cout << confMat( r, c ) << " ";
        }
        std::cout << std::endl;
    }
    overallTrue /= sumAll;

    // binary classification metrics
    double precision, recall, f1score = -1.0;
    if ( confMat.rows() == 2 )
    {
        precision = (double)confMat(1,1) / (double)(confMat(1,1)+confMat(0,1));
        recall = (double)confMat(1,1) / (double)(confMat(1,1)+confMat(1,0));
        f1score = 2.0*(precision*recall)/(precision+recall);
    }

    // normalizing confMat using rows
    for ( int r = 0 ; r < (int) confMat.rows() ; r++ )
    {
        double sum = 0.0;

        for ( int c = 0 ; c < (int) confMat.cols() ; c++ )
            sum += confMat ( r, c );

        if ( std::fabs ( sum ) > 1e-4 )
            for ( int c = 0 ; c < (int) confMat.cols() ; c++ )
                confMat ( r, c ) /= sum;
    }

    // get average recognition rate
    double avgTrue = 0.0;
    int classesTrained = 0;
    for ( int r = 0 ; r < (int) confMat.rows() ; r++ )
    {
        if ( classNames.existsClassno ( r )
             && ( forbiddenClasses.find ( r ) == forbiddenClasses.end() ) )
        {
            avgTrue += confMat ( r, r );
            double lsum = 0.0;
            for ( int r2 = 0; r2 < ( int ) confMat.rows(); r2++ )
                lsum += confMat ( r,r2 );

            if ( lsum != 0.0 )
                classesTrained++;
        }
    }

    // print classification statistics
    std::cout << "\nOverall Recogntion Rate: " << overallTrue;
    std::cout << "\nAverage Recogntion Rate: " << avgTrue / ( classesTrained );
    std::cout << "\nLower Bound: " << 1.0 /(double)classesTrained;
    std::cout << "\nPrecision: " << precision;
    std::cout << "\nRecall: " << recall;
    std::cout << "\nF1Score: " << f1score;

    std::cout <<"\n\nClasses:" << std::endl;
    for ( int r = 0 ; r < (int) confMat.rows() ; r++ )
    {
        if ( classNames.existsClassno ( r )
             && ( forbiddenClasses.find ( r ) == forbiddenClasses.end() ) )
        {
            std::string cname = classNames.text ( r );
            std::cout << cname.c_str() << ": " << confMat ( r, r ) << std::endl;
        }
    }

}

void SemSegTools::saveResultsToImageFile(
        const Config *conf,
        const string &section,
        const ColorImage &orig,
        const ColorImage &gtruth,
        const ColorImage &segment,
        const string &file )
{
    std::string resultDir = conf->gS ( section, "resultdir", "." );
    std::string outputType = conf->gS ( section, "output_type", "ppm" );
    std::string outputPostfix = conf->gS ( section, "output_postfix", "" );

    NICE::ColorImage overlaySegment, overlayGTruth;

    NICE::Image* origGrey = orig.getChannel(1);
    segmentToOverlay( origGrey, segment, overlaySegment );
    segmentToOverlay( origGrey, gtruth, overlayGTruth );

    std::stringstream out;
    out << resultDir << "/" << file << outputPostfix;

#ifdef DEBUG
    std::cout << "Writing to file " << out.str() << "_*." << outputType << std::endl;
#endif

    orig.write ( out.str() + "_orig." + outputType );
    segment.write ( out.str() + "_result." + outputType );
    gtruth.write ( out.str() + "_groundtruth." + outputType );
    overlaySegment.write ( out.str() + "_overlay_res." + outputType );
    overlayGTruth.write ( out.str() + "_overlay_gt." + outputType );

}

void SemSegTools::collectTrainingExamples ( 
        const Config * conf,
        const std::string & section,
        const LabeledSet & train,
        const ClassNames & cn,
        Examples & examples,
        vector<CachedExample *> & imgexamples )
{
    assert ( train.count() > 0 );
    examples.clear();
    imgexamples.clear();

    int grid_size_x = conf->gI(section, "grid_size_x", 5 );
    int grid_size_y = conf->gI(section, "grid_size_y", 5 );
    int grid_border_x = conf->gI(section, "grid_border_x", 20 );
    int grid_border_y = conf->gI(section, "grid_border_y", 20 );

    std::string selection = conf->gS(section, "train_selection" );

    set<int> classnoSelection;
    cn.getSelection ( selection, classnoSelection );
    
    bool useExcludedAsBG = conf->gB(section, "use_excluded_as_background", false );

    int backgroundClassNo = 0;
    
    if ( useExcludedAsBG ) 
    {
        backgroundClassNo = cn.classno("various");
        assert ( backgroundClassNo >= 0 );
    }

    LOOP_ALL_S (train)
    {
        EACH_INFO(image_classno,imgInfo);
        std::string imgfn = imgInfo.img();

        if ( ! imgInfo.hasLocalizationInfo() ) {
            std::cerr << "WARNING: NO localization info found for "
                      << imgfn << " !" << std::endl;
            continue;
        }

        int xsize, ysize;
        CachedExample *ce = new CachedExample ( imgfn );
        ce->getImageSize ( xsize, ysize );
        imgexamples.push_back ( ce );

        const LocalizationResult *locResult = imgInfo.localization();
        if ( locResult->size() <= 0 ) {
            std::cerr << "WARNING: NO ground truth polygons found for "
                      << imgfn << " !" << std::endl;
            continue;
        }

        std::cerr << "SemSegTools: Collecting pixel examples from localization info: "
                  << imgfn << std::endl;

        NICE::Image pixelLabels (xsize, ysize);
        pixelLabels.set(0);
        locResult->calcLabeledImage ( pixelLabels, cn.getBackgroundClass() );

    #ifdef DEBUG_LOCALIZATION
        NICE::Image img (imgfn);
        showImage(img);
        showImage(pixelLabels);
    #endif

        Example pce ( ce, 0, 0 );
        for ( int x = 0 ; x < xsize ; x += grid_size_x )
            for ( int y = 0 ; y < ysize ; y += grid_size_y )
            {
                if ( (x >= grid_border_x) &&
                    ( y >= grid_border_y ) && ( x < xsize - grid_border_x ) &&
                    ( y < ysize - grid_border_x ) )
                {
                    pce.x = x; pce.y = y;
                    int classno = pixelLabels.getPixel(x,y);

                    if ( classnoSelection.find(classno) != classnoSelection.end() ) {
                    examples.push_back ( pair<int, Example> (
                        classno,
                        pce // FIXME: offset handling
                    ) );
                    } else if ( useExcludedAsBG ) {
                    examples.push_back ( pair<int, Example> (
                        backgroundClassNo,
                        pce // FIXME: offset handling
                    ) );
                    }
                }
            }
    }

    std::cerr << "total number of examples: " << (int)examples.size() << std::endl;
}