/*
 *  main.cpp
 *
 *  Copyright (c) 2006 Machine Perception Laboratory
 *  University of California San Diego.
 *
 *	Authors: Andrew Salamon, Josh Susskind
 *
 * Please read the disclaimer and notes about redistribution
 * at the end of this file.
 *
 */


#include <iostream> // for cout and cerr
#include <string>

#include <boost/filesystem/path.hpp> // For system independent path support
#include <boost/filesystem/exception.hpp>
#include <boost/filesystem/operations.hpp>
#include <boost/filesystem/convenience.hpp> // For basename and extension
namespace bf = boost::filesystem;

#include "ImageDirIterator.h"

//#define USE_IMAGEMAGICK
#ifdef USE_IMAGEMAGICK
#include <Magick++.h>
#endif

#ifdef WIN32
#ifndef WINVER
#define WINVER 0x0502
#endif
#include "afxwin.h"
#endif

#include "rimage.h"
#include "rimage_ser.h"
#include "mp_CERT.h"

#ifdef USE_MLR
#include "mp_MLR.h"
mp_MLR *setupMLR();
mp_MLR *setupMLR2();
#endif

void printResults( std::vector<double> &results, bool tabDelim )
{
    std::cout.precision(8);
    for( unsigned int i=0; i < results.size(); ++i )
    {
        std::cout << results[i];
        if( i < (results.size() - 1) )
        {
            if( tabDelim )
                std::cout << '\t';
            else
                std::cout << ", ";
        }
    }
    if( !tabDelim )
        std::cout << std::endl;
}

std::string getsvmweightspath();
std::string getimagepath();
bool loadRImage(RImage<float> &pixels, const std::string &path, int argc);
void processPath( MP_CERT &cert, RImage<float> &pixels, const std::string &path, int argc, bool tabDelim );
bool processImage( MP_CERT &cert, RImage<float> &pixels, const std::string &path, int argc, bool tabDelim );
void processDir( MP_CERT &cert, RImage<float> &pixels, const std::string &path, int argc, bool tabDelim );
bool processArgs( int argc, char *argv[], std::vector< std::string > &paths, bool &tabDelim );

// Command line options
bool doEmotions = true; // if this is set to true, emotion weights initialization code in MP_CERT also needs to be uncommented.
bool inputIsFaces = false; // if this is set to true, faces have have already been found and cropped
bool outputPatches = false; // if this is set to true, write face patches out to files
std::string patchDir("/tmp");

#ifdef USE_MLR
    mp_MLR *mlr = NULL;
    mp_MLR *mlr2 = NULL;
#endif

int main (int argc, char *argv[])
{
    bool tabDelim(false);
    std::vector< std::string > paths;

    if( !processArgs( argc, argv, paths, tabDelim ) )
        return -1;

	MP_CERT cert(getsvmweightspath());

#ifdef USE_FEATURES
    {
#ifdef WIN32
        std::string fdDir( "../../featuredetector/frozenData/" );
#else
        std::string fdDir( "./featuredetector/frozenData/" );
#endif

        try {
            cert.loadFeatureDetector( fdDir );
        }
        catch( const std::string error )
        {
            std::cerr << error << std::endl;
            return -2;
        }

        cert.setUseFeatures( true );
    }
#endif

#ifdef USE_MLR
    mlr = setupMLR();
    mlr2 = setupMLR2();
#endif

    cert.setDoEmotions( doEmotions );

	std::cout << "Finished loading data files" << std::endl;

	RImage<float> pixels;

    if( paths.size() > 0 )
    {
        for( unsigned int index = 0; index < paths.size(); ++index )
        {
            processPath( cert, pixels, paths[index], argc, tabDelim );
        }
    }
    else
    {
        do {
            std::string path = getimagepath();

            if( "quit" == path )
                break;

            processPath( cert, pixels, path, argc, tabDelim );
        } while( true );
    }

#ifdef USE_MLR
    delete mlr;
    delete mlr2;
#endif
}

/*************************************************************/

