// This file contains my implementation for pose esimtation
#include "marker.hpp"
#include <Eigen/SVD>
#include <Eigen/Dense>
#include <glm/gtc/matrix_transform.hpp>
#include <opencv2/opencv.hpp>
#include <cmath>
#include <limits>
#include <iostream>

// refer to: https://www.dropbox.com/s/qkulg4j64lyn0qa/2018_proj_geo_for_cv_projcv_assignment.pdf?dl=0
// update computed rotation matrix
void updateRotationSVD(glm::mat4x3& M)
{
    Eigen::Matrix3f R;
    R << M[0][0], M[1][0], M[2][0],
         M[0][1], M[1][1], M[2][1],
         M[0][2], M[1][2], M[2][2];
    Eigen::JacobiSVD<Eigen::MatrixXf> svdSolver(R, Eigen::ComputeFullV | Eigen::ComputeFullU);
    R = svdSolver.matrixU() * svdSolver.matrixV().transpose();
    M[0][0] = R.col(0)[0];
    M[0][1] = R.col(0)[1];
    M[0][2] = R.col(0)[2];
    M[1][0] = R.col(1)[0];
    M[1][1] = R.col(1)[1];
    M[1][2] = R.col(1)[2];
    M[2][0] = R.col(2)[0];
    M[2][1] = R.col(2)[1];
    M[2][2] = R.col(2)[2];
}

// refer to https://citeseerx.ist.psu.edu/viewdoc/download?doi=10.1.1.895.2324&rep=rep1&type=pdf
// update to Tait-Bryan rotation matrix
void updateRotationTaitBryan(glm::mat4x3& M)
{
    float theta = -std::asin(M[2][0]);
    float thetaS = std::sin(theta);
    float thetaC = std::cos(theta);
    float psi = std::asin(M[1][0] / thetaC);
    float psiS = std::sin(psi);
    float psiC = std::cos(psi);
    float phi = std::asin(M[2][1] / thetaC);
    float phiS = std::sin(phi);
    float phiC = std::cos(phi);

    M[0][0] = thetaC * psiC;
    M[1][0] = thetaC * psiS;
    M[2][0] = -thetaS;

    M[0][1] = phiS * thetaS * psiC - phiC * psiS;
    M[1][1] = phiS * thetaS * psiS + phiC * psiC;
    M[2][1] = phiS * thetaC;

    M[0][2] = phiC * thetaS * psiC + phiS * psiS;
    M[1][2] = phiC * thetaS * psiS - phiS * psiC;
    M[2][2] = phiC * thetaC;
}

// refer to: https://stackoverflow.com/questions/12463487/obtain-rotation-axis-from-rotation-matrix-and-translation-vector-in-opencv
void updateRotationEuclidean(glm::mat4x3& M)
{
    // obtain rotation angle
    float angle = std::acos(0.5f * (M[0][0] + M[1][1] + M[2][2] - 1.0f));
    // obtain axis
    glm::vec3 axis = glm::vec3(M[1][2] - M[2][1], M[2][0] - M[0][2], M[0][1] - M[1][0]);
    axis = glm::normalize(axis);
    glm::mat4 rot = glm::rotate(glm::mat4(1.0f), glm::degrees(angle), axis);
    M[0][0] = rot[0][0]; M[1][0] = rot[1][0]; M[2][0] = rot[2][0];
    M[0][1] = rot[0][1]; M[1][1] = rot[1][1]; M[2][1] = rot[2][1];
    M[0][2] = rot[0][2]; M[1][2] = rot[1][2]; M[2][2] = rot[2][2];
}

// compute reprojection error
float reprojectionError(
    const glm::mat3& cameraK, const glm::mat4x3& M,
    const std::vector<glm::vec2>& objPoints,
    const std::vector<glm::vec2>& imgPoints
)
{
    float error = 0.0f;
    for(size_t i = 0; i < 4; i++)
    {
        // auto projected = M * glm::vec4(objPoints[i], 0.0f, 1.0f);
        auto projected = cameraK * M * glm::vec4(objPoints[i], 0.0f, 1.0f);
        float dx = projected.x - imgPoints[i].x;
        float dy = projected.y - imgPoints[i].y;
        error += dx * dx + dy * dy;
    }
    return std::sqrt(error * 0.125f);
}

// validate solution based on Zhang's method
bool validateSolutionZhang(
    const Eigen::Vector3f& tstar, const Eigen::Vector3f& nstar, const Eigen::MatrixXf& H,
    const Eigen::MatrixXf& mstarT, Eigen::Matrix3f& outR, Eigen::Vector3f& outT
)
{
    // prepare rotation matrix and translation vector
    // R = H (I + t* n^T)^-1
    outR = H * (Eigen::Matrix3f::Identity() + tstar * nstar.transpose()).inverse();
    if(outR.determinant() < 0.0f)
    {
        // flip surface
        outR *= -1.0f;
    }
    outT = outR * tstar;
    // 1 + n^T R^T t
    float positive = 1.0f + (nstar.transpose() * outR.transpose() * outT)[0];
    if(positive <= 0.0f)
        return false;
    // m*^T n*
    Eigen::Vector4f positives = mstarT * nstar;
    if(positives[0] <= 0.0f || positives[1] <= 0.0f ||
       positives[2] <= 0.0f || positives[3] <= 0.0f)
       return false;
    return true;
}

// refer to https://hal.inria.fr/inria-00174036v3/document
// refer to https://github.com/opencv/opencv/blob/4.x/modules/calib3d/src/homography_decomp.cpp#L217
// decompose homography matrix
void decomposeHomoMatrixZhang(
    const glm::mat3& cameraK, const glm::mat3& cameraInvK, const glm::mat3x4& mstarTransposed,
    const glm::mat3& G, const std::vector<glm::vec2>& objPoints,
    const std::vector<glm::vec2>& imgPoints, glm::mat4x3& outputM)
{
    // normalize H
    // glm::mat3 hHat = cameraInvK * G * cameraK;
    glm::mat3 hHat = cameraInvK * G;
    // retrieve homography matrix H
    Eigen::Matrix<float, 3, 3> H;
    H << hHat[0][0], hHat[1][0], hHat[2][0],
         hHat[0][1], hHat[1][1], hHat[2][1],
         hHat[0][2], hHat[1][2], hHat[2][2];
    Eigen::EigenSolver<Eigen::MatrixXf> eigenSolver;
    // eigenSolver.compute(H, false);
    // float gamma = eigenSolver.eigenvalues().real()[1];
    // H = H / gamma; // scale by lambda2
    // H^TH = V A V^T and solve for eigenvalues and eigenvectors
    // H = H.transpose() * H;
    eigenSolver.compute(H, true);
    Eigen::VectorXf lambdas = eigenSolver.eigenvalues().real();
    Eigen::MatrixXf vs = eigenSolver.eigenvectors().real();
    // verify the eigenvectors
    if (vs.rows() < 3 || vs.cols() < 3) return;
    float lambda1 = lambdas[0];
    float lambda3 = lambdas[2];
    // float lambda1 = lambdas[0] >= lambdas[2] ? lambdas[0] : lambdas[2];
    // float lambda3 = lambdas[0] >= lambdas[2] ? lambdas[2] : lambdas[0];
    float lambda1x3 = lambda1 * lambda3;
    float lambda1m3 = lambda1 - lambda3;
    float lambda1m3_2 = lambda1m3 * lambda1m3;
    // prepare zeta
    float tmp1 = 1.0f / (2.0f * lambda1x3);
    float tmp2 = std::sqrt(std::abs(1.0f + 4.0f * lambda1x3 / lambda1m3_2));
    float tmp1x2 = tmp1 * tmp2;
    float zeta1 = -tmp1 + tmp1x2;
    float zeta3 = -tmp1 - tmp1x2;
    float zeta1_2 = zeta1 * zeta1;
    float zeta3_2 = zeta3 * zeta3;
    float zeta1m3 = zeta1 - zeta3;
    float zeta1m3_inv = 1.0f / zeta1m3;
    // compute norms of v'1,3
    float nv1p = std::sqrt(std::abs(zeta1_2 * lambda1m3_2 + 2.0f * zeta1 * (lambda1x3 - 1.0f) + 1.0f));
    float nv3p = std::sqrt(std::abs(zeta3_2 * lambda1m3_2 + 2.0f * zeta3 * (lambda1x3 - 1.0f) + 1.0f));
    // compute v1' and v3'
    Eigen::Vector3f v1p, v3p;
    v1p << vs.col(0)[0]*nv1p, vs.col(0)[1]*nv1p, vs.col(0)[2]*nv1p;
    v3p << vs.col(2)[0]*nv3p, vs.col(2)[1]*nv3p, vs.col(2)[2]*nv3p;
    // 8 solutions in total:
    Eigen::Vector3f tstar[4], nstar[4];
    // solution 1
    tstar[0] =  (v1p - v3p) * zeta1m3_inv;
    tstar[1] = -(v1p - v3p) * zeta1m3_inv;
    nstar[0] =  (zeta1 * v3p - zeta3 * v1p) * zeta1m3_inv;
    nstar[1] = -(zeta1 * v3p - zeta3 * v1p) * zeta1m3_inv;
    // solution 2
    tstar[2] =  (v1p + v3p) * zeta1m3_inv;
    tstar[3] = -(v1p + v3p) * zeta1m3_inv;
    nstar[2] =  (zeta1 * v3p + zeta3 * v1p) * zeta1m3_inv;
    nstar[3] = -(zeta1 * v3p + zeta3 * v1p) * zeta1m3_inv;
    // prepare mstar transposed (for validation)
    Eigen::Matrix<float, 4, 3> mstarT;
    mstarT << mstarTransposed[0][0], mstarTransposed[1][0], mstarTransposed[2][0],
              mstarTransposed[0][1], mstarTransposed[1][1], mstarTransposed[2][1],
              mstarTransposed[0][2], mstarTransposed[1][2], mstarTransposed[2][2],
              mstarTransposed[0][3], mstarTransposed[1][3], mstarTransposed[2][3];
    // validate solutions, collect the first one
    Eigen::Matrix3f R; // rotation
    Eigen::Vector3f T; // translation
    std::vector<glm::mat4x3> candidates;
    for(int i = 0; i < 2; i++)
    {
        for(int j = 0; j < 2; j++)
        {
            if(validateSolutionZhang(tstar[i], nstar[j], H, mstarT, R, T))
            {
                candidates.push_back(glm::mat4x3(
                    glm::vec3(R.col(0)[0], R.col(0)[1], R.col(0)[2]),
                    glm::vec3(R.col(1)[0], R.col(1)[1], R.col(1)[2]),
                    glm::vec3(R.col(2)[0], R.col(2)[1], R.col(2)[2]),
                    glm::vec3(T[0], T[1], T[2])
                ));
            }
            if(validateSolutionZhang(tstar[i+2], nstar[j+2], H, mstarT, R, T))
            {
                candidates.push_back(glm::mat4x3(
                    glm::vec3(R.col(0)[0], R.col(0)[1], R.col(0)[2]),
                    glm::vec3(R.col(1)[0], R.col(1)[1], R.col(1)[2]),
                    glm::vec3(R.col(2)[0], R.col(2)[1], R.col(2)[2]),
                    glm::vec3(T[0], T[1], T[2])
                ));
            }
        }
    }

    if(candidates.size() == 0)
    {
        outputM = glm::mat4x3(0.0f);
        return;
    }
    // compute reprojection error for each
    float minError = std::numeric_limits<float>::max();
    size_t selectedIdx = 0;
    for(size_t i = 0; i < candidates.size(); i++)
    {
        float error = reprojectionError(cameraK, candidates[i], objPoints, imgPoints);
        if(error < minError)
        {
            minError = error;
            selectedIdx = i;
        }
    }
    // std::cout << "Reprojection Error: " << minError << std::endl;
    // set pose matrix
    outputM = candidates[selectedIdx];
}

// refer to: https://github.com/opencv/opencv/blob/master/modules/calib3d/src/undistort.dispatch.cpp#L384
// undistort image points
void undistortPoints(
    const glm::mat3& cameraK, const glm::vec3& cameraDistK, const glm::vec2& cameraDistP,
    glm::vec2& p1, glm::vec2& p2, glm::vec2& p3, glm::vec2& p4)
{
    glm::vec2 src[4] = {p1, p2, p3, p4};
    glm::vec2 dst[4];
    float fx = cameraK[0][0], fy = cameraK[1][1];
    float invfx = 1.0f / fx, invfy = 1.0f / fy;
    float cx = cameraK[2][0], cy = cameraK[2][1];
    // prepare coefficient K
    // k1,k2,p1,p2,k3
    float k[5] = {cameraDistK.x, cameraDistK.y, cameraDistP.x, cameraDistP.y, cameraDistK.z};
    for(int i = 0; i < 4; i++)
    {
        float x = src[i].x, y = src[i].y;
        float x0, y0, u, v;
        u = x;
        v = y;
        x = x0 = (x - cx) * invfx;
        y = y0 = (y - cy) * invfy;
        // init error
        float error = std::numeric_limits<float>::max();
        // iteratively reduce distortion error
        const int MAX_ITER = 10;
        const float EPS = 0.0005f;
        int j = 0;
        for(; j < MAX_ITER; j++)
        {
            if(error < EPS) break;
            float r2 = x*x + y*y;
            float icdist = 1.0f / (
                1.0f + ((k[4] * r2 + k[1]) * r2 + k[0]) * r2
            );
            // if distortion is negative, reset
            if(icdist < 0.0f)
            {
                x = (u - cx) * invfx;
                y = (v - cy) * invfy;
                break;
            }
            float deltaX = 2.0f * k[2] * x * y + k[3] * (r2 + 2.0f * x * x);
            float deltaY = k[2] * (r2 + 2 * y * y) + 2.0f * k[3] * x * y;
            x = (x0 - deltaX) * icdist;
            y = (y0 - deltaY) * icdist;
            // compute error
            {
                float r4, r6;
                float a1, a2, a3;
                float cdist;
                float xd, yd;
                r2 = x * x + y * y;
                r4 = r2 * r2;
                r6 = r4 * r2;
                a1 = 2.0f * x * y;
                a2 = r2 + 2.0f * x * x;
                a3 = r2 + 2.0f * y * y;
                cdist = 1.0f + k[0] * r2 + k[1] * r4 + k[4] * r6;
                xd = x * cdist + k[2] * a1 + k[3] * a2;
                yd = y * cdist + k[2] * a3 + k[3] * a1;
                float xproj, yproj;
                xproj = xd * fx + cx;
                yproj = yd * fy + cy;
                error = std::sqrt((xproj - u) * (xproj - u) + (yproj - v) * (yproj - v));
            }
        }
        // std::cout << j << "," << error << std::endl;
        dst[i].x = x * fx + cx;
        dst[i].y = y * fy + cy;
    }
    // retrieve new positions
    p1 = dst[0];
    p2 = dst[1];
    p3 = dst[2];
    p4 = dst[3];
}

