#include "mylib/commen_include.h"
#include "mylib/g2o_opt_type.h"
#include "mylib/types_six_dof_expmap_ours.h"
#include "mylib/so3.h"
#include "mylib/QuatPlane.h"
#include "ceres/ceres.h"
#include "glog/logging.h"

//openCV
#include <opencv2/opencv.hpp>
#include <opencv2/core/eigen.hpp>

//g2o
#include "g2o/g2o/types/types_sba.h"
#include "g2o/g2o/types/types_six_dof_expmap.h"
#include "g2o/g2o/core/optimization_algorithm_levenberg.h"
#include "g2o/g2o/core/sparse_optimizer.h"
#include "g2o/g2o/core/block_solver.h"
//#include "g2o/g2o/solvers/linear_solver_eigen.h"
#include "g2o/g2o/solvers/linear_solver_dense.h"


/*
using ceres::AutoDiffCostFunction;
using ceres::CostFunction;
using ceres::Problem;
using ceres::Solve;
using ceres::Solver;
// A templated cost functor that implements the residual r = 10 -
// x. The method operator() is templated so that we can then use an
// automatic differentiation wrapper around it to generate its
// derivatives.
struct CostFunctor {
    template <typename T>
    bool operator()(const T* const x, T* residual) const {
        residual[0] = 10.0 - x[0];
        return true;
    }
};*/

struct sFrame
{
    EIGEN_MAKE_ALIGNED_OPERATOR_NEW;
    sFrame(Eigen::Matrix4d Tcw): Tcw_(Tcw){}

    Eigen::Matrix4d Tcw_;

    std::unordered_map<int, Eigen::Vector3d> per_plane_feature_;


};

double fx = 1;//750;
double fy = 1;//500;
double cx = 0;//250;
double cy = 0;//250;

const int nCameras = 5;
const int nMapplanes = 2;  // now plane just 3

std::vector<std::vector<Eigen::Vector3d>> readPlaneFile();
Eigen::Vector4d CreatePlanebyPoints(Eigen::Vector4d pt1_,Eigen::Vector4d pt2_,Eigen::Vector4d pt3_,Eigen::Vector4d pt4_);
void CreateCameras(std::vector<sFrame>& vFrames, std::vector<Sophus::SE3d>& gtpose, int nCams);
void CreateSceneFeatures(const std::vector<sFrame>& vFrames, std::vector<Eigen::Vector3d>& vpGeometryFeatures);
void FeatureDetection(const Eigen::Matrix3d& K, const std::vector<Eigen::Vector3d> &scenefeatures, std::vector<sFrame> &vframes, bool is_add_noise);
void AddNoise(std::vector<sFrame>& vFrames, std::vector<Eigen::Vector3d>& vpGeometryFeatures, double gs_noise, double cam_noise);
//int cerestest(int argc, char** argv);

Eigen::Quaterniond ExpUpDateQuat(const Eigen::Vector3d &omg,const Eigen::Quaterniond q);
Eigen::Vector3d logQuat(Eigen::Quaterniond q);
void SolveBA(std::vector<sFrame> &vNoisyCameras,std::vector<Eigen::Vector3d>& vpGeometryFeatures);