std::string getsvmweightspath() {
	std::string defaultpath = "../../mp_auCoder/SVMWeights";
#ifndef WIN32
    defaultpath = "SVMWeights";
#endif
//	std::string path = "";
//	std::cout << "Enter path to svmweights directory [" << defaultpath << "]: ";
//	std::cin >> path;
//	if ( path == "" ) path = defaultpath;
	std::cout << "Loading data files... " << std::endl;
	return(defaultpath);
}

/*************************************************************/

std::string getimagepath() {
	std::string path;
	std::cout << "Enter path to image";
#ifdef WIN32
#ifndef USE_IMAGEMAGICK
	std::cout << " (Windows demo only accepts 24 bit bitmaps)";
#endif
#endif
	std::cout << " or ctrl-c to exit" << std::endl;
	std::cout << "--> ";
	std::cin >> path;
    if( std::cin.eof() )
        path = "quit";
	return(path);
}

/*************************************************************/

bool loadRImage(RImage<float> &pixels, const std::string &path, int argc)
{
#ifdef USE_IMAGEMAGICK
    // Before trying to read the image we need to make sure the path is valid,
    // otherwise ImageMagick terminates the program.
    Magick::Image image;

    try {
        image.read( path.c_str() );
        /*
            Or we could create an image from in memory data (e.g. video frame capture)
            Magick::Blob blob( inMemoryImageDataPtr, inMemoryImageDataLength );
            Magick::Image image( blob );
        */
        if (image.columns() == 0 && image.rows() == 0)
        {
            std::cerr << "Image contains 0 rows and 0 columns! Exiting..." << std::endl;
            return false;
        }

        image.quantizeColors(256); // convert to 256 colors
        image.quantize(); // convert to grayscale

        pixels.setSize( image.columns(), image.rows() );
        image.write( 0, 0, image.columns(), image.rows(), "I", Magick::FloatPixel, pixels.array );

        return true;
    }
    catch( Magick::ErrorBlob &error )
    {
        std::cerr << "Caught ImageMagick IO exception: " << error.what() << std::endl;
        return false;
    }
    catch( Magick::Exception &error )
    {
        std::cerr << "Caught ImageMagick exception: " << error.what() << std::endl;
        return false;
    }
    catch( std::exception &error )
    {
        std::cerr << "Caught standard exception: " << error.what() << std::endl;
        return false;
    }
#elif WIN32
	// Load bitmap
	HBITMAP hBitmap = (HBITMAP) LoadImage(NULL, path.c_str(), IMAGE_BITMAP, 0, 0, LR_LOADFROMFILE | LR_CREATEDIBSECTION);
	if (hBitmap)
	{
		CBitmap bmpBitmap;
		// Create a CBitmap
		bmpBitmap.Attach(hBitmap);
		BITMAP bm;
		// Get the Bitmap
		bmpBitmap.GetBitmap(&bm);

		if(bm.bmBitsPixel==24)
		// if it's 24 bits ...
		{
			// a bit tricky ...
			// align within 32 bits
			int align = (4-((bm.bmWidth)%4)) & 3;
			int nbr_byte = bm.bmHeight * bm.bmWidth + align * bm.bmHeight/3;
			BYTE* pBits= (BYTE*) bm.bmBits + (bm.bmHeight-1)*(3*bm.bmWidth+align);
			pixels.setSize(bm.bmWidth, bm.bmHeight);
			// read bitmap pixels, convert to grayscale, and copy into RImage for analysis
			for (int l=0;l<bm.bmHeight;l++)
			{
				for (int c=0;c<bm.bmWidth;c++)
				{
					int b=*(pBits);
					int g=*(pBits+1);
					int r=*(pBits+2);
					int theGrayValue = ((306 * (r & 0xff)) +
										(601 * (g & 0xff)) +
										(117 * (b & 0xff))) >> 10;
					pixels.setPixel( c, l, theGrayValue/255.0f ); // copy to RImage
					pBits=pBits+3;
				}
				pBits=pBits-align-2*3*bm.bmWidth;
			}
            return true;
		}
		else
		{
            if( !restoreRImageFromFile( pixels, path ) )
            {
                std::cerr << "Unable to read image file. Exiting..." << std::endl;
                return false;
            }
            else
            {
                return true;
            }
        }
	}
	else {
		std::cerr << "Unable to read image file. Exiting..." << std::endl;
		return false;
	}
#else
	std::cerr << "Unable to read image file. Exiting..." << std::endl;
	return false;
#endif
}

