#include <seq2map/geometry_problems.hpp>
#include <boost/thread.hpp>
#include <random>

using namespace seq2map;

//==[ AlignmentObjective ]====================================================//

//==[ AlignmentObjective::InlierSelector ]====================================//

bool AlignmentObjective::InlierSelector::operator() (const EuclideanTransform& x, IndexList& inliers, IndexList& outliers) const
{
    inliers.clear();
    outliers.clear();

    //try
    //{
        metre.Start();
        cv::Mat error = (*objective)(x);
        metre.Stop(error.rows);

        if (error.type() != CV_64F)
        {
            error.convertTo(error, CV_64F);
        }

        const double* e = error.ptr<double>();

        for (size_t i = 0; i < error.total(); i++)
        {
            if (!std::isfinite(e[i]))
            {
                E_ERROR << "entry " << i << " is not a finite number";
                return false;
            }

            if (e[i] > threshold)
            {
                outliers.push_back(i);
            }
            else
            {
                inliers.push_back(i);
            }
        }
    //}
    //catch (std::exception& ex)
    //{
    //    E_ERROR << "error selecting inliers : " << ex.what();
    //    return false;
    //}

    return true;
}

bool AlignmentObjective::InlierSelector::operator() (const EuclideanTransform& x, IndexList& inliers) const
{
    IndexList outliers;
    return (*this)(x, inliers, outliers);
}

//==[ EpipolarObjective ]=====================================================//

bool EpipolarObjective::SetData(const GeometricMapping& data, ProjectionModel::ConstOwn& src, ProjectionModel::ConstOwn& dst)
{
    if (!data.IsConsistent()) return false;

    const size_t d0 = data.src.GetDimension();
    const size_t d1 = data.dst.GetDimension();

    if (d0 != 2 && d0 != 3)
    {
        E_ERROR << "source geometry has to be either 2D Euclidean or 3D homogeneous (d=" << d0 << ")";
        return false;
    }

    if (d1 != 2 && d1 != 3)
    {
        E_ERROR << "target geometry has to be either 2D Euclidean or 3D homogeneous (d=" << d1 << ")";
        return false;
    }

    if (!src)
    {
        E_ERROR << "missing source projection model";
        return false;
    }

    if (!dst)
    {
        E_ERROR << "missing target projection model";
        return false;
    }

    m_data.indices = data.indices;
    m_data.src = src->Backproject(data.src); //.Reshape(Geometry::ROW_MAJOR);
    m_data.dst = dst->Backproject(data.dst); //.Reshape(Geometry::ROW_MAJOR);
    m_data.metric = data.metric ? data.metric : Metric::Own(new EuclideanMetric());

    return true;
}

AlignmentObjective::Own EpipolarObjective::GetSubObjective(const IndexList& indices) const
{
    EpipolarObjective* sub = new EpipolarObjective(m_src, m_dst, m_distType);
    sub->m_data = m_data[indices]; // avoid calling SetData() as the normalised data should not be normalised again

    return AlignmentObjective::Own(sub);
}

cv::Mat EpipolarObjective::operator() (const EuclideanTransform& tform) const
{
    const Metric& d = *m_data.metric;

    cv::Mat x0 = m_data.src.mat;
    cv::Mat x1 = m_data.dst.mat;

    const EuclideanTransform M = M0.GetInverse() >> tform >> M1;
    const cv::Mat F = M.ToEssentialMatrix();

    const cv::Mat Fx0 = x0 * F.t();
    const cv::Mat Fx00 = Fx0.col(0);
    const cv::Mat Fx01 = Fx0.col(1);
    const cv::Mat Fx02 = Fx0.col(2);

    const cv::Mat xFx = Fx00.mul(x1.col(0)) + Fx01.mul(x1.col(1)) + Fx02;
    //            xFx = Fx00.mul(x1.col(0)) + Fx01.mul(x1.col(1)) + Fx02.mul(x1.col(2));
    //                                                         should be one ^^^^^^^^^

    if (m_distType == ALGEBRAIC)
    {
        // algebraic epipolar distance
        // x1' * F * x0
        Geometry err(m_data.src.shape, xFx);
        return d(err).mat;
    }

    cv::Mat Fx1 = x1 * F;
    cv::Mat Fx10 = Fx1.col(0);
    cv::Mat Fx11 = Fx1.col(1);
    //cv::Mat Fx12 = Fx1.col(2);

    cv::Mat nn0 = Fx00.mul(Fx00) + Fx01.mul(Fx01);
    cv::Mat nn1 = Fx10.mul(Fx10) + Fx11.mul(Fx11);

    Geometry err(m_data.src.shape);

    if (m_distType == SAMPSON)
    {
        // Sampson approximation:
        //
        //              x1' * F * x0
        // ---------------------------------------
        // sqrt(Fx00^2 + Fx01^2 + Fx10^2 + Fx11^2)

        cv::Mat nn; cv::sqrt(nn0 + nn1, nn);
        cv::divide(xFx, nn, err.mat);
    }
    else
    {
        // symmetric geometric error
        //
        //        x1' * F * x0              x1' * F * x0
        // ( ---------------------- , ---------------------- )
        //    sqrt(Fx00^2 + Fx01^2)    sqrt(Fx10^2 + Fx11^2)

        err.mat = cv::Mat(xFx.rows, 2, xFx.type());
        cv::sqrt(nn0, nn0);
        cv::sqrt(nn1, nn1);

        cv::divide(xFx, nn0, err.mat.col(0));
        cv::divide(xFx, nn1, err.mat.col(1));
    }

    return d(err).mat;
}

//==[ ProjectionObjective ]===================================================//

AlignmentObjective::Own ProjectionObjective::GetSubObjective(const IndexList& indices) const
{
    ProjectionObjective* sub = new ProjectionObjective(m_proj, m_forward);
    sub->m_data = m_data[indices];

    return AlignmentObjective::Own(sub);
}

bool ProjectionObjective::SetData(const GeometricMapping& data)
{
    if (!data.IsConsistent()) return false;

    const size_t d0 = data.src.GetDimension();
    const size_t d1 = data.dst.GetDimension();

    if (d0 != 3 && d0 != 4)
    {
        E_ERROR << "source geometry has to be either 3D Euclidean or 4D homogeneous (d=" << d0 << ")";
        return false;
    }

    if (d1 != 2 && d1 != 3)
    {
        E_ERROR << "target geometry has to be either 2D Euclidean or 3D homogeneous (d=" << d1 << ")";
        return false;
    }

    m_data.indices = data.indices;
    m_data.src = (d0 == 3 ? Geometry::MakeHomogeneous(data.src) : data.src); //.Reshape(Geometry::ROW_MAJOR);
    m_data.dst = (d1 == 3 ? Geometry::FromHomogeneous(data.dst) : data.dst); //.Reshape(Geometry::ROW_MAJOR);
    m_data.metric = data.metric ? data.metric : Metric::Own(new EuclideanMetric());

    return true;
}

cv::Mat ProjectionObjective::operator() (const EuclideanTransform& f) const
{
    if (!m_proj)
    {
        E_ERROR << "missing projection model";
        return cv::Mat();
    }

    const EuclideanTransform tf = m_forward ? f : f.GetInverse();

    Geometry x = m_data.src;

    // m_metres.proj.Start();
    Geometry y = m_proj->Project(tf(x, true), ProjectionModel::EUCLIDEAN_2D);
    // m_metres.proj.Stop(x.GetElements());

    // m_metres.jaco.Start();
    Geometry jac = m_proj->GetJacobian(x, y);
    // m_metres.jaco.Stop(x.GetElements());

    // m_metres.tfrm.Start();
    Metric::ConstOwn d = m_data.metric->Transform(tf, jac);
    // m_metres.tfrm.Stop(jac.GetElements());

    // m_metres.eval.Start();
    cv::Mat e = (*d)(y, m_data.dst).mat;
    // m_metres.eval.Stop(y.GetElements());

    // E_TRACE << m_metres.proj.GetElapsedSeconds();
    // E_TRACE << m_metres.jaco.GetElapsedSeconds();
    // E_TRACE << m_metres.tfrm.GetElapsedSeconds();
    // E_TRACE << m_metres.eval.GetElapsedSeconds();

    return e;
}

//==[ PhotometricObjective ]==================================================//

AlignmentObjective::Own PhotometricObjective::GetSubObjective(const IndexList& indices) const
{
    PhotometricObjective* sub = new PhotometricObjective(m_proj, m_dst, m_type, m_interp);
    sub->m_data = m_data[indices];

    return AlignmentObjective::Own(sub);
}

bool PhotometricObjective::SetData(const GeometricMapping& data)
{
    if (!data.IsConsistent()) return false;

    const size_t d0 = data.src.GetDimension();
    const size_t d1 = data.dst.GetDimension();
    const size_t dI = static_cast<size_t>(m_dst.channels());

    if (d0 != 3 && d0 != 4)
    {
        E_ERROR << "source geometry has to be either 3D Euclidean or 4D homogeneous";
        return false;
    }

    if (d1 != dI)
    {
        E_ERROR << "dimensionality mismatches between the referenced image (d=" << dI << ") and the target (d=" << d1 << ")";
        return false;
    }

    m_data.indices = data.indices;
    m_data.src = (d0 == 3 ? Geometry::MakeHomogeneous(data.src) : data.src);
    m_data.dst = data.dst;
    m_data.metric = data.metric ? data.metric : Metric::Own(new EuclideanMetric());

    if (m_data.dst.mat.depth() != m_type)
    {
        m_data.dst.mat.convertTo(m_data.dst.mat, m_type);
    }

    return true;
}

bool PhotometricObjective::SetData(const Geometry& g, const Geometry& p, const cv::Mat& src, bool snapped, const Metric::Own metric, const Indices& indices)
{
    if (g.GetElements() != p.GetElements())
    {
        E_ERROR << "inconsistent number of elements, given " << g.GetElements() << " 3D point(s) and " << p.GetElements() << " 2D point(s)";
        return false;
    }

    GeometricMapping data;
    cv::Mat src32F = src;

    if (src.depth() != m_type)
    {
        src.convertTo(src32F, m_type);
    }

    data.src = g;
    data.dst = Geometry(Geometry::PACKED, interp(src32F, p.mat, snapped ? cv::INTER_NEAREST : m_interp));
    data.metric = metric;
    data.indices = indices;

    return SetData(data);
}

bool PhotometricObjective::SetData(const Geometry& g, const cv::Mat& src, const Metric::Own metric, const Indices& indices)
{
    if (g.shape == Geometry::PACKED && g.mat.rows == src.rows && g.mat.cols == src.cols)
    {
        Points3D xyz;
        std::vector<cv::Point_<short>> sub;

        for (int i = 0; i < g.mat.rows; i++)
        {
            for (int j = 0; j < g.mat.cols; j++)
            {
                Point3D pt = g.mat.at<Point3D>(i, j);
                if (pt.z > 0)
                {
                    xyz.push_back(pt);
                    sub.push_back(cv::Point_<short>(j, i));
                }
            }
        }

        return SetData(
            Geometry(Geometry::PACKED, cv::Mat(xyz, false)),
            Geometry(Geometry::PACKED, cv::Mat(sub, false)),
            src, true, metric, indices);
    }

    if (!m_proj)
    {
        E_ERROR << "missing projection model";
        return false;
    }

    Geometry p = m_proj->Project(g, ProjectionModel::EUCLIDEAN_2D).Reshape(Geometry::PACKED);
    p.mat.convertTo(p.mat, CV_32F);

    return SetData(g, p, src, false, metric, indices);
}

