//=============================================================================================================
/**
* @file     hpifit.cpp
* @author   Lorenz Esch <Lorenz.Esch@tu-ilmenau.de>;
*           Matti Hamalainen <msh@nmr.mgh.harvard.edu>
* @version  1.0
* @date     March, 2017
*
* @section  LICENSE
*
* Copyright (C) 2017, Lorenz Esch and Matti Hamalainen. All rights reserved.
*
* Redistribution and use in source and binary forms, with or without modification, are permitted provided that
* the following conditions are met:
*     * Redistributions of source code must retain the above copyright notice, this list of conditions and the
*       following disclaimer.
*     * Redistributions in binary form must reproduce the above copyright notice, this list of conditions and
*       the following disclaimer in the documentation and/or other materials provided with the distribution.
*     * Neither the name of MNE-CPP authors nor the names of its contributors may be used
*       to endorse or promote products derived from this software without specific prior written permission.
*
* THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND ANY EXPRESS OR IMPLIED
* WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A
* PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT,
* INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,
* PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION)
* HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING
* NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE
* POSSIBILITY OF SUCH DAMAGE.
*
*
* @brief    HPIFit class defintion.
*
*/

//*************************************************************************************************************
//=============================================================================================================
// INCLUDES
//=============================================================================================================

#include "hpifit.h"

#include <fiff/fiff_dig_point_set.h>

#include <utils/ioutils.h>

#include <iostream>
#include <fiff/fiff_cov.h>
#include <fstream>

//*************************************************************************************************************
//=============================================================================================================
// Eigen INCLUDES
//=============================================================================================================

#include <Eigen/SVD>
#include <Eigen/Dense>

//*************************************************************************************************************
//=============================================================================================================
// QT INCLUDES
//=============================================================================================================

#include <QFuture>
#include <QtConcurrent/QtConcurrent>

//*************************************************************************************************************
//=============================================================================================================
// USED NAMESPACES
//=============================================================================================================

using namespace Eigen;
using namespace INVERSELIB;
using namespace FIFFLIB;

//*************************************************************************************************************
//=============================================================================================================
// DEFINE GLOBAL METHODS
//=============================================================================================================

Eigen::MatrixXd pinv(Eigen::MatrixXd a)
{
    double epsilon = std::numeric_limits<double>::epsilon();
    Eigen::JacobiSVD< Eigen::MatrixXd > svd(a ,Eigen::ComputeThinU | Eigen::ComputeThinV);
    double tolerance = epsilon * std::max(a.cols(), a.rows()) * svd.singularValues().array().abs()(0);
    return svd.matrixV() * (svd.singularValues().array().abs() > tolerance).select(svd.singularValues().array().inverse(),0).matrix().asDiagonal() * svd.matrixU().adjoint();
}

/*********************************************************************************
 * magnetic_dipole leadfield for a magnetic dipole in an infinite medium
 * The function has been compared with matlab magnetic_dipole and it gives same output
 *********************************************************************************/

Eigen::MatrixXd magnetic_dipole(Eigen::MatrixXd pos, Eigen::MatrixXd pnt, Eigen::MatrixXd ori) {

    double u0 = 1e-7;
    int nchan;
    Eigen::MatrixXd r, r2, r5, x, y, z, mx, my, mz, Tx, Ty, Tz, lf;

    nchan = pnt.rows();

    // Shift the magnetometers so that the dipole is in the origin
    pnt.array().col(0) -= pos(0);
    pnt.array().col(1) -= pos(1);
    pnt.array().col(2) -= pos(2);

    r = pnt.array().square().rowwise().sum().sqrt();

    r2 = r5 = x = y = z = mx = my = mz = Tx = Ty = Tz = lf = Eigen::MatrixXd::Zero(nchan,3);

    for(int i = 0;i < nchan;i++) {
        r2.row(i).array().fill(pow(r(i),2));
        r5.row(i).array().fill(pow(r(i),5));
    }

    for(int i = 0;i < nchan;i++) {
        x.row(i).array().fill(pnt(i,0));
        y.row(i).array().fill(pnt(i,1));
        z.row(i).array().fill(pnt(i,2));
    }

    mx.col(0).array().fill(1);
    my.col(1).array().fill(1);
    mz.col(2).array().fill(1);

    Tx = 3 * x.cwiseProduct(pnt) - mx.cwiseProduct(r2);
    Ty = 3 * y.cwiseProduct(pnt) - my.cwiseProduct(r2);
    Tz = 3 * z.cwiseProduct(pnt) - mz.cwiseProduct(r2);

    for(int i = 0;i < nchan;i++) {
        lf(i,0) = Tx.row(i).dot(ori.row(i));
        lf(i,1) = Ty.row(i).dot(ori.row(i));
        lf(i,2) = Tz.row(i).dot(ori.row(i));
    }

    for(int i = 0;i < nchan;i++) {
        for(int j = 0;j < 3;j++) {
            lf(i,j) = u0 * lf(i,j)/(4 * M_PI * r5(i,j));
        }
    }

    return lf;
}

/*********************************************************************************
 * compute_leadfield computes a forward solution for a dipole in a a volume
 * conductor model. The forward solution is expressed as the leadfield
 * matrix (Nchan*3), where each column corresponds with the potential or field
 * distributions on all sensors for one of the x,y,z-orientations of the dipole.
 * The function has been compared with matlab ft_compute_leadfield and it gives
 * same output
 *********************************************************************************/

Eigen::MatrixXd compute_leadfield(const Eigen::MatrixXd& pos, const struct SensorInfo& sensors)
{

    Eigen::MatrixXd pnt, ori, lf;

    pnt = sensors.coilpos; // position of each coil
    ori = sensors.coilori; // orientation of each coil

    lf = magnetic_dipole(pos, pnt, ori);

    lf = sensors.tra * lf;

    return lf;
}

/*********************************************************************************
 * dipfitError computes the error between measured and model data
 * and can be used for non-linear fitting of dipole position.
 * The function has been compared with matlab dipfit_error and it gives
 * same output
 *********************************************************************************/

DipFitError dipfitError(const Eigen::MatrixXd& pos, const Eigen::MatrixXd& data, const struct SensorInfo& sensors, const Eigen::MatrixXd& matProjectors)
{
    // Variable Declaration
    struct DipFitError e;
    Eigen::MatrixXd lf, dif;

    // Compute lead field for a magnetic dipole in infinite vacuum
    lf = compute_leadfield(pos, sensors);

    e.moment = pinv(lf) * data;

    //dif = data - lf * e.moment;
    dif = data - matProjectors * lf * e.moment;

    e.error = dif.array().square().sum()/data.array().square().sum();

    e.numIterations = 0;

    return e;
}

/*********************************************************************************
 * Compare function for sorting
 *********************************************************************************/

bool compare(HPISortStruct a, HPISortStruct b)
{
    return (a.base_arr < b.base_arr);
}

/*********************************************************************************
 * fminsearch Multidimensional unconstrained nonlinear minimization (Nelder-Mead).
 * X = fminsearch(X0, maxiter, maxfun, display, data, sensors) starts at X0 and
 * attempts to find a local minimizer
 *********************************************************************************/