void processPath( MP_CERT &cert, RImage<float> &pixels, const std::string &path, int argc, bool tabDelim )
{
    bf::path bfPath( path, bf::native );
    if( bf::exists( bfPath ) )
    {
        if( bf::is_directory( bfPath ) )
            processDir( cert, pixels, path, argc, tabDelim );
        else
            processImage( cert, pixels, path, argc, tabDelim );
    }
}

bool processImage( MP_CERT &cert, RImage<float> &pixels, const std::string &path, int argc, bool tabDelim )
{
    if( loadRImage(pixels, path, argc) ) // handles loading an image from file into an RImage
    {
        std::vector<double> aus;

        if( inputIsFaces )
            aus = cert.calcAUsFromFace( pixels );
        else
            aus = cert.calcAUs( pixels );

	if( cert.didFindFace() )
	{
		if( !tabDelim )
		    std::cout << "AU vector:" << std::endl;

		printResults( aus, tabDelim );

		if( doEmotions )
		{
		    std::vector<double> emotions = cert.getEmotions();
		    if( !tabDelim )
			std::cout << std::endl << "Emotion Detectors: " << std::endl;
		    std::cout << "\t";
		    printResults( emotions, tabDelim );
		}

#ifdef USE_MLR
		std::map< std::string, mp_AUResults > freqs = cert.getFrequencies();
		try {
			double res = mlr->MLR( freqs );
			if( !tabDelim )
				std::cout << std::endl << "MLR 1+2+4: " << res << std::endl;
			else
				std::cout << '\t' << res;
		}
		catch( std::logic_error &err )
		{
			std::cerr << "Failed to run MLR 1+2+4: " << err.what() << std::endl;
			if( !tabDelim )
				std::cout << std::endl << "MLR 1+2+4: Failed" << std::endl;
			else
				std::cout << '\t' << "NaN";
		}

		try {
			double res = mlr2->MLR( freqs );
			if( !tabDelim )
				std::cout << std::endl << "MLR 1-2: " << res << std::endl;
			else
				std::cout << '\t' << res;
		}
		catch( std::logic_error &err )
		{
			std::cerr << "Failed to run MLR 1-2: " << err.what() << std::endl;
			if( !tabDelim )
				std::cout << std::endl << "MLR 1-2: Failed" << std::endl;
			else
				std::cout << '\t' << "NaN";
		}
#endif

		std::vector<float> eyes = cert.eyes();
		std::cout << "\tEyes: {" << eyes[0] << ", " << eyes[1] << "} {" << eyes[2] << ", " << eyes[3] << "}" << "\t";
	}
	else
	{
		if( !tabDelim )
			std::cout << "No face found." << std::endl;
	}

        if( outputPatches && cert.didFindFace() )
        {
            RImage<float> patch = cert.getPatch();

            // generate a filename with _face appended.
            bf::path bfPath(patchDir);
            bf::path fname(path, bf::native);

            bfPath /= bf::basename(fname) + "_face.rimage";
            //bfPath = bf::change_extension( bfPath, bf::extension(fname) );

            if( !saveRImageToFile( patch, bfPath.native_file_string() ) )
            {
                std::cerr << "Unable to write face to file: " << bfPath.native_file_string() << std::endl;
            }
        }

        return true;
    }

    return false;
}

void processDir( MP_CERT &cert, RImage<float> &pixels, const std::string &path, int argc, bool tabDelim )
{
    ImageDirIterator images( path );

    while( ++images )
    {
        std::string fileName = *images;
        std::cout << fileName << '\t';
        processImage( cert, pixels, fileName, argc, tabDelim );
        std::cout << std::endl;
    }
}