// refer to book: "Augmented Reality: Principles and Practice"
// decompose homography matrix
void decomposeHomoMatrixARBook(
    const glm::mat3& cameraK, const glm::mat3& cameraInvK, glm::mat3 H, glm::mat4x3& outputM
)
{
    H = cameraInvK * H;
    // recover rotation matrix and translation vector
    float d = 1.0f / std::sqrt(glm::length(H[0]) * glm::length(H[1]));
    // set translation
    outputM[3] = d * H[2];
    glm::vec3 h1 = H[0];
    glm::vec3 h2 = H[1];
    glm::vec3 h12 = glm::normalize(h1 + h2);
    glm::vec3 h21 = glm::normalize(glm::cross(h12, glm::cross(h1, h2)));
    // set rotations
    d = 1.0f / std::sqrt(2.0f);
    outputM[0] = (h12 + h21) * d; // set R1
    outputM[1] = (h12 - h21) * d; // set R2
    outputM[2] = glm::cross(outputM[0], outputM[1]); // set R3
}

// refer to https://stackoverflow.com/questions/8927771/computing-camera-pose-with-homography-matrix-based-on-4-coplanar-points
// decompose homography matrix
void decomposeHomoMatrixInternet(
    const glm::mat3& cameraK, const glm::mat3& cameraInvK, glm::mat3 H, glm::mat4x3& outputM
)
{
    H = cameraInvK * H;
    float tnorm = 2.0f / (glm::length(H[0]) + glm::length(H[1]));
    // set translation
    outputM[3] = tnorm * H[2];
    // set rotations
    outputM[0] = glm::normalize(H[0]);
    outputM[1] = glm::normalize(H[1]);
    outputM[2] = glm::cross(outputM[0], outputM[1]);
}

// refer to https://gist.github.com/inspirit/740979/97f54a63eb5f61f8f2eb578d60eb44839556ff3f
void decomposeHomoMatrixInternet2(
    const glm::mat3& cameraK, const glm::mat3& cameraInvK, glm::mat3 H, glm::mat4x3& outputM
)
{
    H = cameraInvK * H;
    float lambda = 1.0f / glm::length(H[0]);
    // set translation
    outputM[3] = lambda * H[2];
    // set rotations
    outputM[0] = lambda * H[0];
    outputM[1] = lambda * H[1];
    outputM[2] = glm::cross(outputM[0], outputM[1]);
}

// refer to https://courses.cs.duke.edu//spring22/compsci527/notes/n_10_reconstruction.pdf
// decompose homography matrix
void decomposeHomoMatrixDuke(
    const glm::mat3& cameraK, const glm::mat3& cameraInvK,
    const std::vector<glm::vec2>& objPoints, const std::vector<glm::vec2>& imgPoints,
    glm::mat3 H, glm::mat4x3& outputM
)
{
    H = cameraInvK * H;
    // prepare E
    Eigen::Matrix<float, 3, 3> E;
    E << H[0][0], H[1][0], H[2][0],
         H[0][1], H[1][1], H[2][1],
         H[0][2], H[1][2], H[2][2];
    // decompose E
    Eigen::JacobiSVD<Eigen::MatrixXf> svdSolver;
    svdSolver.compute(E, Eigen::ComputeFullV | Eigen::ComputeFullU);
    // estimate two possible t
    auto& U = svdSolver.matrixU();
    auto& V = svdSolver.matrixV();
    Eigen::Vector3f T1 =  V.col(2);
    Eigen::Vector3f T2 = -V.col(2);
    // prepare alpha, beta
    Eigen::Matrix3f ab1 = U.col(0) * V.col(0).transpose();
    Eigen::Matrix3f ab2 = U.col(1) * V.col(1).transpose();
    Eigen::Matrix3f ab3 = U.col(2) * V.col(2).transpose();
    // compute Q
    Eigen::Matrix3f Q1 = ab1 + ab2 + ab3;
    Eigen::Matrix3f Q2 = ab1 + ab2 - ab3;
    // compute R
    Eigen::Matrix3f R1 = Q1 * Q1.determinant();
    Eigen::Matrix3f R2 = Q2 * Q2.determinant();
    std::vector<glm::mat4x3> candidates;
    candidates.push_back(glm::mat4x3(
        glm::vec3(R1.col(0)[0], R1.col(0)[1], R1.col(0)[2]),
        glm::vec3(R1.col(1)[0], R1.col(1)[1], R1.col(1)[2]),
        glm::vec3(R1.col(2)[0], R1.col(2)[1], R1.col(2)[2]),
        glm::vec3(T1[0], T1[1], T1[2])
    ));
    candidates.push_back(glm::mat4x3(
        glm::vec3(R1.col(0)[0], R1.col(0)[1], R1.col(0)[2]),
        glm::vec3(R1.col(1)[0], R1.col(1)[1], R1.col(1)[2]),
        glm::vec3(R1.col(2)[0], R1.col(2)[1], R1.col(2)[2]),
        glm::vec3(T2[0], T2[1], T2[2])
    ));
    candidates.push_back(glm::mat4x3(
        glm::vec3(R2.col(0)[0], R2.col(0)[1], R2.col(0)[2]),
        glm::vec3(R2.col(1)[0], R2.col(1)[1], R2.col(1)[2]),
        glm::vec3(R2.col(2)[0], R2.col(2)[1], R2.col(2)[2]),
        glm::vec3(T1[0], T1[1], T1[2])
    ));
    candidates.push_back(glm::mat4x3(
        glm::vec3(R2.col(0)[0], R2.col(0)[1], R2.col(0)[2]),
        glm::vec3(R2.col(1)[0], R2.col(1)[1], R2.col(1)[2]),
        glm::vec3(R2.col(2)[0], R2.col(2)[1], R2.col(2)[2]),
        glm::vec3(T2[0], T2[1], T2[2])
    ));
    float minError = std::numeric_limits<float>::max();
    size_t selectedIdx = 0;
    for(size_t i = 0; i < candidates.size(); i++)
    {
        float error = reprojectionError(cameraK, candidates[i], objPoints, imgPoints);
        if(error < minError)
        {
            minError = error;
            selectedIdx = i;
        }
    }
    // set pose matrix
    outputM = candidates[selectedIdx];
}

// refer to https://stanford.edu/class/ee267/notes/ee267_notes_tracking.pdf
// decompose homography matrix
void decomposeHomoMatrixStanford(
    const glm::mat3& cameraK, const glm::mat3& cameraInvK, glm::mat3 H, glm::mat4x3& outputM
)
{
    H = cameraInvK * H;
    float normCol0 = glm::length(H[0]);
    float normCol1 = glm::length(H[1]);
    float s = 2.0f / (normCol0 + normCol1);
    // set translation
    outputM[3] = s * H[2];
    // set rotations
    outputM[0] = H[0] / normCol0;
    outputM[1] = H[1] - outputM[0] * glm::dot(outputM[0], H[1]);
    outputM[1] = glm::normalize(outputM[1]);
    outputM[2] = glm::cross(outputM[0], outputM[1]);
}

// scale the estimated projection matrix
float scalePoseM(glm::mat4x3& M)
{
    float scale = std::abs(M[3][2]);
    M /= scale;
    return scale;
}

void testPoseM(const glm::mat3& cameraK, glm::mat4x3& M, const glm::vec2& q)
{
    auto tmp = cameraK * M * glm::vec4(q, 0.0f, 1.0f);
    std::cout << tmp.x << "," << tmp.y << "," << tmp.z << std::endl;
}

// estimate pose from homography (SVD method)
void Marker::estimatePoseSVD(
    const glm::mat3& cameraK, const glm::mat3& cameraInvK,
    const glm::vec3& cameraDistK, const glm::vec2& cameraDistP
)
{
    if(!_new_marker) return;
    if(_marker_borderp1p2.x <= 0.0f)
    {
        _poseM = glm::mat4x3(0.0f);
        _poseMRefined = glm::mat4x3(0.0f);
        return;
    }
    // prepare p
    glm::vec2 p1 = glm::vec2(_marker_borderp1p2.x, _marker_borderp1p2.y);
    glm::vec2 p2 = glm::vec2(_marker_borderp1p2.z, _marker_borderp1p2.w);
    glm::vec2 p3 = glm::vec2(_marker_borderp3p4.x, _marker_borderp3p4.y);
    glm::vec2 p4 = glm::vec2(_marker_borderp3p4.z, _marker_borderp3p4.w);
    // undistort points
    undistortPoints(
        cameraK, cameraDistK, cameraDistP,
        p1, p2, p3, p4
    );
    // glm::mat3x4 mstarT = glm::transpose(
    //     cameraInvK * glm::mat4x3(
    //         glm::vec3(p1, 1.0f),
    //         glm::vec3(p2, 1.0f),
    //         glm::vec3(p3, 1.0f),
    //         glm::vec3(p4, 1.0f)
    //     )
    // );
    // prepare q
    // const glm::vec2 q1 = glm::vec2(0.0f, 0.0f);
    // const glm::vec2 q2 = glm::vec2(0.0f, 1.0f);
    // const glm::vec2 q4 = glm::vec2(1.0f, 0.0f);
    // const glm::vec2 q3 = glm::vec2(1.0f, 1.0f);
    const glm::vec2 q1 = glm::vec2(-1.0f, -1.0f);
    const glm::vec2 q2 = glm::vec2(-1.0f,  1.0f);
    const glm::vec2 q3 = glm::vec2( 1.0f, -1.0f);
    const glm::vec2 q4 = glm::vec2( 1.0f,  1.0f);
    // set up matrix A
    Eigen::Matrix<float, 8, 9> A;
    A << q1.x, q1.y, 1.0f, 0.0f, 0.0f, 0.0f, -p1.x*q1.x, -p1.x*q1.y, -p1.x,
         0.0f, 0.0f, 0.0f, q1.x, q1.y, 1.0f, -p1.y*q1.x, -p1.y*q1.y, -p1.y,

         q2.x, q2.y, 1.0f, 0.0f, 0.0f, 0.0f, -p2.x*q2.x, -p2.x*q2.y, -p2.x,
         0.0f, 0.0f, 0.0f, q2.x, q2.y, 1.0f, -p2.y*q2.x, -p2.y*q2.y, -p2.y,

         q3.x, q3.y, 1.0f, 0.0f, 0.0f, 0.0f, -p3.x*q3.x, -p3.x*q3.y, -p3.x,
         0.0f, 0.0f, 0.0f, q3.x, q3.y, 1.0f, -p3.y*q3.x, -p3.y*q3.y, -p3.y,

         q4.x, q4.y, 1.0f, 0.0f, 0.0f, 0.0f, -p4.x*q4.x, -p4.x*q4.y, -p4.x,
         0.0f, 0.0f, 0.0f, q4.x, q4.y, 1.0f, -p4.y*q4.x, -p4.y*q4.y, -p4.y;
    // solve SVD for A
    Eigen::JacobiSVD<Eigen::MatrixXf> svdSolver(A, Eigen::ComputeFullV);
    auto& matrixV = svdSolver.matrixV();
    auto& h = matrixV.col(matrixV.cols() - 1);
    glm::mat3 H = glm::mat3(
        glm::vec3(h[0], h[3], h[6]),
        glm::vec3(h[1], h[4], h[7]),
        glm::vec3(h[2], h[5], h[8])
    );
    std::vector<glm::vec2> objPoints = {q1, q2, q3, q4};
    std::vector<glm::vec2> imgPoints = {p1, p2, p3, p4};
    // decomposeHomoMatrixZhang(cameraK, cameraInvK, mstarT, H, objPoints, imgPoints, _poseM);
    // decomposeHomoMatrixARBook(cameraK, cameraInvK, H, _poseM);
    decomposeHomoMatrixInternet(cameraK, cameraInvK, H, _poseM);
    // decomposeHomoMatrixInternet2(cameraK, cameraInvK, H, _poseM);
    // decomposeHomoMatrixStanford(cameraK, cameraInvK, H, _poseM);
    _err_scale = scalePoseM(_poseM);
    refinePoseM(cameraInvK, objPoints, imgPoints);
    _err_reproj = reprojectionError(cameraK, _poseMRefined, objPoints, imgPoints);
}