int main(int argc, char** argv) {
    //cerestest(argc, argv);
/********************** prepare data for opt part**********************************************************************/
    std::vector<sFrame> vGTCameras;
    //std::vector<Eigen::Vector3d> vGTGeomeryFeatures; instead -->PlaneFeature
    std::vector<Sophus::SE3d> gt_cam_pose;
    std::vector<Eigen::Vector3d> PlaneFeature;
    CreateCameras(vGTCameras, gt_cam_pose, nCameras);
    CreateSceneFeatures(vGTCameras, PlaneFeature);
    cv::Mat cv_K = (cv::Mat_<double>(3,3) << 1, 0,  1, 0, 1, 1,0,  0,  1);
    Eigen::Matrix3d eigen_K;
    cv::cv2eigen(cv_K, eigen_K);

    FeatureDetection(eigen_K,PlaneFeature, vGTCameras, false);

    std::vector<Eigen::Vector3d> vNoisyGeomeryPlaneFeatures = PlaneFeature;
    std::vector<sFrame> vNoisyCameras = vGTCameras;
    double cam_noise = 0.0;
    double gs_noise = 0.0;
    AddNoise(vNoisyCameras, vNoisyGeomeryPlaneFeatures, gs_noise, cam_noise);
/************************* error of structure before opt **************************************************************/
    double error_structure = 0.0;
    for(int i = 0; i < vNoisyGeomeryPlaneFeatures.size(); i++)
    {
        Eigen::Vector3d gs = vNoisyGeomeryPlaneFeatures[i];
        Eigen::Vector3d gt_plane = PlaneFeature[i];
        Eigen::Quaterniond  q(1,0,0,0);
        Eigen::Quaterniond  structure_err = ExpUpDateQuat(gs,q)*ExpUpDateQuat(gt_plane,q).inverse();
        error_structure += logQuat(structure_err).norm();
    }

    std::cout << "RMSE for the structure is:                " << error_structure / (double)nMapplanes << std::endl;

/************************* error of camera before opt *****************************************************************/
    double error_cam_rotation = 0.0;
    double error_cam_translation = 0.0;
    for(int i = 0; i < vNoisyCameras.size(); i++)
    {
        Eigen::Matrix4d Tcw = vNoisyCameras[i].Tcw_;
        Eigen::Matrix3d Rcw = Tcw.block(0,0,3,3);
        Eigen::Vector3d tcw = Tcw.block(0,3,3,1);

        Sophus::SE3d noise_cam = Sophus::SE3d(Rcw, tcw);
        Sophus::SE3d gt_cam = gt_cam_pose[i];

        Sophus::SE3d pose_err = gt_cam.inverse() * noise_cam;
        error_cam_rotation += pose_err.so3().log().norm();
        error_cam_translation += pose_err.translation().norm();
    }
    std::cout << "RMSE for the rotation of camera is:       " << error_cam_rotation / (double)nCameras << std::endl;
    std::cout << "RMSE for the translation of camera is:    " << error_cam_translation / (double)nCameras << std::endl;

/************************** opt part using g2o ************************************************************************/
    SolveBA(vNoisyCameras,vNoisyGeomeryPlaneFeatures);

    return 0;


}
const double SMALL_EPS = 1e-5;
Eigen::Quaterniond ExpUpDateQuat(const Eigen::Vector3d &omg,const Eigen::Quaterniond q){
    Eigen::Quaterniond res;
    double theta,half_theta;
    theta = omg.norm();//M_PI / 6;
    half_theta = 0.5*theta;
    EIGEN_USING_STD(sin)
    EIGEN_USING_STD(cos)
    // if(theta<SMALL_EPS) TODO
    if(theta<SMALL_EPS){
        res.vec() = omg;
        res.w() = 1;
    }else{
        res.vec() = sin(half_theta)/theta*omg;
        res.w() = cos(half_theta);
    }

    return res;
}
// inverse of the exponential map:log map
Eigen::Vector3d logQuat(Eigen::Quaterniond q){
/*    Eigen::Vector3d w;
    Eigen::Vector3d q_vec = q.vec();
    double q_norm = q.vec().norm();
    double q_w = q.w();
    if (q_norm<SMALL_EPS){
        w = q_vec;
    }else {
        w = 2 * acos(q_w) / q_norm * q_vec;
    }
    return w;*/
    q.normalize();
    Eigen::Vector3d w;
    double q_norm = q.vec().norm();
    Eigen::Vector3d q_vec = q.vec();
    double q_w = q.w();
    if (q_norm < SMALL_EPS) {
        w = q_vec;
    } else {
        double q_n_3 = q_norm * q_norm * q_norm;
        double q_w_2 = q_w * q_w;
        w = q_vec / q_norm * (1 - 1 / 3 * q_n_3 / q_w_2);
    }
    return w;
}
/***********************************************************************************************************************/