#ifdef USE_MLR
mp_MLR *setupMLR()
{
    std::string weightsPath("../../MLR/BHAT_1_2_4.xml");
    inputAUVector inputs;
    // These input values are listed in ../MLR/bhat_requirements.txt
    inputs.push_back( std::pair<std::string,unsigned int>( "EmotionWeights", 0 ) );
    inputs.push_back( std::pair<std::string,unsigned int>( "SVMWeights", 0 ) );
    inputs.push_back( std::pair<std::string,unsigned int>( "SVMWeights", 1 ) );
    inputs.push_back( std::pair<std::string,unsigned int>( "SVMWeights", 2 ) );
    mp_MLR *mlr = new mp_MLR( weightsPath, inputs );

    return mlr;
}

mp_MLR *setupMLR2()
{
    std::string weightsPath("../../MLR/BHAT_1-2.xml");
    inputAUVector inputs;
    // These input values are listed in ../MLR/bhat_requirements.txt
    inputs.push_back( std::pair<std::string,unsigned int>( "EmotionWeights", 1 ) );
    inputs.push_back( std::pair<std::string,unsigned int>( "SVMWeights", 0 ) );
    inputs.push_back( std::pair<std::string,unsigned int>( "SVMWeights", 1 ) );
    inputs.push_back( std::pair<std::string,unsigned int>( "SVMWeights", 2 ) );
    inputs.push_back( std::pair<std::string,unsigned int>( "EmotionWeights", 0 ) );
    mp_MLR *mlr = new mp_MLR( weightsPath, inputs );

    return mlr;
}
#endif

void usage( char *prog )
{
    std::cout << "Usage: CERTDemo [-h] [-f] [-p <patchDirectory>] [-t] [<images>|<imageDirectories>]" << std::endl;
    std::cout << "       -h Output this help message then quit." << std::endl;
    std::cout << "       -f Input images are already cropped and rotated faces." << std::endl;
    std::cout << "       -p Generate cropped and rotated face patches in the given directory. Filenames will have '_face' appended to them." << std::endl;
    std::cout << "       -t Output simple tab delimited results: <filename><tab><First AU value><tab><Second AU value><etc>." << std::endl;
    std::cout << "       Note: If images or directories are given on the command line they will all be processed then the application will quit." << std::endl;
}

bool processArgs( int argc, char *argv[], std::vector< std::string > &paths, bool &tabDelim )
{
    bool argsOK = true;
    for( int argIndex = 1; argIndex < argc; ++argIndex )
    {
        std::string arg1( argv[argIndex] );
        if( "-f" == arg1)
        {
            inputIsFaces = true;
        }
        else if( "-p" == arg1 )
        {
            if( argIndex >= (argc - 1) )
            {
                std::cerr << "A directory is required for the '-p' option." << std::endl;
                argsOK = false;
            }
            else
            {
                // get the next arg and if it's a directory, put the patches there
                char *nextArg = argv[argIndex+1];
                bf::path bfArg( nextArg, bf::native );
                if( bf::exists( bfArg ) && bf::is_directory( bfArg ) )
                {
                    outputPatches = true;
                    patchDir = nextArg;
                    ++argIndex;
                }
                else
                {
                    std::cerr << "Invalid directory for the '-p' option: " << nextArg << "." << std::endl;
                    argsOK = false;
                }
            }
        }
        else if( "-t" == arg1 )
        {
            tabDelim = true;
        }
        else if( "-h" == arg1 )
        {
            usage( argv[0] );
            return false;
        }
        else
        {
            paths.push_back( argv[argIndex] );
        }
    }

    if( !argsOK )
        usage( argv[0] );

    return argsOK;
}
/*
 *
 * Redistribution and use in source and binary forms, with or without modification, are permitted provided that the following conditions are met:
 *
 *    1. Redistributions of source code must retain the above copyright notice, this list of conditions and the following disclaimer.
 *    2. Redistributions in binary form must reproduce the above copyright notice, this list of conditions and the following disclaimer in the documentation and/or other materials provided with the distribution.
 *    3. The name of the author may not be used to endorse or promote products derived from this software without specific prior written permission.
 *
 * THIS SOFTWARE IS PROVIDED BY THE AUTHOR ``AS IS'' AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE AUTHOR BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
 *
 */