// reference: https://franklinta.com/2014/09/08/computing-css-matrix3d-transforms/
// estimate pose from homography (linear equation method)
void Marker::estimatePoseLinear(
    const glm::mat3& cameraK, const glm::mat3& cameraInvK,
    const glm::vec3& cameraDistK, const glm::vec2& cameraDistP
)
{
    if(!_new_marker) return;
    if(_marker_borderp1p2.x <= 0.0f)
    {
        _poseM = glm::mat4x3(0.0f);
        _poseMRefined = glm::mat4x3(0.0f);
        return;
    }
    // prepare p
    glm::vec2 p1 = glm::vec2(_marker_borderp1p2.x, _marker_borderp1p2.y);
    glm::vec2 p2 = glm::vec2(_marker_borderp1p2.z, _marker_borderp1p2.w);
    glm::vec2 p3 = glm::vec2(_marker_borderp3p4.x, _marker_borderp3p4.y);
    glm::vec2 p4 = glm::vec2(_marker_borderp3p4.z, _marker_borderp3p4.w);
    // undistort points
    undistortPoints(
        cameraK, cameraDistK, cameraDistP,
        p1, p2, p3, p4
    );
    // prepare q
    static glm::vec2 q1 = glm::vec2(-1.0f, -1.0f);
    static glm::vec2 q2 = glm::vec2(-1.0f,  1.0f);
    static glm::vec2 q3 = glm::vec2( 1.0f, -1.0f);
    static glm::vec2 q4 = glm::vec2( 1.0f,  1.0f);
    // set up matrix A
    Eigen::Matrix<float, 8, 8> A;
    A << q1.x, q1.y, 1.0f, 0.0f, 0.0f, 0.0f, -p1.x*q1.x, -p1.x*q1.y,
         0.0f, 0.0f, 0.0f, q1.x, q1.y, 1.0f, -p1.y*q1.x, -p1.y*q1.y,

         q2.x, q2.y, 1.0f, 0.0f, 0.0f, 0.0f, -p2.x*q2.x, -p2.x*q2.y,
         0.0f, 0.0f, 0.0f, q2.x, q2.y, 1.0f, -p2.y*q2.x, -p2.y*q2.y,

         q3.x, q3.y, 1.0f, 0.0f, 0.0f, 0.0f, -p3.x*q3.x, -p3.x*q3.y,
         0.0f, 0.0f, 0.0f, q3.x, q3.y, 1.0f, -p3.y*q3.x, -p3.y*q3.y,

         q4.x, q4.y, 1.0f, 0.0f, 0.0f, 0.0f, -p4.x*q4.x, -p4.x*q4.y,
         0.0f, 0.0f, 0.0f, q4.x, q4.y, 1.0f, -p4.y*q4.x, -p4.y*q4.y;

    Eigen::Vector<float, 8> b;
    b << p1.x, p1.y, p2.x, p2.y, p3.x, p3.y, p4.x, p4.y;
    auto h = A.colPivHouseholderQr().solve(b);
    glm::mat3 H = glm::mat3(
        glm::vec3(h[0], h[3], h[6]),
        glm::vec3(h[1], h[4], h[7]),
        glm::vec3(h[2], h[5], 1.0f)
    );

    std::vector<glm::vec2> objPoints = {q1, q2, q3, q4};
    std::vector<glm::vec2> imgPoints = {p1, p2, p3, p4};
    // decomposeHomoMatrixStanford(cameraK, cameraInvK, H, _poseM);
    // decomposeHomoMatrixARBook(cameraK, cameraInvK, H, _poseM);
    decomposeHomoMatrixInternet(cameraK, cameraInvK, H, _poseM);
    // decomposeHomoMatrixInternet2(cameraK, cameraInvK, H, _poseM);
    // decomposeHomoMatrixDuke(cameraK, cameraInvK, objPoints, imgPoints, H, _poseM);
    _err_scale = scalePoseM(_poseM);
    refinePoseM(cameraInvK, objPoints, imgPoints);
    _err_reproj = reprojectionError(cameraK, _poseMRefined, objPoints, imgPoints);
}

// use opencv as reference
// estimate pose (OpenCV method)
void Marker::estimatePoseOpenCV(
    const glm::mat3& cameraK, const glm::mat3& cameraInvK,
    const glm::vec3& cameraDistK, const glm::vec2& cameraDistP
)
{
    if(!_new_marker) return;
    if(_marker_borderp1p2.x <= 0.0f)
    {
        _poseM = glm::mat4x3(0.0f);
        _poseMRefined = glm::mat4x3(0.0f);
        return;
    }
    std::vector<cv::Point2d> imgPoints = {
        cv::Point2d(_marker_borderp1p2.x, _marker_borderp1p2.y),
        cv::Point2d(_marker_borderp1p2.z, _marker_borderp1p2.w),
        cv::Point2d(_marker_borderp3p4.x, _marker_borderp3p4.y),
        cv::Point2d(_marker_borderp3p4.z, _marker_borderp3p4.w),
    };
    std::vector<cv::Point3d> objPoints = {
        cv::Point3d(-1.0, -1.0, 0.0),
        cv::Point3d(-1.0,  1.0, 0.0),
        cv::Point3d( 1.0, -1.0, 0.0),
        cv::Point3d( 1.0,  1.0, 0.0),
    };
    // std::vector<cv::Point3d> objPoints = {
    cv::Mat3d cameraMat = (
        cv::Mat_<cv::Vec<double,3>>(3,3) << cameraK[0][0], cameraK[1][0], cameraK[2][0],
            cameraK[0][1], cameraK[1][1], cameraK[2][1],
            cameraK[0][2], cameraK[1][2], cameraK[2][2]
    );
    std::vector<double> cameraDist = {
        cameraDistK.x, cameraDistK.y, cameraDistP.x, cameraDistP.y, cameraDistK.z
    };
    cv::Mat rvec, tvec;
    if(!cv::solvePnP(objPoints, imgPoints, cameraMat, cameraDist, rvec, tvec, false, cv::SOLVEPNP_IPPE))
    {
        _poseM = glm::mat4x3(0.0f);
        return;
    }
    cv::Mat rotMat;
    cv::Rodrigues(rvec, rotMat);
    rotMat.convertTo(rotMat, CV_32F);
    tvec.convertTo(tvec, CV_32F);
    _poseM = glm::mat4x3(
        glm::vec3(rotMat.at<float>(0,0), rotMat.at<float>(1,0), rotMat.at<float>(2,0)),
        glm::vec3(rotMat.at<float>(0,1), rotMat.at<float>(1,1), rotMat.at<float>(2,1)),
        glm::vec3(rotMat.at<float>(0,2), rotMat.at<float>(1,2), rotMat.at<float>(2,2)),
        glm::vec3(tvec.at<float>(0,0), tvec.at<float>(0,1), tvec.at<float>(0,2))
    );
    _err_scale = scalePoseM(_poseM);
}


#include <eigen3/Eigen/Dense>
#include <fstream>
#include <sstream>
#include <cmath>
#include <iostream>
#include <boost/function.hpp>
#include <boost/bind.hpp>

#include "ros/ros.h"
#include "std_msgs/Float64.h"
#include "std_msgs/Float64MultiArray.h"
#include "sensor_msgs/JointState.h"
#include "gazebo_msgs/LinkStates.h"
#include "define.h"

#include <typeinfo>

#define DESIRED_VEL 40  // RW_qdot_des [rad/s]
#define NUM_OF_MEASUREMENTS 1000

typedef Eigen::Matrix<float, Eigen::Dynamic, Eigen::Dynamic> Matrix;

template<typename MatType>
using PseudoInverseType = Eigen::Matrix<typename MatType::Scalar, MatType::ColsAtCompileTime, MatType::RowsAtCompileTime>;

template<typename MatType>
PseudoInverseType<MatType> pseudoInverse(const MatType &a, double epsilon = std::numeric_limits<double>::epsilon())
{
    using WorkingMatType = Eigen::Matrix<typename MatType::Scalar, Eigen::Dynamic, Eigen::Dynamic, 0, MatType::MaxRowsAtCompileTime, MatType::MaxColsAtCompileTime>;
    Eigen::BDCSVD<WorkingMatType> svd(a, Eigen::ComputeThinU | Eigen::ComputeThinV);
    svd.setThreshold(epsilon*std::max(a.cols(), a.rows()));
    Eigen::Index rank = svd.rank();
    Eigen::Matrix<typename MatType::Scalar, Eigen::Dynamic, MatType::RowsAtCompileTime,
    0, Eigen::BDCSVD<WorkingMatType>::MaxDiagSizeAtCompileTime, MatType::MaxRowsAtCompileTime>
    tmp = svd.matrixU().leftCols(rank).adjoint();
    tmp = svd.singularValues().head(rank).asDiagonal().inverse() * tmp;
    return svd.matrixV().leftCols(rank) * tmp;
}

bool reachedVel = false;

float q1;       // angle of first joint [rad]
float q2;       // angle of second joint [rad]
float q1dot;    // rate of first joint [rad/s]
float q2dot;    // rate of second joint [rad/s]
float omega0;   // base angular velocity [rad/s]
float RW_vel;   // reaction wheel velocity [rad/s]

void velocityCheckCallback(const sensor_msgs::JointState::ConstPtr& msg) {

    q1      = msg->position[0];
    q2      = msg->position[1];
    q1dot   = msg->velocity[0];
    q2dot   = msg->velocity[1];
    RW_vel  = msg->velocity[4];

    // ROS_INFO("RW_vel: %.5f | q1: %.5f | q2: %.5f | q1dot: %.5f | q2dot: %.5f", RW_vel, q1, q2, q1dot, q2dot);

    if (!reachedVel && RW_vel >= DESIRED_VEL)
        reachedVel = true;
    else if (reachedVel) {
        // ROS_INFO("Measurement number: %d");
    }
}

void positionCheckCallback(const gazebo_msgs::LinkStates::ConstPtr& msg) {

    omega0 = msg->twist[3].angular.z;
    // ROS_INFO("angular twist z: %.5f", omega0);
}

int main(int argc, char **argv) {

    // calculations
    double EulerConstant = std::exp(1.0);

    double M   = M0+M1+M2;
    double I0  = ((double)1/2)*(M0*pow(RH, (double)2));
    double I1  = 3.46*pow(EulerConstant, (double)-4);
    double I2  = 3.46*pow(EulerConstant, (double)-4);
    // double R0X = 0.1954*cos(27.9*PI/180);
    // double R0Y = 0.1954*sin(27.9*PI/180);

    double frequency = (float)1/DT;

    double hrw = Irw * DESIRED_VEL;

    /* define the robot 8 inertial parameters */
    double pi1 = M0*R0X*((M1+M2)*L1+M2*R1)/M;
    double pi2 = M0*R0X*M2*L2/M;
    double pi3 = M0*R0Y*((M1+M2)*L1+M2*R1)/M;
    double pi4 = M0*R0Y*M2*L2/M;
    double pi5 = I0+M0*(M1+M2)*(pow(R0X, (double)2) + pow(R0Y, (double)2))/M;
    double pi6 = M2*L2*(M0*L1+(M0+M1)*R1)/M;
    double pi7 = I1+(M0*(M1+M2)*pow(L1, (double)2)+2*M0*M2*L1*R1+M2*(M0+M1)*pow(R1, (double)2))/M;
    double pi8 = I2+(M2*(M0+M1)*pow(L2, (double)2))/M;

    Eigen::Matrix<double, 8, 1> robotInertialParameters;
    robotInertialParameters(0, 0) = pi1;
    robotInertialParameters(1, 0) = pi2;
    robotInertialParameters(2, 0) = pi3;
    robotInertialParameters(3, 0) = pi4;
    robotInertialParameters(4, 0) = pi5;
    robotInertialParameters(5, 0) = pi6;
    robotInertialParameters(6, 0) = pi7;
    robotInertialParameters(7, 0) = pi8;

    q1 = q2 = q1dot = q2dot = omega0 = 0.0;

    /* Eigen Matrix */
    Matrix Y;
    Y.resize(NUM_OF_MEASUREMENTS, 8);

    /* Define Hrw matrix as a Nx1 column vector and all components equal to hrw */
    Eigen::Matrix<float, NUM_OF_MEASUREMENTS, 1> Hcm;
    for (int i = 0; i < NUM_OF_MEASUREMENTS; ++i)
        Hcm(i, 0) = hrw;

    /* ros init */
    ros::init(argc, argv, "cepheus_controller_node");
    ros::NodeHandle n;

    /* Create publishers */
    ros::Publisher RW_velocity_pub = n.advertise<std_msgs::Float64>("/cepheus/reaction_wheel_velocity_controller/command", 1);
    ros::Publisher LE_position_pub = n.advertise<std_msgs::Float64>("/cepheus/left_elbow_position_controller/command", 1);
    ros::Publisher LS_position_pub = n.advertise<std_msgs::Float64>("/cepheus/left_shoulder_position_controller/command", 1);

    /* messages to publish */
    std_msgs::Float64 msg_RW;
    std_msgs::Float64 msg_LE;
    std_msgs::Float64 msg_LS;

    /* init messages */
    msg_RW.data = 0.1;
    msg_LE.data = 0.1;
    msg_LS.data = 0.1;

    int currentMeasurement = 0;

    /* Create subscribers */
    ros::Subscriber RW_velocity_sub = n.subscribe<sensor_msgs::JointState>("/cepheus/joint_states", 1, velocityCheckCallback);
    ros::Subscriber position_sub = n.subscribe<gazebo_msgs::LinkStates>("/gazebo/link_states", 1, positionCheckCallback);

    ros::Rate loop_rate(frequency);

    while (ros::ok()) {

        RW_velocity_pub.publish(msg_RW);
        LE_position_pub.publish(msg_LE);
        LS_position_pub.publish(msg_LS);

        ros::spinOnce();

        // arm joins sinusoidal movement
        msg_LE.data = 3 * sin(ros::Time::now().toSec());
        msg_LS.data = -1 * sin(ros::Time::now().toSec());

        if (reachedVel && (currentMeasurement < NUM_OF_MEASUREMENTS)) {

            /* Keep RW desired velocity while taking measurements */

            msg_RW.data = DESIRED_VEL;

            ROS_INFO("-----------------------------------------------------------------");
            ROS_INFO("current measurement number: %d", currentMeasurement+1);
            ROS_INFO("q1: %.3f, q2: %.3f, q1dot: %.3f, q2dot: %.3f, omega0: %.3f", q1, q2, q1dot, q2dot, omega0);

            Y(currentMeasurement, 0) = (2*omega0 + q1dot) * cos(q1);
            Y(currentMeasurement, 1) = (2*omega0 + q1dot + q2dot) * cos(q1+q2);
            Y(currentMeasurement, 2) = (2*omega0 + q1dot) * sin(q1);
            Y(currentMeasurement, 3) = (2*omega0 + q1dot + q2dot) * sin(q1+q2);
            Y(currentMeasurement, 4) = omega0;
            Y(currentMeasurement, 5) = (2*omega0 + 2*q1dot + q2dot) * cos(q2);
            Y(currentMeasurement, 6) = omega0 + q1dot;
            Y(currentMeasurement, 7) = omega0 + q1dot + q2dot;

            currentMeasurement++;
        }
        else if (currentMeasurement >= NUM_OF_MEASUREMENTS) {

            /* Measurements completed */

            RW_velocity_sub.shutdown();
            position_sub.shutdown();
            // std::cout << Y << std::endl;

            Eigen::ColPivHouseholderQR<Eigen::MatrixXf> qr_decomp(Y);
            // Eigen::FullPivLU<Matrix> lu_decomp(Y);
            // auto fs = Y.colPivHouseholderQr();
            auto rank = qr_decomp.rank();
            ROS_INFO("rank = %ld", rank);
            if (rank == 8) {
                // std::cout << Y << std::endl;
                auto pinv = pseudoInverse(Y);
                // auto pinv = Y.completeOrthogonalDecomposition().pseudoInverse();
                // std::cout << Y.rows() << " " << Y.cols() << " " << pinv.rows() << " " << pinv.cols() <<'\n';

                auto pi_est = pinv*Hcm;
                std::ofstream file;
                file.open("/home/pelekoudas/cepheus_simulator/ls_pi_est.txt");
                if (file.is_open()) {
                    double e = 0.0;
                    for (int i = 0; i < 8; ++i) {
                        e = (robotInertialParameters(i, 0) - pi_est(i, 0))/robotInertialParameters(i, 0)*100;
                        file << "Robot Parameter " << i << ": " << robotInertialParameters(i, 0) << ", LS estimation: " << pi_est(i, 0) << ", error(%): "<< e << '\n';
                    }
                    file.close();
                } else {
                    std::cout << pi_est << std::endl;
                }
                break;
            }

            // break;
        }
        else {
            msg_RW.data += 0.1;
        }

        loop_rate.sleep();
    }

    return 0;
}