std::vector<std::vector<Eigen::Vector3d>> readPlaneFile(){
    std::vector<std::vector<Eigen::Vector3d>> planesPoints;

    std::ifstream planesFile("../res/planes.txt");

    if(!planesFile.is_open()){
        assert("Error - planes file not open");//throw is not from std
    }

    while(planesFile.good()){
        std::vector<Eigen::Vector3d> curPlane;
        int id;
        planesFile >> id;
        for(int v = 0; v < 4; ++v){ //four points to describe a plane
            Eigen::Vector3d curPoint;
            for(int c = 0; c < 3; ++c){
                double val;
                planesFile >> val;
                if(planesFile.good()){
                    curPoint[c] = val;
                }
            }
            if(planesFile.good()){
                //std::cout << "curPoint = " << curPoint.transpose() << std::endl;
                curPlane.push_back(curPoint);
            }
        }
        if(planesFile.good()){
//            std::cout << "curPlane.size() = " << curPlane.size() << std::endl;
            planesPoints.push_back(curPlane);
        }
    }
//    std::cout << "planesPoints.size() = " << planesPoints.size() << std::endl;
    return planesPoints;
}

Eigen::Vector4d CreatePlanebyPoints(Eigen::Vector4d pt1_,Eigen::Vector4d pt2_,Eigen::Vector4d pt3_,Eigen::Vector4d pt4_)
{
    Eigen::Vector3d pt1=pt1_.head(3);Eigen::Vector3d pt2=pt2_.head(3);Eigen::Vector3d pt3=pt3_.head(3);Eigen::Vector3d pt4=pt4_.head(3);
    Eigen::Vector4d plane_homo;
    plane_homo.head(3) =  (pt1-pt3).cross((pt2-pt3));
    plane_homo(3) = -pt3.transpose()*(pt1.cross(pt2));
    Eigen::Vector3d normal;double dis;
    normal =plane_homo.head(3)/plane_homo.head(3).norm();//.normalized();
    dis = -plane_homo(3)/plane_homo.head(3).norm();
    Eigen::Vector4d HessePlane;
    HessePlane.head(3) = normal;
    HessePlane(3) = dis;
    return HessePlane;
}


void CreateCameras(std::vector<sFrame>& vFrames, std::vector<Sophus::SE3d>& gtpose, int nCams) {

        std::ifstream trajFile("../res/groundtruth.txt");

        if (!trajFile.is_open()) {
            assert("Error - groundtruth file not open");
        }

        std::vector<Eigen::Matrix<double, 7, 1>> gtPoses_7;
        while (trajFile.good()) {
            Eigen::Matrix<double, 7, 1> curPose;
            int id;
            trajFile >> id;
            for (int v = 0; v < 7; ++v) {
                double val;
                trajFile >> val;
                if (trajFile.good()) {
                    curPose[v] = val;
                }
            }
            if (trajFile.good()) {
                //std::cout << "curPose = " << curPose.transpose() << std::endl;
                gtPoses_7.push_back(curPose);
            }
        }
        for (int i=0;i<nCams;++i){
            Eigen::Matrix<double, 7, 1> camPose;
            camPose = gtPoses_7[i];
            Eigen::Vector3d t = camPose.head(3);
            Eigen::Quaterniond q;
            q.x() = camPose(3);q.y() = camPose(4);q.z() = camPose(5);q.w() = camPose(6);
            q.normalize();
            Eigen::Matrix3d R(q);
            Eigen::Matrix4d T;
            T<<1,0,0,0, 0,1,0,0, 0,0,1,0,0,0,0,1;
            T.block(0,0,3,3) = R;
            T.block(0,3,3,1) = t;
            sFrame frame(T);
            vFrames.push_back(frame);
            Sophus::SE3d SE3_qt_cam(q, t);
            gtpose.push_back(SE3_qt_cam);
        }

}

