#include "optimization.h"

Optimization::~Optimization()
{

}

void Optimization::init()
{

}


double Optimization::deg2rad(double angDeg){
    return angDeg / 180.0 * PI;
}

void Optimization::setMeasurementMatrix(cv::Mat inW)
{
    nbParams = -1;

    Win = inW.clone();
    //AR_Assert(!(Win->empty()));

    W.release();
    tm.clear();

    cv::Mat Wmean;
    computeWmean(Win, Wmean, tm);

    //util::makeNonHomogenious(Wmean);
    W = Wmean.clone();

    nbCams = W.rows / 2;
    nbPts  = W.cols ;

    paramConst   = cv::Mat::ones(nbCams,PARAMS_PER_CAM,CV_8UC1); // 1 for varying parameter
    paramBounds  = PI/2*cv::Mat::ones(nbCams,PARAMS_PER_CAM,cv::DataType<double>::type);
    paramOffset  = cv::Mat::zeros(nbCams,PARAMS_PER_CAM,cv::DataType<double>::type);
    paramInitial = cv::Mat::zeros(nbCams,PARAMS_PER_CAM,cv::DataType<double>::type);

    // Define default constant parameters:
    paramConst.col(K)     = cv::Mat::zeros(nbCams,1,paramConst.type());
    paramConst.col(ALPHA) = cv::Mat::zeros(nbCams,1,paramConst.type());
    paramConst.col(S)     = cv::Mat::zeros(nbCams,1,paramConst.type());
    paramConst.row(0).colRange(RT1,RT2+1) = cv::Mat::zeros(1,3,paramConst.type());

    // Define default offset parameters:
    paramOffset.col(K)     = 1 * cv::Mat::ones(nbCams,1,paramOffset.type());
    paramOffset.col(ALPHA) = 1 * cv::Mat::ones(nbCams,1,paramOffset.type());

    // Define default bounds (considered symmetric):

    // Define initial values. Actual value of parameter: offset + initial
    paramInitial.rowRange(1,nbCams).col(RR) = deg2rad(5.0) * cv::Mat::ones(nbCams-1,1,paramInitial.type());

    std::cout << paramInitial;
}


void Optimization::run()
{
    _minf = INFINITY;
    _bestObjValue = INFINITY;

    std::vector<double> x0,lb,ub;
    this->getInitialConditionsAndBounds(x0, lb, ub);
    this->setObjectiveFunction(Optimization::OBJFUNC_Rectification);  
    //this->setObjectiveFunction(Optimization::OBJFUNC_Distance);
   

    nlopt::opt optimizer(nlopt::GN_CRS2_LM, this->getParametersNumber()); // 1
    optimizer.set_lower_bounds(lb);
    optimizer.set_upper_bounds(ub);
    optimizer.set_min_objective( *Optimization::wrap, this);
    optimizer.set_ftol_rel(1e-30);
    optimizer.set_maxtime(_maxTimeStep1);

    optimizer.optimize(x0, _minf);

    cout << "Found minimum: " << _minf <<endl;

    cv::Mat paramX = 0 * paramOffset.clone();
    copyVectorToMatElements(x0, paramVarIdx, paramX);
    paramResult = paramOffset.clone() + paramX;
}

double Optimization::operator() (const std::vector<double> &x, std::vector<double> &grad)
{
    (void)(grad); // remove warning for unused variable
    switch (_objFuncMethod) {
        case OBJFUNC_DistanceSquared: return distanceSquared(x); break;
        case OBJFUNC_Distance       : return testFunction(x); break;
        case OBJFUNC_PseudoInverse  : return distancePseudoInverse(x); break;
        case OBJFUNC_Rectification  : return distanceWminusMMW(x); break;
        default: break;
    }
    return 0;
}


double Optimization::testFunction (const std::vector<double> &x){
   return (x[0]*x[0] + x[1]*x[1] + 20*sin(x[0]) + 20*sin(x[1]));
}

MatCell Optimization::getMfromParams(const std::vector<double> &x){
    MatCell M(nbCams);

    cv::Mat A = cv::Mat::eye(2, 2, cv::DataType<double>::type);
    A.at<double>(0,0) = x[0] * x[1]; // k*E
    A.at<double>(0,1) = x[0] * x[2];    // k*s

    M[0] = A * cv::Mat::eye(2, 3, cv::DataType<double>::type);

    cv::Mat angles = cv::Mat::eye(3, nbCams-1, cv::DataType<double>::type);
    //angles = params(4 : 4 + 3*(nCams-1)-1);
    //angles = reshape(angles, [3 length(angles)/3]);
    for(int i = 0 ; i < angles.cols ; i++){
        angles.at<double>(0,i) = x[3*i+3];
        angles.at<double>(1,i) = x[3*i+4];
        angles.at<double>(2,i) = x[3*i+5];
    }
    //std::cout << angles << std::endl;

    for (int i = 1 ; i < nbCams ; i++){
        double p = angles.at<double>(0,i-1);
        double t = angles.at<double>(1,i-1);
        double k = angles.at<double>(2,i-1);

        // R = Rz(k)*Ry(t)*Rx(p)
        cv::Mat R = (cv::Mat_<double>(2,3) << cos(k)*cos(t), cos(k)*sin(p)*sin(t) - cos(p)*sin(k), sin(k)*sin(p) + cos(k)*cos(p)*sin(t),
                     cos(t)*sin(k), cos(k)*cos(p) + sin(k)*sin(p)*sin(t), cos(p)*sin(k)*sin(t) - cos(k)*sin(p));

        M[i] = A * R ;
    }
    return M;
}