// Copyright Matt Overby 2021.
// Distributed under the MIT License.

#ifndef MCL_NORMAL_HPP
#define MCL_NORMAL_HPP 1

#include <Eigen/Dense>

namespace mcl
{

template <typename T>
static inline Eigen::Matrix<T,3,1> triangle_normal(
	const Eigen::Matrix<T,3,1> &a,
	const Eigen::Matrix<T,3,1> &b,
	const Eigen::Matrix<T,3,1> &c,
	bool normalize = true)
{
	Eigen::Matrix<T,3,1> n = (b-a).cross(c-a);
	if (normalize) { n.stableNormalize(); }
	return n;
}

template <typename T>
static inline Eigen::Matrix<T,2,1> edge_normal(
	const Eigen::Matrix<T,2,1> &p0,
	const Eigen::Matrix<T,2,1> &p1,
	bool normalize = true)
{
	Eigen::Matrix<T,2,1> n(p1[1]-p0[1], -(p1[0]-p0[0]));
	if (normalize) { n.stableNormalize(); }
	return n;
}

} // ns mcl

#endif


#include <stan/math/fwd/scal.hpp>
#include <gtest/gtest.h>
#include <boost/math/special_functions/hypot.hpp>
#include <math/fwd/scal/fun/nan_util.hpp>

TEST(AgradFwdHypot, Fvar) {
  using boost::math::hypot;
  using stan::math::fvar;
  using std::isnan;

  fvar<double> x(0.5, 1.0);
  fvar<double> y(2.3, 2.0);

  fvar<double> a = hypot(x, y);
  EXPECT_FLOAT_EQ(hypot(0.5, 2.3), a.val_);
  EXPECT_FLOAT_EQ((0.5 * 1.0 + 2.3 * 2.0) / hypot(0.5, 2.3), a.d_);

  fvar<double> z(0.0, 1.0);
  fvar<double> w(-2.3, 2.0);
  fvar<double> b = hypot(x, z);

  EXPECT_FLOAT_EQ(0.5, b.val_);
  EXPECT_FLOAT_EQ(1.0, b.d_);

  fvar<double> c = hypot(x, w);
  isnan(c.val_);
  isnan(c.d_);

  fvar<double> d = hypot(z, x);
  EXPECT_FLOAT_EQ(0.5, d.val_);
  EXPECT_FLOAT_EQ(1.0, d.d_);
}

TEST(AgradFwdHypot, FvarFvarDouble) {
  using boost::math::hypot;
  using stan::math::fvar;

  fvar<fvar<double> > x;
  x.val_.val_ = 3.0;
  x.val_.d_ = 1.0;

  fvar<fvar<double> > y;
  y.val_.val_ = 6.0;
  y.d_.val_ = 1.0;

  fvar<fvar<double> > a = hypot(x, y);

  EXPECT_FLOAT_EQ(hypot(3.0, 6.0), a.val_.val_);
  EXPECT_FLOAT_EQ(3.0 / hypot(3.0, 6.0), a.val_.d_);
  EXPECT_FLOAT_EQ(6.0 / hypot(3.0, 6.0), a.d_.val_);
  EXPECT_FLOAT_EQ(-0.059628479, a.d_.d_);
}

struct hypot_fun {
  template <typename T0, typename T1>
  inline typename boost::math::tools::promote_args<T0, T1>::type operator()(
      const T0 arg1, const T1 arg2) const {
    return hypot(arg1, arg2);
  }
};

TEST(AgradFwdHypot, nan_0) {
  hypot_fun hypot_;
  test_nan_fwd(hypot_, 3.0, 5.0, false);
}


/*
 * COPYRIGHT AND PERMISSION NOTICE
 * Penn Software MSCKF_VIO
 * Copyright (C) 2017 The Trustees of the University of Pennsylvania
 * All rights reserved.
 */

#include <iostream>
#include <Eigen/Dense>
#include <gtest/gtest.h>
#include <gtsam_vio/math_utils.hpp>

using namespace std;
using namespace Eigen;
using namespace gtsam_vio;

TEST(MathUtilsTest, skewSymmetric) {
  Vector3d w(1.0, 2.0, 3.0);
  Matrix3d w_hat = skewSymmetric(w);
  Vector3d zero_vector = w_hat * w;

  FullPivLU<Matrix3d> lu_helper(w_hat);
  EXPECT_EQ(lu_helper.rank(), 2);
  EXPECT_DOUBLE_EQ(zero_vector.norm(), 0.0);
  return;
}

TEST(MathUtilsTest, quaternionNormalize) {
  Vector4d q(1.0, 1.0, 1.0, 1.0);
  quaternionNormalize(q);

  EXPECT_DOUBLE_EQ(q.norm(), 1.0);
  return;
}

TEST(MathUtilsTest, quaternionToRotation) {
  Vector4d q(0.0, 0.0, 0.0, 1.0);
  Matrix3d R = quaternionToRotation(q);
  Matrix3d zero_matrix = R - Matrix3d::Identity();

  FullPivLU<Matrix3d> lu_helper(zero_matrix);
  EXPECT_EQ(lu_helper.rank(), 0);
  return;
}

TEST(MathUtilsTest, rotationToQuaternion) {
  Vector4d q1(0.0, 0.0, 0.0, 1.0);
  Matrix3d I = Matrix3d::Identity();
  Vector4d q2 = rotationToQuaternion(I);
  Vector4d zero_vector = q1 - q2;

  EXPECT_DOUBLE_EQ(zero_vector.norm(), 0.0);
  return;
}

TEST(MathUtilsTest, quaternionMultiplication) {
  Vector4d q1(2.0, 2.0, 1.0, 1.0);
  Vector4d q2(1.0, 2.0, 3.0, 1.0);
  q1 = q1 / q1.norm();
  q2 = q2 / q2.norm();
  Vector4d q_prod = quaternionMultiplication(q1, q2);

  Matrix3d R1 = quaternionToRotation(q1);
  Matrix3d R2 = quaternionToRotation(q2);
  Matrix3d R_prod = R1 * R2;
  Matrix3d R_prod_cp = quaternionToRotation(q_prod);

  Matrix3d zero_matrix = R_prod - R_prod_cp;

  EXPECT_NEAR(zero_matrix.sum(), 0.0, 1e-10);
  return;
}

int main(int argc, char** argv) {
  testing::InitGoogleTest(&argc, argv);
  return RUN_ALL_TESTS();
}


#include "gtest/gtest.h"

#include "conex/exponential_map_pade.h"

#include <Eigen/Dense>
#include <unsupported/Eigen/MatrixFunctions>

#include "conex/debug_macros.h"

namespace conex {

using Eigen::Map;
using Eigen::MatrixXd;
using Eigen::VectorXd;

GTEST_TEST(ExponentialMapPadeApproximation, CompareWithEigen) {
  int n = 4;
  MatrixXd A(n, n);
  // clang-format off
  A << 3, 1, 0, 1,
       1, 3, 1, 0,
       0, 1, 4, 1,
       1, 0, 1, 5;
  // clang-format on
  A = A / A.trace();

  MatrixXd reference = A.exp();
  MatrixXd calculated(n, n);

  Map<MatrixXd, Eigen::Aligned> map(calculated.data(), n, n);
  ExponentialMapPadeApproximation(A, &map);
  for (int i = 0; i < n; i++) {
    for (int j = 0; j < n; j++) {
      EXPECT_NEAR(reference(i, j), calculated(i, j), 1e-7);
    }
  }
}

}  // namespace conex


//
// $Id$
//
//
// Original author: Darren Kessner <darren@proteowizard.org>
//
// Copyright 2006 Louis Warschaw Prostate Cancer Center
//   Cedars Sinai Medical Center, Los Angeles, California  90048
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
//

#include "MagnitudeLorentzian.hpp"
#include "MagnitudeLorentzianTestData.hpp"
#include "pwiz/data/misc/FrequencyData.hpp"
#include "pwiz/utility/misc/unit.hpp"
#include <boost/filesystem/operations.hpp>
#include "pwiz/utility/misc/Std.hpp"
#include <cstring>

using namespace pwiz::util;
using namespace pwiz::frequency;
using namespace pwiz::data;

ostream* os_ = 0;
double epsilon_ = numeric_limits<double>::epsilon();

void testBasic()
{
    MagnitudeLorentzian m(1,0,1); // m(x) = 1/sqrt(x^2+1)
    unit_assert_equal(m(0), 1, epsilon_);
    unit_assert_equal(m(1), 1/sqrt(2.), epsilon_);
    unit_assert_equal(m(2), 1/sqrt(5.), epsilon_);
    unit_assert_equal(m(3), 1/sqrt(10.), epsilon_);
    unit_assert_equal(m(4), 1/sqrt(17.), epsilon_);

    // center == 0, alpha == 2*pi, tau == 1/(2*pi)
    unit_assert_equal(m.center(), 0, epsilon_);
    unit_assert_equal(m.alpha(), 2*M_PI, epsilon_);
    unit_assert_equal(m.tau(), 1/(2*M_PI), epsilon_);

    if (os_) *os_ << "testBasic(): success!\n";
}

void testFit()
{
    MagnitudeLorentzian ref(1,0,1);

    // choose sample values near 1!
    // weighting pow(y,6) gives big roundoff errors

    vector< pair<double,double> > samples;
    for (int i=-2; i<3; i++)
        samples.push_back(make_pair(i/10.,ref(i/10.)));

    MagnitudeLorentzian m(samples);

    if (os_)
    {
        *os_ << "coefficients: " << setprecision(14);
        copy(m.coefficients().begin(), m.coefficients().end(), ostream_iterator<double>(*os_, " "));
        *os_ << endl;

        *os_ << "error: " << m(0)-1 << endl;

        for (int i=0; i<5; i++)
            *os_ << i << ", " << m(i) << endl;
    }

    unit_assert_equal(m(0), 1, epsilon_*100);
    unit_assert_equal(m(1), 1/sqrt(2.), epsilon_*100);
    unit_assert_equal(m(2), 1/sqrt(5.), epsilon_*100);
    unit_assert_equal(m(3), 1/sqrt(10.), epsilon_*100);
    unit_assert_equal(m(4), 1/sqrt(17.), epsilon_*100);
    if (os_) *os_ << "testFit(): success!\n";
}

