/*
 *  Gaussian Kernel Regression model class
 *
 *  Author: Andrew Salamon
 *
 *  Copyright (c) 2008 Machine Perception Laboratory
 *  University of California San Diego.
 *
 * 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.
 *
 */

#include "GKR_Model.h"

#include <iostream>
#include <fstream>
#include <sstream>
#include <limits>

#ifdef WIN32
#define _USE_MATH_DEFINES
#include <math.h>
#else
#include <cmath>
#endif

#include <boost/serialization/version.hpp>
#include <boost/archive/xml_oarchive.hpp>
#include <boost/archive/xml_iarchive.hpp>
#include <boost/serialization/map.hpp>
#include <boost/serialization/vector.hpp>

namespace GKR
{

Model::Model() : sqrt2pi( std::sqrt( 2.0 * M_PI ) )
{
    initTemporalValues();
    initLogisticValues();
}

Model::Model( const std::string &_ID ) : ID( _ID ), sqrt2pi( std::sqrt( 2.0 * M_PI ) )
{
    initTemporalValues();
    initLogisticValues();
}

Model::Model( const Model &other ) : ID( other.ID ), sqrt2pi( std::sqrt( 2.0 * M_PI ) )
{
    Gmean.resize( boost::extents[other.Gmean.shape()[0]][other.Gmean.shape()[1]] );
    Gmean = other.Gmean;
    Gsigma.resize( boost::extents[other.Gsigma.shape()[0]][other.Gsigma.shape()[1]] );
    Gsigma = other.Gsigma;
    prediction.resize( boost::extents[other.prediction.shape()[0]][other.prediction.shape()[1]] );
    prediction = other.prediction;
    c.resize( boost::extents[other.c.shape()[0]][other.c.shape()[1]] );
    c = other.c;
    initE();
    initTemporalValues();
    initLogisticValues();
}

void
Model::initE()
{
    if( Evec.size() != Gmean.shape()[0] )
    {
        Evec = std::vector<float>( Gmean.shape()[0], 1.0 );
        Emat.resize( boost::extents[Gmean.shape()[0]][1] );
        Emat.assign( Evec.begin(), Evec.end() );
    }
}

void
Model::initTemporalValues( float _S_x, float _S_y, float _S_0, float _Y_0 )
{
    S_x = _S_x;
    S_y = _S_y;
    S_0 = _S_0;
    Y_0 = _Y_0;

    S_t = S_0; // filter variance
    // K_t = ; // filter gain
    // Y_t; // filter output

    firstT = true;;
}

void
Model::initLogisticValues( float _alpha, float _beta )
{
    alpha = _alpha;
    beta = _beta;
}

std::vector<float>
Model::getCKFValues()
{
// Order is: alpha, beta, S_x, S_y, S_0, Y_0,
    std::vector<float> vals;
    vals.push_back( alpha );
    vals.push_back( beta );
    vals.push_back( S_x );
    vals.push_back( S_y );
    vals.push_back( S_0 );
    vals.push_back( Y_0 );

    return vals;
}

void
Model::minAndMax( Matrix2d &poses, std::vector<float> &min, std::vector<float> &max )
{
    unsigned int height = poses.shape()[0];
    unsigned int width = poses.shape()[1];
    min = std::vector<float>( width, std::numeric_limits<float>::max() );
    max = std::vector<float>( width, std::numeric_limits<float>::min() );
    for( unsigned int i = 0; i < height; ++i )
    {
        for( unsigned int subi = 0; subi < width; ++subi )
        {
            float val = poses[i][subi];
            if( val < min[subi] )
                min[subi] = val;
            if( val > max[subi] )
                max[subi] = val;
        }
    }
}

bool
Model::train( std::vector<float> &au, Matrix2d &poses, int nGM )
{
    std::vector<float> Xm; // Minimums for each pose dimension, yaw, pitch, roll
    std::vector<float> Xr; // Range for each pose dimension

    minAndMax( poses, Xm, Xr );
    for( unsigned int i = 0; i < Xr.size(); ++i )
    {
        Xr[i] -= Xm[i];
    }

    Gmean.resize( boost::extents[nGM][Xr.size()] ); // Make sure Gmean is the right size
    initE();
    const std::vector<float> &E = Evec;

    for( int i = 0; i < nGM; ++i )
    {
        float GMratio = static_cast<float>(i + 1) / static_cast<float>(nGM);
        for( unsigned int subi = 0; subi < Xr.size(); ++subi )
        {
            Gmean[i][subi] = (GMratio * Xr[subi]) + (E[i] * Xm[subi]);
        }
    }

    // Fill in the Gsigma vector, and also calculate two constant subexpressions for use in the next step
    Gsigma.resize( boost::extents[Xr.size()][1] ); // Make sure Gsigma is the right size
    std::vector<float> initSigmas( Xr.size(), 0.0 );
    std::vector< std::vector<float> > tmpSigma1( E.size(), initSigmas ); // (2*E*(Gsigma.^2)); 8x3
    std::vector< std::vector<float> > tmpSigma2( E.size(), initSigmas ); // (E*Gsigma * sqrt(2*pi)); 8x3
    for( unsigned int subi = 0; subi < Xr.size(); ++subi )
    {
        float tmp = Xr[subi] / nGM;
        Gsigma[subi][0] = tmp;

        for( unsigned int i = 0; i < E.size(); ++i )
        {
            tmpSigma1[i][subi] = 2.0 * E[i] * (tmp * tmp);
            tmpSigma2[i][subi] = E[i] * tmp * sqrt2pi;
        }
    }

    int trXp_width = nGM * Xr.size();
    Matrix2d trXp_matrix( boost::extents[poses.shape()[0]][trXp_width] );

//    std::cout << "trXp (before):\n";
//    printVecVec( trXp );

    for( unsigned int i = 0; i < trXp_matrix.size(); ++i )
    {
        // Matlab code:
        // f=exp(-(E*trX(i,:) - Gmean).^2 ./ (2*E*Gsigma.^2)) ./ (E*Gsigma * sqrt(2*pi));
        // trXp(i,:) = f(:)';

        int trXp_ind = 0;
        for( unsigned int col = 0; col < Gmean[0].size(); ++col )
        {
            for( unsigned int row = 0; row < Gmean.size(); ++row )
            {
                float val = E[row] * poses[i][col] - Gmean[row][col];
//                std::cout << "poses[" << i << "][" << col << "] = " << poses[i][col] << "    ";
//                std::cout << "Gmean[" << row << "][" << col << "] = " << Gmean[row][col] << "  -> ";
//                std::cout << trXp_ind << std::endl;
                val = -(val * val);
                val /= tmpSigma1[row][col];
                val = std::exp( val ) / tmpSigma2[row][col];

                trXp_matrix[i][ trXp_ind++ ] = val;
            }
        }

    }

//    std::cout << "Xm: ";
//    printVec( Xm );
//    std::cout << "Xr: ";
//    printVec( Xr );
//
//    std::cout << std::endl;
//    printMatrix( Gmean, "Gmean" );
//
//    std::cout << "Gsigma:\n";
//    printVec( Gsigma );
//
//    std::cout << "tmpSigma1:\n";
//    printVecVec( tmpSigma1 );
//
//    std::cout << "tmpSigma2:\n";
//    printVecVec( tmpSigma2 );

//    std::cout << "trXp:\n";
//    printVecVec( trXp );

//    std::cout << "poses:\n";
//    printVecVec( poses );

    //linear regression
    // c = trXp \ trY; % Same as inv(trXp' * trXp) * trXp' * trY

    // copy au into boost_multi_array and use routines from simplematrix to do the left multiply (linear regression)
    Matrix2d trY_matrix( boost::extents[au.size()][1] );
    Matrix2d trXp_trans;
    Matrix2d tmp;

    for( unsigned int row = 0; row < au.size(); ++row )
    {
        trY_matrix[row][0] = au[row];
    }

    SimpleMatrixOps::transpose( trXp_matrix, trXp_trans );
    SimpleMatrixOps::multiply( trXp_trans, trXp_matrix, tmp );
    SimpleMatrixOps::inverse( tmp, c );
    SimpleMatrixOps::multiply( c, trXp_trans, tmp );
//    printMatrix( tmp );
    SimpleMatrixOps::multiply( tmp, trY_matrix, c );
//    std::cout << "size(c) = " << c.shape()[0] << "x" << c.shape()[1] << std::endl;

//    printMatrix( trY_matrix );
//    printMatrix( c );

//    printVec( au );

    return true;
}

bool
Model::predict( const Matrix2d &poses )
{
    int frames = poses.shape()[0];
    int ttXp_width = Gmean.shape()[0] * Gmean.shape()[1];
    Matrix2d ttXp( boost::extents[frames][ttXp_width] );

    SimpleMatrixOps::setAllZero( ttXp );

    Matrix2d tmpSigma1( boost::extents[Gmean.shape()[0]][Gmean.shape()[1]] ); // (2*E*(Gsigma.^2)); 8x3
    Matrix2d tmpSigma2( tmpSigma1 ); // (E*Gsigma * sqrt(2*pi)); 8x3
    for( unsigned int subi = 0; subi < Gmean.shape()[1]; ++subi )
    {
        float tmp = Gsigma[subi][0];

        for( unsigned int i = 0; i < Gmean.shape()[0]; ++i )
        {
            tmpSigma1[i][subi] = 2.0 * Evec[i] * (tmp * tmp);
            tmpSigma2[i][subi] = Evec[i] * tmp * sqrt2pi;
        }
    }

//    printMatrix( tmpSigma1, "tmpSignal1" );
//    printMatrix( tmpSigma2, "tmpSignal2" );

    for( unsigned int i = 0; i < ttXp.shape()[0]; ++i )
    {
        Matrix2d tmp;
        Matrix2d tmp2( boost::extents[1][poses.shape()[1]] );
        for( unsigned int pi = 0; pi < poses.shape()[1]; ++pi )
            tmp2[0][pi] = poses[i][pi];

        SimpleMatrixOps::multiply( Emat, tmp2, tmp );
        SimpleMatrixOps::subtract( tmp, Gmean, tmp2 );
        SimpleMatrixOps::elementwisePower( tmp2, 2.0, tmp);
        SimpleMatrixOps::multiply( tmp, -1.0 );

        // For each element of tmp:
        //    divide by the corresponding element of tmpSigma1
        //    exp
        //    divide by corresponding element of tmpSigma2
        int ttXp_ind = 0;
        for( unsigned int col = 0; col < tmp.shape()[1]; ++col )
        {
            for( unsigned int row = 0; row < tmp.shape()[0]; ++row )
            {
                float val = std::exp( tmp[row][col] / tmpSigma1[row][col] ) / tmpSigma2[row][col];
                ttXp[i][ttXp_ind++] = val;
            }
        }
    }

//    printMatrix( ttXp, "after all" );

    SimpleMatrixOps::multiply( ttXp, c, prediction );
//    printMatrix( prediction, "Prediction" );

    return true;
}

/*
   Returns the filtered AU value, plus the confidence
 */
std::pair<float,float>
Model::CKF( bool foundFace, float au )
{
    float X_tp1 = au;

    if( firstT )
    {
        S_t = S_0;
        Y_t = X_tp1;
        firstT = false;
    }
    else
    {
        updateVariance( foundFace );
    }
    updateGain( foundFace );

    float Y_tp1 = Y_t + (K_t * ( X_tp1 - Y_t ));
/*
Y_{t+1} = Y_t + K_{t+1} ( X_{t+1} - Y_t)

where X_{t+1} is the (baseline corrected) CERT output at time t+1.
*/
    Y_t = Y_tp1;
    return std::pair<float,float>( Y_tp1, std::sqrt(S_t) );
}

void
Model::updateVariance( bool foundFace )
{
    if( foundFace )
    {
// S_{t+1} = (( S_t + S_y)^{-1} + S_x^{-1})^{-1}
        S_t = 1.0 / ( (1.0 / (S_t + S_y)) + (1.0 / S_x) );
//        test = S_x / ((S_x / (S_t + S_y)) + 1.0);
    }
    else
    {
// S_{t+1} = S_t + S_y
        S_t = S_t + S_y;
    }

}

void
Model::updateGain( bool foundFace )
{
    if( foundFace )
    {
// K_{t+1} = S_{t+1} S_x^{-1}
        K_t = S_t / S_x; // S_t * (1.0 / S_x);
    }
    else
    {
        K_t = 0.0;
    }
}

float
Model::logistic( float au )
{
    // f(au) = 1 / (1 + exp(alpha * au + beta));
    au = 1.0 / (1.0 + std::exp( (alpha * au) + beta ) );
    return au;
}

void
Model::clear()
{
	Gmean.resize( boost::extents[0][0] );
	Gsigma.resize( boost::extents[0][0] );
	c.resize( boost::extents[0][0] );
	prediction.resize( boost::extents[0][0] );
}
	
bool
Model::operator==( const Model &other ) const
{
    if( ID != other.ID ) return false;
    if( Gmean != other.Gmean ) return false;
    if( Gsigma != other.Gsigma ) return false;
    if( c != other.c ) return false;
    if( std::abs( sqrt2pi - other.sqrt2pi ) > 0.0001 ) return false;
    if( Emat != other.Emat ) return false;
    if( Evec != other.Evec ) return false;

//    std::vector<float> Evec;
    return true;
}

void
Model::operator=( const Model &other )
{
    ID = other.ID;
    Gmean.resize( boost::extents[other.Gmean.shape()[0]][other.Gmean.shape()[1]] );
    Gmean = other.Gmean;
    Gsigma.resize( boost::extents[other.Gsigma.shape()[0]][other.Gsigma.shape()[1]] );
    Gsigma = other.Gsigma;
    prediction.resize( boost::extents[other.prediction.shape()[0]][other.prediction.shape()[1]] );
    prediction = other.prediction;
    c.resize( boost::extents[other.c.shape()[0]][other.c.shape()[1]] );
    c = other.c;
    initE();
}

void printMatrix( const Matrix2d &matrix, std::string title )
{
    if( title.length() > 0 )
        std::cout << title << ":" << std::endl;
    for( unsigned int row = 0; row < matrix.shape()[0]; ++row )
    {
        for( unsigned int col = 0; col < matrix.shape()[1]; ++col )
        {
            if( title.length() > 0 )
                std::cout << "  ";
            std::cout << matrix[row][col];
            if( col < (matrix.shape()[1] - 1) )
                std::cout << " ";
        }
        std::cout << std::endl;
    }
}

} // end namespace GKR

BOOST_CLASS_VERSION( GKR::Model, 1 );