cv::Mat PhotometricObjective::operator() (const EuclideanTransform& tf) const
{
    if (!m_proj)
    {
        E_ERROR << "missing projection model";
        return cv::Mat();
    }

    // source 3D geometry
    Geometry x = m_data.src;

    // project to the 2D image plane
    Geometry p = m_proj->Project(tf(x, true), ProjectionModel::EUCLIDEAN_2D);

    // convert the projected image coordinates to packed points for interpolation
    cv::Mat p32f;
    p.Reshape(Geometry::PACKED).mat.convertTo(p32f, CV_32F);

    // find mapped pixel values
    Geometry y = Geometry(Geometry::PACKED, interp(m_dst, p32f, m_interp));

    const MahalanobisMetric* m = dynamic_cast<const MahalanobisMetric*>(m_data.metric.get());

    if (!m)
    {
        return (*m_data.metric)(m_data.dst, y).mat;
    }

    // error propagation for Mahalanobis metric
    Geometry jac = m_proj->GetJacobian(x, p);
    cv::Mat dx = interp(m_gradient.Ix, p32f, m_interp);
    cv::Mat dy = interp(m_gradient.Iy, p32f, m_interp);
    cv::Mat dI = cv::Mat(jac.mat.rows, 3, jac.mat.depth());

    // propagate 3D -> 2D error metric to image manifold (2D -> 1D)
    if (dx.depth() != dI.depth()) dx.convertTo(dx, dI.depth());
    if (dy.depth() != dI.depth()) dy.convertTo(dy, dI.depth());

    for (int j = 0; j < 3; j++)
    {
        const int k = j * 2;
        cv::add(jac.mat.col(k).mul(dx), jac.mat.col(k+1).mul(dy), dI.col(j));
    }

    jac.mat = dI; // the new metric will be in 1D space

    Metric::ConstOwn d = m_data.metric->Transform(tf, jac);

    /////////////////////////////////////////////////////////////////////////////////////////
    // PersistentMat(d->ToMahalanobis()->GetCovariance().mat.clone()).Store(Path("jac.bin"));
    // PersistentMat(x.mat).Store(Path("x.bin"));
    // PersistentMat(y.mat).Store(Path("y.bin"));
    // PersistentMat(m_data.dst.mat.clone()).Store(Path("y0.bin"));
    /////////////////////////////////////////////////////////////////////////////////////////

    return (*d)(m_data.dst, y).mat;
}

//==[ PhotometricObjective::GradientImage ]===================================//

PhotometricObjective::GradientImage::GradientImage(const cv::Mat& im, int depth)
{
    cv::Scharr(im, Ix, depth, 1, 0);
    cv::Scharr(im, Iy, depth, 0, 1);
}

//==[ RigidObjective ]========================================================//

AlignmentObjective::Own RigidObjective::GetSubObjective(const IndexList& indices) const
{
    RigidObjective* sub = new RigidObjective();
    sub->m_data = m_data[indices];

    return AlignmentObjective::Own(sub);
}

bool RigidObjective::SetData(const GeometricMapping& data)
{
    if (!data.IsConsistent()) return false;

    const size_t d0 = data.src.GetDimension();
    const size_t d1 = data.dst.GetDimension();

    if (d0 != 3 && d0 != 4)
    {
        E_ERROR << "source geometry has to be either 3D Euclidean or 4D homogeneous (d=" << d0 << ")";
        return false;
    }

    if (d1 != 3 && d1 != 4)
    {
        E_ERROR << "target geometry has to be either 3D Euclidean or 4D homogeneous (d=" << d1 << ")";
        return false;
    }

    m_data.indices = data.indices;
    m_data.src = (d0 == 3 ? Geometry::MakeHomogeneous(data.src) : data.src);
    m_data.dst = (d1 == 4 ? Geometry::FromHomogeneous(data.dst) : data.dst);
    m_data.metric = data.metric;

    return true;
}

cv::Mat RigidObjective::operator() (const EuclideanTransform& f) const
{
    Metric::ConstOwn d = m_data.metric->Transform(f);

    Geometry x = m_data.src;
    Geometry y = f(x, true);

    return (*d)(m_data.dst, y).mat;
}

//==[ PoseEstimator ]=========================================================//

//==[ EssentialMatrixDecomposer ]=============================================//

bool EssentialMatrixDecomposer::operator() (const GeometricMapping& mapping, Estimate& estimate) const
{
    if (!mapping.Check(2, 2))
    {
        E_ERROR << "invalid mapping";
        return false;
    }

    if (mapping.GetSize() < GetMinPoints())
    {
        E_ERROR << "insufficient mapping size (" << mapping.GetSize() << " < " << GetMinPoints() << ")";
        return false;
    }

    if (!m_srcProj || !m_dstProj)
    {
        E_ERROR << "missing projeciton model(s)";
        return false;
    }

    GeometricMapping bpm;
    bpm.src.mat = Geometry::FromHomogeneous(m_srcProj->Backproject(mapping.src)).mat;
    bpm.dst.mat = Geometry::FromHomogeneous(m_dstProj->Backproject(mapping.dst)).mat;

    cv::Mat E;

    try
    {
        E = cv::findEssentialMat(bpm.src.mat, bpm.dst.mat, cv::Mat::eye(3, 3, CV_64F));
    }
    catch (std::exception& ex)
    {
        E_ERROR << "cv::findEssentialMat failed : " << ex.what();
        return false;
    }

    return estimate.pose.FromEssentialMatrix(E, bpm);
}

//==[ PerspevtivePoseEstimator ]==============================================//

bool PerspevtivePoseEstimator::operator() (const GeometricMapping& mapping, Estimate& estimate) const
{
    if (!mapping.Check(3, 2))
    {
        E_ERROR << "invalid mapping";
        return false;
    }

    if (mapping.GetSize() < GetMinPoints())
    {
        E_ERROR << "insufficient mapping size (" << mapping.GetSize() << " < " << GetMinPoints() << ")";
        return false;
    }

    if (!m_proj)
    {
        E_ERROR << "missing projection model";
        return false;
    }

    const cv::Mat opts = mapping.src.mat.reshape(3);
    const cv::Mat ipts = Geometry::FromHomogeneous(m_proj->Backproject(mapping.dst)).mat.reshape(2);

    cv::Mat K = cv::Mat::eye(3, 3, CV_64F);
    cv::Mat D = cv::Mat();
    cv::Mat rvec, tvec;

    bool success = cv::solvePnP(opts, ipts, K, D, rvec, tvec, false, cv::SOLVEPNP_EPNP);

    if (!success)
    {
        E_ERROR << "cv::solvePnP failed";
        return false;
    }

    if (!estimate.pose.GetRotation().FromVector(rvec))
    {
        E_ERROR << "error setting rotation vector";
        return false;
    }

    if (!estimate.pose.SetTranslation(tvec))
    {
        E_ERROR << "error setting translation vector";
        return false;
    }

    return true;
}

//==[ QuatAbsOrientationSolver ]==============================================//

bool QuatAbsOrientationSolver::operator() (const GeometricMapping& mapping, Estimate& estimate) const
{
    throw std::logic_error("not implemented");
    return false;
}

//==[ DummyPoseEstimator ]====================================================//

bool DummyPoseEstimator::operator() (const GeometricMapping& mapping, Estimate& estimate) const
{
    estimate.pose = m_pose;
    return true;
}

//==[ InversePoseEstimator ]=================================================//

bool InversePoseEstimator::operator() (const GeometricMapping& mapping, Estimate& estimate) const
{
    if (!m_estimator || !(*m_estimator)(mapping, estimate))
    {
        return false;
    }

    // invert the solution
    estimate.pose = estimate.pose.GetInverse();

    // TODO: need to invert the metric as well
    // ...
    // ..
    // .

    return true;
}

//==[ ConsensusPoseEstimator ]================================================//

bool ConsensusPoseEstimator::operator() (const GeometricMapping& mapping, Estimate& estimate) const
{
    IndexLists inliers, outliers;
    return (*this)(mapping, estimate, inliers, outliers);
}