void testData()
{
    string filename = "MagnitudeLorentizianTest.cfd.temp.txt";
    ofstream temp(filename.c_str());
    temp << sampleData_;
    temp.close();

    FrequencyData fd(filename);
    boost::filesystem::remove(filename);

    FrequencyData::const_iterator max = fd.max();
    if (os_) *os_ << "max: (" << max->x << ", " << abs(max->y) << ")\n";

    // fit MagnitudeLorentzian to 3 points on unnormalized data

    vector< pair<double,double> > samples1;
    transform(fd.max()-1, fd.max()+2, back_inserter(samples1), FrequencyData::magnitudeSample);

    if (os_)
    {
        *os_ << "raw data:\n";
        for (unsigned int i=0; i<samples1.size(); i++)
            *os_ << "sample " << i << ": (" << samples1[i].first << ", " << samples1[i].second << ")\n";
    }

    const MagnitudeLorentzian m1(samples1);

    if (os_)
    {
        *os_ << "m1: ";
        copy(m1.coefficients().begin(), m1.coefficients().end(), ostream_iterator<double>(*os_, " "));
        *os_ << endl;
        *os_ << "error: " << scientific << m1.leastSquaresError() << endl;

        for (unsigned int i=0; i<samples1.size(); i++)
            *os_ << "m1(" << i << ") == " << m1(samples1[i].first) << endl;
    }

    // now on normalized data

    fd.normalize();

    vector< pair<double,double> > samples2;
    transform(fd.max()-1, fd.max()+2, back_inserter(samples2), FrequencyData::magnitudeSample);

    if (os_)
    {
        *os_ << "normalized: \n";
        for (unsigned int i=0; i<samples2.size(); i++)
            *os_ << "sample " << i << ": (" << samples2[i].first << ", " << samples2[i].second << ")\n";
    }

    const MagnitudeLorentzian m2(samples2);

    if (os_)
    {
        *os_ << "m2: ";
        copy(m2.coefficients().begin(), m2.coefficients().end(), ostream_iterator<double>(*os_, " "));
        *os_ << endl;
        *os_ << "error: " << scientific << m2.leastSquaresError() << endl;

        for (unsigned int i=0; i<samples2.size(); i++)
            *os_ << "m2(" << i << ") == " << m2(samples2[i].first) << " [" << fd.scale()*m2(samples2[i].first) << "]\n";
    }

    unit_assert_equal(m2.leastSquaresError(), 0, 1e-15);
}

int main(int argc, char* argv[])
{
    TEST_PROLOG(argc, argv)

    try
    {
        if (argc>1 && !strcmp(argv[1],"-v")) os_ = &cout;
        if (os_) *os_ << "MagnitudeLorentzianTest\n";
        testBasic();
        testFit();
        testData();
    }
    catch (exception& e)
    {
        TEST_FAILED(e.what())
    }
    catch (...)
    {
        TEST_FAILED("Caught unknown exception.")
    }

    TEST_EPILOG
}


#include <iostream>
#include <Eigen/Core>
#include <Eigen/Geometry>

using namespace std;
int main()
{
    Eigen::Quaterniond q1(0.35,0.2,0.3,0.1);
    q1=q1.normalized();
    Eigen::Quaterniond p(0,0.1,0.2,0.3);
    cout<<"p"<<endl<<p.coeffs()<<endl;
    Eigen::Quaterniond p1=q1*p*q1.inverse();
    cout<<"p after rotate:"<<endl<<p1.coeffs()<<endl;

    return 0;
}


/* Copyright (C) 2012-2019 IBM Corp.
 * This program is Licensed under the Apache License, Version 2.0
 * (the "License"); you may not use this file except in compliance
 * with the License. You may obtain a copy of the License at
 *   http://www.apache.org/licenses/LICENSE-2.0
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License. See accompanying LICENSE file.
 */

/* Test_Permutations.cpp - Applying plaintext permutation to encrypted vector
 */
#include <NTL/ZZ.h>
NTL_CLIENT

#include <helib/NumbTh.h>
#include <helib/timing.h>
#include <helib/permutations.h>
#include <helib/EncryptedArray.h>
#include <helib/ArgMap.h>

using namespace helib;

static bool noPrint = true;

void testCtxt(long m, long p, long widthBound=0, long L=0, long r=1);

// OLD CODE
//void usage(char *prog)
//{
//  cout << "Usage: "<<prog<<" [test=? [optional parameters...]]\n";
//  cout << "  optional parameters have the form 'attr1=val1 attr2=val2 ...'\n";
//  cout << "  e.g, 'test=1 m=108 p=2 r=1\n";
//  cout << "  test is either 0 (plaintext) or 1 (ciphertext)[default=1]\n\n";
//  cout << "test=0, permuting plaintext hypercubes (dimension upto 4):\n";
//  cout << "  ord1,ord2,ord3,ord4 size of dimensions 1..4 [default ord1=30, ord2,3,4=0]\n";
//  cout << "  good1,good2,good3,good4 native rotation flags (0/1) [default=1]\n";
//  cout << "  depth bounds the depth of permutation network [default=5]\n";
//  cout << "\ntest=1, permuting ciphertext slots:\n";
//  cout << "  m is the cyclotomic field [default=4369]\n";
//  cout << "  p,r define the plaintext space p^r [default p=2,r=1]\n";
//  cout << "  depth bounds the depth of permutation network [default=5]\n";
//  cout << "  L is number of bits in chain [default=30*depth]\n\n";
//  cout << "dry=1 for dry run [default=0]\n";
//  exit(0);
//}

void testCube(Vec<GenDescriptor>& vec, long widthBound)
{
  GeneratorTrees trees;
  long cost = trees.buildOptimalTrees(vec, widthBound);
  if (!noPrint) {
    cout << "@TestCube: trees=" << trees << endl;
    cout << " cost =" << cost << endl;
  }
  Vec<long> dims;
  trees.getCubeDims(dims);
  CubeSignature sig(dims);

  for (long cnt=0; cnt<3; cnt++) {
    Permut pi;
    randomPerm(pi, trees.getSize());

    PermNetwork net;
    net.buildNetwork(pi, trees);

    HyperCube<long> cube1(sig), cube2(sig);
    for (long i=0; i<cube1.getSize(); i++) cube1[i] = i;
    HyperCube<long> cube3 = cube1;
    applyPermToVec(cube2.getData(), cube1.getData(), pi); // direct application
    net.applyToCube(cube3); // applying permutation netwrok
    if (cube2==cube3) cout << "GOOD\n";
    else {
      cout << "BAD\n";
      if (cube1.getSize()<100 && !noPrint) {
	cout << "in="<<cube1.getData() << endl;
	cout << "out1="<<cube2.getData()<<", out2="
	     << cube3.getData()<<endl<<endl;
      }
    }
  }
}

void testCtxt(long m, long p, long widthBound, long L, long r)
{
  if (!noPrint)
    cout << "@testCtxt(m="<<m<<",p="<<p<<",depth="<<widthBound<< ",r="<<r<<")";

  Context context(m,p,r);
  EncryptedArray ea(context); // Use G(X)=X for this ea object

  // Some arbitrary initial plaintext array
  vector<long> in(ea.size());
  for (long i=0; i<ea.size(); i++) in[i] = i % p;

  // Setup generator-descriptors for the PAlgebra generators
  Vec<GenDescriptor> vec(INIT_SIZE, ea.dimension());
  for (long i=0; i<ea.dimension(); i++)
    vec[i] = GenDescriptor(/*order=*/ea.sizeOfDimension(i),
			   /*good=*/ ea.nativeDimension(i), /*genIdx=*/i);

  // Some default for the width-bound, if not provided
  if (widthBound<=0) widthBound = 1+log2((double)ea.size());

  // Get the generator-tree structures and the corresponding hypercube
  GeneratorTrees trees;
  long cost = trees.buildOptimalTrees(vec, widthBound);
  if (!noPrint) {
    context.zMStar.printout();
    cout << ": trees=" << trees << endl;
    cout << " cost =" << cost << endl;
  }
  //  Vec<long> dims;
  //  trees.getCubeDims(dims);
  //  CubeSignature sig(dims);

  // 1/2 prime per level should be more or less enough, here we use 1 per layer
  if (L<=0) L = (1+trees.numLayers())*context.BPL();
  buildModChain(context, /*nLevels=*/L, /*nDigits=*/3);
  if (!noPrint) cout << "**Using "<<L<<" and "
		     << context.ctxtPrimes.card() << " Ctxt-primes\n";

  // Generate a sk/pk pair
  SecKey secretKey(context);
  const PubKey& publicKey = secretKey;
  secretKey.GenSecKey(); // A +-1/0 secret key
  Ctxt ctxt(publicKey);

  for (long cnt=0; cnt<3; cnt++) {
    resetAllTimers();
    // Choose a random permutation
    Permut pi;
    randomPerm(pi, trees.getSize());

    // Build a permutation network for pi
    PermNetwork net;
    net.buildNetwork(pi, trees);

    // make sure we have the key-switching matrices needed for this network
    addMatrices4Network(secretKey, net);

    // Apply the permutation pi to the plaintext
    vector<long> out1(ea.size());
    vector<long> out2(ea.size());
    applyPermToVec(out1, in, pi); // direct application

    // Encrypt plaintext array, then apply permutation network to ciphertext
    ea.encrypt(ctxt, publicKey, in);
    if (!noPrint)
      cout << "  ** applying permutation network to ciphertext... " << flush;
    double t = GetTime();
    net.applyToCtxt(ctxt, ea); // applying permutation netwrok
    t = GetTime() -t;
    if (!noPrint)
      cout << "done in " << t << " seconds" << endl;
    ea.decrypt(ctxt, secretKey, out2);

    if (out1==out2) cout << "GOOD\n";
    else {
      cout << "************ BAD\n";
    }
    // printAllTimers();
  }
}

/* m = 31, p = 2, phi(m) = 30
  ord(p)=5
  generator 6 has order (== Z_m^*) of 6
  T = [1 6 5 30 25 26 ]

  m = 61, p = 3, phi(m) = 60
  ord(p)=10
  generator 13 has order (== Z_m^*) of 3
  generator 2 has order (!= Z_m^*) of 2
  T = [1 2 13 26 47 33 ]

  m = 683, p = 2, phi(m) = 682
  ord(p)=22
  generator 3 has order (== Z_m^*) of 31

  m = 47127, p = 2, phi(m) = 30008
  ord(p)=22
  generator 5 has order (== Z_m^*) of 682
  generator 13661 has order (== Z_m^*) of 2
*/

int main(int argc, char *argv[])
{
  long test = 1;
  long p = 2;
  long r = 1;
  long m = 4369;
  long depth = 5;
  long L = 0;

  long ord1 = 30;
  long ord2 = 0;
  long ord3 = 0;
  long ord4 = 0;
  long good1 = 1;
  long good2 = 1;
  long good3 = 1;
  long good4 = 1;

  bool dry = 0;
  noPrint = 1;

  ArgMap amap;
  amap.arg("test", test);
  amap.arg("p", p);
  amap.arg("r", r);
  amap.arg("m", m);
  amap.arg("depth", depth);
  amap.arg("L", L);
  amap.arg("ord1", ord1);
  amap.arg("ord2", ord2);
  amap.arg("ord3", ord3);
  amap.arg("ord4", ord4);
  amap.arg("good1", good1);
  amap.arg("good2", good2);
  amap.arg("good3", good3);
  amap.arg("good4", good4);
  amap.arg("dry", dry);
  amap.arg("noPrint", noPrint);
  amap.parse(argc, argv);

  // get parameters from the command line
  //if (!parseArgs(argc, argv, argmap)) usage(argv[0]);

  setDryRun(dry);
  if (test==0 || dry!=0) {
    Vec<GenDescriptor> vec;
    long nGens;
    if (ord2<=1) nGens=1;
    else if (ord3<=1) nGens=2;
    else if (ord4<=1) nGens=3;
    else nGens=4;
    vec.SetLength(nGens);

    switch (nGens) {
    case 4:  vec[3] = GenDescriptor(ord4, good4, /*genIdx=*/3);
    case 3:  vec[2] = GenDescriptor(ord3, good3, /*genIdx=*/2);
    case 2:  vec[1] = GenDescriptor(ord2, good2, /*genIdx=*/1);
    default: vec[0] = GenDescriptor(ord1, good1, /*genIdx=*/0);
    }
    if (!noPrint) {
      cout << "***Testing ";
      if (isDryRun()) cout << "(dry run) ";
      for (long i=0; i<vec.length(); i++)
	cout << "("<<vec[i].order<<","<<vec[i].good<<")";
      cout << ", depth="<<depth<<"\n";
    }
    testCube(vec, depth);
  }
  else {
    setTimersOn();
    if (!noPrint)
      cout << "***Testing m="<<m<<", p="<<p<<", depth="<<depth<< endl;
    testCtxt(m,p,depth,L,r);
  }
}

#if 0
  cout << "***Testing m=31, p=2, width=3\n"; // (6 good)
  testCtxt(/*m=*/31, /*p=*/2, /*width=*/3);

  cout << "\n***Testing m=61, p=3, width=3\n"; // (3 good), (2, bad)
  testCtxt(/*m=*/61, /*p=*/3, /*width=*/3);

  cout << "\n***Testing m=683, p=2, width=5\n"; // (31, good)
  testCtxt(/*m=*/683, /*p=*/2, /*width=*/5);

  //  cout << "\n***Testing m=47127, p=2, width=11\n"; // (682,good),(2,good)
  //  testCtxt(/*m=*/47127, /*p=*/2, /*width=*/11);

  // Test 1: a single good small prime-order generator (3)
  {
  Vec<GenDescriptor> vec(INIT_SIZE, 1);
  vec[0] = GenDescriptor(/*order=*/3, /*good=*/true, /*genIdx=*/0);
  cout << "***Testing (3,good), width=1\n";
  testCube(vec, /*width=*/1);
  }

  // Test 2: a single bad larger prime-order generator (31)
  {
  Vec<GenDescriptor> vec(INIT_SIZE, 1);
  vec[0] = GenDescriptor(/*order=*/31, /*good=*/false, /*genIdx=*/0);
  cout << "\n***Testing (31,bad), width=5\n";
  testCube(vec, /*width=*/5);
  }

  // Test 3: two generators with small prime orders (2,3), both bad
  {
  Vec<GenDescriptor> vec(INIT_SIZE, 2);
  vec[0] = GenDescriptor(/*order=*/2, /*good=*/false, /*genIdx=*/0);
  vec[1] = GenDescriptor(/*order=*/3, /*good=*/false, /*genIdx=*/1);
  cout << "\n***Testing [(2,bad),(3,bad)], width=3\n";
  testCube(vec, /*width=*/3);
  }

  // Test 4: two generators with small prime orders (2,3), one good
  {
  Vec<GenDescriptor> vec(INIT_SIZE, 2);
  vec[0] = GenDescriptor(/*order=*/3, /*good=*/true, /*genIdx=*/0);
  vec[1] = GenDescriptor(/*order=*/2, /*good=*/false, /*genIdx=*/1);
  cout << "\n***Testing [(3,good),(2,bad)], width=3\n";
  testCube(vec, /*width=*/3);
  }

  // Test 5: a single good composite-order generator (6)
  {
  Vec<GenDescriptor> vec(INIT_SIZE, 1);
  vec[0] = GenDescriptor(/*order=*/6, /*good=*/true, /*genIdx=*/0);
  cout << "\n***Testing (6,good), width=3\n";
  testCube(vec, /*width=*/3);
  }

  // Test 6: (6,good),(2,bad)
  {
  Vec<GenDescriptor> vec(INIT_SIZE, 2);
  vec[0] = GenDescriptor(/*order=*/6,/*good=*/true, /*genIdx=*/0);
  vec[1] = GenDescriptor(/*order=*/ 2, /*good=*/false,/*genIdx=*/1);
  cout << "\n**Testing [(6,good),(2,bad)], width=5\n";
  testCube(vec, /*width=*/5);
  }

  // Test 7: the "general case", (682,good),(2,bad)
  {
  Vec<GenDescriptor> vec(INIT_SIZE, 2);
  vec[0] = GenDescriptor(/*order=*/682,/*good=*/true, /*genIdx=*/0);
  vec[1] = GenDescriptor(/*order=*/ 2, /*good=*/false,/*genIdx=*/1);
  cout << "\n**Testing [(682,good),(2,bad)], width=11\n";
  testCube(vec, /*width=*/11);
  }