cv::Mat Optimization::getPoints3DfromParams(const std::vector<double> &x){
    cv::Mat X = cv::Mat_<double>(3, nbPts);
    for (int i = 0 ; i < nbPts ; i++){
        X.at<double>(0,i) = x[3 + 3*(nbCams-1) + 3*i];
        X.at<double>(1,i) = x[3 + 3*(nbCams-1) + 3*i + 1];
        X.at<double>(2,i) = x[3 + 3*(nbCams-1) + 3*i + 2];
    }
    return X;
}


double Optimization::distanceSquared(const std::vector<double> &x)
{
    _iterCnt++;

    MatCell M = getMfromParams(x);
    cv::Mat X = getPoints3DfromParams(x);

    cv::Mat reprojErrorMat = W - M*X;
    double reprojError = cv::sum(reprojErrorMat.mul(reprojErrorMat))[0];
    double s = x[2];

    double objValue = std::abs( reprojError + s );

    if (_iterCnt % nbParams * 500 == 0){
        if ( objValue < _bestObjValue) _bestObjValue = objValue ;
        std::stringstream stream;
        cout <<   "  *  " << std::setw(15) << _iterCnt
                  << "  *  " << std::setw(12) << std::scientific << std::setprecision(5) << _bestObjValue
                  << "  *  " << std::setw(12) << std::scientific << std::setprecision(5) << objValue << endl;
    }

    return objValue;
}

double Optimization::distancePseudoInverse(const std::vector<double> &x)
{
    _iterCnt++;

    MatCell M = getMfromParams(x);
    cv::Mat X = getPoints3DfromParams(x);

    cv::Mat Minv;
    cv::invert( (cv::Mat) M,Minv,cv::DECOMP_SVD);
    cv::Mat reprojErrorMat = W - M*Minv*W;
    double reprojError = cv::sum(reprojErrorMat.mul(reprojErrorMat))[0];
    double s = x[2];

    double objValue = std::abs( reprojError + s );

    if (_iterCnt % nbParams * 500 == 0){
        if ( objValue < _bestObjValue) _bestObjValue = objValue ;
        std::stringstream stream;
        cout <<   "  :  " << std::setw(15) << _iterCnt
                  << "  :  " << std::setw(12) << std::scientific << std::setprecision(5) << _bestObjValue
                  << "  :  " << std::setw(12) << std::scientific << std::setprecision(5) << objValue << endl;
    }

    return objValue;
}

double Optimization::distanceWminusMMW(const std::vector<double> &x)
{

    cv::Mat paramX = 0 * paramOffset.clone();
    util::copyVectorToMatElements( x, paramVarIdx, paramX);
    cv::Mat paramTable = paramOffset.clone() + paramX;

    double k     = paramTable.at<double>(0,K);
    double alpha = paramTable.at<double>(0,ALPHA);
    double s     = paramTable.at<double>(0,S);
    cv::Mat A = k * (cv::Mat_<double>(2,2) << alpha, s, 0, 1.);

    std::vector<cv::Mat> Marray(nbCams);
    std::vector<cv::Mat> Rarray(nbCams);

    for (auto i = 0; i < nbCams; i++){
        double t1  = paramTable.at<double>(i,RT1);
        double rho = paramTable.at<double>(i,RR);
        double t2  = paramTable.at<double>(i,RT2);

        Rarray[i] = rotmat::fromEulerZYZ(t1,rho,t2);
        if (i != 0) Rarray[i] = Rarray[i] * Rarray[i-1];

        Marray[i] = A * Rarray[i].rowRange(0,0+2);
    }

    cv::Mat M = util::concatenateMat(Marray, util::CONCAT_VERTICAL);
    cv::Mat pinvM;
    cv::invert(M, pinvM, cv::DECOMP_SVD);

    cv::Mat errorMat = W - M*pinvM*W;
    double objValue = cv::sum(errorMat.mul(errorMat))[0];

    _iterCnt++;
    if ( _iterCnt % 10000 == 0){
        if ( objValue < _bestObjValue) _bestObjValue = objValue ;
        std::stringstream stream;
        cout <<   "  :  " << std::setw(15) << _iterCnt
                  << "  :  " << std::setw(12) << std::scientific << std::setprecision(5) << _bestObjValue
                  << "  :  " << std::setw(12) << std::scientific << std::setprecision(5) << objValue << endl;
    }

    return objValue;
}