bool ConsensusPoseEstimator::operator() (const GeometricMapping& mapping, Estimate& estimate, IndexLists& inliers, IndexLists& outliers) const
{
    if (!m_solver)
    {
        E_ERROR << "missing inner pose estimator";
        return false;
    }

    if (m_selectors.empty())
    {
        E_ERROR << "missing selector";
        return false;
    }

    const PoseEstimator& f = *m_solver;
    const size_t m = mapping.GetSize();
    const size_t n = f.GetMinPoints();

    if (m < n)
    {
        E_ERROR << "insufficient number of correspondences";
        return false;
    }

    const size_t population = GetPopulation();
    const size_t minInliers = static_cast<size_t>(population * m_minInlierRatio);
    const double logOneMinusConfidence = std::log(1 - m_confidence);
    const size_t iter = std::min(m_maxIter, n > 0 ? static_cast<size_t>(std::ceil(logOneMinusConfidence / std::log(1 - std::pow(m_minInlierRatio, n)))) : 1);
    size_t numInliers = 0;

    Speedometre solve, eval;

    if (m_verbose)
    {
        E_TRACE << std::setw(80) << std::setfill('=') << "";
        E_TRACE << std::setw(6)  << std::right << "Trial"
                << std::setw(19) << std::right << "Inliers"
                << std::setw(19) << std::right << "Best"
                << std::setw(18) << std::right << "Solve Time"
                << std::setw(18) << std::right << "Eval. Time";
        E_TRACE << std::setw(80) << std::setfill('=') << "";
    }

    struct Result
    {
        IndexLists inliers;
        IndexLists outliers;
    };

    //for (size_t k = 0, iter = m_maxIter; k < iter; k++)
    for (size_t k = 0; k < iter; k++)
    {
        IndexList idx = DrawSamples(m, n);
        GeometricMapping samples = mapping[idx];
        PoseEstimator::Estimate trial;
        Result result;
        size_t hits = 0;

        solve.Start();
        if (!f(samples, trial))
        {
            E_ERROR << "inner pose estimator failed";
            return false;
        }
        solve.Stop(1);

        eval.Start();
        if (!m_threaded)
        {
            BOOST_FOREACH(const AlignmentObjective::InlierSelector& g, m_selectors)
            {
                EvalResult rs;

                if (!g(trial.pose, rs.accepted, rs.rejected))
                {
                    E_ERROR << "selector failed";
                    return false;
                }

                result.inliers .push_back(rs.accepted);
                result.outliers.push_back(rs.rejected);

                hits += rs.accepted.size();
            }
        }
        else
        {
            boost::thread_group threads;
            std::vector<EvalResult> results(m_selectors.size());

            for (size_t i = 0; i < m_selectors.size(); i++)
            {
                threads.add_thread(
                    new boost::thread(
                        ConsensusPoseEstimator::EvalThread,
                        boost::cref(m_selectors[i]),
                        boost::cref(trial.pose),
                        boost::ref(results[i])
                    )
                );
            }

            threads.join_all();

            BOOST_FOREACH (const EvalResult& rs, results)
            {
                if (!rs.success)
                {
                    E_ERROR << "selector failed";
                    return false;
                }

                result.inliers .push_back(rs.accepted);
                result.outliers.push_back(rs.rejected);

                hits += rs.accepted.size();
            }
        }
        eval.Stop(1);

        if (m_verbose)
        {
            E_TRACE << std::setw(6)  << std::right << (k + 1)
                    << std::setw(12) << std::right << hits       << " (" << std::setw(3) << std::right << (100*hits/population)       << "%)"
                    << std::setw(12) << std::right << numInliers << " (" << std::setw(3) << std::right << (100*numInliers/population) << "%)"
                    << std::setw(15) << std::right << (solve.GetElapsedSeconds() * 1000) << " ms"
                    << std::setw(15) << std::right << (eval .GetElapsedSeconds() * 1000) << " ms";
        }

        if (hits < numInliers) continue;

        // accept the trial
        numInliers = hits;
        estimate = trial;
        inliers  = result.inliers;
        outliers = result.outliers;

        if (hits < minInliers) continue;

        // convergence control
        //iter = std::min(iter, static_cast<size_t>(std::ceil(logOneMinusConfidence / std::log(1 - std::pow(hits / population, n)))));
    }

    const bool success = numInliers >= minInliers;

    if (!success)
    {
        const int inlierRate = (int) (std::round(100 * numInliers / population));
        const int targetRate = (int) (std::round(100 * minInliers / population));

        E_WARNING << "unable to find enough inliers in " << iter << " iteration(s)";
        E_WARNING << "the best trial achieves " << inlierRate << "% inliers while " << targetRate << "% is required";
    }

    // post-estimation non-linear optimisation
    if (success && m_optimisation)
    {
        MultiObjectivePoseEstimation refinement;
        refinement.SetDifferentiationStep(1e-2);
        refinement.SetPose(estimate.pose);
        std::vector<Indices> inliersOptim;

        size_t i = 0;
        inliersOptim.resize(m_selectors.size());

        BOOST_FOREACH (const AlignmentObjective::InlierSelector& g, m_selectors)
        {
            const IndexList& idx = inliers[i];
            AlignmentObjective::Own sub = g.objective->GetSubObjective(idx);

            if (!sub)
            {
                E_WARNING << "error building sub-objective for model " << i;
            }
            else if (sub->GetData().GetSize() > 0)
            {
                refinement.AddObjective(AlignmentObjective::ConstOwn(sub));
                inliersOptim[i].insert(inliersOptim[i].end(), idx.begin(), idx.end());
            }

            i++;
        }

        LeastSquaresSolver::State state;
        LevenbergMarquardtAlgorithm levmar;

        levmar.SetVervbose(m_verbose);
        levmar.SetInitialDamp(1e-2);

        if (!levmar.Solve(refinement, state))
        {
            E_ERROR << "non-linear optimisation failed";
            return false;
        }

        //////////////////////////////////////////////////////////////
        // VectorisableD::Vec x; refinement.GetPose().Store(x);
        // PersistentMat(cv::Mat(refinement(x))).Store(Path("y.bin"));
        // PersistentMat(state.hessian).Store(Path("hessian.bin"));
        //////////////////////////////////////////////////////////////

        if (state.hessian.empty())
        {
            E_ERROR << "empty Hessian returned";
            return false;
        }

        if (!checkPositiveDefinite(state.hessian))
        {
            E_ERROR << "returned Hessian is not positive definite";
            E_ERROR << mat2string(state.hessian, "H");

            return false;
        }

        EuclideanTransform initPose = estimate.pose;
        cv::Mat cov; cv::invert(state.hessian, cov);

        estimate.pose = refinement.GetPose();
        estimate.metric = Metric::Own(new MahalanobisMetric(MahalanobisMetric::ANISOTROPIC_ROTATED, 6, symmat(cov)));
        estimate.valid = true;

        // do evaluation again
        inliers.clear();
        outliers.clear();
        numInliers = 0;

        BOOST_FOREACH (const AlignmentObjective::InlierSelector& g, m_selectors)
        {
            IndexList accepted, rejected;

            g(estimate.pose, accepted, rejected);

            inliers.push_back(accepted);
            outliers.push_back(rejected);

            numInliers += accepted.size();
        }

        if (numInliers < minInliers)
        {
            E_WARNING << "the optimised model results in too few inliers (" << numInliers << " < " << minInliers << ")";
            E_WARNING << "begin of analysis..";

            size_t i = 0;
            BOOST_FOREACH(const AlignmentObjective::InlierSelector& g, m_selectors)
            {
                cv::Mat err = (*g.objective)(estimate.pose);
                const Indices& idx = inliersOptim[i];

                E_WARNING << "objective " << i << " : " << idx.size() << " -> " << inliers[i].size() << " / " << outliers[i].size() << " (rmse=" << rms(err) << ")";

                const GeometricMapping& data = g.objective->GetData();
                std::stringstream ss; ss << "obj" << i;

                PersistentMat(data.src.mat.clone()).Store(Path(ss.str() + ".src.bin"));
                PersistentMat(data.dst.mat.clone()).Store(Path(ss.str() + ".dst.bin"));
                PersistentMat(err).Store(Path(ss.str() + ".err.bin"));

                if (!idx.empty())
                {
                    PersistentMat(cv::Mat(idx)).Store(Path(ss.str() + ".idx.bin"));
                }

                const WeightedEuclideanMetric* we = dynamic_cast<const WeightedEuclideanMetric*>(data.metric.get());

                if (we)
                {
                    PersistentMat(we->GetWeights().clone()).Store(Path(ss.str() + ".we.bin"));
                }

                i++;
            }

            PersistentMat(initPose.GetTransformMatrix()).Store(Path("pose.init.bin"));
            PersistentMat(estimate.pose.GetTransformMatrix()).Store(Path("pose.optm.bin"));

            E_WARNING << "..end of analysis";
        }
    }

    return estimate.valid = numInliers >= minInliers;
}

size_t ConsensusPoseEstimator::GetPopulation() const
{
    size_t n = 0;

    BOOST_FOREACH (const AlignmentObjective::InlierSelector& selector, m_selectors)
    {
        n += selector.objective ? selector.objective->GetData().GetSize() : 0;
    }

    return n;
}

IndexList ConsensusPoseEstimator::DrawSamples(size_t population, size_t samples)
{
    if (population < samples)
    {
        E_ERROR << "insufficient population (n=" << population << ") for " << samples << " sample(s)";
        return IndexList();
    }

    if (samples == 0)
    {
        return IndexList();
    }

    std::vector<size_t> idx(population);
    for (size_t i = 0; i < population; i++)
    {
        idx[i] = i;
    }

    std::random_shuffle(idx.begin(), idx.end());

    return IndexList(idx.begin(), std::next(idx.begin(), samples));
}

void ConsensusPoseEstimator::EvalThread(const AlignmentObjective::InlierSelector& g, const EuclideanTransform& tf, EvalResult& result)
{
    result.success = g(tf, result.accepted, result.rejected);
}

//==[ MultiObjectivePoseEstimation ]==========================================//

bool MultiObjectivePoseEstimation::Initialise(VectorisableD::Vec& x)
{
    m_conds = GetConds();
    return m_conds >= 6 && m_tform.Store(x);
}

VectorisableD::Vec MultiObjectivePoseEstimation::operator() (const VectorisableD::Vec& x) const
{
    EuclideanTransform tform(m_tform.GetRotation().GetParameterisation());

    if (!tform.Restore(x))
    {
        E_ERROR << "error devectorising transform";
        return VectorisableD::Vec();
    }

    VectorisableD::Vec y;
    size_t m = 0;

    y.reserve(m_conds);

    BOOST_FOREACH (const AlignmentObjective::ConstOwn& obj, m_objectives)
    {
        if (!obj) continue;

        cv::Mat yi = (*obj)(tform);
        size_t  mi = static_cast<size_t>(yi.total());

        assert(m + mi <= m_conds);

        if (yi.type() != CV_64F)
        {
            yi.convertTo(yi, CV_64F);
        }

        y.insert(y.end(), (double*)yi.datastart, (double*)yi.dataend);
        m += mi;
    }

    assert(m == m_conds);

    return y;
}

size_t MultiObjectivePoseEstimation::GetConds() const
{
    size_t m = 0;

    BOOST_FOREACH (const AlignmentObjective::ConstOwn& obj, m_objectives)
    {
        if (!obj) continue;
        m += obj->GetData().dst.GetElements();
    }

    return m;
}

//==[ StructureEstimation ]===================================================//

//==[ StructureEstimation::Estimate ]=========================================//

StructureEstimation::Estimate StructureEstimation::Estimate::operator[] (const IndexList& indices) const
{
    Estimate estimate(structure.shape);
    estimate.structure = structure[indices];
    estimate.metric    = metric ? (*metric)[indices] : Metric::Own();

    return estimate;
}

StructureEstimation::Estimate StructureEstimation::Estimate::operator+ (const Estimate& estimate) const
{
    Estimate clone(Geometry(structure), metric ? metric->Clone() : Metric::Own());
    return clone += estimate;
}

StructureEstimation::Estimate& StructureEstimation::Estimate::operator+= (const Estimate& estimate)
{
    if (!metric || !estimate.metric)
    {
        E_ERROR << "fusion failed due to missing metric(es)";
        return *this;
    }

    bool native;

    boost::shared_ptr</***/ MahalanobisMetric> m0 = metric->ToMahalanobis(native);
    boost::shared_ptr<const MahalanobisMetric> m1 = estimate.metric->ToMahalanobis();

    if (!m0 || !m1)
    {
        E_ERROR << "fusion failed due to missing Mahalanobis metric(es)";
        return *this;
    }

    // PersistentMat(m0->GetCovariance().mat.clone()).Store(Path("C0.bin"));
    // PersistentMat(m1->GetCovariance().mat.clone()).Store(Path("C1.bin"));

    Geometry g0  = structure.Reshape(Geometry::ROW_MAJOR);
    Geometry g10 = estimate.structure - g0;

    // PersistentMat(g0.mat).Store(Path("g0.bin"));
    // PersistentMat(estimate.structure.mat.clone()).Store(Path("g1.bin"));

    if (g10.IsEmpty())
    {
        E_WARNING << "the subtraction of structure geometry is empty";
        return *this;
    }

    cv::Mat K = m0->Update(*m1);

    if (K.empty())
    {
        E_ERROR << "error updating metric, invalid Kalman gain returned";
        return *this;
    }

    // PersistentMat(K).Store(Path("K.bin"));

    switch (m0->type)
    {
    case MahalanobisMetric::ISOTROPIC:

        for (int j = 0; j < g0.mat.cols; j++)
        {
            g0.mat.col(j) += K.mul(g10.mat.col(j));
        }

        break;

    case MahalanobisMetric::ANISOTROPIC_ORTHOGONAL:

        for (int j = 0; j < g0.mat.cols; j++)
        {
            g0.mat.col(j) += K.col(j).mul(g10.mat.col(j));
        }

        break;

    case MahalanobisMetric::ANISOTROPIC_ROTATED:

        for (int j0 = 0; j0 < g0.mat.cols; j0++)
        {
            for (int j1 = 0; j1 < g10.mat.cols; j1++)
            {
                const int j = j0 * g10.mat.cols + j1;
                g0.mat.col(j0) += K.col(j).mul(g10.mat.col(j1));
            }
        }

        break;
    }

    // PersistentMat(g0.mat).Store(Path("G2.bin"));

    // just in case : structure.shape == Geometry::COL_MAJOR
    structure = g0;

    if (!native)
    {
        boost::shared_ptr<WeightedEuclideanMetric> we =
            boost::dynamic_pointer_cast<WeightedEuclideanMetric, Metric>(metric);

        if (!we)
        {
            E_WARNING << "metric fusion impossible";
            return *this;
        }

        if (!we->FromMahalanobis(*m0))
        {
            E_ERROR << "error converting Mahalanobis distance to weighted Euclidean";
        }
    }

    return *this;
}

//==[ TwoViewTriangulation ]==================================================//