#endif


// Copyright (c) 2021 FRC Team 3512. All Rights Reserved.

#pragma once

#include <algorithm>

#include <Eigen/Cholesky>
#include <Eigen/Core>
#include <Eigen/Eigenvalues>
#include <Eigen/QR>
#include <drake/math/discrete_algebraic_riccati_equation.h>
#include <frc/StateSpaceUtil.h>
#include <frc/system/Discretization.h>
#include <frc/system/LinearSystem.h>
#include <units/time.h>
#include <wpi/array.h>

#include "controllers/LQR.hpp"

namespace frc3512 {

/**
 * Contains the controller coefficients and logic for an implicit model
 * follower.
 *
 * Implicit model following lets us design a feedback controller that erases the
 * dynamics of our system and makes it behave like some other system. This can
 * be used to make a drivetrain more controllable during teleop driving by
 * making it behave like a slower or more benign drivetrain.
 *
 * For more on the underlying math, read appendix C.3 in
 * https://file.tavsys.net/control/controls-engineering-in-frc.pdf.
 */
template <int States, int Inputs>
class ImplicitModelFollower {
public:
    /**
     * Constructs a controller with the given coefficients and plant.
     *
     * @param plant    The plant being controlled.
     * @param plantRef The plant whose dynamics should be followed.
     * @param Qelems   The maximum desired error tolerance for each state.
     * @param Relems   The maximum desired control effort for each input.
     * @param dt       Discretization timestep.
     */
    template <int Outputs>
    ImplicitModelFollower(
        const frc::LinearSystem<States, Inputs, Outputs>& plant,
        const frc::LinearSystem<States, Inputs, Outputs>& plantRef,
        const wpi::array<double, States>& Qelems,
        const wpi::array<double, Inputs>& Relems, units::second_t dt)
        : ImplicitModelFollower<States, Inputs>(plant.A(), plant.B(),
                                                plantRef.A(), plantRef.B(),
                                                Qelems, Relems, dt) {}

    /**
     * Constructs a controller with the given coefficients and plant.
     *
     * @param A      Continuous system matrix of the plant being controlled.
     * @param B      Continuous input matrix of the plant being controlled.
     * @param Aref   Continuous system matrix whose dynamics should be followed.
     * @param Bref   Continuous input matrix whose dynamics should be followed.
     * @param Qelems The maximum desired error tolerance for each state.
     * @param Relems The maximum desired control effort for each input.
     * @param dt     Discretization timestep.
     */
    ImplicitModelFollower(const Eigen::Matrix<double, States, States>& A,
                          const Eigen::Matrix<double, States, Inputs>& B,
                          const Eigen::Matrix<double, States, States>& Aref,
                          const Eigen::Matrix<double, States, States>& Bref,
                          const wpi::array<double, States>& Qelems,
                          const wpi::array<double, Inputs>& Relems,
                          units::second_t dt) {
        // Discretize real dynamics
        Eigen::Matrix<double, States, States> discA;
        Eigen::Matrix<double, States, Inputs> discB;
        frc::DiscretizeAB<States, Inputs>(A, B, dt, &discA, &discB);

        // Discretize desired dynamics
        Eigen::Matrix<double, States, States> discAref;
        Eigen::Matrix<double, States, Inputs> discBref;
        frc::DiscretizeAB<States, Inputs>(Aref, Bref, dt, &discAref, &discBref);

        // Find initial Q and R weights
        Eigen::Matrix<double, States, States> Q = frc::MakeCostMatrix(Qelems);
        Eigen::Matrix<double, Inputs, Inputs> R = frc::MakeCostMatrix(Relems);

        Eigen::Matrix<double, States, States> Adiff = discA - discAref;

        Eigen::Matrix<double, States, States> Qimf =
            Adiff.transpose() * Q * Adiff;
        Eigen::Matrix<double, Inputs, Inputs> Rimf =
            discB.transpose() * Q * discB + R;
        Eigen::Matrix<double, States, Inputs> Nimf =
            Adiff.transpose() * Q * discB;

        // Nudge eigenvalues of Qimf slightly more positive if Q <= 0, since
        // this is usually caused by numerical imprecision
        Eigen::SelfAdjointEigenSolver<decltype(Qimf)> Qeigen{Qimf};
        if (!std::all_of(Qeigen.eigenvalues().data(),
                         Qeigen.eigenvalues().data() + States,
                         [](const auto& elem) { return elem >= 0.0; })) {
            Qimf += decltype(Qimf)::Identity() * 1e-10;
        }

        m_K = LQR<States, Inputs>(discA, discB, Qimf, Rimf, Nimf);

        // Find u_imf that makes real model match reference model.
        //
        // x_k+1 = Ax_k + Bu_imf
        // z_k+1 = Aref z_k + Bref u_k
        //
        // Let x_k = z_k.
        //
        // x_k+1 = z_k+1
        // Ax_k + Bu_imf = Aref x_k + Bref u_k
        // Bu_imf = Aref x_k - Ax_k + Bref u_k
        // Bu_imf = (Aref - A)x_k + Bref u_k
        // u_imf = B^+ ((Aref - A)x_k + Bref u_k)
        // u_imf = -B^+ (A - Aref)x_k + B^+ Bref u_k

        // The first term makes the open-loop poles that of the reference
        // system, and the second term makes the input behave like that of the
        // reference system.
        m_B = discB.householderQr().solve(discBref);

        Reset();
    }

    /**
     * Returns the controller matrix K.
     */
    const Eigen::Matrix<double, Inputs, States>& K() const { return m_K; }

    /**
     * Returns an element of the controller matrix K.
     *
     * @param i Row of K.
     * @param j Column of K.
     */
    double K(int i, int j) const { return m_K(i, j); }

    /**
     * Returns the control input vector u.
     *
     * @return The control input.
     */
    const Eigen::Matrix<double, Inputs, 1>& U() const { return m_u; }

    /**
     * Returns an element of the control input vector u.
     *
     * @param i Row of u.
     *
     * @return The row of the control input vector.
     */
    double U(int i) const { return m_u(i, 0); }

    /**
     * Resets the controller.
     */
    void Reset() { m_u.setZero(); }

    /**
     * Returns the next output of the controller.
     *
     * @param x The current state x.
     * @param u The current input for the original model.
     */
    Eigen::Matrix<double, Inputs, 1> Calculate(
        const Eigen::Matrix<double, States, 1>& x,
        const Eigen::Matrix<double, Inputs, 1>& u) {
        m_u = -m_K * x + m_B * u;
        return m_u;
    }

private:
    // Computed controller output
    Eigen::Matrix<double, Inputs, 1> m_u;

    // Controller gain
    Eigen::Matrix<double, Inputs, States> m_K;

    // Input space conversion gain
    Eigen::Matrix<double, Inputs, Inputs> m_B;
};

}  // namespace frc3512


#include "lib/header.h"
#include "lib/random.h"

#include <armadillo/armadillo>

#include "glog/logging.h"

struct Kalman {
   public:
    Kalman(size_t n_out, size_t n_states) : n_out_(n_out) {
        x_ = arma::colvec(n_states, arma::fill::zeros);

        a_ = arma::mat(n_states, n_states, arma::fill::zeros);
        for (size_t i = 0; i < n_states; ++i) {
            for (size_t j = i; j < n_states; ++j) {
                a_(i, j) = 1.0 / std::pow(10.0, j - i);
            }
        }

        // LOG(INFO) << "A:\n" << a_;

        i_ = arma::eye(n_states, n_states);

        h_ = arma::mat(n_out, n_states, arma::fill::zeros);
        for (size_t i = 0; i < n_out; ++i) {
            h_(i, i) = 1.0;
        }

        q_ = 10.0 * arma::eye(n_states, n_states) + 0.1 * arma::ones(n_states, n_states);
        r_ = 40.0 * arma::eye(n_out, n_out) + 0.1 * arma::ones(n_out, n_out);
        p_ = arma::zeros(n_states, n_states);
    }

    void update(const DoubleVector& z) {
        x_hat_ = a_ * x_;
        p_ = a_ * p_ * a_.t() + q_;
        const arma::colvec y = arma::colvec(z) - h_ * x_hat_;
        const arma::mat s = h_ * p_ * h_.t() + r_;
        k_ = p_ * h_.t() * s.i();
        x_hat_ += k_ * y;
        p_ = (i_ - k_ * h_) * p_;
        x_ = x_hat_;
        // LOG(INFO) << "Z:\n" << z;
        // LOG(INFO) << "X:\n" << x_;
    }

    DoubleVector state() const {
        DoubleVector result(n_out_);
        const arma::vec s = h_ * x_;
        for (size_t i = 0; i < n_out_; ++i) {
            result[i] = s(i);
        }
        return result;
    }

   private:
    size_t n_out_;
    arma::colvec x_;
    arma::colvec x_hat_;
    arma::mat a_;
    arma::mat h_;
    arma::mat q_;
    arma::mat r_;
    arma::mat k_;
    arma::mat p_;
    arma::mat i_;
};

int main() {
    Kalman k(1, 3);
    double sum1 = 0;
    double sum2 = 0;
    for (size_t i = 0; i < 100000; ++i) {
        const double value = std::sin(static_cast<double>(i) / 1000.);
        const double noise_value = value + 0.01 * randNorm01<double>();
        k.update({noise_value});
        sum1 += std::abs(noise_value - value);
        sum2 += std::abs(k.state()[0] - value);
        std::cout << (noise_value - value) << " " << (k.state()[0] - value) << std::endl;
    }
    std::cout << sum2 / sum1 << std::endl;
    return 0;
}


//===- gemm.cc ------------------------------------------------------------===//
//
//                       The CIM Hardware Simulator Project
//
// See LICENSE.TXT for details.
//
//===----------------------------------------------------------------------===//
#include "gemm.hh"
#include "diagnostic/msgHandling.hh"
#include <Eigen/Dense>

namespace cimHW {

cimHWGemmOp::cimHWGemmOp(
  const std::string &configFile,
  const_element_type* input_A_, const_dim_type input_A_ndim_, const_dim_type* input_A_dims_,
  const_element_type* input_B_, const_dim_type input_B_ndim_, const_dim_type* input_B_dims_,
  const_element_type* input_C_, const_dim_type input_C_ndim_, const_dim_type* input_C_dims_,
  element_type* output_Y_, const_dim_type output_Y_ndim_, const_dim_type* output_Y_dims_,
  const_element_type alpha_, const_element_type beta_,
  const_dim_type transA_, const_dim_type transB_)
	: cimHWOp("Gemm")
  , m_input_A(input_A_)
  , m_input_B(input_B_)
  , m_input_C(input_C_)
  , m_output_Y(output_Y_)
  , m_alpha(alpha_), m_beta(beta_)
  , m_cimCU(configFile)
{
  verbose1(opName());

  // input data validation check
  if(input_A_ndim_ != 2 || input_B_ndim_ != 2)
    error("Rank of inputs should be 2");

  // handle transition of matrix, finally the matrix should look like...
  // (m_M, m_K) x (m_K, m_N)
  if(!transA_)
  {
    m_M = input_A_dims_[0];
    m_K = input_A_dims_[1];
  }
  else
  {
    m_M = input_A_dims_[1];
    m_K = input_A_dims_[0];
  }
  if(!transB_)
  {
    m_N = input_B_dims_[1];
    if(m_K != input_B_dims_[0])
      error("The dimension of matrics to be gemm should be consistent");
  }
  else
  {
    m_N = input_B_dims_[0];
    if(m_K != input_B_dims_[1])
      error("The dimension of matrics to be gemm should be consistent");
  }

  // map to proper view of matrix according to transision
  m_matrix_A = (!transA_)?
              Eigen::Map<const MatrixXfRowMajor>(m_input_A, m_M, m_K).eval() :
              Eigen::Map<const MatrixXfRowMajor>(m_input_A, m_K, m_M).transpose();
  m_matrix_B = (!transB_)?
              Eigen::Map<const MatrixXfRowMajor>(m_input_B, m_K, m_N).eval() :
              Eigen::Map<const MatrixXfRowMajor>(m_input_B, m_N, m_K).transpose();

  // transform C's shape and size
  m_C_shape = std::vector<dim_type>(input_C_dims_, input_C_dims_ + input_C_ndim_);
  m_C_size = std::accumulate(m_C_shape.cbegin(), m_C_shape.cend(), 1, std::multiplies<dim_type>());
}

void cimHWGemmOp::simulate()
{
  // take matrixA as kernel and matrixB as data to use ComputeUnit
  MatrixXfRowMajor result = MatrixXfRowMajor::Zero(m_M, m_N);
  for(int i=0; i<m_matrix_A.rows(); i++)
    result.row(i) = m_cimCU.compute(m_matrix_B, m_matrix_A.row(i).transpose());

  // multiply alpha
  result *= m_alpha;
  // Broadcast the bias as needed if bias is given
  if (m_beta != 0 && m_input_C != nullptr) {
    MatrixXfRowMajor bias_mat(m_M, m_N);
    if (m_C_size == 1) {
      // C is (), (1,) or (1, 1), set the scalar
      bias_mat.setConstant(*m_input_C);
    } else if (m_C_shape.size() == 1 || m_C_shape[0] == 1) {
      // C is (N,) or (1, N)
      bias_mat.rowwise() = Eigen::Map<const RowVectorXf>(m_input_C, m_N);
    } else if (m_C_shape[1] == 1) {
      // C is (M, 1)
      bias_mat.colwise() = Eigen::Map<const Eigen::VectorXf>(m_input_C, m_M);
    } else {
      // C is (M, N), no broadcast needed.
      bias_mat = Eigen::Map<const MatrixXfRowMajor>(m_input_C, m_M, m_N);
    }
    // multiply apply beta
    result += bias_mat * m_beta;
  }

  // copy back to output
  memcpy(m_output_Y, result.data(), sizeof(numType) * m_M * m_N);
  verbose1("\n");
}

} // namespace cimHW