Eigen::MatrixXd fminsearch(const Eigen::MatrixXd& pos,
                           int maxiter,
                           int maxfun,
                           int display,
                           const Eigen::MatrixXd& data,
                           const Eigen::MatrixXd& matProjectors,
                           const struct SensorInfo& sensors,
                           int &simplex_numitr)
{
    double tolx, tolf, rho, chi, psi, sigma, func_evals, usual_delta, zero_term_delta, temp1, temp2;
    std::string header, how;
    int n, itercount, prnt;
    Eigen::MatrixXd onesn, two2np1, one2n, v, y, v1, tempX1, tempX2, xbar, xr, x, xe, xc, xcc, xin, posCopy;
    std::vector <double> fv, fv1;
    std::vector <int> idx;

    DipFitError tempdip, fxr, fxe, fxc, fxcc;

    //tolx = tolf = 1e-4;
    // Seok
    tolx = tolf = 1e-9;

    switch(display) {
        case 0:
            prnt = 0;
            break;
        default:
            prnt = 1;
    }

    header = " Iteration   Func-count     min f(x) Procedure";

    posCopy = pos;

    n = posCopy.cols();

    // Initialize parameters
    rho = 1; chi = 2; psi = 0.5; sigma = 0.5;
    onesn = Eigen::MatrixXd::Ones(1,n);
    two2np1 = one2n = Eigen::MatrixXd::Zero(1,n);

    for(int i = 0;i < n;i++) {
        two2np1(i) = 1 + i;
        one2n(i) = i;
    }

    v = v1 = Eigen::MatrixXd::Zero(n, n+1);
    fv.resize(n+1);
    idx.resize(n+1);
    fv1.resize(n+1);

    for(int i = 0;i < n; i++) {
        v(i,0) = posCopy(i);
    }

    tempdip = dipfitError(posCopy, data, sensors, matProjectors);
    fv[0] = tempdip.error;

    func_evals = 1;
    itercount = 0;
    how = "";

    // Continue setting up the initial simplex.
    // Following improvement suggested by L.Pfeffer at Stanford
    usual_delta = 0.05;             // 5 percent deltas for non-zero terms
    zero_term_delta = 0.00025;      // Even smaller delta for zero elements of x
    xin = posCopy.transpose();

    for(int j = 0;j < n;j++) {
        y = xin;

        if(y(j) != 0) {
            y(j) = (1 + usual_delta) * y(j);
        } else {
            y(j) = zero_term_delta;
        }

        v.col(j+1).array() = y;
        posCopy = y.transpose();
        tempdip = dipfitError(posCopy, data, sensors, matProjectors);
        fv[j+1] = tempdip.error;
    }

    // Sort elements of fv
    std::vector<HPISortStruct> vecSortStruct;

    for (int i = 0; i < fv.size(); i++) {
        HPISortStruct structTemp;
        structTemp.base_arr = fv[i];
        structTemp.idx = i;
        vecSortStruct.push_back(structTemp);
    }

    sort (vecSortStruct.begin(), vecSortStruct.end(), compare);

    for (int i = 0; i < vecSortStruct.size(); i++) {
        idx[i] = vecSortStruct[i].idx;
    }

    for (int i = 0;i < n+1;i++) {
        v1.col(i) = v.col(idx[i]);
        fv1[i] = fv[idx[i]];
    }

    v = v1;fv = fv1;

    how = "initial simplex";
    itercount = itercount + 1;
    func_evals = n + 1;

    tempX1 = Eigen::MatrixXd::Zero(1,n);

    while ((func_evals < maxfun) && (itercount < maxiter)) {

        for (int i = 0;i < n;i++) {
            tempX1(i) = abs(fv[0] - fv[i+1]);
        }

        temp1 = tempX1.maxCoeff();

        tempX2 = Eigen::MatrixXd::Zero(n,n);

        for(int i = 0;i < n;i++) {
            tempX2.col(i) = v.col(i+1) -  v.col(0);
        }

        tempX2 = tempX2.array().abs();

        temp2 = tempX2.maxCoeff();

        if((temp1 <= tolf) && (temp2 <= tolx)) {
            break;
        }

        xbar = v.block(0,0,n,n).rowwise().sum();
        xbar /= n;

        xr = (1+rho) * xbar - rho * v.block(0,n,v.rows(),1);

        x = xr.transpose();
        //std::cout << "Iteration Count: " << itercount << ":" << x << std::endl;

        fxr = dipfitError(x, data, sensors, matProjectors);

        func_evals = func_evals+1;

        if (fxr.error < fv[0]) {
            // Calculate the expansion point
            xe = (1 + rho * chi) * xbar - rho * chi * v.col(v.cols()-1);
            x = xe.transpose();
            fxe = dipfitError(x, data, sensors, matProjectors);
            func_evals = func_evals+1;

            if(fxe.error < fxr.error) {
                v.col(v.cols()-1) = xe;
                fv[n] = fxe.error;
                how = "expand";
            } else {
                v.col(v.cols()-1) = xr;
                fv[n] = fxr.error;
                how = "reflect";
            }
        }
        else {
            if(fxr.error < fv[n-1]) {
                v.col(v.cols()-1) = xr;
                fv[n] = fxr.error;
                how = "reflect";
            } else { // fxr.error >= fv[:,n-1]
                // Perform contraction
                if(fxr.error < fv[n]) {
                    // Perform an outside contraction
                    xc = (1 + psi * rho) * xbar - psi * rho * v.col(v.cols()-1);
                    x = xc.transpose();
                    fxc = dipfitError(x, data, sensors, matProjectors);
                    func_evals = func_evals + 1;

                    if(fxc.error <= fxr.error) {
                        v.col(v.cols()-1) = xc;
                        fv[n] = fxc.error;
                        how = "contract outside";
                    } else {
                        // perform a shrink
                        how = "shrink";
                    }
                } else {
                    xcc = (1 - psi) * xbar + psi * v.col(v.cols()-1);
                    x = xcc.transpose();
                    fxcc = dipfitError(x, data, sensors, matProjectors);
                    func_evals = func_evals+1;
                    if(fxcc.error < fv[n]) {
                        v.col(v.cols()-1) = xcc;
                        fv[n] = fxcc.error;
                        how = "contract inside";
                    } else {
                        // perform a shrink
                        how = "shrink";
                    }
                }

                if(how.compare("shrink") == 0) {
                    for(int j = 1;j < n+1;j++) {
                        v.col(j).array() = v.col(0).array() + sigma * (v.col(j).array() - v.col(0).array());
                        x = v.col(j).array().transpose();
                        tempdip = dipfitError(x,data, sensors, matProjectors);
                        fv[j] = tempdip.error;
                    }
                }
            }
        }

        // Sort elements of fv
        vecSortStruct.clear();

        for (int i = 0; i < fv.size(); i++) {
            HPISortStruct structTemp;
            structTemp.base_arr = fv[i];
            structTemp.idx = i;
            vecSortStruct.push_back(structTemp);
        }

        sort (vecSortStruct.begin(), vecSortStruct.end(), compare);

        for (int i = 0; i < vecSortStruct.size(); i++) {
            idx[i] = vecSortStruct[i].idx;
        }

        for (int i = 0;i < n+1;i++) {
            v1.col(i) = v.col(idx[i]);
            fv1[i] = fv[idx[i]];
        }

        v = v1;
        fv = fv1;
        itercount = itercount + 1;
    }

    x = v.col(0).transpose();

    // Seok
    simplex_numitr = itercount;

    return x;
}

/*********************************************************************************
 * dipfit function is adapted from Fieldtrip Software. It has been
 * heavily edited for use with MNE Scan Software
 *********************************************************************************/

void doDipfitConcurrent(FittingCoilData& lCoilData)
{
    // Initialize variables
    Eigen::RowVectorXd currentCoil = lCoilData.coilPos;
    Eigen::VectorXd currentData = lCoilData.sensorData;
    SensorInfo currentSensors = lCoilData.sensorPos;

    int display = 0;
    int maxiter = 500;
    int simplex_numitr = 0;

    lCoilData.coilPos = fminsearch(currentCoil,
                                       maxiter,
                                       2 * maxiter * currentCoil.cols(),
                                       display,
                                       currentData,
                                       lCoilData.matProjector,
                                       currentSensors,
                                       simplex_numitr);

    lCoilData.errorInfo = dipfitError(currentCoil, currentData, currentSensors, lCoilData.matProjector);
    lCoilData.errorInfo.numIterations = simplex_numitr;
}

//*************************************************************************************************************
//=============================================================================================================
// DEFINE MEMBER METHODS
//=============================================================================================================

HPIFit::HPIFit()
{

}

//*************************************************************************************************************