StructureEstimation::Estimate TwoViewTriangulation::operator() (const GeometricMapping& m) const
{
    if (!m.Check(2, 2))
    {
        E_ERROR << "ill-shaped mapping";
        return Estimate(m.src.shape);
    }

    Geometry g = (*this)(
        m.src.Reshape(Geometry::ROW_MAJOR),
        m.dst.Reshape(Geometry::ROW_MAJOR),
        M01.pose).Reshape(m.src.shape);

    // error propogation
    const int ROWS = static_cast<int>(g.GetElements());
    const int COLS = static_cast<int>(MahalanobisMetric::GetCovMatCols(MahalanobisMetric::ANISOTROPIC_ROTATED, 3));
    cv::Mat jac = GetJacobian(m.src, m.dst, g).mat;
    cv::Mat cov = cv::Mat(ROWS, COLS, g.mat.depth());
    boost::shared_ptr<const MahalanobisMetric> metric = m.metric ? m.metric->ToMahalanobis() : 0;

    if (metric)
    {
        for (int i = 0; i < ROWS; i++)
        {
            cv::Mat Ci = metric->GetFullCovMat(static_cast<size_t>(i)); // C, 4-by-4
            cv::Mat Ji = jac.row(i).reshape(1, 4); // J', 4-by-3

            Ci = Ji.t() * Ci * Ji; // J * C * J', 3-by-3

            // if (!checkPositiveDefinite(Ci))
            // {
            //   const cv::Mat C0 = metric->GetFullCovMat(static_cast<size_t>(i));
            //
            //   E_WARNING << "the Jacobian-transformed covariance of entry " << i << " is not positive definite";
            //   E_WARNING << mat2string(C0, "C0");
            //   E_WARNING << mat2string(Ci, "C1");
            //   E_WARNING << mat2string(Ji, "J");
            //
            //   Ci.setTo(0);
            // }

            for (int d0 = 0; d0 < 3; d0++)
            {
                for (int d1 = d0; d1 < 3; d1++)
                {
                    cov.at<double>(i, sub2symind(d0, d1, 3)) = Ci.at<double>(d0, d1);
                }
            }
        }
    }
    else // identity covaraince assumed
    {
        for (int i = 0; i < ROWS; i++)
        {
            cv::Mat Ji = jac.row(i).reshape(1, 4); // 4-by-3
            cv::Mat Ci = Ji.t() * Ji; // 3-by-3

            for (int d0 = 0; d0 < 3; d0++)
            {
                for (int d1 = d0; d1 < 3; d1++)
                {
                    cov.at<double>(i, sub2symind(d0, d1, 3)) = Ci.at<double>(d0, d1);
                }
            }
        }
    }

    // do error propagation from motion estimate
    if (M01.metric)
    {
        Geometry jacpos = GetPoseJacobian(m.src, m.dst, g);

        if (!jacpos.IsEmpty())
        {
            boost::shared_ptr<const MahalanobisMetric> pmetric =
                M01.metric->Transform(EuclideanTransform::Identity, jacpos)->ToMahalanobis();

            if (pmetric)
            {
                cv::Mat covpos = pmetric->GetFullCovMat();
                cv::add(cov, covpos, cov); // the pose-derived covariance is independent to the image-derived one

                for (int i = 0; i < covpos.rows; i++)
                {
                    if (covpos.at<double>(i, 0) == 0)
                    {
                        // E_TRACE << "nullified entry " << i;
                        cov.row(i).setTo(0.0f);
                    }
                }
            }
            else
            {
                E_WARNING << "error obtaining pose metric as a Mahalanobis metric";
            }
        }
        else
        {
            E_WARNING << "error evaluating Jacobian with respect to the pose";
        }
    }

    ////////////////////////////////////////////////////////////////////////
    // cv::Mat det(cov.rows, 1, cov.depth());
    // for (int i = 0; i < ROWS; i++)
    // {
    //     det.at<double>(i) = cv::determinant(symmat(cov.row(i), 3));
    // }
    // PersistentMat(m.src.mat.clone()).Store(Path("x0.bin"));
    // PersistentMat(m.dst.mat.clone()).Store(Path("x1.bin"));
    // PersistentMat(g.mat).Store(Path("g.bin"));
    // PersistentMat(cov).Store(Path("cov.bin"));
    // PersistentMat(M01.pose.GetTransformMatrix()).Store(Path("M01.bin"));
    // PersistentMat(det).Store(Path("det.bin"));
    ////////////////////////////////////////////////////////////////////////

    // remove points behind
    for (int i = 0; i < ROWS; i++)
    {
        if (g.mat.at<double>(i, 2) < 0)
        {
            // E_TRACE << "lmk=" << m.indices.at(i) << ",pos=" << g.mat.row(i);

            g.mat.row(i).setTo(0.0f);
            cov.row(i).setTo(0.0f);
        }
    }

    return StructureEstimation::Estimate(
        g, Metric::Own(new MahalanobisMetric(MahalanobisMetric::ANISOTROPIC_ROTATED, 3, cov))
    );
}

Geometry TwoViewTriangulation::GetJacobian(const Geometry& x0, const Geometry& x1, const Geometry& g) const
{
    const double dx = 1e-2;
    cv::Mat J(static_cast<int>(g.GetElements()), g.mat.cols * 4, g.mat.type());

    for (size_t i = 0; i < 4; i++)
    {
        Geometry gi = (i < 2) ? (*this)(x0.Step(i, dx), x1, M01.pose) : (*this)(x0, x1.Step(i - 2, dx), M01.pose);
        Geometry dy = g - gi;

        const int j = static_cast<int>(i) * g.mat.cols;
        J.colRange(j, j + g.mat.cols) = dy.Reshape(Geometry::ROW_MAJOR).mat / dx;
    }

    return Geometry(Geometry::ROW_MAJOR, J);
}

Geometry TwoViewTriangulation::GetPoseJacobian(const Geometry& x0, const Geometry& x1, const Geometry& g) const
{
    const double dx = 1e-2;
    const VectorisableD::Vec v = M01.pose.ToVector();

    if (v.empty())
    {
        E_ERROR << "error vectorising pose";
        return Geometry(Geometry::ROW_MAJOR);
    }

    cv::Mat J(static_cast<int>(g.GetElements()), g.mat.cols * static_cast<int>(v.size()), g.mat.type());

    for (size_t i = 0; i < v.size(); i++)
    {
        VectorisableD::Vec vi = v;
        vi[i] += dx;

        EuclideanTransform mi(M01.pose.GetRotation().GetParameterisation());
        if (!mi.FromVector(vi))
        {
            E_ERROR << "error devectorising pose for numerical differentiation of pose parameter " << i;
            return Geometry(Geometry::ROW_MAJOR);
        }

        Geometry gi = (*this)(x0, x1, mi);
        Geometry dy = g - gi;

        const int j = static_cast<int>(i) * g.mat.cols;
        J.colRange(j, j + g.mat.cols) = dy.Reshape(Geometry::ROW_MAJOR).mat / dx;
    }

    return Geometry(Geometry::ROW_MAJOR, J);
}

//==[ OptimalTriangulation ]==================================================//

Geometry OptimalTriangulation::operator() (const Geometry& x0, const Geometry& x1, const EuclideanTransform& tform) const
{
    const PosedProjection Q1(tform, P1.Clone());

    Geometry p0 = Geometry::FromHomogeneous(P0.Backproject(x0)).Reshape(Geometry::PACKED);
    Geometry p1 = Geometry::FromHomogeneous(Q1.Backproject(x1)).Reshape(Geometry::PACKED);
    Geometry g(Geometry::COL_MAJOR);

    // type casting to fit cv::triangulatePoints
    if (p0.mat.depth() != CV_32F) p0.mat.convertTo(p0.mat, CV_32F);
    if (p1.mat.depth() != CV_32F) p1.mat.convertTo(p1.mat, CV_32F);

    cv::Mat pmat0 = P0.pose.GetTransformMatrix(false, true, CV_32F); // I * (R0 | t0)
    cv::Mat pmat1 = Q1.pose.GetTransformMatrix(false, true, CV_32F); // I * (R1 | t1)

    cv::triangulatePoints(pmat0, pmat1, p0.mat, p1.mat, g.mat);

    // cast structure to the type of the given geometry data, if neccessary
    if (g.mat.type() != x0.mat.type())
    {
        g.mat.convertTo(g.mat, x0.mat.type());
    }

    return Geometry::FromHomogeneous(g).Reshape(x0.shape);
}

//==[ MidPointTriangulation ]=================================================//

void MidPointTriangulation::DecomposeProjMatrix(const cv::Mat& P, cv::Mat& KRinv, cv::Mat& c)
{
    cv::Mat KR = P.rowRange(0, 3).colRange(0, 3);
    KRinv = KR.inv();
    c = -KRinv * P.rowRange(0, 3).col(3);
}

Geometry MidPointTriangulation::operator() (const Geometry& x0, const Geometry& x1, const EuclideanTransform& tform) const
{
    // const PosedProjection Q1(tform, P1.Clone());
    const EuclideanTransform PQ = tform >> P1.pose;
    const PosedProjection Q1(PQ, P1.proj->Clone());

    cv::Mat c0 = P0.pose.GetInverse().GetTranslation();
    cv::Mat c1 = Q1.pose.GetInverse().GetTranslation();

    cv::Mat s = c0 + c1;
    cv::Mat t = c1 - c0;
    cv::Mat y0 = P0.Backproject(x0).mat;
    cv::Mat y1 = Q1.Backproject(x1).mat;

    if (y1.type() != y0.type())
    {
        y1.convertTo(y1, y0.type());
    }

    cv::Mat g = cv::Mat(y0.rows, 3, y0.type());
    cv::Mat w = cv::Mat(y0.rows, 1, y0.type());

    for (int i = 0; i < g.rows; i++)
    {
        cv::Mat Bt = cv::Mat(2, 3, y0.type());
        y0.row(i).copyTo(Bt.row(0));
        y1.row(i).copyTo(Bt.row(1));

        cv::Mat At = Bt.clone();
        At.row(1) = -Bt.row(1);

        cv::Mat A = At.t();
        cv::Mat B = Bt.t();
        cv::Mat k = (At * A).inv() * At * t;
        cv::Mat x = (B * k + s) * 0.5f;
        cv::Mat d = (A * k - t);

        g.row(i) = x.t();
        w.row(i) = 1.0f / cv::norm(d);
    }

    return Geometry(Geometry::ROW_MAJOR, g);
}

//==[ MultipleViewBlockMatcher ]==============================================//

// PinholeModel MultipleViewBlockMatcher::RefView::NullProjection;


/* ----------------------------------------------------------------------------

 * GTSAM Copyright 2010, Georgia Tech Research Corporation,
 * Atlanta, Georgia 30332-0415
 * All Rights Reserved
 * Authors: Frank Dellaert, et al. (see THANKS for the full author list)

 * See LICENSE for the license information

 * -------------------------------------------------------------------------- */

/**
 * @file    DoglegOptimizer.cpp
 * @brief
 * @author  Richard Roberts
 * @date   Feb 26, 2012
 */

#include <gtsam/nonlinear/DoglegOptimizer.h>
#include <gtsam/nonlinear/DoglegOptimizerImpl.h>
#include <gtsam/nonlinear/internal/NonlinearOptimizerState.h>
#include <gtsam/linear/GaussianBayesTree.h>
#include <gtsam/linear/GaussianBayesNet.h>
#include <gtsam/linear/GaussianFactorGraph.h>
#include <gtsam/linear/VectorValues.h>

#include <boost/algorithm/string.hpp>