#include <iostream>
#include <fstream>
#include <string>
#include <Eigen/Dense>
#include <vector>
#include <type_traits>

using namespace std;
using namespace Eigen;

MatrixXd open(string fileToOpen)
{
    vector<double> MtxEntries;
    ifstream MtxDataFile(fileToOpen);
    string MtxRowString;
    string MtxEntry;

    int MtxRowNumber = 0;

    while (getline(MtxDataFile, MtxRowString))
    {
        stringstream MtxRowStringStream(MtxRowString);

        while (getline(MtxRowStringStream, MtxEntry, ','))
        {
            MtxEntries.push_back(stod(MtxEntry));
        }
        MtxRowNumber++;
    }

    return Map<Matrix<double, Dynamic, Dynamic, RowMajor>>(matrixEntries.data(), matrixRowNumber, matrixEntries.size() / matrixRowNumber);
 }

void save(string fileName, MatrixXd  matrix)
{
    const static IOFormat CSVFormat(FullPrecision, DontAlignCols, " ", "\n");
    ofstream file(fileName);
    if (file.is_open())
    {
        file << matrix.format(CSVFormat);
        file.close();
    }
}

void Print(MatrixXd a, int n, int m)
{
     for (int i = 0; i < n; i++)
{
	 for (int j = n; j < m; j++)
{
            printf("%.3f  ", a(i, j));
        }
        fprintf(output_file,"\n");
    }

     return;
}

void GaussJordan(MatrixXd mtx){

   MatrixXd I_ = MatrixXd::Identity(mtx.rows(),mtx.cols());

   MatrixXd C(mtx.rows(), mtx.cols()+I_.cols());

   C << mtx, I_;

   C.row(0).swap(C.row(mtx.rows() - 1));

   for (int i = 0; i < mtx.rows(); i++) {

        for (int j = 0; j < mtx.rows(); j++)
{
            if (j != i) {

                float temp = C(j,i) / C(i,i);
                for (int k = 0; k < 2*mtx.rows(); k++)
{
                    C(j,k) -= C(i,k)*temp;
                }
            }
        }
    }

    for (int i = 0; i < mtx.rows(); i++)
{
        float temp_ = C(i,i);

        for (int j = 0; j < 2 * mtx.rows(); j++)
{
            C(i,j) = C(i,j)/temp_;
        }
      }

    Print(C, mtx.cols(), 2*mtx.cols());
}


/*
 * Copyright (c) 2015, The Regents of the University of California (Regents).
 * All rights reserved.
 *
 * 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. Neither the name of the copyright holder nor the names of its
 *       contributors may be used to endorse or promote products derived
 *       from this software without specific prior written permission.
 *
 * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS 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 COPYRIGHT HOLDER OR CONTRIBUTORS 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.
 *
 * Please contact the author(s) of this library if you have any questions.
 * Authors: Erik Nelson            ( eanelson@eecs.berkeley.edu )
 *          David Fridovich-Keil   ( dfk@eecs.berkeley.edu )
 */

#include "rotation.h"

#include <cmath>
#include <Eigen/LU>
#include <glog/logging.h>

namespace bsfm {

// Convert from Euler angles to a rotation matrix. Phi, theta, and psi define the
// angles of the intermediate rotations about x (R_x), y (R_y), and z (R_z)
// respectively. See https://en.wikipedia.org/wiki/Rotation_matrix.
Matrix3d EulerAnglesToMatrix(double phi, double theta, double psi) {
  double c1 = std::cos(phi);
  double c2 = std::cos(theta);
  double c3 = std::cos(psi);
  double s1 = std::sin(phi);
  double s2 = std::sin(theta);
  double s3 = std::sin(psi);

  Matrix3d R;
  R(0, 0) = c2*c3;
  R(0, 1) = c3*s1*s2 - c1*s3;
  R(0, 2) = s1*s3 + c1*c3*s2;
  R(1, 0) = c2*s3;
  R(1, 1) = c1*c3 + s1*s2*s3;
  R(1, 2) = c1*s2*s3 - c3*s1;
  R(2, 0) = -s2;
  R(2, 1) = c2*s1;
  R(2, 2) = c1*c2;

  return R;
}

// Same thing as above, but where phi, theta, and psi are specified as a vector.
Matrix3d EulerAnglesToMatrix(const Vector3d& euler_angles) {
  return EulerAnglesToMatrix(euler_angles(0), euler_angles(1), euler_angles(2));
}

// Convert from a rotation matrix to Euler angles.
// From: http://staff.city.ac.uk/~sbbh653/publications/euler.pdf
// Note that the solution that is returned is only unique when phi, theta, and
// psi are all <= 0.5 * PI. If this is not the case, they will still be correct,
// but may not be unique!
Vector3d MatrixToEulerAngles(const Matrix3d& R) {
  // Make sure R is actually a rotation matrix.
  if (std::abs(R.determinant() - 1) > 1e-4) {
    LOG(WARNING) << "R does not have a determinant of 1.";
    return Vector3d::Zero();
  }

  double theta = -std::asin(R(2, 0));

  if (std::abs(cos(theta)) < 1e-8) {
    LOG(WARNING) << "Theta is approximately +/- PI/2, which yields a "
                    "singularity. Cannot decompose matrix into Euler angles.";
    return Vector3d(theta, 0.0, 0.0);
  }

  double phi = std::atan2(R(2, 1), R(2, 2));
  double psi = std::atan2(R(1, 0) / std::cos(theta), R(0, 0) / std::cos(theta));

  return Vector3d(phi, theta, psi);
}

// Get roll angle from a rotation matrix.
// Just like above, the solution will only be unique if roll < 0.5 * PI.
double Roll(const Matrix3d& R) {
  double theta = -std::asin(R(2, 0));
  if (std::abs(std::cos(theta)) < 1e-8)
    return 0.0;
  return std::atan2(R(2, 1), R(2, 2));
}

// Get pitch angle from a rotation matrix.
// This solution is unique.
double Pitch(const Matrix3d& R) {
  return -std::asin(R(2, 0));
}

// Get yaw angle from a rotation matrix.
// Just like above, the solution will only be unique if yaw < 0.5 * PI.
double Yaw(const Matrix3d& R) {
  double theta = -std::asin(R(2, 0));
  if (std::abs(std::cos(theta)) < 1e-8)
    return 0.0;
  return std::atan2(R(1, 0) / std::cos(theta), R(0, 0) / std::cos(theta));
}

// Unroll an angle to be \in [0, 2*PI)
double Unroll(double angle) {
  angle = fmod(angle, 2.0 * M_PI);
  if (angle < 0)
    angle += 2.0 * M_PI;
  return angle;
}

// Normalize an angle to be \in [-PI, PI)
double Normalize(double angle) {
  angle = fmod(angle + M_PI, 2.0 * M_PI);
  if (angle < 0)
    angle += 2.0 * M_PI;
  return angle - M_PI;
}

// Computes the shortest distance between two angles on S^1.
// Found by manipulating the first answer on:
// stackoverflow.com/questions/1878907/the-smallest-difference-between-2-angles
double S1Distance(double from, double to) {
  double d = Unroll(Unroll(to) - Unroll(from));
  if (d > M_PI)
    d -= 2.0*M_PI;
  return Normalize(d);
}

// Convert from degrees to radians.
double D2R(double angle) {
  return angle * M_PI / 180.0;
}

// Convert from radians to degrees.
double R2D(double angle) {
  return angle * 180.0 / M_PI;
}

// An error metric between two rotation matrices on SO3.
double SO3Error(const Matrix3d& R1, const Matrix3d& R2) {
  const Matrix3d R_error = R1.transpose()*R2 - R2.transpose()*R1;

  // 'vee' is the inverse of the hat operator, and extracts a vector from a
  // cross-product matrix.
  const Vector3d R_vee(R_error(2,1), R_error(0,2), R_error(1,0));
  return (0.5 * R_vee).norm();
}

}  //\namespace bsfm


//==================================================================================================
/*!
  @file

  @copyright 2016 NumScale SAS

  Distributed under the Boost Software License, Version 1.0.
  (See accompanying file LICENSE.md or copy at http://boost.org/LICENSE_1_0.txt)
*/
//==================================================================================================
#ifndef BOOST_SIMD_FUNCTION_GAMMALN_HPP_INCLUDED
#define BOOST_SIMD_FUNCTION_GAMMALN_HPP_INCLUDED

#if defined(DOXYGEN_ONLY)
namespace boost { namespace simd
{

 /*!

    @ingroup group-euler
    Function object implementing gammaln capabilities

    Natural logarithm of the absolute value of the Gamma function
     \f$\displaystyle \log |\Gamma(x)|\f$

    @par Semantic:

    For every parameter of floating type T

    @code
    T r = gammaln(x);
    @endcode

    is similar to:

    @code
    T r = log(abs(gamma(x))));
    @endcode

    @par Notes

    - The accuracy of the function is not uniformly good for negative entries
      The algorithm used is currently an adapted vesion of the cephes one.
      For better accuracy in the negative entry case one can use the extern
      boost_math gammaln functor but at a loss of speed.

      However, as stated in boost math:

      "While the relative errors near the positive roots of lgamma are very low,
       the  function has an infinite number of irrational roots for negative arguments:
       very close to these negative roots only a low absolute error can be guaranteed."

    - The call gammaln(x, sgn) also returns the sign of gamma in the output parameter sgn.

       Be aware that POSIX version of lgamma is not thread-safe: each execution of the function
       stores the sign of the gamma function of x in the static external variable signgam.

       boost.simd also provides @ref signgam which independantly computes the sign.

    @par Decorators

    std_ for floating entries  provides access to @c std::lgamma

    @see gamma, signgam

  **/
  Value gammaln(Value const & v0);
} }
#endif

#include <boost/simd/function/scalar/gammaln.hpp>
#include <boost/simd/function/simd/gammaln.hpp>

#endif


#include <Eigen/Dense>
#include <complex>
#include <iostream>

#include <fft/fft2.hpp>

using namespace std;

int main(int argc, char *argv[])
{
  typedef std::complex<double> cdouble;

  typedef Eigen::Array<cdouble, Eigen::Dynamic, Eigen::Dynamic, Eigen::RowMajor> complex_array_t;
  typedef Eigen::Array<double, Eigen::Dynamic, Eigen::Dynamic, Eigen::RowMajor> array_t;

  complex_array_t f_tilde_hat(4, 4);

  f_tilde_hat << 0, 0, 0, 0, 0, 1, 1, 1, 0, 1, 1, 1, 0, 1, 1, 1;
  array_t f;

  FFT fft;

  std::cout << "f_tilde_hat:\n" << f_tilde_hat << "\n";

  fft.ifft2(f, f_tilde_hat);
  std::cout << "f:\n" << f << "\n";

  return 0;
}


#include <OpenTissue/core/math/math_vector3.h>
#include <util_angle.h>

#define BOOST_AUTO_TEST_MAIN
#include <boost/test/auto_unit_test.hpp>
#include <boost/test/unit_test_suite.hpp>
#include <boost/test/floating_point_comparison.hpp>
#include <boost/test/test_tools.hpp>

typedef double                        T;
typedef OpenTissue::math::Vector3<T>  V;
typedef V::value_traits               VT;

BOOST_AUTO_TEST_SUITE(util_angle);

BOOST_AUTO_TEST_CASE(test1)
{

  V const v0 = V(0.0, 1.0, 0.0);
  V const v1 = V(0.0, 0.0, 0.0);
  V const v2 = V(1.0, 0.0, 0.0);

  T const angle = util::angle(v0,v1,v2);

  BOOST_CHECK_CLOSE(angle, VT::pi_half(), 0.01 );
}

BOOST_AUTO_TEST_CASE(test2)
{

  V const v0 = V(0.0, 1.0, 0.0);
  V const v1 = V(0.0, 0.0, 0.0);
  V const v2 = V(1.0, 1.0, 0.0);

  T const angle = util::angle(v0,v1,v2);

  BOOST_CHECK_CLOSE(angle, VT::pi_quarter(), 0.01 );
}

BOOST_AUTO_TEST_SUITE_END();