void CreateSceneFeatures(const std::vector<sFrame>& vFrames, std::vector<Eigen::Vector3d>& PlaneFeature)
{/*
    std::mt19937 gen(std::chrono::system_clock::now().time_since_epoch().count());

    std::uniform_real_distribution<double> d_x{0.0, 4.0};
    std::uniform_real_distribution<double> d_y{0.0, 5.0};
    std::uniform_real_distribution<double> d_z{0.0, 5.0};
    std::uniform_real_distribution<double> d_h{0.1, 5.0};
*/
    int cnt = 0;
    Eigen::Vector4d pt1_,pt2_,pt3_,pt4_;
    std::vector<Eigen::Vector3d> curPlane;
    std::vector<std::vector<Eigen::Vector3d>> planesPoints = readPlaneFile();
    //std::vector<Eigen::Vector3d> PlaneFeature;  //3 planes
    //std::vector<std::vector<Eigen::Vector3d>>::iterator ele;
    for (int i = 0;i<planesPoints.size();++i) {
        curPlane = planesPoints[i];
        pt1_.head(3) = curPlane[0];pt1_(3) = 1;
        pt2_.head(3) = curPlane[1];pt2_(3) = 1;
        pt3_.head(3) = curPlane[2];pt3_(3) = 1;
        pt4_.head(3) = curPlane[3];pt4_(3) = 1;
        Eigen::Vector4d curPlane_hesse = CreatePlanebyPoints(pt1_,pt2_,pt3_,pt4_);
        // std::cout <<"curPlane_hesse is "<<curPlane_hesse.transpose()<<std::endl;
        // PlaneFeature.push_back(curPlane_hesse.head(3)*curPlane_hesse(3));
        Eigen::Vector4d homo_para_w;
        double n_norm = 1;
        Eigen::Vector3d norm = curPlane_hesse.head(3);
        double dis = curPlane_hesse(3);
        norm.normalized();
        homo_para_w[0] = norm[0] * n_norm;
        homo_para_w[1] = norm[1] * n_norm;
        homo_para_w[2] = norm[2] * n_norm;
        homo_para_w[3] = -1 * dis * n_norm;
        int n_obervation_times = 0;
        //PlaneFeature.push_back(curPlane_hesse.head(3)*curPlane_hesse(3));}

            for(int j = 0; j < nCameras; j++)
            {
                //Eigen::Matrix3d Rcw = vFrames[j].Tcw_.block(0, 0, 3, 3);
                //Eigen::Vector3d tcw = vFrames[j].Tcw_.block(0, 3, 3, 1);
                // Eigen::Matrix3d Rcw = Rwc.transpose();
                // Eigen::Vector3d tcw = -Rcw * twc;

                Eigen::Matrix4d camera_T_world = vFrames[j].Tcw_;
                Eigen::Matrix4d transM =camera_T_world.transpose().inverse();
                Eigen::Vector4d homo_para_c = transM* homo_para_w;

                double d = -homo_para_c[3] / std::sqrt(homo_para_c[0] * homo_para_c[0] + homo_para_c[1] * homo_para_c[1] + homo_para_c[2] * homo_para_c[2]);
                if( d < 0)
                    continue;
                n_obervation_times ++;
            }

            if(n_obervation_times == nCameras)
            {
                PlaneFeature.push_back(curPlane_hesse.head(3)*curPlane_hesse(3));
                cnt ++;
            }

            if(cnt == nMapplanes)
                break;
    }
}


void FeatureDetection(const Eigen::Matrix3d& K, const std::vector<Eigen::Vector3d> &scenefeatures, std::vector<sFrame> &vframes, bool is_add_noise)
{
    // detect point feature
    std::mt19937 gen{12345};
    const double pixel_sigma = 0.0;
    std::normal_distribution<> d{0.0, pixel_sigma};

    int planeID = 0;
    for(auto mIter = scenefeatures.begin(); mIter != scenefeatures.end(); ++mIter)
    {
        Eigen::Vector3d Pw = (*mIter);
        double dis = Pw.norm();
        Eigen::Vector3d normal = Pw/dis;
        Eigen::Vector4d homo_para_w;
        double n_norm = 1;
        normal.normalized();
        homo_para_w[0] = normal[0] * n_norm;
        homo_para_w[1] = normal[1] * n_norm;
        homo_para_w[2] = normal[2] * n_norm;
        homo_para_w[3] = -1 * dis * n_norm;

        for(auto fIter = vframes.begin(); fIter != vframes.end(); ++fIter)
        {
            Eigen::Matrix3d Rcw = (*fIter).Tcw_.block(0,0,3,3);
            Eigen::Vector3d tcw = (*fIter).Tcw_.block(0,3,3,1);

            Eigen::Vector4d homo_para_c = (*fIter).Tcw_.transpose().inverse()*homo_para_w;
            //QuatPlane::HomogeneousToHesse(const Eigen::Vector4d &homo_para, Eigen::Vector3d &norm, double &dis)
            Eigen::Vector3d normh;
            double dish;
            QuatPlane::HomogeneousToHesse(homo_para_c,normh,dish);
            Eigen::Vector3d ob = normh*dish;

            (*fIter).per_plane_feature_.insert(std::make_pair(planeID, ob));


        }
        planeID++;
    }


}