namespace gtsam {

/* ************************************************************************* */
DoglegParams::VerbosityDL DoglegParams::verbosityDLTranslator(const std::string &verbosityDL) const {
  std::string s = verbosityDL;  boost::algorithm::to_upper(s);
  if (s == "SILENT") return DoglegParams::SILENT;
  if (s == "VERBOSE") return DoglegParams::VERBOSE;

  /* default is silent */
  return DoglegParams::SILENT;
}

/* ************************************************************************* */
std::string DoglegParams::verbosityDLTranslator(VerbosityDL verbosityDL) const {
  std::string s;
  switch (verbosityDL) {
  case DoglegParams::SILENT:  s = "SILENT"; break;
  case DoglegParams::VERBOSE: s = "VERBOSE"; break;
  default:                    s = "UNDEFINED"; break;
  }
  return s;
}

/* ************************************************************************* */
namespace internal {
struct DoglegState : public NonlinearOptimizerState {
  const double delta;

  DoglegState(const Values& values, double error, double delta, unsigned int iterations = 0)
      : NonlinearOptimizerState(values, error, iterations), delta(delta) {}
};
}

typedef internal::DoglegState State;

/* ************************************************************************* */
DoglegOptimizer::DoglegOptimizer(const NonlinearFactorGraph& graph, const Values& initialValues,
                                 const DoglegParams& params)
    : NonlinearOptimizer(
          graph, std::unique_ptr<State>(
                     new State(initialValues, graph.error(initialValues), params.deltaInitial))),
      params_(ensureHasOrdering(params, graph)) {}

DoglegOptimizer::DoglegOptimizer(const NonlinearFactorGraph& graph, const Values& initialValues,
                                 const Ordering& ordering)
    : NonlinearOptimizer(graph, std::unique_ptr<State>(
                                    new State(initialValues, graph.error(initialValues), 1.0))) {
  params_.ordering = ordering;
}

double DoglegOptimizer::getDelta() const {
  return static_cast<const State*>(state_.get())->delta;
}

/* ************************************************************************* */
GaussianFactorGraph::shared_ptr DoglegOptimizer::iterate(void) {

  // Linearize graph
  GaussianFactorGraph::shared_ptr linear = graph_.linearize(state_->values);

  // Pull out parameters we'll use
  const bool dlVerbose = (params_.verbosityDL > DoglegParams::SILENT);

  // Do Dogleg iteration with either Multifrontal or Sequential elimination
  DoglegOptimizerImpl::IterationResult result;

  if ( params_.isMultifrontal() ) {
    GaussianBayesTree bt = *linear->eliminateMultifrontal(*params_.ordering, params_.getEliminationFunction());
    VectorValues dx_u = bt.optimizeGradientSearch();
    VectorValues dx_n = bt.optimize();
    result = DoglegOptimizerImpl::Iterate(getDelta(), DoglegOptimizerImpl::ONE_STEP_PER_ITERATION,
      dx_u, dx_n, bt, graph_, state_->values, state_->error, dlVerbose);
  }
  else if ( params_.isSequential() ) {
    GaussianBayesNet bn = *linear->eliminateSequential(*params_.ordering, params_.getEliminationFunction());
    VectorValues dx_u = bn.optimizeGradientSearch();
    VectorValues dx_n = bn.optimize();
    result = DoglegOptimizerImpl::Iterate(getDelta(), DoglegOptimizerImpl::ONE_STEP_PER_ITERATION,
      dx_u, dx_n, bn, graph_, state_->values, state_->error, dlVerbose);
  }
  else if ( params_.isIterative() ) {
    throw std::runtime_error("Dogleg is not currently compatible with the linear conjugate gradient solver");
  }
  else {
    throw std::runtime_error("Optimization parameter is invalid: DoglegParams::elimination");
  }

  // Maybe show output
  if(params_.verbosity >= NonlinearOptimizerParams::DELTA) result.dx_d.print("delta");

  // Create new state with new values and new error
  state_.reset(new State(state_->values.retract(result.dx_d), result.f_error, result.delta,
                         state_->iterations + 1));
  return linear;
}

/* ************************************************************************* */
DoglegParams DoglegOptimizer::ensureHasOrdering(DoglegParams params, const NonlinearFactorGraph& graph) const {
  if (!params.ordering)
    params.ordering = Ordering::Create(params.orderingType, graph);
  return params;
}

} /* namespace gtsam */


// Copyright (c) 2018-present Baidu, Inc. All Rights Reserved.
//
// 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 "internal_functions.h"
#include <openssl/md5.h>
#include "hll_common.h"
#include "datetime.h"
#include <boost/date_time/gregorian/gregorian.hpp>
#include <cctype>
#include <cmath>
#include <algorithm>