void HPIFit::fitHPI(const MatrixXd& t_mat,
                        const Eigen::MatrixXd& t_matProjectors,
                        FiffCoordTrans& transDevHead,
                        const QVector<int>& vFreqs,
                        QVector<double>& vGof,
                        FiffDigPointSet& fittedPointSet,
                        FiffInfo::SPtr pFiffInfo,
                        bool bDoDebug,
                        const QString& sHPIResourceDir)
{
    //Check if data was passed
    if(t_mat.rows() == 0 || t_mat.cols() == 0 ) {
        std::cout<<std::endl<< "HPIFit::fitHPI - No data passed. Returning.";
    }

    //Check if projector was passed
    if(t_matProjectors.rows() == 0 || t_matProjectors.cols() == 0 ) {
        std::cout<<std::endl<< "HPIFit::fitHPI - No projector passed. Returning.";
    }

    vGof.clear();

    struct SensorInfo sensors;
    struct CoilParam coil;
    int numCh = pFiffInfo->nchan;
    int samF = pFiffInfo->sfreq;
    int samLoc = t_mat.cols(); // minimum samples required to localize numLoc times in a second

    //Get HPI coils from digitizers and set number of coils
    int numCoils = 0;
    QList<FiffDigPoint> lHPIPoints;

    for(int i = 0; i < pFiffInfo->dig.size(); ++i) {
        if(pFiffInfo->dig[i].kind == FIFFV_POINT_HPI) {
            numCoils++;
            lHPIPoints.append(pFiffInfo->dig[i]);
        }
    }

    //Set coil frequencies
    Eigen::VectorXd coilfreq(numCoils);

    if(vFreqs.size() >= numCoils) {
        for(int i = 0; i < numCoils; ++i) {
            coilfreq[i] = vFreqs.at(i);
            //std::cout<<std::endl << coilfreq[i] << "Hz";
        }
    } else {
        std::cout<<std::endl<< "HPIFit::fitHPI - Not enough coil frequencies specified. Returning.";
        return;
    }

    // Initialize HPI coils location and moment
    coil.pos = Eigen::MatrixXd::Zero(numCoils,3);
    coil.mom = Eigen::MatrixXd::Zero(numCoils,3);
    coil.dpfiterror = Eigen::VectorXd::Zero(numCoils);
    coil.dpfitnumitr = Eigen::VectorXd::Zero(numCoils);

    // Generate simulated data
    Eigen::MatrixXd simsig(samLoc,numCoils*2);
    Eigen::VectorXd time(samLoc);

    for (int i = 0; i < samLoc; ++i) {
        time[i] = i*1.0/samF;
    }

    for(int i = 0; i < numCoils; ++i) {
        for(int j = 0; j < samLoc; ++j) {
            simsig(j,i) = sin(2*M_PI*coilfreq[i]*time[j]);
            simsig(j,i+numCoils) = cos(2*M_PI*coilfreq[i]*time[j]);
        }
    }

    // Create digitized HPI coil position matrix
    Eigen::MatrixXd headHPI(numCoils,3);

    // check the pFiffInfo->dig information. If dig is empty, set the headHPI is 0;
    if (lHPIPoints.size() > 0) {
        for (int i = 0; i < lHPIPoints.size(); ++i) {
            headHPI(i,0) = lHPIPoints.at(i).r[0];
            headHPI(i,1) = lHPIPoints.at(i).r[1];
            headHPI(i,2) = lHPIPoints.at(i).r[2];
        }
    } else {
        for (int i = 0; i < numCoils; ++i) {
            headHPI(i,0) = 0;
            headHPI(i,1) = 0;
            headHPI(i,2) = 0;
        }
    }

    // Get the indices of inner layer channels and exclude bad channels.
    //TODO: Only supports babymeg and vectorview gradiometeres for hpi fitting.
    QVector<int> innerind(0);

    for (int i = 0; i < numCh; ++i) {
        if(pFiffInfo->chs[i].chpos.coil_type == FIFFV_COIL_BABY_MAG ||
                pFiffInfo->chs[i].chpos.coil_type == FIFFV_COIL_VV_PLANAR_T1 ||
                pFiffInfo->chs[i].chpos.coil_type == FIFFV_COIL_VV_PLANAR_T2 ||
                pFiffInfo->chs[i].chpos.coil_type == FIFFV_COIL_VV_PLANAR_T3) {
            // Check if the sensor is bad, if not append to innerind
            if(!(pFiffInfo->bads.contains(pFiffInfo->ch_names.at(i)))) {
                innerind.append(i);
            }
        }
    }

    //Create new projector based on the excluded channels, first exclude the rows then the columns
    MatrixXd matProjectorsRows(innerind.size(),t_matProjectors.cols());
    MatrixXd matProjectorsInnerind(innerind.size(),innerind.size());

    for (int i = 0; i < matProjectorsRows.rows(); ++i) {
        matProjectorsRows.row(i) = t_matProjectors.row(innerind.at(i));
    }

    for (int i = 0; i < matProjectorsInnerind.cols(); ++i) {
        matProjectorsInnerind.col(i) = matProjectorsRows.col(innerind.at(i));
    }

    //UTILSLIB::IOUtils::write_eigen_matrix(matProjectorsInnerind, "matProjectorsInnerind.txt");
    //UTILSLIB::IOUtils::write_eigen_matrix(t_matProjectors, "t_matProjectors.txt");

    // Initialize inner layer sensors
    sensors.coilpos = Eigen::MatrixXd::Zero(innerind.size(),3);
    sensors.coilori = Eigen::MatrixXd::Zero(innerind.size(),3);
    sensors.tra = Eigen::MatrixXd::Identity(innerind.size(),innerind.size());

    for(int i = 0; i < innerind.size(); i++) {
        sensors.coilpos(i,0) = pFiffInfo->chs[innerind.at(i)].chpos.r0[0];
        sensors.coilpos(i,1) = pFiffInfo->chs[innerind.at(i)].chpos.r0[1];
        sensors.coilpos(i,2) = pFiffInfo->chs[innerind.at(i)].chpos.r0[2];
        sensors.coilori(i,0) = pFiffInfo->chs[innerind.at(i)].chpos.ez[0];
        sensors.coilori(i,1) = pFiffInfo->chs[innerind.at(i)].chpos.ez[1];
        sensors.coilori(i,2) = pFiffInfo->chs[innerind.at(i)].chpos.ez[2];
    }

    Eigen::MatrixXd topo(innerind.size(), numCoils*2);
    Eigen::MatrixXd amp(innerind.size(), numCoils);
    Eigen::MatrixXd ampC(innerind.size(), numCoils);

    // Get the data from inner layer channels
    Eigen::MatrixXd innerdata(innerind.size(), t_mat.cols());

    for(int j = 0; j < innerind.size(); ++j) {
        innerdata.row(j) << t_mat.row(innerind[j]);
    }

    // Calculate topo
    topo = innerdata * pinv(simsig).transpose(); // topo: # of good inner channel x 8

    // Select sine or cosine component depending on the relative size
    amp  = topo.leftCols(numCoils); // amp: # of good inner channel x 4
    ampC = topo.rightCols(numCoils);

    for(int j = 0; j < numCoils; ++j) {
       float nS = 0.0;
       float nC = 0.0;
       for(int i = 0; i < innerind.size(); ++i) {
           nS += amp(i,j)*amp(i,j);
           nC += ampC(i,j)*ampC(i,j);
       }

       if(nC > nS) {
         for(int i = 0; i < innerind.size(); ++i) {
           amp(i,j) = ampC(i,j);
         }
       }
    }

    //Find good seed point/starting point for the coil position in 3D space
    //Find biggest amplitude per pickup coil (sensor) and store corresponding sensor channel index
    VectorXi chIdcs(numCoils);

    for (int j = 0; j < numCoils; j++) {
        double maxVal = 0;
        int chIdx = 0;

        for (int i = 0; i < amp.rows(); ++i) {
            if(abs(amp(i,j)) > maxVal) {
                maxVal = abs(amp(i,j));

                if(chIdx < innerind.size()) {
                    chIdx = innerind.at(i);
                }
            }
        }

        chIdcs(j) = chIdx;
    }

    //Generate seed point by projection the found channel position 3cm inwards
    Eigen::MatrixXd coilPos = Eigen::MatrixXd::Zero(numCoils,3);

    for (int j = 0; j < chIdcs.rows(); ++j) {
        int chIdx = chIdcs(j);

        if(chIdx < pFiffInfo->chs.size()) {
            double x = pFiffInfo->chs.at(chIdcs(j)).chpos.r0[0];
            double y = pFiffInfo->chs.at(chIdcs(j)).chpos.r0[1];
            double z = pFiffInfo->chs.at(chIdcs(j)).chpos.r0[2];

            coilPos(j,0) = -1 * pFiffInfo->chs.at(chIdcs(j)).chpos.ez[0] * 0.03 + x;
            coilPos(j,1) = -1 * pFiffInfo->chs.at(chIdcs(j)).chpos.ez[1] * 0.03 + y;
            coilPos(j,2) = -1 * pFiffInfo->chs.at(chIdcs(j)).chpos.ez[2] * 0.03 + z;
        }

        //std::cout << "HPIFit::fitHPI - Coil " << j << " max value index " << chIdx << std::endl;
    }

    coil.pos = coilPos;

    coil = dipfit(coil, sensors, amp, numCoils, matProjectorsInnerind);

    Eigen::Matrix4d trans = computeTransformation(headHPI, coil.pos);
    //Eigen::Matrix4d trans = computeTransformation(coil.pos, headHPI);

    // Store the final result to fiff info
    // Set final device/head matrix and its inverse to the fiff info
    transDevHead.from = 1;
    transDevHead.to = 4;

    for(int r = 0; r < 4; ++r) {
        for(int c = 0; c < 4 ; ++c) {
            transDevHead.trans(r,c) = trans(r,c);
        }
    }

    // Also store the inverse
    transDevHead.invtrans = transDevHead.trans.inverse();

    //Calculate GOF
    MatrixXd temp = coil.pos;
    temp.conservativeResize(coil.pos.rows(),coil.pos.cols()+1);

    temp.block(0,3,numCoils,1).setOnes();
    temp.transposeInPlace();

    MatrixXd testPos = trans * temp;
    MatrixXd diffPos = testPos.block(0,0,3,numCoils) - headHPI.transpose();

    for(int i = 0; i < diffPos.cols(); ++i) {
        vGof.append(diffPos.col(i).norm());
    }

    //Generate final fitted points and store in digitizer set
    for(int i = 0; i < coil.pos.rows(); ++i) {
        FiffDigPoint digPoint;
        digPoint.kind = FIFFV_POINT_EEG;
        digPoint.ident = i;
        digPoint.r[0] = coil.pos(i,0);
        digPoint.r[1] = coil.pos(i,1);
        digPoint.r[2] = coil.pos(i,2);

        fittedPointSet << digPoint;
    }

    if(bDoDebug) {
        // DEBUG HPI fitting and write debug results
        std::cout << std::endl << std::endl << "HPIFit::fitHPI - dpfiterror" << coil.dpfiterror << std::endl << coil.pos << std::endl;
        std::cout << std::endl << std::endl << "HPIFit::fitHPI - Initial seed point for HPI coils" << std::endl << coil.pos << std::endl;
        std::cout << std::endl << std::endl << "HPIFit::fitHPI - temp" << std::endl << temp << std::endl;
        std::cout << std::endl << std::endl << "HPIFit::fitHPI - testPos" << std::endl << testPos << std::endl;
        std::cout << std::endl << std::endl << "HPIFit::fitHPI - Diff fitted - original" << std::endl << diffPos << std::endl;
        std::cout << std::endl << std::endl << "HPIFit::fitHPI - dev/head trans" << std::endl << trans << std::endl;

        QString sTimeStamp = QDateTime::currentDateTime().toString("yyMMdd_hhmmss");

        if(!QDir(sHPIResourceDir).exists()) {
            QDir().mkdir(sHPIResourceDir);
        }

        UTILSLIB::IOUtils::write_eigen_matrix(coilPos, QString("%1/%2_coilPosSeed_mat").arg(sHPIResourceDir).arg(sTimeStamp));

        UTILSLIB::IOUtils::write_eigen_matrix(coil.pos, QString("%1/%2_coilPos_mat").arg(sHPIResourceDir).arg(sTimeStamp));

        UTILSLIB::IOUtils::write_eigen_matrix(headHPI, QString("%1/%2_headHPI_mat").arg(sHPIResourceDir).arg(sTimeStamp));

        MatrixXd testPosCut = testPos.transpose();//block(0,0,3,4);
        UTILSLIB::IOUtils::write_eigen_matrix(testPosCut, QString("%1/%2_testPos_mat").arg(sHPIResourceDir).arg(sTimeStamp));

        MatrixXi idx_mat(chIdcs.rows(),1);
        idx_mat.col(0) = chIdcs;
        UTILSLIB::IOUtils::write_eigen_matrix(idx_mat, QString("%1/%2_idx_mat").arg(sHPIResourceDir).arg(sTimeStamp));

        MatrixXd coilFreq_mat(coilfreq.rows(),1);
        coilFreq_mat.col(0) = coilfreq;
        UTILSLIB::IOUtils::write_eigen_matrix(coilFreq_mat, QString("%1/%2_coilFreq_mat").arg(sHPIResourceDir).arg(sTimeStamp));

        UTILSLIB::IOUtils::write_eigen_matrix(diffPos, QString("%1/%2_diffPos_mat").arg(sHPIResourceDir).arg(sTimeStamp));

        UTILSLIB::IOUtils::write_eigen_matrix(amp, QString("%1/%2_amp_mat").arg(sHPIResourceDir).arg(sTimeStamp));
    }
}