void AddNoise(std::vector<sFrame>& vFrames, std::vector<Eigen::Vector3d>& vpGeometryFeatures, double gs_noise , double cam_noise )
{
    cv::RNG rng ( cv::getTickCount() );

    // add noise for structure
    for(auto fIter = vpGeometryFeatures.begin(); fIter != vpGeometryFeatures.end(); ++fIter)
    {

        double noise_x = rng.gaussian ( gs_noise );
        double noise_y = rng.gaussian ( gs_noise );
        double noise_z = rng.gaussian ( gs_noise );

        Eigen::Vector3d error = Eigen::Vector3d(noise_x, noise_y, noise_z);

        Eigen::Vector3d plane_3d = (*fIter);
        plane_3d += error;
        (*fIter) = plane_3d;

    }

    std::normal_distribution<> noise_cam{0.0, cam_noise};
    Eigen::Matrix3d Rx, Ry, Rz;
    Eigen::Matrix3d R;
    Eigen::Vector3d t;
    Eigen::Matrix4d T;
    for(auto cIter = vFrames.begin(); cIter != vFrames.end(); ++cIter)
    {
        // skip the first camera.
        if(cIter == vFrames.begin())
            continue;

        double tz = rng.gaussian ( cam_noise );
        double ty = rng.gaussian ( cam_noise );
        double tx = rng.gaussian ( cam_noise );

        Rz << cos ( tz ), -sin ( tz ), 0.0,
                sin ( tz ), cos ( tz ), 0.0,
                0.0, 0.0, 1.0;
        Ry << cos ( ty ), 0.0, sin ( ty ),
                0.0, 1.0, 0.0,
                -sin ( ty ), 0.0, cos ( ty );
        Rx << 1.0, 0.0, 0.0,
                0.0, cos ( tx ), -sin ( tx ),
                0.0, sin ( tx ), cos ( tx );
        R = Rz * Ry * Rx;

        // translation.
        double x = rng.gaussian ( cam_noise );
        double y = rng.gaussian ( cam_noise );
        double z = rng.gaussian ( cam_noise );
        t << x, y, z;

        // SE3
        Sophus::SE3d cam_noise ( R, t );

        Eigen::Matrix3d Rcw = (*cIter).Tcw_.block(0,0,3,3);
        Eigen::Vector3d tcw = (*cIter).Tcw_.block(0,3,3,1);
        Sophus::SE3d cam_temp(Rcw, tcw);
        Sophus::SE3d newcamera = cam_temp * cam_noise;
        (*cIter).Tcw_.block(0,0,3,3) = newcamera.rotationMatrix();
        (*cIter).Tcw_.block(0,3,3,1) = newcamera.translation();

    }
}
/* using SO3
void SolveBA(std::vector<sFrame> &vNoisyCameras,std::vector<Eigen::Vector3d>& vpNoiseGeometryFeatures){
    g2o::SparseOptimizer optimizer;
    g2o::BlockSolverX::LinearSolverType *linearSolver;
    linearSolver = new g2o::LinearSolverDense<g2o::BlockSolverX::PoseMatrixType>();
    g2o::BlockSolverX *solver_ptr = new g2o::BlockSolverX(linearSolver);
    g2o::OptimizationAlgorithmLevenberg *solver = new g2o::OptimizationAlgorithmLevenberg(solver_ptr);
    optimizer.setAlgorithm(solver);
    // add camera vertex
    for (int i=0;i<nCameras;i++){
        g2o::VertexSE3Expmap *v = new g2o::VertexSE3Expmap();
        Eigen::Matrix4d T_cw = vNoisyCameras[i].Tcw_;
        Eigen::Matrix3d R_cw = T_cw.block(0,0,3,3);
        Eigen::Vector3d t_cw = T_cw.block(0,3,3,1);
        v ->setId(i);
        v ->setFixed(i==0);
        g2o::SE3Quat est = g2o::SE3Quat(R_cw,t_cw);
        v ->setEstimate(est);
        optimizer.addVertex(v);
    }
    // add plane feature vertex
    for (int i = 0;i < vpNoiseGeometryFeatures.size();i++){
        g2o::VertexSO3Expmap *vp = new g2o::VertexSO3Expmap();
        vp ->setId(i+nCameras);
        vp ->setFixed(false);
        vp ->setEstimate(g2o::SO3::exp(vpNoiseGeometryFeatures[i]));
        optimizer.addVertex(vp);
        for(int j = 0; j < vNoisyCameras.size(); j++)
        {
            g2o::EdgeSE3ProjectPlaneSO3* e = new g2o::EdgeSE3ProjectPlaneSO3();

            e->setVertex(0, dynamic_cast<g2o::OptimizableGraph::Vertex*> (optimizer.vertex(i + nCameras)));
            e->setVertex(1,dynamic_cast<g2o::OptimizableGraph::Vertex*> (optimizer.vertex(j)));
            e->setParams(g2o::SO3::exp(vNoisyCameras[j].per_plane_feature_.find(i)->second));
            e->setInformation(Eigen::Matrix3d::Identity());

            optimizer.addEdge(e);
        }
    }

    optimizer.setVerbose(true);
    optimizer.initializeOptimization();
    optimizer.optimize(10);
}
*/