namespace baikaldb {
static const int32_t DATE_FORMAT_LENGTH = 128;
static const std::vector<std::string> day_names = {
        "Sunday", "Monday", "Tuesday", "Wednesday",
        "Thursday", "Friday", "Saturday"};
static const std::vector<std::string> month_names = {
        "January", "February", "March", "April", "May",
        "June", "July", "August", "September",
        "October", "November", "December"};
ExprValue round(const std::vector<ExprValue>& input) {
    if (input.size() == 0 || input[0].is_null()) {
        return ExprValue::Null();
    }
    int bits = input.size() == 2 ? input[1].get_numberic<int>() : 0;
    double base = std::pow(10, bits);
    double orgin = input[0].get_numberic<double>();
    ExprValue tmp(pb::DOUBLE);
    if (base > 0) {
        if (orgin < 0) {
            tmp._u.double_val = -::round(-orgin * base) / base;
        } else {
            tmp._u.double_val = ::round(orgin * base) / base;
        }
    }
    return tmp;
}

ExprValue floor(const std::vector<ExprValue>& input) {
    if (input.size() == 0 || input[0].is_null()) {
        return ExprValue::Null();
    }
    ExprValue tmp(pb::INT64);
    tmp._u.int64_val = ::floor(input[0].get_numberic<double>());
    return tmp;
}

ExprValue ceil(const std::vector<ExprValue>& input) {
    if (input.size() == 0 || input[0].is_null()) {
        return ExprValue::Null();
    }
    ExprValue tmp(pb::INT64);
    tmp._u.int64_val = ::ceil(input[0].get_numberic<double>());
    return tmp;
}

ExprValue abs(const std::vector<ExprValue>& input) {
    if (input.size() == 0 || input[0].is_null()) {
        return ExprValue::Null();
    }
    ExprValue tmp(pb::DOUBLE);
    tmp._u.double_val = ::abs(input[0].get_numberic<double>());
    return tmp;
}

ExprValue sqrt(const std::vector<ExprValue>& input) {
    if (input.size() == 0 || input[0].is_null()) {
        return ExprValue::Null();
    }
    double val = input[0].get_numberic<double>();
    if (val < 0) {
        return ExprValue::Null();
    }
    ExprValue tmp(pb::DOUBLE);
    tmp._u.double_val = std::sqrt(val);
    return tmp;
}

ExprValue mod(const std::vector<ExprValue>& input) {
    if (input.size() < 2 || input[0].is_null() || input[1].is_null()) {
        return ExprValue::Null();
    }
    double rhs = input[1].get_numberic<double>();
    if (float_equal(rhs, 0)) {
        return ExprValue::Null();
    }
    double lhs = input[0].get_numberic<double>();

    ExprValue tmp(pb::DOUBLE);
    tmp._u.double_val = std::fmod(lhs, rhs);
    return tmp;
}

ExprValue rand(const std::vector<ExprValue>& input) {
    ExprValue tmp(pb::DOUBLE);
    tmp._u.double_val = butil::fast_rand_double();
    return tmp;
}

ExprValue sign(const std::vector<ExprValue>& input) {
    if (input.size() < 1 || input[0].is_null()) {
        return ExprValue::Null();
    }
    ExprValue tmp(pb::INT64);
    double val = input[0].get_numberic<double>();
    tmp._u.int64_val = val > 0 ? 1 : (val < 0 ? -1 : 0);
    return tmp;
}

ExprValue sin(const std::vector<ExprValue>& input) {
    if (input.size() < 1 || input[0].is_null()) {
        return ExprValue::Null();
    }
    ExprValue tmp(pb::DOUBLE);
    tmp._u.double_val = std::sin(input[0].get_numberic<double>());
    return tmp;
}
ExprValue asin(const std::vector<ExprValue>& input) {
    if (input.size() < 1 || input[0].is_null()) {
        return ExprValue::Null();
    }
    ExprValue tmp(pb::DOUBLE);
    double  val = input[0].get_numberic<double>();
    if (val < -1 || val > 1) {
        return ExprValue::Null();
    }
    tmp._u.double_val = std::asin(val);
    return tmp;
}

ExprValue cos(const std::vector<ExprValue>& input) {
    if (input.size() < 1 || input[0].is_null()) {
        return ExprValue::Null();
    }
    ExprValue tmp(pb::DOUBLE);
    tmp._u.double_val = std::cos(input[0].get_numberic<double>());
    return tmp;
}

ExprValue acos(const std::vector<ExprValue>& input) {
    if (input.size() < 1 || input[0].is_null()) {
        return ExprValue::Null();
    }
    double  val = input[0].get_numberic<double>();
    if (val < -1 || val > 1) {
        return ExprValue::Null();
    }
    ExprValue tmp(pb::DOUBLE);
    tmp._u.double_val = std::acos(val);
    return tmp;
}

ExprValue tan(const std::vector<ExprValue>& input) {
    if (input.size() < 1 || input[0].is_null()) {
        return ExprValue::Null();
    }
    ExprValue tmp(pb::DOUBLE);
    tmp._u.double_val = std::tan(input[0].get_numberic<double>());
    return tmp;
}

ExprValue cot(const std::vector<ExprValue>& input) {
    if (input.size() < 1 || input[0].is_null()) {
        return ExprValue::Null();
    }
    ExprValue tmp(pb::DOUBLE);
    double val = input[0].get_numberic<double>();
    double sin_val = std::sin(val);
    double cos_val = std::cos(val);
    if (float_equal(sin_val, 0)) {
        return ExprValue::Null();
    }
    tmp._u.double_val = cos_val/sin_val;
    return tmp;
}

ExprValue atan(const std::vector<ExprValue>& input) {
    if (input.size() < 1 || input[0].is_null()) {
        return ExprValue::Null();
    }
    ExprValue tmp(pb::DOUBLE);
    tmp._u.double_val = std::atan(input[0].get_numberic<double>());
    return tmp;
}

ExprValue ln(const std::vector<ExprValue>& input) {
    if (input.size() < 1 || input[0].is_null()) {
        return ExprValue::Null();
    }
    double val = input[0].get_numberic<double>();
    if (val <= 0) {
        return ExprValue::Null();
    }
    ExprValue tmp(pb::DOUBLE);
    tmp._u.double_val = std::log(input[0].get_numberic<double>());
    return tmp;
}

ExprValue log(const std::vector<ExprValue>& input) {
    if (input.size() < 2 || input[0].is_null() || input[1].is_null()) {
        return ExprValue::Null();
    }
    double base = input[0].get_numberic<double>();
    double val = input[1].get_numberic<double>();
    if (base <= 0 || val <= 0 || base == 1) {
        return ExprValue::Null();
    }
    ExprValue tmp(pb::DOUBLE);
    tmp._u.double_val = std::log(val) / std::log(base);
    return tmp;
}

ExprValue pow(const std::vector<ExprValue>& input) {
    if (input.size() < 2 || input[0].is_null() || input[1].is_null()) {
        return ExprValue::Null();
    }
    double base = input[0].get_numberic<double>();
    double exp = input[1].get_numberic<double>();
    ExprValue tmp(pb::DOUBLE);
    tmp._u.double_val = std::pow(base, exp);
    return tmp;
}

ExprValue pi(const std::vector<ExprValue>& input) {
    ExprValue tmp(pb::DOUBLE);
    tmp._u.double_val = M_PI;
    return tmp;
}

ExprValue greatest(const std::vector<ExprValue>& input) {
    bool find_flag = false;
    double ret = std::numeric_limits<double>::lowest();
    for (const auto& item : input) {
        if (item.is_null()) {
            return ExprValue::Null();
        } else {
            double val = item.get_numberic<double>();
            if (!find_flag) {
                find_flag = true;
                ret = val;
            } else {
                if (val > ret) {
                    ret = val;
                }
            }
        }
    }
    if (find_flag) {
        ExprValue tmp(pb::DOUBLE);
        tmp._u.double_val = ret;
        return tmp;
    } else {
        return ExprValue::Null();
    }
}

ExprValue least(const std::vector<ExprValue>& input) {
    bool find_flag = false;
    double ret = std::numeric_limits<double>::max();
    for (const auto& item : input) {
        if (item.is_null()) {
            return ExprValue::Null();
        } else {
            double val = item.get_numberic<double>();
            if (!find_flag) {
                find_flag = true;
                ret = val;
            } else {
                if (val < ret) {
                    ret = val;
                }
            }
        }
    }
    if (find_flag) {
        ExprValue tmp(pb::DOUBLE);
        tmp._u.double_val = ret;
        return tmp;
    } else {
        return ExprValue::Null();
    }
}

ExprValue length(const std::vector<ExprValue>& input) {
    if (input.size() == 0 || input[0].is_null()) {
        return ExprValue::Null();
    }
    ExprValue tmp(pb::UINT32);
    tmp._u.uint32_val = input[0].get_string().size();
    return tmp;
}

ExprValue bit_length(const std::vector<ExprValue>& input) {
    if (input.size() == 0 || input[0].is_null()) {
        return ExprValue::Null();
    }
    ExprValue tmp(pb::UINT32);
    tmp._u.uint32_val = input[0].get_string().size() * 8;
    return tmp;
}

ExprValue lower(const std::vector<ExprValue>& input) {
    if (input.size() == 0 || input[0].is_null()) {
        return ExprValue::Null();
    }
    ExprValue tmp(pb::STRING);
    tmp.str_val = input[0].get_string();
    std::transform(tmp.str_val.begin(), tmp.str_val.end(), tmp.str_val.begin(), ::tolower);
    return tmp;
}

ExprValue lower_gbk(const std::vector<ExprValue>& input) {
    if (input.size() == 0 || input[0].is_null()) {
        return ExprValue::Null();
    }
    ExprValue tmp(pb::STRING);
    tmp.str_val = input[0].get_string();
    std::string& literal = tmp.str_val;
    size_t idx = 0;
    while (idx < literal.size()) {
        if ((literal[idx] & 0x80) != 0) {
            idx += 2;
        } else {
            literal[idx] = tolower(literal[idx]);
            idx++;
        }
    }
    return tmp;
}

ExprValue upper(const std::vector<ExprValue>& input) {
    if (input.size() == 0 || input[0].is_null()) {
        return ExprValue::Null();
    }
    ExprValue tmp(pb::STRING);
    tmp.str_val = input[0].get_string();
    std::transform(tmp.str_val.begin(), tmp.str_val.end(), tmp.str_val.begin(), ::toupper);
    return tmp;
}

ExprValue concat(const std::vector<ExprValue>& input) {
    ExprValue tmp(pb::STRING);
    for (auto& s : input) {
        if (s.is_null()) {
            return ExprValue::Null();
        }
        tmp.str_val += s.get_string();
    }
    return tmp;
}

ExprValue substr(const std::vector<ExprValue>& input) {
    if (input.size() < 2) {
        return ExprValue::Null();
    }
    for (auto& s : input) {
        if (s.is_null()) {
            return ExprValue::Null();
        }
    }
    std::string str = input[0].get_string();
    ExprValue tmp(pb::STRING);
    int pos = input[1].get_numberic<int>();
    if (pos < 0) {
        pos = str.size() + pos;
    } else {
        --pos;
    }
    if (pos < 0 || pos >= (int)str.size()) {
        return tmp;
    }
    int len = -1;
    if (input.size() == 3) {
        len = input[2].get_numberic<int>();
        if (len <= 0) {
            return tmp;
        }
    }
    tmp.str_val = str;
    if (len == -1) {
        tmp.str_val = tmp.str_val.substr(pos);
    } else {
        tmp.str_val = tmp.str_val.substr(pos, len);
    }
    return tmp;
}

ExprValue left(const std::vector<ExprValue>& input) {
    if (input.size() < 2) {
        return ExprValue::Null();
    }
    for (auto& s : input) {
        if (s.is_null()) {
            return ExprValue::Null();
        }
    }
    ExprValue tmp(pb::STRING);
    int len = input[1].get_numberic<int>();
    if (len <= 0) {
        return tmp;
    }
    tmp.str_val = input[0].str_val;
    tmp.str_val = tmp.str_val.substr(0, len);
    return tmp;
}

ExprValue right(const std::vector<ExprValue>& input) {
    if (input.size() < 2) {
        return ExprValue::Null();
    }
    for (auto& s : input) {
        if (s.is_null()) {
            return ExprValue::Null();
        }
    }
    ExprValue tmp(pb::STRING);
    int len = input[1].get_numberic<int>();
    if (len <= 0) {
        return tmp;
    }
    int pos = input[0].str_val.size() - len;
    if (pos < 0) {
        pos = 0;
    }
    tmp.str_val = input[0].str_val;
    tmp.str_val = tmp.str_val.substr(pos);
    return tmp;
}

ExprValue trim(const std::vector<ExprValue>& input) {
    if (input.size() != 1) {
        return ExprValue::Null();
    }

    ExprValue tmp(pb::STRING);
    tmp.str_val = input[0].str_val;
    tmp.str_val.erase(0, tmp.str_val.find_first_not_of(" "));
    tmp.str_val.erase(tmp.str_val.find_last_not_of(" ") + 1);

    return tmp;
}

ExprValue ltrim(const std::vector<ExprValue>& input) {
    if (input.size() != 1) {
        return ExprValue::Null();
    }

    ExprValue tmp(pb::STRING);
    tmp.str_val = input[0].str_val;
    tmp.str_val.erase(0, tmp.str_val.find_first_not_of(" "));

    return tmp;
}

ExprValue rtrim(const std::vector<ExprValue>& input) {
    if (input.size() != 1) {
        return ExprValue::Null();
    }

    ExprValue tmp(pb::STRING);
    tmp.str_val = input[0].str_val;
    tmp.str_val.erase(tmp.str_val.find_last_not_of(" ") + 1);

    return tmp;
}

ExprValue concat_ws(const std::vector<ExprValue>& input) {
    if (input.size() < 2) {
        return ExprValue::Null();
    }

    if (input[0].is_null()) {
        return ExprValue::Null();
    }

    ExprValue tmp(pb::STRING);

    bool first_push = false;
    for (int i = 1; i < input.size(); i++) {
        if (!input[i].is_null()) {
            if (!first_push) {
                first_push = true;
                tmp.str_val = input[i].get_string();
            } else {
                tmp.str_val += input[0].get_string() + input[i].get_string();
            }
        }
    }

    if (!first_push) {
        return ExprValue::Null();
    }

    return tmp;
}

ExprValue ascii(const std::vector<ExprValue>& input) {
    if (input.size() < 1) {
        return ExprValue::Null();
    }

    if (input[0].is_null()) {
        return ExprValue::Null();
    }

    ExprValue tmp(pb::INT32);

    if (input[0].str_val.empty()) {
        tmp._u.int32_val = 0;
    } else {
        tmp._u.int32_val = static_cast<int32_t>(input[0].str_val[0]);
    }

    return tmp;
}

ExprValue strcmp(const std::vector<ExprValue>& input) {
    if (input.size() != 2) {
        return ExprValue::Null();
    }

    ExprValue tmp(pb::INT32);

    int64_t ret = input[0].compare(input[1]);
    if (ret < 0) {
        tmp._u.int32_val = -1;
    } else if (ret > 0) {
        tmp._u.int32_val = 1;
    } else {
        tmp._u.int32_val = 0;
    }

    return tmp;
}

ExprValue insert(const std::vector<ExprValue>& input) {
    if (input.size() != 4) {
        return ExprValue::Null();
    }

    for (auto s : input) {
        if (s.is_null()) {
            return ExprValue::Null();
        }
    }

    int pos = input[1].get_numberic<int>();
    if (pos < 0) {
        return input[0];
    }

    int len = input[2].get_numberic<int>();
    if (len <= 0) {
        return input[0];
    }

    ExprValue tmp(pb::STRING);
    tmp.str_val = input[0].str_val;
    tmp.str_val.replace(pos, len, input[2].str_val);

    return tmp;
}

ExprValue replace(const std::vector<ExprValue>& input) {
    if (input.size() != 3) {
        return ExprValue::Null();
    }

    if (input[0].is_null()) {
        return ExprValue::Null();
    }

    if (input[1].str_val.empty()) {
        return input[0];
    }

    ExprValue tmp(pb::STRING);
    tmp.str_val = input[0].str_val;

    for (auto pos = 0; pos != std::string::npos; pos += input[2].str_val.length()) {
        pos = tmp.str_val.find(input[1].str_val, pos);
        if (pos != std::string::npos) {
            tmp.str_val.replace(pos, input[1].str_val.length(), input[2].str_val);
        } else {
            break;
        }
    }

    return tmp;
}

ExprValue repeat(const std::vector<ExprValue>& input) {
    if (input.size() != 2) {
        return ExprValue::Null();
    }

    if (input[0].is_null()) {
        return ExprValue::Null();
    }

    int len = input[1].get_numberic<int>();
    if (len <= 0) {
        return ExprValue::Null();
    }

    std::string val = input[0].get_string();
    ExprValue tmp(pb::STRING);
    tmp.str_val.reserve(val.size() * len);

    for (int i = 0; i < len; i++) {
        tmp.str_val += val;
    }

    return tmp;
}

ExprValue reverse(const std::vector<ExprValue>& input) {
    if (input.size() != 1) {
        return ExprValue::Null();
    }

    if (input[0].is_null()) {
        return ExprValue::Null();
    }

    ExprValue tmp(pb::STRING);
    tmp.str_val = input[0].get_string();
    std::reverse(tmp.str_val.begin(), tmp.str_val.end());

    return tmp;
}

ExprValue locate(const std::vector<ExprValue>& input) {
    if (input.size() < 2 || input.size() > 3) {
        return ExprValue::Null();
    }

    for (auto s : input) {
        if (s.is_null()) {
            return ExprValue::Null();
        }
    }

    int begin_pos = 0;
    if (input.size() == 3) {
        begin_pos = input[2].get_numberic<int>();
    }

    ExprValue tmp(pb::INT32);
    auto pos = input[1].str_val.find(input[0].str_val, begin_pos);
    if (pos != std::string::npos) {
        tmp._u.int32_val = pos;
    } else {
        tmp._u.int32_val = 0;
    }

    return tmp;
}

ExprValue unix_timestamp(const std::vector<ExprValue>& input) {
    ExprValue tmp(pb::UINT32);
    if (input.size() == 0) {
        tmp._u.uint32_val = time(NULL);
    } else {
        if (input[0].is_null()) {
            return ExprValue::Null();
        }
        ExprValue in = input[0];
        if (in.type == pb::INT64) {
            in.cast_to(pb::STRING);
        }
        tmp._u.uint32_val = in.cast_to(pb::TIMESTAMP)._u.uint32_val;
    }
    return tmp;
}

ExprValue from_unixtime(const std::vector<ExprValue>& input) {
    if (input.size() == 0 || input[0].is_null()) {
        return ExprValue::Null();
    }
    ExprValue tmp(pb::TIMESTAMP);
    tmp._u.uint32_val = input[0].get_numberic<uint32_t>();
    return tmp;
}

ExprValue now(const std::vector<ExprValue>& input) {
    return ExprValue::Now();
}
ExprValue date_format(const std::vector<ExprValue>& input) {
    if (input.size() != 2) {
        return ExprValue::Null();
    }
    for (auto& s : input) {
        if (s.is_null()) {
            return ExprValue::Null();
        }
    }
    ExprValue tmp = input[0];
    time_t t = tmp.cast_to(pb::TIMESTAMP)._u.uint32_val;
    struct tm t_result;
    localtime_r(&t, &t_result);
    char s[DATE_FORMAT_LENGTH];
    strftime(s, sizeof(s), input[1].str_val.c_str(), &t_result);
    ExprValue format_result(pb::STRING);
    format_result.str_val = s;
    return format_result;
}
ExprValue timediff(const std::vector<ExprValue>& input) {
    if (input.size() < 2) {
        return ExprValue::Null();
    }
    for (auto& s : input) {
        if (s.is_null()) {
            return ExprValue::Null();
        }
    }
    ExprValue arg1 = input[0];
    ExprValue arg2 = input[1];
    int32_t seconds = arg1.cast_to(pb::TIMESTAMP)._u.uint32_val -
        arg2.cast_to(pb::TIMESTAMP)._u.uint32_val;
    ExprValue ret(pb::TIME);
    ret._u.int32_val = seconds_to_time(seconds);
    return ret;
}
ExprValue timestampdiff(const std::vector<ExprValue>& input) {
    if (input.size() < 3) {
        return ExprValue::Null();
    }
    for (auto& s : input) {
        if (s.is_null()) {
            return ExprValue::Null();
        }
    }

    ExprValue arg2 = input[1];
    ExprValue arg3 = input[2];
    int32_t seconds = arg3.cast_to(pb::TIMESTAMP)._u.uint32_val -
        arg2.cast_to(pb::TIMESTAMP)._u.uint32_val;
    ExprValue ret(pb::INT64);
    if (input[0].str_val == "second") {
        ret._u.int64_val = seconds;
    } else if (input[0].str_val == "minute") {
        ret._u.int64_val = seconds / 60;
    } else if (input[0].str_val == "hour") {
        ret._u.int64_val = seconds / 3600;
    } else if (input[0].str_val == "day") {
        ret._u.int64_val = seconds / (24 * 3600);
    } else {
        // un-support
        return ExprValue::Null();
    }
    return ret;
}

ExprValue curdate(const std::vector<ExprValue>& input) {
    ExprValue tmp(pb::DATE);
    uint64_t datetime = timestamp_to_datetime(time(NULL));
    tmp._u.uint32_val = (datetime >> 41) & 0x3FFFFF;
    return tmp;
}
ExprValue current_date(const std::vector<ExprValue>& input) {
    return curdate(input);
}
ExprValue curtime(const std::vector<ExprValue>& input) {
    ExprValue tmp(pb::TIME);
    uint64_t datetime = timestamp_to_datetime(time(NULL));
    tmp._u.int32_val = datetime_to_time(datetime);
    return tmp;
}
ExprValue current_time(const std::vector<ExprValue>& input) {
    return curtime(input);
}
ExprValue day(const std::vector<ExprValue>& input) {
    if (input.size() == 0 || input[0].is_null()) {
        return ExprValue::Null();
    }
    ExprValue in = input[0];
    if (in.type == pb::INT64) {
        in.cast_to(pb::STRING);
    }
    ExprValue tmp(pb::UINT32);
    time_t t = in.cast_to(pb::TIMESTAMP)._u.uint32_val;
    struct tm tm;
    localtime_r(&t, &tm);
    tmp._u.uint32_val = tm.tm_mday;
    return tmp;
}
ExprValue dayname(const std::vector<ExprValue>& input) {
    if (input.size() == 0 || input[0].is_null()) {
        return ExprValue::Null();
    }
    ExprValue tmp(pb::STRING);
    uint32_t week_num = dayofweek(input)._u.uint32_val;
    if (week_num <= day_names.size()) {
        tmp.str_val = day_names[week_num - 1];
    }
    return tmp;
}
ExprValue monthname(const std::vector<ExprValue>& input) {
    if (input.size() == 0 || input[0].is_null()) {
        return ExprValue::Null();
    }
    uint32_t month_num = month(input)._u.uint32_val;
    ExprValue tmp(pb::STRING);
    if (month_num <= month_names.size()) {
        tmp.str_val = month_names[month_num - 1];
    }
    return tmp;
}
ExprValue dayofweek(const std::vector<ExprValue>& input) {
    if (input.size() == 0 || input[0].is_null()) {
        return ExprValue::Null();
    }
    ExprValue in = input[0];
    if (in.type == pb::INT64) {
        in.cast_to(pb::STRING);
    }
    ExprValue tmp(pb::UINT32);
    time_t t = in.cast_to(pb::TIMESTAMP)._u.uint32_val;
    struct tm tm;
    localtime_r(&t, &tm);
    boost::gregorian::date today(tm.tm_year + 1900, ++tm.tm_mon, tm.tm_mday);
    /*
      DAYOFWEEK(d) 函数返回 d 对应的一周中的索引（位置）。1 表示周日，2 表示周一，……，7 表示周六
    */
    tmp._u.uint32_val = today.day_of_week() + 1;
    return tmp;
}
ExprValue dayofmonth(const std::vector<ExprValue>& input) {
    return day(input);
}
ExprValue month(const std::vector<ExprValue>& input) {
    if (input.size() == 0 || input[0].is_null()) {
        return ExprValue::Null();
    }
    ExprValue in = input[0];
    if (in.type == pb::INT64) {
        in.cast_to(pb::STRING);
    }
    ExprValue tmp(pb::UINT32);
    time_t t = in.cast_to(pb::TIMESTAMP)._u.uint32_val;
    struct tm tm;
    localtime_r(&t, &tm);
    tmp._u.uint32_val = ++tm.tm_mon;
    return tmp;
}
ExprValue year(const std::vector<ExprValue>& input) {
    if (input.size() == 0 || input[0].is_null()) {
        return ExprValue::Null();
    }
    ExprValue in = input[0];
    if (in.type == pb::INT64) {
        in.cast_to(pb::STRING);
    }
    ExprValue tmp(pb::UINT32);
    time_t t = in.cast_to(pb::TIMESTAMP)._u.uint32_val;
    struct tm tm;
    localtime_r(&t, &tm);
    tmp._u.uint32_val = tm.tm_year + 1900;
    return tmp;
}
ExprValue time_to_sec(const std::vector<ExprValue>& input) {
    if (input.size() == 0 || input[0].is_null()) {
        return ExprValue::Null();
    }
    ExprValue tmp(pb::INT32);
    ExprValue in = input[0];
    time_t time = in.cast_to(pb::TIME)._u.int32_val;
    bool minus = false;
    if (time < 0) {
        minus = true;
        time = -time;
    }
    uint32_t hour = (time >> 12) & 0x3FF;
    uint32_t min = (time >> 6) & 0x3F;
    uint32_t sec = time & 0x3F;
    uint32_t sec_sum = hour * 3600 + min * 60 + sec;
    if (!minus) {
        tmp._u.int32_val = sec_sum;
    } else {
        tmp._u.int32_val = -sec_sum;
    }
    return tmp;
}
ExprValue sec_to_time(const std::vector<ExprValue>& input) {
    if (input.size() == 0 || input[0].is_null()) {
        return ExprValue::Null();
    }
    ExprValue in = input[0];
    if (in.type == pb::STRING) {
        in.cast_to(pb::INT32);
    }
    int32_t secs = in._u.int32_val;
    bool minus = false;
    if (secs < 0) {
        minus = true;
        secs = - secs;
    }
    ExprValue tmp(pb::TIME);

    uint32_t hour = secs / 3600;
    uint32_t min = (secs - hour * 3600) / 60;
    uint32_t sec = secs % 60;
    int32_t time = 0;
    time |= sec;
    time |= (min << 6);
    time |= (hour << 12);
    if (minus) {
        time = -time;
    }
    tmp._u.int32_val = time;
    return tmp;
}
ExprValue dayofyear(const std::vector<ExprValue>& input) {
    if (input.size() == 0 || input[0].is_null()) {
        return ExprValue::Null();
    }
    ExprValue in = input[0];
    if (in.type == pb::INT64) {
        in.cast_to(pb::STRING);
    }
    ExprValue tmp(pb::UINT32);
    time_t t = in.cast_to(pb::TIMESTAMP)._u.uint32_val;
    struct tm tm;
    localtime_r(&t, &tm);
    boost::gregorian::date today(tm.tm_year += 1900, ++tm.tm_mon, tm.tm_mday);
    tmp._u.uint32_val = today.day_of_year();
    return tmp;
}
ExprValue weekday(const std::vector<ExprValue>& input) {
    if (input.size() == 0 || input[0].is_null()) {
        return ExprValue::Null();
    }
    ExprValue in = input[0];
    if (in.type == pb::INT64) {
        in.cast_to(pb::STRING);
    }
    ExprValue one = input[0];
    time_t t = one.cast_to(pb::TIMESTAMP)._u.uint32_val;
    struct tm tm;
    localtime_r(&t, &tm);
    boost::gregorian::date today(tm.tm_year + 1900, ++tm.tm_mon, tm.tm_mday);
    ExprValue tmp(pb::UINT32);
    uint32_t day_of_week = today.day_of_week();
    if (day_of_week >= 1) {
        tmp._u.uint32_val = day_of_week - 1;
    } else {
        tmp._u.uint32_val = 6;
    }
    return tmp;
}
ExprValue week(const std::vector<ExprValue>& input) {
    if (input.size() == 0 || input[0].is_null()) {
        return ExprValue::Null();
    }
    ExprValue one = input[0];
    time_t t = one.cast_to(pb::TIMESTAMP)._u.uint32_val;
    struct tm tm;
    localtime_r(&t, &tm);
    boost::gregorian::date today(tm.tm_year += 1900, ++tm.tm_mon, tm.tm_mday);
    uint32_t week_number = today.week_number();
    ExprValue tmp(pb::UINT32);
    if (input.size() > 1) {
        ExprValue two = input[1];
        uint32_t mode = two.cast_to(pb::UINT32)._u.uint32_val;
        if (!mode) {
            week_number -= 1;
        }
    }
    tmp._u.uint32_val = week_number;
    return tmp;
}

ExprValue datediff(const std::vector<ExprValue>& input) {
    if (input.size() == 0 || input[0].is_null() || input[1].is_null()) {
        return ExprValue::Null();
    }
    ExprValue left = input[0];
    if (left.type == pb::INT64) {
        left.cast_to(pb::STRING);
    }
    ExprValue right = input[1];
    if (right.type == pb::INT64) {
        right.cast_to(pb::STRING);
    }
    time_t t1 = left.cast_to(pb::TIMESTAMP)._u.uint32_val;
    time_t t2 = right.cast_to(pb::TIMESTAMP)._u.uint32_val;
    ExprValue tmp(pb::INT32);
    tmp._u.int32_val = (t1 - t2) / (3600 * 24);
    return tmp;
}

ExprValue hll_add(const std::vector<ExprValue>& input) {
    if (input.size() == 0) {
        return ExprValue::Null();
    }
    if (!input[0].is_hll()) {
        (ExprValue&)input[0] = hll::hll_init();
    }
    for (size_t i = 1; i < input.size(); i++) {
        if (!input[i].is_null()) {
            hll::hll_add((ExprValue&)input[0], input[i].hash());
        }
    }
    return input[0];
}

ExprValue hll_init(const std::vector<ExprValue>& input) {
    if (input.size() == 0) {
        return ExprValue::Null();
    }
    ExprValue hll_init = hll::hll_init();
    for (size_t i = 0; i < input.size(); i++) {
        if (!input[i].is_null()) {
            hll::hll_add(hll_init, input[i].hash());
        }
    }
    return hll_init;
}

ExprValue hll_merge(const std::vector<ExprValue>& input) {
    if (input.size() == 0) {
        return ExprValue::Null();
    }
    if (!input[0].is_hll()) {
        (ExprValue&)input[0] = hll::hll_init();
    }
    for (size_t i = 1; i < input.size(); i++) {
        if (input[i].is_hll()) {
            hll::hll_merge((ExprValue&)input[0], (ExprValue&)input[i]);
        }
    }
    return input[0];
}

ExprValue hll_estimate(const std::vector<ExprValue>& input) {
    if (input.size() == 0) {
        return ExprValue::Null();
    }
    ExprValue tmp(pb::INT64);
    if (input[0].is_null()) {
        tmp._u.int64_val = 0;
        return tmp;
    }
    tmp._u.int64_val = hll::hll_estimate(input[0]);
    return tmp;
}

ExprValue case_when(const std::vector<ExprValue>& input) {
    for (size_t i = 0; i < input.size() / 2; ++i) {
        auto if_index = i * 2;
        auto then_index = i * 2 + 1;
        if (input[if_index].get_numberic<bool>()) {
            return input[then_index];
        }
    }
    //没有else分支, 返回null
    if (input.size() % 2 == 0) {
        return ExprValue();
     } else {
        return input[input.size() - 1];
    }
}

ExprValue case_expr_when(const std::vector<ExprValue>& input) {
    for (size_t i = 0; i < (input.size() - 1) / 2; ++i) {
        auto if_index = i * 2 + 1;
        auto then_index = i * 2 + 2;
        if (input[0].compare(input[if_index]) == 0) {
            return input[then_index];
        }
    }
    //没有else分支, 返回null
    if ((input.size() - 1) % 2 == 0) {
        return ExprValue();
     } else {
        return input[input.size() - 1];
    }
}

ExprValue if_(const std::vector<ExprValue>& input) {
    if (input.size() != 3) {
        return ExprValue::Null();
    }
    return input[0].get_numberic<bool>() ? input[1] : input[2];
}

ExprValue ifnull(const std::vector<ExprValue>& input) {
    if (input.size() != 2) {
        return ExprValue::Null();
    }
    return input[0].is_null() ? input[1] : input[0];
}

ExprValue nullif(const std::vector<ExprValue>& input) {
    if (input.size() != 2) {
        return ExprValue::Null();
    }
    ExprValue arg1 = input[0];
    ExprValue arg2 = input[1];
    if (arg1.compare_diff_type(arg2) == 0) {
        return ExprValue::Null();
    } else {
        return input[0];
    }
}
ExprValue murmur_hash(const std::vector<ExprValue>& input) {
    if (input.size() == 0) {
        return ExprValue::Null();
    }
    ExprValue tmp(pb::UINT64);
    if (input.size() == 0) {
        tmp._u.uint64_val = 0;
    } else {
        tmp._u.uint64_val = make_sign(input[0].str_val);
    }
    return tmp;
}

ExprValue md5(const std::vector<ExprValue>& input) {
    if (input.size() == 0) {
        return ExprValue::Null();
    }
    ExprValue orig = input[0];
    if (orig.type != pb::STRING) {
        orig.cast_to(pb::STRING);
    }
    ExprValue tmp(pb::STRING);
    unsigned char md5_str[16] = {0};
    MD5_CTX ctx;
    MD5_Init(&ctx);
    MD5_Update(&ctx, orig.str_val.c_str(), orig.str_val.size());
    MD5_Final(md5_str, &ctx);
    tmp.str_val.resize(32);

    int j = 0;
    static char const zEncode[] = "0123456789abcdef";
    for (int i = 0; i < 16; i ++) {
        int a = md5_str[i];
        tmp.str_val[j++] = zEncode[(a >> 4) & 0xf];
        tmp.str_val[j++] = zEncode[a & 0xf];
    }
    return tmp;
}

ExprValue sha1(const std::vector<ExprValue>& input) {
    if (input.size() == 0) {
        return ExprValue::Null();
    }
    ExprValue orig = input[0];
    if (orig.type != pb::STRING) {
        orig.cast_to(pb::STRING);
    }
    ExprValue tmp(pb::STRING);
    tmp.str_val = butil::SHA1HashString(orig.str_val);
    return tmp;
}

ExprValue sha(const std::vector<ExprValue>& input) {
    return sha1(input);
}
}