//*************************************************************************************************************

CoilParam HPIFit::dipfit(struct CoilParam coil, struct SensorInfo sensors, const Eigen::MatrixXd& data, int numCoils, const Eigen::MatrixXd& t_matProjectors)
{
    //Do this in conncurrent mode
    //Generate QList structure which can be handled by the QConcurrent framework
    QList<FittingCoilData> lCoilData;

    for(qint32 i = 0; i < numCoils; ++i) {
        FittingCoilData coilData;
        coilData.coilPos = coil.pos.row(i);
        coilData.sensorData = data.col(i);
        coilData.sensorPos = sensors;
        coilData.matProjector = t_matProjectors;

        lCoilData.append(coilData);
    }
    //Do the concurrent filtering
    if(!lCoilData.isEmpty()) {
        //Do sequential
        for(int l = 0; l < lCoilData.size(); ++l) {
            doDipfitConcurrent(lCoilData[l]);
        }

//        //Do concurrent
//        QFuture<void> future = QtConcurrent::map(lCoilData,
//                                             doDipfitConcurrent);
//        future.waitForFinished();

        //Transform results to final coil information
        for(qint32 i = 0; i < lCoilData.size(); ++i) {
            coil.pos.row(i) = lCoilData.at(i).coilPos;
            coil.mom = lCoilData.at(i).errorInfo.moment.transpose();
            coil.dpfiterror(i) = lCoilData.at(i).errorInfo.error;
            coil.dpfitnumitr(i) = lCoilData.at(i).errorInfo.numIterations;

            //std::cout<<std::endl<< "HPIFit::dipfit - Itr steps for coil " << i << " =" <<coil.dpfitnumitr(i);
        }
    }

    return coil;
}

//*************************************************************************************************************

Eigen::Matrix4d HPIFit::computeTransformation(Eigen::MatrixXd NH, Eigen::MatrixXd BT)
{
    Eigen::MatrixXd xdiff, ydiff, zdiff, C, Q;
    Eigen::Matrix4d transFinal = Eigen::Matrix4d::Identity(4,4);
    Eigen::Matrix4d Rot = Eigen::Matrix4d::Zero(4,4);
    Eigen::Matrix4d Trans = Eigen::Matrix4d::Identity(4,4);
    double meanx,meany,meanz,normf;

    for(int i = 0; i < 15; ++i) {
        // Calcualte mean translation for all points -> centroid of both data sets
        xdiff = NH.col(0) - BT.col(0);
        ydiff = NH.col(1) - BT.col(1);
        zdiff = NH.col(2) - BT.col(2);

        meanx = xdiff.mean();
        meany = ydiff.mean();
        meanz = zdiff.mean();

        // Apply translation -> bring both data sets to the same center location
        for (int j = 0; j < BT.rows(); ++j) {
            BT(j,0) = BT(j,0) + meanx;
            BT(j,1) = BT(j,1) + meany;
            BT(j,2) = BT(j,2) + meanz;
        }

        // Estimate rotation component
        C = BT.transpose() * NH;

        Eigen::JacobiSVD< Eigen::MatrixXd > svd(C ,Eigen::ComputeThinU | Eigen::ComputeThinV);

        Q = svd.matrixU() * svd.matrixV().transpose();

        //Handle special reflection case
        if(Q.determinant() < 0) {
            Q(0,2) = Q(0,2) * -1;
            Q(1,2) = Q(1,2) * -1;
            Q(2,2) = Q(2,2) * -1;
        }

        // Apply rotation on translated points
        BT = BT * Q;

        // Calculate GOF
        normf = (NH.transpose()-BT.transpose()).norm();

        // Store rotation part to transformation matrix
        Rot(3,3) = 1;
        for(int j = 0; j < 3; ++j) {
            for(int k = 0; k < 3; ++k) {
                Rot(j,k) = Q(k,j);
            }
        }

        // Store translation part to transformation matrix
        Trans(0,3) = meanx;
        Trans(1,3) = meany;
        Trans(2,3) = meanz;

        // Safe rotation and translation to final matrix for next iteration step
        transFinal = Rot * Trans * transFinal;
    }

    return transFinal;
}


#pragma once

#include <cmath>
#include <cstdint>
#include <functional>
#include <fstream>
#include <iostream>
#include <memory>
#include <stdexcept>
#include <string>
#include <unordered_map>
#include <utility>
#include <vector>

#include <Eigen/Geometry>
#include <common_robotics_utilities/maybe.hpp>
#include <common_robotics_utilities/serialization.hpp>
#include <common_robotics_utilities/voxel_grid.hpp>
#include <common_robotics_utilities/zlib_helpers.hpp>

namespace voxelized_geometry_tools
{
/// Using declaration to make naming clearer and easier to read.
using EstimateDistanceQuery = common_robotics_utilities::OwningMaybe<double>;

/// This is similar to std::optional<Eigen::Vector4d>, but lets us enforce
/// specific behavior in the contained Vector4d.
class GradientQuery
{
private:
  Eigen::Vector4d gradient_ = Eigen::Vector4d(0.0, 0.0, 0.0, 0.0);
  bool has_value_ = false;

public:
  explicit GradientQuery(const Eigen::Vector4d& gradient)
      : gradient_(gradient), has_value_(true)
  {
    if (gradient_(3) != 0.0)
    {
      throw std::invalid_argument("gradient(3) != 0.0");
    }
  }

  GradientQuery(const double x, const double y, const double z)
      : gradient_(Eigen::Vector4d(x, y, z, 0.0)), has_value_(true) {}

  GradientQuery() : has_value_(false) {}

  const Eigen::Vector4d& Value() const
  {
    if (HasValue())
    {
      return gradient_;
    }
    else
    {
      throw std::runtime_error("GradientQuery does not have value");
    }
  }

  bool HasValue() const { return has_value_; }

  explicit operator bool() const { return HasValue(); }
};

/// This is equivalent to std::optional<Eigen::Vector4d>, but kept separate here
/// so as to not require C++17 support with a working std::optional<T>. This
/// also allows us to enforce specific behavior in the contained Vector4d.
class ProjectedPosition
{
private:
  const Eigen::Vector4d gradient_ = Eigen::Vector4d(0.0, 0.0, 0.0, 1.0);
  const bool has_value_ = false;

public:
  explicit ProjectedPosition(const Eigen::Vector4d& gradient)
      : gradient_(gradient), has_value_(true)
  {
    if (gradient_(3) != 1.0)
    {
      throw std::invalid_argument("gradient(3) != 1.0");
    }
  }

  ProjectedPosition(const double x, const double y, const double z)
      : gradient_(Eigen::Vector4d(x, y, z, 1.0)), has_value_(true) {}

  ProjectedPosition() : has_value_(false) {}

  const Eigen::Vector4d& Value() const
  {
    if (HasValue())
    {
      return gradient_;
    }
    else
    {
      throw std::runtime_error("ProjectedPosition does not have value");
    }
  }

  bool HasValue() const { return has_value_; }