#include <Eigen/Core>
#include <iostream>
using namespace Eigen;
using namespace std;
void PolygonToEquations(const MatrixX2d& pts, MatrixX2d& ab, VectorXd& c) {
  // ax + by + c <= 0
  // assume polygon is convex

  Vector2d p0 = pts.row(0);

  for (int i=0; i < pts.rows(); ++i) {
    int i1 = (i+1) % pts.rows();
    double x0 = pts(i,0),
        y0 = pts(i,1),
        x1 = pts(i1,0),
        y1 = pts(i1,1);
    ab(i,0) = -(y1 - y0);
    ab(i,1) = x1 - x0;
    ab.row(i).normalize();
    c(i) = -ab.row(i).dot(pts.row(i));
  }

}

int main() {
  MatrixX2d m(4,2), ab(4,2);
  VectorXd c(4);
  m << 0,0,
      0,1,
      1,1,
      1,0;
  PolygonToEquations(m, ab, c);
  cout << "ab: " << ab << endl;
  cout << "c: " << c.transpose() << endl;
}


#include "gaussianprocessregressor.h"
#include <iostream>
#include <cmath>
#include <ctime>
#include <Eigen/LU>
#include "nloptutility.h"
#include "utility.h"

//#define VERBOSE
//#define NOISELESS

using Eigen::MatrixXd;
using Eigen::VectorXd;

namespace
{

const bool   useLogNormalPrior     = true;
const double a_prior_mu            = std::log(0.500);
const double a_prior_sigma_squared = 0.10;
#ifdef NOISELESS
const double b_fixed               = 1e-06;
#else
const double b_prior_mu            = std::log(0.001);
const double b_prior_sigma_squared = 0.10;
#endif
const double r_prior_mu            = std::log(0.500);
const double r_prior_sigma_squared = 0.10;

double calc_grad_a_prior(const double a)
{
    return (a_prior_mu - a_prior_sigma_squared - std::log(a)) / (a_prior_sigma_squared * a);
}

#ifndef NOISELESS
double calc_grad_b_prior(const double b)
{
    return (b_prior_mu - b_prior_sigma_squared - std::log(b)) / (b_prior_sigma_squared * b);
}
#endif

double calc_grad_r_i_prior(const Eigen::VectorXd &r, const int index)
{
    return (r_prior_mu - r_prior_sigma_squared - std::log(r(index))) / (r_prior_sigma_squared * r(index));
}

double calc_a_prior(const double a)
{
    return std::log(Utility::log_normal(a, a_prior_mu, a_prior_sigma_squared));
}

#ifndef NOISELESS
double calc_b_prior(const double b)
{
    return std::log(Utility::log_normal(b, b_prior_mu, b_prior_sigma_squared));
}
#endif

double calc_r_i_prior(const Eigen::VectorXd &r, const int index)
{
    return std::log(Utility::log_normal(r(index), r_prior_mu, r_prior_sigma_squared));
}

double calc_grad_a(const MatrixXd& X, const MatrixXd& C_inv, const VectorXd& y, const double a, const double b, const VectorXd& r)
{
    const MatrixXd C_grad_a = Regressor::calc_C_grad_a(X, a, b, r);
    const double term1 = + 0.5 * y.transpose() * C_inv * C_grad_a * C_inv * y;
    const double term2 = - 0.5 * (C_inv * C_grad_a).trace();
    return term1 + term2 + (useLogNormalPrior ? calc_grad_a_prior(a) : 0.0);
}

#ifndef NOISELESS
double calc_grad_b(const MatrixXd& X, const MatrixXd& C_inv, const VectorXd& y, const double a, const double b, const VectorXd& r)
{
    const MatrixXd C_grad_b = Regressor::calc_C_grad_b(X, a, b, r);
    const double term1 = + 0.5 * y.transpose() * C_inv * C_grad_b * C_inv * y;
    const double term2 = - 0.5 * (C_inv * C_grad_b).trace();
    return term1 + term2 + (useLogNormalPrior ? calc_grad_b_prior(b) : 0.0);
}
#endif

double calc_grad_r_i(const MatrixXd& X, const MatrixXd& C_inv, const VectorXd& y, const double a, const double b, const VectorXd& r, const int index)
{
    const MatrixXd C_grad_r_i = Regressor::calc_C_grad_r_i(X, a, b, r, index);
    const double term1 = + 0.5 * y.transpose() * C_inv * C_grad_r_i * C_inv * y;
    const double term2 = - 0.5 * (C_inv * C_grad_r_i).trace();
    return term1 + term2 + (useLogNormalPrior ? calc_grad_r_i_prior(r, index) : 0.0);
}

VectorXd calc_grad(const MatrixXd& X, const MatrixXd& C_inv, const VectorXd& y, const double a, const double b, const VectorXd& r)
{
    const unsigned D = X.rows();

    VectorXd grad(D + 2);
    grad(0) = calc_grad_a(X, C_inv, y, a, b, r);
#ifdef NOISELESS
    grad(1) = 0.0;
#else
    grad(1) = calc_grad_b(X, C_inv, y, a, b, r);
#endif

    for (unsigned i = 2; i < D + 2; ++ i)
    {
        const unsigned index = i - 2;
        grad(i) = calc_grad_r_i(X, C_inv, y, a, b, r, index);
    }

    return grad;
}

struct Data
{
    Data(const MatrixXd& X, const VectorXd& y) : X(X), y(y)
    {
    }

    const MatrixXd X;
    const VectorXd y;
};

// For counting the number of function evaluations
unsigned count;

// Log likelihood that will be maximized
double objective(const std::vector<double> &x, std::vector<double>& grad, void* data)
{
    // For counting the number of function evaluations
    ++ count;

    const MatrixXd& X = static_cast<const Data*>(data)->X;
    const VectorXd& y = static_cast<const Data*>(data)->y;

    const unsigned N = X.cols();

    const double   a = x[0];
#ifdef NOISELESS
    const double   b = b_fixed;
#else
    const double   b = x[1];
#endif
    const VectorXd r = [&x]() { std::vector<double> _x = x; return Eigen::Map<VectorXd>(&_x[2], _x.size() - 2); }();

    const MatrixXd C     = Regressor::calc_C(X, a, b, r);
    const MatrixXd C_inv = C.inverse();

    // When the algorithm is gradient-based, compute the gradient vector
    if (grad.size() == x.size())
    {
        const VectorXd g = calc_grad(X, C_inv, y, a, b, r);
        for (unsigned i = 0; i < g.rows(); ++ i) grad[i] = g(i);
    }

    const double term1 = - 0.5 * y.transpose() * C_inv * y;
    const double term2 = - 0.5 * std::log(C.determinant());
    const double term3 = - 0.5 * N * std::log(2.0 * M_PI);

    // Computing the regularization terms from a prior assumptions
    const double a_prior = calc_a_prior(a);
#ifdef NOISELESS
    const double b_prior = 1.0;
#else
    const double b_prior = calc_b_prior(b);
#endif
    const double r_prior = [&r]()
    {
        double sum = 0.0;
        for (unsigned i = 0; i < r.rows(); ++ i) sum += calc_r_i_prior(r, i);
        return sum;
    }();
    const double regularization = useLogNormalPrior ? (a_prior + b_prior + r_prior) : 0.0;

    return term1 + term2 + term3 + regularization;
}

}

GaussianProcessRegressor::GaussianProcessRegressor(const MatrixXd& X, const VectorXd& y)
{
    this->X = X;
    this->y = y;

    compute_MAP();

    C     = calc_C(X, a, b, r);
    C_inv = C.inverse();
}

GaussianProcessRegressor::GaussianProcessRegressor(const Eigen::MatrixXd &X, const Eigen::VectorXd &y, double a, double b, const Eigen::VectorXd &r)
{
    this->X = X;
    this->y = y;
    this->a = a;
    this->b = b;
    this->r = r;

    C     = calc_C(X, a, b, r);
    C_inv = C.inverse();
}

double GaussianProcessRegressor::estimate_y(const VectorXd &x) const
{
    const VectorXd k = calc_k(x, X, a, b, r);
    return k.transpose() * C_inv * y;
}

double GaussianProcessRegressor::estimate_s(const VectorXd &x) const
{
    const VectorXd k = calc_k(x, X, a, b, r);
    return std::sqrt(a + b - k.transpose() * C_inv * k);
}

void GaussianProcessRegressor::compute_MAP()
{
    const unsigned D = X.rows();

    Data data(X, y); // Constant during optimization

#ifdef VERBOSE
    const auto t1 = std::chrono::system_clock::now();
    count = 0;
#endif

    const VectorXd x_ini = VectorXd::Constant(D + 2, 1e+00);
    const VectorXd upper = VectorXd::Constant(D + 2, 5e+01);
    const VectorXd lower = VectorXd::Constant(D + 2, 1e-08);

    const VectorXd x_glo = nloptUtility::compute(x_ini, upper, lower, objective, &data, nlopt::GN_DIRECT, 300);

#ifdef VERBOSE
    const auto t2 = std::chrono::system_clock::now();
    const unsigned g_count = count;
    count = 0;
#endif

    const VectorXd x_loc = nloptUtility::compute(x_glo, upper, lower, objective, &data, nlopt::LD_TNEWTON, 1000);

    a = x_loc(0);
    b = x_loc(1);
    r = x_loc.block(2, 0, D, 1);

#ifdef NOISELESS
    b = b_fixed;
#endif

#ifdef VERBOSE
    const auto t3 = std::chrono::system_clock::now();
    const unsigned l_count = count;

    std::cout << "global optimization: " << std::chrono::duration_cast<std::chrono::milliseconds>(t2 - t1).count() << " ms, " << g_count << " evaluations, obj = " << objective(x_star_global, x_ini, &data) << std::endl;
    std::cout << "local optimization: " << std::chrono::duration_cast<std::chrono::milliseconds>(t3 - t2).count() << " ms, " << l_count << " evaluations, obj = " << objective(x_star_local, x_ini, &data) << std::endl;
    std::cout << "a = " << a << ", b = " << b << ", r = " << r.transpose() << std::endl;
#endif
}


#include <functional>
#include <iostream>
#include <sstream>
#include <stdlib.h>
#include <string.h>

#include <boost/random.hpp>
#include <boost/tuple/tuple.hpp>
#include <boost/tuple/tuple_io.hpp>

#include "biggles/detail/random.hpp"
#include "biggles/sampling/metropolis_hastings.hpp"

extern "C" {
#include <ccan/tap/tap.h>
}

typedef boost::tuples::tuple<float, float> point;

inline float rosenbrock(const point& p)
{
    float x, y;
    boost::tuples::tie(x,y) = p;
    return (1.f-x)*(1.f-x) + 100.f*(y-x*x)*(y-x*x);
}

// a cooked-up PDF which has a maximum where the Rosenbrock function has a minimum.
inline float rosenbrock_log_pdf(const point& p)
{
    return -rosenbrock(p);
}

typedef boost::tuples::tuple<point, float> proposal_result_t;

inline proposal_result_t propose_gaussian(const point& p)
{
    static boost::normal_distribution<float> norm(0.f, 0.2f);
    static boost::variate_generator<boost::mt19937&, boost::normal_distribution<float> > variate(
        biggles::detail::random_generator, norm);

    point new_p(p.get<0>() + variate(), p.get<1>() + variate());

    return proposal_result_t(new_p, 0.f);
}

using namespace biggles::sampling::metropolis_hastings;

int main(int argc, char** argv)
{
    biggles::detail::seed_prng(0xdeadbeef);

    bool verbose = (argc > 1) && (0 == strcmp(argv[1], "-v"));

    plan_tests(4);

    typedef std::pointer_to_unary_function<const point&, float> pdf_t;
    typedef std::pointer_to_unary_function<const point&, proposal_result_t> proposal_func_t;

    // MH sampler for Rosenbrock function. An example of using a function pointer directly.
    sampler<pdf_t, proposal_func_t> s(std::ptr_fun(rosenbrock_log_pdf), std::ptr_fun(propose_gaussian));

    point optimum(s.last_sample());
    float optimum_log_density(s.current_sample_log_density());

    for(int i=0; i<4096; ++i)
    {
        s.draw();

        if(s.current_sample_log_density() > optimum_log_density)
        {
            optimum_log_density = s.current_sample_log_density();
            optimum = s.last_sample();
        }

        if(verbose)
        {
            std::stringstream ss;
            ss << "sample: " << s.last_sample();
            diag("%s", ss.str().c_str());
        }
    }

    {
        std::stringstream ss;
        ss << "optimum: " << optimum << ", log-pdf: " << optimum_log_density
        << ", alpha: " << s.acceptance_rate() << ", rosenbrock: " << rosenbrock(optimum);
        diag("%s", ss.str().c_str());
    }

    diag("acceptance rate = %.2f%%", s.acceptance_rate()*100.f);
    //ok(fabs(s.acceptance_rate() - 0.25f) < 0.2f, "acceptance rate within 0.2 of 0.25");
    ok(fabs(optimum.get<0>() - 1.f) < 0.1f, "x-co-ordinate within 10%% of optimum");
    ok(fabs(optimum.get<1>() - 1.f) < 0.1f, "y-co-ordinate within 10%% of optimum");

    // the same thing using the optimise function
    point optimum2 = optimise(std::ptr_fun(rosenbrock_log_pdf), std::ptr_fun(propose_gaussian));

    {
        std::stringstream ss;
        ss << "optimum via optimise(): " << optimum2 << ", log-pdf: " << rosenbrock_log_pdf(optimum2)
        << ", rosenbrock: " << rosenbrock(optimum2);
        diag("%s", ss.str().c_str());
    }

    ok(fabs(optimum2.get<0>() - 1.f) < 0.1f, "x-co-ordinate within 10%% of optimum");
    ok(fabs(optimum2.get<1>() - 1.f) < 0.1f, "y-co-ordinate within 10%% of optimum");

    return exit_status();
}