/* vim: set ts=4 sw=4 sts=4 tw=100 */


#ifndef MATRIX_HPP
#define MATRIX_HPP
#include <Eigen/Dense>
#include "precision.hpp"
typedef  Eigen::Matrix<double,4,4> Matrix4;

#endif


//
// Created by kerin on 2019-01-08.
//

#ifndef BGEN_PROG_PVE_HPP
#define BGEN_PROG_PVE_HPP

#include "genotype_matrix.hpp"
#include "file_utils.hpp"
#include "parameters.hpp"
#include "eigen_utils.hpp"
#include "mpi_utils.hpp"
#include "data.hpp"
#include "eigen_utils.hpp"
#include "rhe_reg_component.hpp"

#include "tools/eigen3.3/Dense"

#include <boost/iostreams/filtering_stream.hpp>

#include <random>

namespace boost_io = boost::iostreams;

class RHEreg {
public:
	// constants
	long n_draws;
	long n_samples;
	long n_components;
	long n_covar;
	long n_env;
	long n_var;
	double Nglobal;
	int world_rank, world_size;
	std::vector<std::string> env_names;

	std::map<long, int> sample_location;

	std::chrono::duration<double> elapsed_streamBgen, elapsed_solveRHE, elapsed_solveJacknife;
	double elapsed_LM, elapsed_LM_copyData;