cv::Mat Optimization::getCalibrationMatrixFromParamTable(const cv::Mat &paramTable ) const{
    double k     = paramTable.at<double>(0,K);
    double alpha = paramTable.at<double>(0,ALPHA);
    double s     = paramTable.at<double>(0,S);
    cv::Mat A = k * (cv::Mat_<double>(2,2) << alpha, s, 0, 1.);
    return A.clone();
}


int Optimization::getParametersNumber()
{
    return nbParams;
}

void Optimization::copyVectorToMatElements(const std::vector<double> &vec, const cv::Mat &idx, cv::Mat &mat)
{
    int nbIdx = (int) idx.total();

    for (int i = 0; i < nbIdx; i++){
        int col = (int) idx.at<cv::Point>(i).x;
        int row = (int) idx.at<cv::Point>(i).y;

        mat.row(row).col(col) = vec[i];
    }
}

void Optimization::getInitialConditionsAndBounds(std::vector<double> &_x0, std::vector<double> &_lb, std::vector<double> &_ub)
{
    paramVarIdx.release();
    cv::findNonZero(paramConst, paramVarIdx);
    nbParams = (int) paramVarIdx.total();
    cout << "Number of parameters: " << nbParams << endl;

    _x0.clear();
    _lb.clear();
    _ub.clear();

    _x0 = std::vector<double>(nbParams);
    _lb = std::vector<double>(nbParams);
    _ub = std::vector<double>(nbParams);

    cv::Mat lbMat = -paramBounds.clone();
    cv::Mat ubMat =  paramBounds.clone();

    lbMat.at<double>(1,RR) = 0;

    util::copyMatElementsToVector( paramInitial, paramVarIdx, _x0);
    util::copyMatElementsToVector(       lbMat , paramVarIdx, _lb);
    util::copyMatElementsToVector(       ubMat , paramVarIdx, _ub);
}


std::vector<cv::Mat> Optimization::getCameraMatrices() const
{
    std::vector<cv::Mat> Parray(nbCams);

    cv::Mat A = getCalibrationMatrixFromParamTable(paramResult);

    std::vector<cv::Mat> Rarray(nbCams);

    for (auto i = 0; i < nbCams; i++){
        double t1  = paramResult.at<double>(i,RT1);
        double rho = paramResult.at<double>(i,RR);
        double t2  = paramResult.at<double>(i,RT2);

        Rarray[i] = rotmat::fromEulerZYZ(t1,rho,t2);
        if (i != 0) Rarray[i] = Rarray[i] * Rarray[i-1];

        cv::Mat M = A * Rarray[i].rowRange(0,0+2);

        cv::Matx34d P = cv::Matx34d::zeros();

        P(0,0) = M.at<double>(0,0);
        P(0,1) = M.at<double>(0,1);
        P(0,2) = M.at<double>(0,2);

        P(1,0) = M.at<double>(1,0);
        P(1,1) = M.at<double>(1,1);
        P(1,2) = M.at<double>(1,2);

        P(0,3) = tm[i].at<double>(0,0);
        P(1,3) = tm[i].at<double>(1,0);
        P(2,3) = 1;

        Parray[i] = ((cv::Mat)P).clone();
    }

    return Parray;
}

cv::Mat Optimization::getCalibrationMatrix() const
{
    cv::Mat A = getCalibrationMatrixFromParamTable(paramResult);
    return A.clone();
}

cv::Mat Optimization::getRotationAnglesDeg() const
{
    //return cv::Mat(paramResult.colRange(RT1,RT2+1)).clone() / PI * 180.0;
    return cv::Mat(paramResult.col(RR)).clone() / PI * 180.0;
}

void Optimization::computeWmean(const cv::Mat & W, cv::Mat & Wmean, std::vector<cv::Mat> & t)
{

    W.copyTo(Wmean);
    int nCams = W.rows / 2;
    t.clear();

    for (int i = 0; i < nCams ; i++)
    {
        cv::Mat tC = cv::Mat::zeros(2,1,cv::DataType<double>::type);
        tC.at<double>(0,0) = cv::mean(W.row( 2*i     ))(0);
        tC.at<double>(1,0) = cv::mean(W.row( 2*i + 1 ))(0);
        t.push_back(tC);
        for (int j = 0; j < W.cols ; j++)
        {
            Wmean.at<double>(2*i    ,j) = Wmean.at<double>(2*i    ,j) - tC.at<double>(0,0);
            Wmean.at<double>(2*i + 1,j) = Wmean.at<double>(2*i + 1,j) - tC.at<double>(1,0);
        }
    }

}