//using own node edge #include "mylib/g2o_opt_type.h"   VertexSBAPlaneomg and EdgeSE3ProjectPlane
void SolveBA(std::vector<sFrame> &vNoisyCameras,std::vector<Eigen::Vector3d>& vpNoiseGeometryFeatures){
    g2o::SparseOptimizer optimizer;
    g2o::BlockSolverX::LinearSolverType *linearSolver;
    linearSolver = new g2o::LinearSolverDense<g2o::BlockSolverX::PoseMatrixType>();
    g2o::BlockSolverX *solver_ptr = new g2o::BlockSolverX(linearSolver);
    g2o::OptimizationAlgorithmLevenberg *solver = new g2o::OptimizationAlgorithmLevenberg(solver_ptr);
    optimizer.setAlgorithm(solver);
    // add camera vertex
    for (int i=0;i<nCameras;i++){
        g2o::VertexSE3Expmap *v = new g2o::VertexSE3Expmap();
        Eigen::Matrix4d T_cw = vNoisyCameras[i].Tcw_;
        Eigen::Matrix3d R_cw = T_cw.block(0,0,3,3);
        Eigen::Vector3d t_cw = T_cw.block(0,3,3,1);
        v ->setId(i);
        v ->setFixed(i==0);
        g2o::SE3Quat est = g2o::SE3Quat(R_cw,t_cw);
        v ->setEstimate(est);
        optimizer.addVertex(v);
    }
    // add plane feature vertex
    for (int i = 0;i < vpNoiseGeometryFeatures.size();i++){
        g2o::VertexSBAPlaneomg *vp = new g2o::VertexSBAPlaneomg();
        vp ->setId(i+nCameras);
        vp ->setFixed(false);
        vp ->setEstimate(vpNoiseGeometryFeatures[i]);
        optimizer.addVertex(vp);
        for(int j = 0; j < vNoisyCameras.size(); j++)
        {
            g2o::EdgeSE3ProjectPlane* e = new g2o::EdgeSE3ProjectPlane();

            e->setVertex(0, dynamic_cast<g2o::OptimizableGraph::Vertex*> (optimizer.vertex(i + nCameras)));
            e->setVertex(1,dynamic_cast<g2o::OptimizableGraph::Vertex*> (optimizer.vertex(j)));
            e->setMeasurement(vNoisyCameras[j].per_plane_feature_.find(i)->second);
            e->setInformation(Eigen::Matrix3d::Identity());

            optimizer.addEdge(e);
        }
    }

    optimizer.setVerbose(true);
    optimizer.initializeOptimization();
    optimizer.optimize(10);
}