	// RHEreg-NLS
	Eigen::VectorXd nls_env_weights;
	std::vector<Eigen::MatrixXd> ytEXXtEys;

	const parameters& p;
	const GenotypeMatrix& X;

	Eigen::VectorXd eta;
	Eigen::VectorXd Y;
	Eigen::MatrixXd& C, E;
	Eigen::MatrixXd CtC_inv;
	Eigen::ArrayXd sigmas;
	Eigen::ArrayXXd sigmas_jack, h2_jack;
	Eigen::ArrayXd h2, h2_se_jack, h2_bias_corrected;
	Eigen::ArrayXd n_var_jack;

	EigenDataMatrix zz;
	const std::unordered_map<long, bool>& sample_is_invalid;
	Data& data;

	// std::vector<std::string> components;
	std::vector<RHEreg_Component> components;
	std::vector<std::string> SNPGROUPS_snpid, SNPGROUPS_group, all_SNPGROUPS;

	RHEreg(Data& dat,
	       Eigen::VectorXd& myY,
	       Eigen::MatrixXd& myC,
	       Eigen::MatrixXd myE) : p(dat.p), X(dat.G), E(myE), Y(myY), C(myC),
		sample_is_invalid(dat.sample_is_invalid),
		sample_location(dat.sample_location), data(dat),
		env_names(dat.env_names) {
		n_samples = data.n_samples;
		Nglobal = mpiUtils::mpiReduce_inplace(n_samples);
		n_draws = p.n_pve_samples;

		n_covar = C.cols();
		n_env = E.cols();

		MPI_Comm_rank(MPI_COMM_WORLD, &world_rank);
		MPI_Comm_size(MPI_COMM_WORLD, &world_size);

		elapsed_LM = 0;
		elapsed_LM_copyData = 0;
	}

	RHEreg(Data& dat,
	       Eigen::VectorXd& myY,
	       Eigen::MatrixXd& myC) : p(dat.p), X(dat.G), Y(myY), C(myC),
		sample_is_invalid(dat.sample_is_invalid),
		sample_location(dat.sample_location), data(dat),
		env_names(dat.env_names)  {
		n_samples = data.n_samples;
		Nglobal = mpiUtils::mpiReduce_inplace(n_samples);
		n_draws = p.n_pve_samples;

		n_covar = C.cols();
		n_env = 0;

		MPI_Comm_rank(MPI_COMM_WORLD, &world_rank);
		MPI_Comm_size(MPI_COMM_WORLD, &world_size);

		elapsed_LM = 0;
		elapsed_LM_copyData = 0;
	}

	void run();

	void standardise_non_genetic_data();

	void initialise_components();

	void compute_RHE_trace_operators();

	void solve_RHE(std::vector<RHEreg_Component>& components);

	Eigen::VectorXd run_RHE_levenburgMarquardt();

	Eigen::VectorXd run_RHE_nelderMead();

	double RHE_nelderMead_obj(Eigen::VectorXd env_weights, void *grad_out) const;

	Eigen::MatrixXd construct_vc_system(const std::vector<RHEreg_Component>& vec_of_components) const;

	Eigen::ArrayXd calc_h2(const Eigen::Ref<const Eigen::MatrixXd>& AA,
	                       const Eigen::Ref<const Eigen::VectorXd>& bb,
	                       const bool &reweight_sigmas) const;

	void process_jacknife_samples();

	Eigen::MatrixXd project_out_covars(Eigen::Ref<Eigen::MatrixXd> rhs);

	void fill_gaussian_noise(unsigned int seed,
	                         Eigen::Ref<Eigen::MatrixXd> zz,
	                         long nn,
	                         long n_draws);

	void to_file(const std::string& file);

	void read_RHE_groups(const std::string& filename);

	double get_jacknife_var(Eigen::ArrayXd jack_estimates);

	double get_jacknife_bias_correct(Eigen::ArrayXd jack_estimates, double full_data_est);
};

#endif //BGEN_PROG_PVE_HPP


//  (C) Copyright Raffi Enficiaud 2014.
//  Distributed under the Boost Software License, Version 1.0.
//  (See accompanying file LICENSE_1_0.txt or copy at
//  http://www.boost.org/LICENSE_1_0.txt)

//  See http://www.boost.org/libs/test for the library home page.

//[example_code
#define BOOST_TEST_MODULE dataset_example59
#include <boost/test/included/unit_test.hpp>
#include <boost/test/data/test_case.hpp>
#include <boost/test/data/monomorphic.hpp>

namespace bdata = boost::unit_test::data;

BOOST_DATA_TEST_CASE( test1, bdata::xrange(5) )
{
  std::cout << "test 1: " << sample << std::endl;
  BOOST_TEST((sample <= 4 && sample >= 0));
}

BOOST_DATA_TEST_CASE(
  test2,
  bdata::xrange<int>( (bdata::begin=1, bdata::end=10, bdata::step=3)) )
{
  std::cout << "test 2: " << sample << std::endl;
  BOOST_TEST((sample <= 4 && sample >= 0));
}
//]