  explicit operator bool() const { return HasValue(); }
};

template<typename ScalarType>
class SignedDistanceField final
    : public common_robotics_utilities::voxel_grid
        ::VoxelGridBase<ScalarType, std::vector<ScalarType>>
{
private:
  using ScalarTypeSerializer
      = common_robotics_utilities::serialization::Serializer<ScalarType>;
  using ScalarTypeDeserializer
      = common_robotics_utilities::serialization::Deserializer<ScalarType>;
  using DeserializedSignedDistanceField
      = common_robotics_utilities::serialization
          ::Deserialized<SignedDistanceField<ScalarType>>;

  std::string frame_;
  bool locked_ = false;

  /// Internal helper used in "Fine" gradient computation.
  inline double ComputeAxisFineGradient(
      const EstimateDistanceQuery& query_point_distance_estimate,
      const EstimateDistanceQuery& minus_axis_distance_estimate,
      const EstimateDistanceQuery& plus_axis_distance_estimate,
      const double query_point_axis_value,
      const double minus_point_axis_value,
      const double plus_point_axis_value) const
  {
    if (query_point_distance_estimate.HasValue()
        && minus_axis_distance_estimate.HasValue()
        && plus_axis_distance_estimate.HasValue())
    {
      const double window_size = plus_point_axis_value
                                 - minus_point_axis_value;
      const double distance_delta = plus_axis_distance_estimate.Value()
                                    - minus_axis_distance_estimate.Value();
      return distance_delta / window_size;
    }
    else if (query_point_distance_estimate.HasValue()
             && minus_axis_distance_estimate.HasValue())
    {
      const double window_size = query_point_axis_value
                                 - minus_point_axis_value;
      const double distance_delta = query_point_distance_estimate.Value()
                                    - minus_axis_distance_estimate.Value();
      return distance_delta / window_size;
    }
    else if (query_point_distance_estimate.HasValue()
             && plus_axis_distance_estimate.HasValue())
    {
      const double window_size = plus_point_axis_value
                                 - query_point_axis_value;
      const double distance_delta = plus_axis_distance_estimate.Value()
                                    - query_point_distance_estimate.Value();
      return distance_delta / window_size;
    }
    else
    {
      throw std::runtime_error(
            "Window size for GetSmoothGradient is too large for SDF");
    }
  }

  /// Bilinear interpolation used in trilinear interpolation.
  template<typename T>
  inline T BilinearInterpolate(const double low_d1,
                               const double high_d1,
                               const double low_d2,
                               const double high_d2,
                               const T query_d1,
                               const T query_d2,
                               const double l1l2_val,
                               const double l1h2_val,
                               const double h1l2_val,
                               const double h1h2_val) const
  {
    Eigen::Matrix<T, 1, 2> d1_offsets;
    d1_offsets(0, 0) = high_d1 - query_d1;
    d1_offsets(0, 1) = query_d1 - low_d1;
    Eigen::Matrix<T, 2, 2> values;
    values(0, 0) = l1l2_val;
    values(0, 1) = l1h2_val;
    values(1, 0) = h1l2_val;
    values(1, 1) = h1h2_val;
    Eigen::Matrix<T, 2, 1> d2_offsets;
    d2_offsets(0, 0) = high_d2 - query_d2;
    d2_offsets(1, 0) = query_d2 - low_d2;
    const T multiplier = 1.0 / ((high_d1 - low_d1) * (high_d2 - low_d2));
    const T bilinear_interpolated
        = multiplier * d1_offsets * values * d2_offsets;
    return bilinear_interpolated;
  }

  /// Wrapper around bilinear interpolation to set arguments more easily.
  template<typename T>
  inline T BilinearInterpolateDistanceXY(
      const Eigen::Vector4d& corner_location,
      const Eigen::Matrix<T, 4, 1>& query_location,
      const double mxmy_dist, const double mxpy_dist,
      const double pxmy_dist, const double pxpy_dist) const
  {
    return BilinearInterpolate(corner_location(0),
                               corner_location(0) + GetResolution(),
                               corner_location(1),
                               corner_location(1) + GetResolution(),
                               query_location(0),
                               query_location(1),
                               mxmy_dist, mxpy_dist,
                               pxmy_dist, pxpy_dist);
  }

  /// Trilinear distance interpolation.
  template<typename T>
  inline T TrilinearInterpolateDistance(
      const Eigen::Vector4d& corner_location,
      const Eigen::Matrix<T, 4, 1>& query_location,
      const double mxmymz_dist, const double mxmypz_dist,
      const double mxpymz_dist, const double mxpypz_dist,
      const double pxmymz_dist, const double pxmypz_dist,
      const double pxpymz_dist, const double pxpypz_dist) const
  {
    // Do bilinear interpolation in the lower XY plane
    const T mz_bilinear_interpolated
        = BilinearInterpolateDistanceXY(corner_location, query_location,
                                        mxmymz_dist, mxpymz_dist,
                                        pxmymz_dist, pxpymz_dist);
    // Do bilinear interpolation in the upper XY plane
    const T pz_bilinear_interpolated
        = BilinearInterpolateDistanceXY(corner_location, query_location,
                                        mxmypz_dist, mxpypz_dist,
                                        pxmypz_dist, pxpypz_dist);
    // Perform linear interpolation/extrapolation between lower and upper planes
    const double inv_resolution = 1.0 / GetResolution();
    const T distance_delta
        = pz_bilinear_interpolated - mz_bilinear_interpolated;
    const T distance_slope = distance_delta * inv_resolution;
    const T query_z_delta = query_location(2) - T(corner_location(2));
    return mz_bilinear_interpolated + (query_z_delta * distance_slope);
  }

  /// Computes "corrected" center distance accounting for the fact that SDF
  /// distance is the distance to the next filled cell center, *not* the
  /// distance to the boundary.
  inline double GetCorrectedCenterDistance(const int64_t x_idx,
                                           const int64_t y_idx,
                                           const int64_t z_idx) const
  {
    const auto query = this->GetImmutable(x_idx, y_idx, z_idx);
    if (query)
    {
      const double nominal_sdf_distance = static_cast<double>(query.Value());
      const double cell_center_distance_offset = GetResolution() * 0.5;
      if (nominal_sdf_distance >= 0.0)
      {
        return nominal_sdf_distance - cell_center_distance_offset;
      }
      else
      {
        return nominal_sdf_distance + cell_center_distance_offset;
      }
    }
    else
    {
      throw std::invalid_argument("Index out of bounds");
    }
  }

  /// Computes the axis lookup indices to use for interpolation.
  template<typename T>
  std::pair<int64_t, int64_t> GetAxisInterpolationIndices(
      const int64_t initial_index,
      const int64_t axis_size,
      const T axis_offset) const
  {
    int64_t lower = initial_index;
    int64_t upper = initial_index;
    if (axis_offset >= 0.0)
    {
      upper = initial_index + 1;
      if (upper >= axis_size)
      {
        upper = initial_index;
        lower = initial_index -1;
        if (lower < 0)
        {
          lower = initial_index;
        }
      }
    }
    else
    {
      lower = initial_index - 1;
      if (lower < 0)
      {
        upper = initial_index + 1;
        lower = initial_index;
        if (upper >= axis_size)
        {
          upper = initial_index;
        }
      }
    }
    return std::make_pair(lower, upper);
  }

  /// Estimates distance via trilinear interpolation of the surrounding 8 cells.
  template<typename T>
  inline T EstimateDistanceInterpolateFromNeighbors(
      const Eigen::Matrix<T, 4, 1>& query_location,
      const int64_t x_idx, const int64_t y_idx, const int64_t z_idx) const
  {
    // Get the query location in grid frame
    const Eigen::Matrix<T, 4, 1> grid_frame_query_location
        = this->GetInverseOriginTransform() * query_location;
    // Switch between all the possible options of where we are
    const Eigen::Vector4d cell_center_location
        = this->GridIndexToLocationInGridFrame(x_idx, y_idx, z_idx);
    const Eigen::Matrix<T, 4, 1> query_offset
        = grid_frame_query_location - cell_center_location.cast<T>();
    // We can't catch the easiest case of being at a cell center, since doing so
    // breaks the ability of Eigen's Autodiff to autodiff this function.
    // Find the best-matching 8 surrounding cell centers
    const std::pair<int64_t, int64_t> x_axis_indices
        = GetAxisInterpolationIndices(
            x_idx, this->GetNumXCells(), query_offset(0));
    const std::pair<int64_t, int64_t> y_axis_indices
        = GetAxisInterpolationIndices(
            y_idx, this->GetNumYCells(), query_offset(1));
    const std::pair<int64_t, int64_t> z_axis_indices
        = GetAxisInterpolationIndices(
            z_idx, this->GetNumZCells(), query_offset(2));
    const Eigen::Vector4d lower_corner_location
        = this->GridIndexToLocationInGridFrame(x_axis_indices.first,
                                               y_axis_indices.first,
                                               z_axis_indices.first);
    const double mxmymz_distance
        = GetCorrectedCenterDistance(x_axis_indices.first,
                                     y_axis_indices.first,
                                     z_axis_indices.first);
    const double mxmypz_distance
        = GetCorrectedCenterDistance(x_axis_indices.first,
                                     y_axis_indices.first,
                                     z_axis_indices.second);
    const double mxpymz_distance
        = GetCorrectedCenterDistance(x_axis_indices.first,
                                     y_axis_indices.second,
                                     z_axis_indices.first);
    const double mxpypz_distance
        = GetCorrectedCenterDistance(x_axis_indices.first,
                                     y_axis_indices.second,
                                     z_axis_indices.second);
    const double pxmymz_distance
        = GetCorrectedCenterDistance(x_axis_indices.second,
                                     y_axis_indices.first,
                                     z_axis_indices.first);
    const double pxmypz_distance
        = GetCorrectedCenterDistance(x_axis_indices.second,
                                     y_axis_indices.first,
                                     z_axis_indices.second);
    const double pxpymz_distance
        = GetCorrectedCenterDistance(x_axis_indices.second,
                                     y_axis_indices.second,
                                     z_axis_indices.first);
    const double pxpypz_distance
        = GetCorrectedCenterDistance(x_axis_indices.second,
                                     y_axis_indices.second,
                                     z_axis_indices.second);
    return TrilinearInterpolateDistance(lower_corner_location,
                                        grid_frame_query_location,
                                        mxmymz_distance, mxmypz_distance,
                                        mxpymz_distance, mxpypz_distance,
                                        pxmymz_distance, pxmypz_distance,
                                        pxpymz_distance, pxpypz_distance);
  }

  /// You *must* provide valid indices to this function, hence why it's private.
  void FollowGradientsToLocalExtremaUnsafe(
      const int64_t x_index, const int64_t y_index, const int64_t z_index,
      common_robotics_utilities::voxel_grid::VoxelGrid<Eigen::Vector3d>&
          watershed_map) const
  {
    // First, check if we've already found the local extrema for the current
    // cell
    const Eigen::Vector3d& stored
        = watershed_map.GetImmutable(x_index, y_index, z_index).Value();
    if (stored.x() != -std::numeric_limits<double>::infinity()
        && stored.y() != -std::numeric_limits<double>::infinity()
        && stored.z() != -std::numeric_limits<double>::infinity())
    {
      // We've already found it for this cell, so we can skip it
      return;
    }
    // Find the local extrema
    GradientQuery raw_gradient
        = GetCoarseGradient(x_index, y_index, z_index, true);
    Eigen::Vector4d gradient_vector = raw_gradient.Value();
    if (GradientIsEffectiveFlat(gradient_vector))
    {
      const Eigen::Vector4d location
          = this->GridIndexToLocationInGridFrame(x_index, y_index, z_index);
      const Eigen::Vector3d local_extrema(
          location(0), location(1), location(2));
      watershed_map.SetValue(x_index, y_index, z_index, local_extrema);
    }
    else
    {
      // Follow the gradient, one cell at a time, until we reach a local maxima
      std::unordered_map<common_robotics_utilities::voxel_grid::GridIndex,
                         int8_t> path;
      common_robotics_utilities::voxel_grid::GridIndex current_index(
            x_index, y_index, z_index);
      path[current_index] = 1;
      Eigen::Vector3d local_extrema(-std::numeric_limits<double>::infinity(),
                                    -std::numeric_limits<double>::infinity(),
                                    -std::numeric_limits<double>::infinity());
      while (true)
      {
        if (path.size() == 10000)
        {
          std::cerr << "Warning, gradient path is long (i.e >= 10000 steps)"
                    << std::endl;
        }
        current_index = GetNextFromGradient(current_index, gradient_vector);
        if (path[current_index] != 0)
        {
          // If we've already been here, then we are done
          const Eigen::Vector4d location
              = this->GridIndexToLocationInGridFrame(current_index);
          local_extrema
              = Eigen::Vector3d(location(0), location(1), location(2));
          break;
        }
        // Check if we've been pushed past the edge
        if (!this->IndexInBounds(current_index))
        {
          // We have the "off the grid" local maxima
          local_extrema
              = Eigen::Vector3d(std::numeric_limits<double>::infinity(),
                                std::numeric_limits<double>::infinity(),
                                std::numeric_limits<double>::infinity());
          break;
        }
        path[current_index] = 1;
        // Check if the new index has already been checked
        const Eigen::Vector3d& new_stored
            = watershed_map.GetImmutable(current_index).Value();
        if (new_stored.x() != -std::numeric_limits<double>::infinity()
            && new_stored.y() != -std::numeric_limits<double>::infinity()
            && new_stored.z() != -std::numeric_limits<double>::infinity())
        {
          // We have the local maxima
          local_extrema = new_stored;
          break;
        }
        else
        {
          raw_gradient = GetCoarseGradient(current_index, true);
          gradient_vector = raw_gradient.Value();
          if (GradientIsEffectiveFlat(gradient_vector))
          {
            // We have the local maxima
            const Eigen::Vector4d location
                = this->GridIndexToLocationInGridFrame(current_index);
            local_extrema
                = Eigen::Vector3d(location(0), location(1), location(2));
            break;
          }
        }
      }
      // Now, go back and mark the entire explored path with the local maxima
      for (auto path_itr = path.begin(); path_itr != path.end(); ++path_itr)
      {
        const common_robotics_utilities::voxel_grid::GridIndex& index
            = path_itr->first;
        watershed_map.SetValue(index, local_extrema);
      }
    }
  }

  bool GradientIsEffectiveFlat(const Eigen::Vector4d& gradient) const
  {
    // A gradient is at a local maxima if the absolute value of all components
    // (x,y,z) are less than 1/2 SDF resolution
    const double step_resolution = GetResolution() * 0.06125;
    if (std::abs(gradient(0)) <= step_resolution
        && std::abs(gradient(1)) <= step_resolution
        && std::abs(gradient(2)) <= step_resolution)
    {
      return true;
    }
    else
    {
      return false;
    }
  }

  common_robotics_utilities::voxel_grid::GridIndex GetNextFromGradient(
      const common_robotics_utilities::voxel_grid::GridIndex& index,
      const Eigen::Vector4d& gradient) const
  {
    // Check if it's inside an obstacle
    const ScalarType stored_distance = this->GetImmutable(index).Value();
    Eigen::Vector4d working_gradient = gradient;
    if (stored_distance < 0.0)
    {
      working_gradient = gradient * -1.0;
    }
    // Given the gradient, pick the "best fit" of the 26 neighboring points
    common_robotics_utilities::voxel_grid::GridIndex next_index = index;
    const double step_resolution = GetResolution() * 0.06125;
    if (working_gradient(0) > step_resolution)
    {
      next_index.X() += 1;
    }
    else if (working_gradient(0) < -step_resolution)
    {
      next_index.X() -= 1;
    }
    if (working_gradient(1) > step_resolution)
    {
      next_index.Y() += 1;
    }
    else if (working_gradient(1) < -step_resolution)
    {
      next_index.Y() -= 1;
    }
    if (working_gradient(2) > step_resolution)
    {
      next_index.Z() += 1;
    }
    else if (working_gradient(2) < -step_resolution)
    {
      next_index.Z() -= 1;
    }
    return next_index;
  }

  /// Implement the VoxelGridBase interface.

  /// We need to implement cloning.
  std::unique_ptr<common_robotics_utilities::voxel_grid
      ::VoxelGridBase<ScalarType, std::vector<ScalarType>>>
  DoClone() const override
  {
    return std::unique_ptr<SignedDistanceField<ScalarType>>(
        new SignedDistanceField<ScalarType>(*this));
  }

  /// We need to serialize the frame and locked flag.
  uint64_t DerivedSerializeSelf(
      std::vector<uint8_t>& buffer,
      const ScalarTypeSerializer& value_serializer) const override
  {
    CRU_UNUSED(value_serializer);
    const uint64_t start_size = buffer.size();
    common_robotics_utilities::serialization::SerializeString(frame_, buffer);
    common_robotics_utilities::serialization::SerializeMemcpyable<uint8_t>(
        static_cast<uint8_t>(locked_), buffer);
    const uint64_t bytes_written = buffer.size() - start_size;
    return bytes_written;
  }

  /// We need to deserialize the frame and locked flag.
  uint64_t DerivedDeserializeSelf(
      const std::vector<uint8_t>& buffer, const uint64_t starting_offset,
      const ScalarTypeDeserializer& value_deserializer) override
  {
    CRU_UNUSED(value_deserializer);
    uint64_t current_position = starting_offset;
    // Deserialize SDF stuff
    const auto frame_deserialized
        = common_robotics_utilities::serialization::DeserializeString<char>(
            buffer, current_position);
    frame_ = frame_deserialized.Value();
    current_position += frame_deserialized.BytesRead();
    const auto locked_deserialized
        = common_robotics_utilities::serialization
            ::DeserializeMemcpyable<uint8_t>(buffer, current_position);
    locked_ = static_cast<bool>(locked_deserialized.Value());
    current_position += locked_deserialized.BytesRead();
    // Figure out how many bytes were read
    const uint64_t bytes_read = current_position - starting_offset;
    return bytes_read;
  }

  /// We do not allow mutable access if the SDF is locked.
  bool OnMutableAccess(const int64_t x_index,
                       const int64_t y_index,
                       const int64_t z_index) override
  {
    CRU_UNUSED(x_index);
    CRU_UNUSED(y_index);
    CRU_UNUSED(z_index);
    return !IsLocked();
  }

public:
  static uint64_t Serialize(
      const SignedDistanceField<ScalarType>& sdf, std::vector<uint8_t>& buffer)
  {
    return sdf.SerializeSelf(
        buffer,
        common_robotics_utilities::serialization
            ::SerializeMemcpyable<ScalarType>);
  }

  static DeserializedSignedDistanceField Deserialize(
      const std::vector<uint8_t>& buffer, const uint64_t starting_offset)
  {
    SignedDistanceField<ScalarType> temp_sdf;
    const uint64_t bytes_read = temp_sdf.DeserializeSelf(
        buffer, starting_offset,
        common_robotics_utilities::serialization
            ::DeserializeMemcpyable<ScalarType>);
    return common_robotics_utilities::serialization::MakeDeserialized(
        temp_sdf, bytes_read);
  }

  static void SaveToFile(
      const SignedDistanceField<ScalarType>& sdf, const std::string& filepath,
      const bool compress)
  {
    std::vector<uint8_t> buffer;
    SignedDistanceField<ScalarType>::Serialize(sdf, buffer);
    std::ofstream output_file(filepath, std::ios::out|std::ios::binary);
    if (compress)
    {
      output_file.write("SDFZ", 4);
      const std::vector<uint8_t> compressed
          = common_robotics_utilities::zlib_helpers::CompressBytes(buffer);
      const size_t serialized_size = compressed.size();
      output_file.write(
          reinterpret_cast<const char*>(compressed.data()),
          static_cast<std::streamsize>(serialized_size));
    }
    else
    {
      output_file.write("SDFR", 4);
      const size_t serialized_size = buffer.size();
      output_file.write(
          reinterpret_cast<const char*>(buffer.data()),
          static_cast<std::streamsize>(serialized_size));
    }
    output_file.close();
  }

  static SignedDistanceField<ScalarType> LoadFromFile(
      const std::string& filepath)
  {
    std::ifstream input_file(
        filepath, std::ios::in | std::ios::binary | std::ios::ate);
    if (input_file.good() == false)
    {
      throw std::invalid_argument("File does not exist");
    }
    const std::streampos end = input_file.tellg();
    input_file.seekg(0, std::ios::beg);
    const std::streampos begin = input_file.tellg();
    const std::streamsize serialized_size = end - begin;
    const std::streamsize header_size = 4;
    if (serialized_size >= header_size)
    {
      // Load the header
      std::vector<uint8_t> file_header(header_size + 1, 0x00);
      input_file.read(reinterpret_cast<char*>(file_header.data()),
                      header_size);
      const std::string header_string(
            reinterpret_cast<const char*>(file_header.data()));
      // Load the rest of the file
      std::vector<uint8_t> file_buffer(
            static_cast<size_t>(serialized_size - header_size), 0x00);
      input_file.read(reinterpret_cast<char*>(file_buffer.data()),
                      serialized_size - header_size);
      // Deserialize
      if (header_string == "SDFZ")
      {
        const std::vector<uint8_t> decompressed
            = common_robotics_utilities::zlib_helpers
                ::DecompressBytes(file_buffer);
        return SignedDistanceField<ScalarType>::Deserialize(
            decompressed, 0).Value();
      }
      else if (header_string == "SDFR")
      {
        return SignedDistanceField<ScalarType>::Deserialize(
            file_buffer, 0).Value();
      }
      else
      {
        throw std::invalid_argument(
              "File has invalid header [" + header_string + "]");
      }
    }
    else
    {
      throw std::invalid_argument("File is too small");
    }
  }

  SignedDistanceField(
      const Eigen::Isometry3d& origin_transform, const std::string& frame,
      const common_robotics_utilities::voxel_grid::GridSizes& sizes,
      const ScalarType default_value)
      : SignedDistanceField(
          origin_transform, frame, sizes, default_value, default_value) {}

  SignedDistanceField(
      const std::string& frame,
      const common_robotics_utilities::voxel_grid::GridSizes& sizes,
      const ScalarType default_value)
      : SignedDistanceField(frame, sizes, default_value, default_value) {}

  SignedDistanceField(
      const Eigen::Isometry3d& origin_transform, const std::string& frame,
      const common_robotics_utilities::voxel_grid::GridSizes& sizes,
      const ScalarType default_value, const ScalarType oob_value)
      : common_robotics_utilities::voxel_grid
            ::VoxelGridBase<ScalarType, std::vector<ScalarType>>(
                origin_transform, sizes, default_value, oob_value),
        frame_(frame), locked_(false)
  {
    if (!this->HasUniformCellSize())
    {
      throw std::invalid_argument("SDF cannot have non-uniform cell sizes");
    }
  }

  SignedDistanceField(
      const std::string& frame,
      const common_robotics_utilities::voxel_grid::GridSizes& sizes,
      const ScalarType default_value, const ScalarType oob_value)
      : common_robotics_utilities::voxel_grid
            ::VoxelGridBase<ScalarType, std::vector<ScalarType>>(
                sizes, default_value, oob_value),
        frame_(frame), locked_(false)
  {
    if (!this->HasUniformCellSize())
    {
      throw std::invalid_argument("SDF cannot have non-uniform cell sizes");
    }
  }

  SignedDistanceField()
      : common_robotics_utilities::voxel_grid
            ::VoxelGridBase<ScalarType, std::vector<ScalarType>>() {}

  bool IsLocked() const { return locked_; }

  void Lock() { locked_ = true; }

  void Unlock() { locked_ = false; }

  double GetResolution() const { return this->GetCellSizes().x(); }

  const std::string& GetFrame() const { return frame_; }

  void SetFrame(const std::string& frame) { frame_ = frame; }

  /// For classical SDF distance queries, see the API of VoxelGridBase for how
  /// to retrieve and set cell values.

  /// "Estimated" distance is computed by performing trilinear interpolation
  /// against the distances in eight cells that surround the query point. This
  /// produces a continuous, but not necessarily smooth distance function that
  /// is more accurate and smoother than the raw distance values inside cells.
  /// Note that it is both more expensive than a simple grid lookup, but also
  /// this rounds inside and outside corners of cells. This can be a problem at
  /// coarse grid resolutions.

  EstimateDistanceQuery EstimateDistance(
      const double x, const double y, const double z) const
  {
      return EstimateDistance4d(Eigen::Vector4d(x, y, z, 1.0));
  }

  EstimateDistanceQuery EstimateDistance3d(
      const Eigen::Vector3d& location) const
  {
    const common_robotics_utilities::voxel_grid::GridIndex index
        = this->LocationToGridIndex3d(location);
    if (this->IndexInBounds(index))
    {
      return EstimateDistanceQuery(
          EstimateDistanceInterpolateFromNeighbors<double>(
              Eigen::Vector4d(location.x(), location.y(), location.z(), 1.0),
              index.X(), index.Y(), index.Z()));
    }
    else
    {
      return EstimateDistanceQuery();
    }
  }

  EstimateDistanceQuery EstimateDistance4d(
      const Eigen::Vector4d& location) const
  {
    const common_robotics_utilities::voxel_grid::GridIndex index
        = this->LocationToGridIndex4d(location);
    if (this->IndexInBounds(index))
    {
      return EstimateDistanceQuery(
          EstimateDistanceInterpolateFromNeighbors<double>(
              location, index.X(), index.Y(), index.Z()));
    }
    else
    {
      return EstimateDistanceQuery();
    }
  }

  /// "Coarse" gradient is computed by retrieving the distance from the
  /// surrounding size cells (+/-x, +/-y, +/-z) and differencing. This is the
  /// fastest method to compute a gradient, and also has the most potential
  /// error and least-smooth behavior.

  GradientQuery GetCoarseGradient(
      const double x, const double y, const double z,
      const bool enable_edge_gradients=false) const
  {
    return GetCoarseGradient4d(
        Eigen::Vector4d(x, y, z, 1.0), enable_edge_gradients);
  }

  GradientQuery GetCoarseGradient3d(
      const Eigen::Vector3d& location,
      const bool enable_edge_gradients=false) const
  {
    const common_robotics_utilities::voxel_grid::GridIndex index
        = this->LocationToGridIndex3d(location);
    if (this->IndexInBounds(index))
    {
      return GetCoarseGradient(index, enable_edge_gradients);
    }
    else
    {
      return GradientQuery();
    }
  }

  GradientQuery GetCoarseGradient4d(
      const Eigen::Vector4d& location,
      const bool enable_edge_gradients=false) const
  {
    const common_robotics_utilities::voxel_grid::GridIndex index
        = this->LocationToGridIndex4d(location);
    if (this->IndexInBounds(index))
    {
      return GetCoarseGradient(index, enable_edge_gradients);
    }
    else
    {
      return GradientQuery();
    }
  }

  GradientQuery GetCoarseGradient(
      const common_robotics_utilities::voxel_grid::GridIndex& index,
      const bool enable_edge_gradients=false) const
  {
    return GetCoarseGradient(
        index.X(), index.Y(), index.Z(), enable_edge_gradients);
  }

  GradientQuery GetCoarseGradient(
      const int64_t x_index, const int64_t y_index, const int64_t z_index,
      const bool enable_edge_gradients=false) const
  {
    const GradientQuery grid_aligned_gradient
        = GetGridAlignedCoarseGradient(x_index, y_index, z_index,
                                       enable_edge_gradients);
    if (grid_aligned_gradient.HasValue())
    {
      const Eigen::Vector4d gradient
          = this->GetOriginTransform() * grid_aligned_gradient.Value();
      return GradientQuery(gradient);
    }
    else
    {
      return grid_aligned_gradient;
    }
  }

  GradientQuery GetGridAlignedCoarseGradient(
      const int64_t x_index, const int64_t y_index, const int64_t z_index,
      const bool enable_edge_gradients=false) const
  {
    // Make sure the index is inside bounds
    if (this->IndexInBounds(x_index, y_index, z_index))
    {
      // See if the index we're trying to query is one cell in from the edge
      if ((x_index > 0) && (y_index > 0) && (z_index > 0)
          && (x_index < (this->GetNumXCells() - 1))
          && (y_index < (this->GetNumYCells() - 1))
          && (z_index < (this->GetNumZCells() - 1)))
      {
        const double inv_twice_resolution = 1.0 / (2.0 * GetResolution());
        const double gx
            = (this->GetImmutable(x_index + 1, y_index, z_index).Value()
               - this->GetImmutable(x_index - 1, y_index, z_index).Value())
              * inv_twice_resolution;
        const double gy
            = (this->GetImmutable(x_index, y_index + 1, z_index).Value()
               - this->GetImmutable(x_index, y_index - 1, z_index).Value())
              * inv_twice_resolution;
        const double gz
            = (this->GetImmutable(x_index, y_index, z_index + 1).Value()
               - this->GetImmutable(x_index, y_index, z_index - 1).Value())
              * inv_twice_resolution;
        return GradientQuery(gx, gy, gz);
      }
      // If we're on the edge, handle it specially
      // TODO: we actually need to handle corners even more carefully,
      // since if the SDF is build with a virtual border, these cells will
      // get zero gradient from this approach!
      else if (enable_edge_gradients)
      {
        // Get the "best" indices we can use
        const int64_t low_x_index
            = std::max(static_cast<int64_t>(0), x_index - 1);
        const int64_t high_x_index
            = std::min(this->GetNumXCells() - 1, x_index + 1);
        const int64_t low_y_index
            = std::max(static_cast<int64_t>(0), y_index - 1);
        const int64_t high_y_index
            = std::min(this->GetNumYCells() - 1, y_index + 1);
        const int64_t low_z_index
            = std::max(static_cast<int64_t>(0), z_index - 1);
        const int64_t high_z_index
            = std::min(this->GetNumZCells() - 1, z_index + 1);
        // Compute the axis increments
        const double x_increment
            = static_cast<double>(high_x_index - low_x_index) * GetResolution();
        const double y_increment
            = static_cast<double>(high_y_index - low_y_index) * GetResolution();
        const double z_increment
            = static_cast<double>(high_z_index - low_z_index) * GetResolution();
        // Compute the gradients for each axis - by default these are zero
        double gx = 0.0;
        double gy = 0.0;
        double gz = 0.0;
        // Only if the increments are non-zero do we compute the axis gradient
        if (x_increment > 0.0)
        {
          const double inv_x_increment = 1.0 / x_increment;
          const double high_x_value
              = this->GetImmutable(high_x_index, y_index, z_index).Value();
          const double low_x_value
              = this->GetImmutable(low_x_index, y_index, z_index).Value();
          // Compute the gradient
          gx = (high_x_value - low_x_value) * inv_x_increment;
        }
        if (y_increment > 0.0)
        {
          const double inv_y_increment = 1.0 / y_increment;
          const double high_y_value
              = this->GetImmutable(x_index, high_y_index, z_index).Value();
          const double low_y_value
              = this->GetImmutable(x_index, low_y_index, z_index).Value();
          // Compute the gradient
          gy = (high_y_value - low_y_value) * inv_y_increment;
        }
        if (z_increment > 0.0)
        {
          const double inv_z_increment = 1.0 / z_increment;
          const double high_z_value
              = this->GetImmutable(x_index, y_index, high_z_index).Value();
          const double low_z_value
              = this->GetImmutable(x_index, y_index, low_z_index).Value();
          // Compute the gradient
          gz = (high_z_value - low_z_value) * inv_z_increment;
        }
        // Assemble and return the computed gradient
        return GradientQuery(gx, gy, gz);
      }
      // Edge gradients disabled, return no gradient
      else
      {
        return GradientQuery();
      }
    }
    // If we're out of bounds, return no gradient
    else
    {
      return GradientQuery();
    }
  }

  /// "Fine" gradient is also computed by differencing, but instead of using the
  /// surrounding cells, this uses EstimateDistance() at +/- half of
  /// "window size" to compute the gradient. This is slower than the "coarse"
  /// gradient, since this needs to call EstimateDistance() six times, but
  /// produces a more accurate and smoother gradient. In the limit as
  /// "window size" -> 0 this produces the true gradient.

  GradientQuery GetFineGradient3d(
      const Eigen::Vector3d& location,
      const double nominal_window_size) const
  {
    return GetFineGradient(location.x(), location.y(), location.z(),
                           nominal_window_size);
  }

  GradientQuery GetFineGradient4d(
      const Eigen::Vector4d& location,
      const double nominal_window_size) const
  {
    return GetFineGradient(location(0), location(1), location(2),
                           nominal_window_size);
  }

  GradientQuery GetFineGradient(
      const double x, const double y, const double z,
      const double nominal_window_size) const
  {
    const double ideal_window_size = std::abs(nominal_window_size);
    if (this->LocationInBounds(x, y, z))
    {
      const double min_x = x - ideal_window_size;
      const double max_x = x + ideal_window_size;
      const double min_y = y - ideal_window_size;
      const double max_y = y + ideal_window_size;
      const double min_z = z - ideal_window_size;
      const double max_z = z + ideal_window_size;
      // Retrieve distance estimates
      const EstimateDistanceQuery point_distance = EstimateDistance(x, y, z);
      const EstimateDistanceQuery mx_distance = EstimateDistance(min_x, y, z);
      const EstimateDistanceQuery px_distance = EstimateDistance(max_x, y, z);
      const EstimateDistanceQuery my_distance = EstimateDistance(x, min_y, z);
      const EstimateDistanceQuery py_distance = EstimateDistance(x, max_y, z);
      const EstimateDistanceQuery mz_distance = EstimateDistance(x, y, min_z);
      const EstimateDistanceQuery pz_distance = EstimateDistance(x, y, max_z);
      // Compute gradient for each axis
      const double gx = ComputeAxisFineGradient(point_distance,
                                                  mx_distance,
                                                  px_distance,
                                                  x, min_x, max_x);
      const double gy = ComputeAxisFineGradient(point_distance,
                                                  my_distance,
                                                  py_distance,
                                                  y, min_y, max_y);
      const double gz = ComputeAxisFineGradient(point_distance,
                                                  mz_distance,
                                                  pz_distance,
                                                  z, min_z, max_z);
      return GradientQuery(gx, gy, gz);
    }
    else
    {
      return GradientQuery();
    }
  }

  GradientQuery GetFineGradient(
      const common_robotics_utilities::voxel_grid::GridIndex& index,
      const double nominal_window_size) const
  {
    return GetFineGradient4d(
        this->GridIndexToLocation(index), nominal_window_size);
  }

  GradientQuery GetFineGradient(
      const int64_t x_index, const int64_t y_index, const int64_t z_index,
      const double nominal_window_size) const
  {
    return GetFineGradient4d(
        this->GridIndexToLocation(x_index, y_index, z_index),
        nominal_window_size);
  }

  /// Project the provided point out of collision.

  ProjectedPosition ProjectOutOfCollision(
      const double x, const double y, const double z,
      const double stepsize_multiplier = 1.0 / 10.0) const
  {
    return ProjectOutOfCollision4d(
        Eigen::Vector4d(x, y, z, 1.0), stepsize_multiplier);
  }

  ProjectedPosition ProjectOutOfCollision3d(
      const Eigen::Vector3d& location,
      const double stepsize_multiplier = 1.0 / 10.0) const
  {
    return ProjectOutOfCollision(
        location.x(), location.y(), location.z(), stepsize_multiplier);
  }

  ProjectedPosition ProjectOutOfCollision4d(
      const Eigen::Vector4d& location,
      const double stepsize_multiplier = 1.0 / 10.0) const
  {
    return ProjectOutOfCollisionToMinimumDistance4d(
        location, 0.0, stepsize_multiplier);
  }

  ProjectedPosition ProjectOutOfCollisionToMinimumDistance(
      const double x, const double y, const double z,
      const double minimum_distance,
      const double stepsize_multiplier = 1.0 / 10.0) const
  {
    return ProjectOutOfCollisionToMinimumDistance4d(
        Eigen::Vector4d(x, y, z, 1.0), minimum_distance, stepsize_multiplier);
  }

  ProjectedPosition ProjectOutOfCollisionToMinimumDistance3d(
      const Eigen::Vector3d& location, const double minimum_distance,
      const double stepsize_multiplier = 1.0 / 10.0) const
  {
    return ProjectOutOfCollisionToMinimumDistance(
        location.x(), location.y(), location.z(),
        minimum_distance, stepsize_multiplier);
  }

  ProjectedPosition ProjectOutOfCollisionToMinimumDistance4d(
      const Eigen::Vector4d& location,
      const double minimum_distance,
      const double stepsize_multiplier = 1.0 / 10.0) const
  {
    // To avoid potential problems with alignment, we need to pass location
    // by reference, so we make a local copy here that we can change.
    // See https://eigen.tuxfamily.org/dox/group__TopicPassingByValue.html
    Eigen::Vector4d mutable_location = location;
    // If we are in bounds, start the projection process,
    // otherwise return the location unchanged
    if (this->LocationInBounds4d(mutable_location))
    {
      // Add a small collision margin to account for rounding and similar
      const double minimum_distance_with_margin
          = minimum_distance + GetResolution() * stepsize_multiplier * 1e-3;
      const double max_stepsize = GetResolution() * stepsize_multiplier;
      const bool enable_edge_gradients = true;
      double sdf_dist = EstimateDistance4d(mutable_location).Value();
      while (sdf_dist <= minimum_distance)
      {
        const GradientQuery gradient
            = GetCoarseGradient4d(mutable_location, enable_edge_gradients);
        if (gradient.HasValue())
        {
          const Eigen::Vector4d& grad_vector = gradient.Value();
          if (grad_vector.norm() > GetResolution() * 0.25) // Sanity check
          {
            // Don't step any farther than is needed
            const double step_distance
                = std::min(max_stepsize,
                           minimum_distance_with_margin - sdf_dist);
            mutable_location += grad_vector.normalized() * step_distance;
            sdf_dist = EstimateDistance4d(mutable_location).Value();
          }
          else
          {
            std::cerr << "Encountered flat gradient - stuck" << std::endl;
            return ProjectedPosition();
          }
        }
        else
        {
          std::cerr << "Failed to compute gradient - out of SDF?" << std::endl;
          return ProjectedPosition();
        }
      }
    }
    return ProjectedPosition(mutable_location);
  }

  /// The following function can be *very expensive* to compute, since it
  /// performs gradient ascent/descent across the SDF.
  common_robotics_utilities::voxel_grid::VoxelGrid<Eigen::Vector3d>
  ComputeLocalExtremaMap() const
  {
    const Eigen::Vector3d default_value(
        -std::numeric_limits<double>::infinity(),
        -std::numeric_limits<double>::infinity(),
        -std::numeric_limits<double>::infinity());
    common_robotics_utilities::voxel_grid::VoxelGrid<Eigen::Vector3d>
        watershed_map(this->GetOriginTransform(), this->GetGridSizes(),
                      default_value);
    for (int64_t x_idx = 0; x_idx < watershed_map.GetNumXCells(); x_idx++)
    {
      for (int64_t y_idx = 0; y_idx < watershed_map.GetNumYCells(); y_idx++)
      {
        for (int64_t z_idx = 0; z_idx < watershed_map.GetNumZCells(); z_idx++)
        {
          // We use an "unsafe" function here because we know all the indices
          // we provide it will be safe
          FollowGradientsToLocalExtremaUnsafe(
              x_idx, y_idx, z_idx, watershed_map);
        }
      }
    }
    return watershed_map;
  }
};
}  // namespace voxelized_geometry_tools

extern template class voxelized_geometry_tools::SignedDistanceField<double>;
extern template class voxelized_geometry_tools::SignedDistanceField<float>;