classdef signedDistanceMap3D < handle
%signedDistanceMap3D Store signed distance over variably sized 3D region
%
%   The signedDistanceMap3D object is used to create and store a Truncated
%   Signed Distance Field over a voxelized 3D space. Voxels within a
%   specified distance from an occupied boundary contain positive distance
%   values and negative values if they lie within occupied space. Voxels
%   that have not been witnessed, or that lie beyond the truncated distance
%   will be given a "background" value.
%
%   SDF = signedDistanceMap3D creates an empty 3D signed distance field
%         with default resolution and raycast settings.
%
%   SDF = signedDistanceMap3D(Name=Value) accepts additional name-value
%         pairs:
%
%           Resolution
%           FullTracing
%           TruncationDistance
%
%   signedDistanceMap3D properties:
%       Resolution          - Grid resolution in cells per meter
%       FullTracing         - Determines whether full or local raycasting occurs
%       TruncationDistance  - Max distance from occupied boundary
%       NumActiveVoxels     - Number of active voxels in the map
%       MapLimits           - Min and max limits of active voxels in XYZ
%
%   signedDistanceMap3D methods:
%       distance            - Compute distance to the zero level set
%       activeVoxels        - Return centers, sizes, and values of all active voxels
%       gradient            - Compute gradient of the signed distance field
%       insertPointCloud    - Integrate new sensor readings into the map
%       show                - Display the signed distance map
%
%   Example:
%
%       % Create empty 3D signed distance field with default settings
%       sdf3D = signedDistanceMap3D;
%
%       % Create empty field with resolution of 10 cells/meter
%       sdf3D = signedDistanceMap3D(Resolution=10);
%
%       % Insert a pointcloud
%       origin = [0 0 0];
%       endPoints = rand(10,3)*10;
%       insertPointCloud(sdf3D,origin,endPoints);
%
%       % Compute distances along one of the rays
%       rayLen = norm(endPoints(1,:),2);
%       L = linspace(0, rayLen + sdf.TruncationDistance*2);
%       signedDistance = distance(sdf3D,endPoints(1,:)/rayLen*L);
%
%       % Compute distance and gradient around the capsule
%       dv = linspace(-5,5,50);
%       [xx,yy,zz] = meshgrid(dv+T(1,end),dv+T(2,end),dv+T(3,end))
%       pts = [xx(:) yy(:) zz(:)];
%       dist = distance(sdf3D,pts);
%       grad = gradient(sdf3D,pts);
%
%   See also: signedDistanceMap

%   Copyright 2023 The MathWorks, Inc.

%#codegen

    properties (Access = private)
        VDBBuiltin
    end

    properties (GetAccess = public)
        %Resolution Grid resolution in cells per meter
        %
        %   Default: 1 cell/meter
        Resolution

        %FullTracing Determines whether full or local raycasting occurs
        %
        %   false - Inserted rays only update voxels within
        %           +/-TruncationDistance of the ray's endpoint
        %
        %              (Default)
        %
        %   true  - All voxels between the ray origin and
        %           TruncationDistance past the ray endpoint are updated
        FullTracing

        %TruncationDistance Max distance from occupied boundary
        %
        %   If not provided, the default TruncationDistance will be set to
        %   the minimum distance required to ensure that the TSDF region is
        %   extends at least 3 voxels thick on either side of the occupied
        %   boundary.
        %
        %   Default: 4 (voxels)
        TruncationDistance
    end

    properties (Dependent)
        %NumActiveVoxels Number of active voxels in the map
        %
        %   Active voxels are those that contain computed distance values
        NumActiveVoxels

        %MapLimits Min and max limits of active voxels in XYZ
        MapLimits
    end

    methods
        function obj = signedDistanceMap3D(nv)
            arguments
                nv.Resolution (1,1) {mustBeNumeric, mustBePositive, mustBeFinite} = 1.0;                    % cell/m
                nv.FullTracing (1,1) {mustBeNumericOrLogical, mustBeMember(nv.FullTracing,[0,1])} = false   % logical
                nv.TruncationDistance (1,1) = nan                                                           % m
            end

            if isnan(nv.TruncationDistance)
                % The minimum truncation distance
                nv.TruncationDistance = 3/nv.Resolution;
            else
                obj.validateTruncationDistance(3/nv.Resolution,nv.TruncationDistance);
            end

            % Construct builtin
            obj.VDBBuiltin = nav.geometry.internal.VDBVolumeBuiltin(single(nv.Resolution), ...
                                                                    single(nv.TruncationDistance),logical(nv.FullTracing));
        end

        function insertPointCloud(obj,origin,pt)
        %insertPointCloud Integrate new sensor readings into the map
        %
        %   This method uses the DDA algorithm to update distance values in
        %   all voxels hit by each ray. If FullTracing is true, then the
        %   ray is traced between the origin and a point lying past the
        %   end-point by TruncationDistance. When FullTracing is false, the
        %   ray will only begin updating cells that are within
        %   TruncationDistance of the endpoint along the ray.
        %
        %   insertPointCloud(OBJ,ORIGIN,PT) casts rays between the [x y z] origin
        %   and each point in the N-by-[X Y Z] PT matrix.
        %
        %   Example:
        %
        %       % Create empty field with resolution of 10 cells/meter
        %       sdf3D = signedDistanceMap3D(Resolution=10);
        %
        %       % Insert a pointcloud
        %       origin = [0 0 0];
        %       endPoints = rand(10,3)*10;
        %       insertPointCloud(sdf3D,origin,endPoints);
        %
        %       % Compute distances along one of the rays
        %       rayLen = norm(endPoints(1,:),2);
        %       L = linspace(0, rayLen + sdf.TruncationDistance*2);
        %       signedDistance = distance(sdf3D,endPoints(1,:)/rayLen*L);
        %
        %   See also: distance, gradient

            arguments
                obj
                origin (1,3) double {mustBeReal}
                pt (:,3) double {mustBeReal}
            end
            goodPts = vecnorm(pt-origin,2,2) > 1e-7;
            if any(goodPts)
                obj.VDBBuiltin.insertPointCloud(origin,pt(goodPts,:));
            end
        end

        function dist = distance(obj,pts,nv)
        %distance Compute distance to the zero level set
        %
        %   DIST = distance(OBJ,PTS) takes in an Nx3 matrix of xyz
        %   query points, PTS, and returns the distance stored in their
        %   corresponding voxel. If the voxel lies outside the truncated or
        %   dilated radius, the distance will be the background value.
        %
        %   DIST = distance(OBJ,PTS,Name,Value) takes in optional
        %   name-value pairs:
        %
        %       InterpolationMethod - Method used for interpolation
        %
        %           nearest     : Returns the distance stored in the xyz's
        %                         corresponding voxel.
        %
        %           linear      : Uses trilinear interpolation to compute
        %                         distance using the 8 voxels surrounding
        %                         the query point.
        %
        %           quadratic   : Uses quadratic interpolation to compute
        %                         distance using the 27 voxels surrounding
        %                         the query point.
        %
        %   Example:
        %
        %       % Create empty field with resolution of 10 cells/meter
        %       sdf3D = signedDistanceMap3D(Resolution=10);
        %
        %       % Insert a pointcloud
        %       origin = [0 0 0];
        %       endPoints = rand(10,3)*10;
        %       sdf3D.insertPointCloud(origin,endPoints);
        %
        %       % Compute distance recorded in active voxels
        %       queryPt = T(1:3,end)' + 0.5;
        %       distNone        = distance(sdf3D,queryPt);
        %       distLinear      = distance(sdf3D,queryPt,InterpolationMethod="linear");
        %       distQuadratic   = distance(sdf3D,queryPt,InterpolationMethod="quadratic");
        %
        %   See also: gradient

            arguments
                obj
                pts (:,3) double {mustBeReal, mustBeFinite}
                nv.InterpolationMethod {mustBeMember(nv.InterpolationMethod,["nearest","linear","quadratic"])} = "nearest"
            end
            switch nv.InterpolationMethod
              case "nearest"
                iMethod = 0;
              case "linear"
                iMethod = 1;
              case "quadratic"
                iMethod = 2;
            end
            dist = obj.VDBBuiltin.distance(pts,iMethod);
        end

        function grad = gradient(obj,pts,nv)
        %gradient Compute gradient of the signed distance field
        %
        %   GRAD = gradient(OBJ,PTS) takes in an Nx3 matrix of xyz
        %   query points, PTS, and returns the Nx3 gradient computed from
        %   voxels in the vicinity.
        %
        %   GRAD = gradient(OBJ,PTS,Name=Value) takes in optional
        %   name-value pairs:
        %
        %       InterpolationMethod - Method used for interpolation
        %
        %           linear      : Computes partial derivatives of the
        %                         trilinearly-interpolated distance field.
        %
        %           quadratic   : Computes partial derivatives of the
        %                         triquadratically-interpolated distance
        %                         field.
        %   Example:
        %
        %       % Create empty field with resolution of 10 cells/meter
        %       sdf3D = signedDistanceMap3D(Resolution=10);
        %
        %       % Insert a pointcloud
        %       origin = [0 0 0];
        %       endPoints = rand(10,3)*10;
        %       sdf3D.insertPointCloud(origin,endPoints);
        %
        %       % Compute gradient in active voxels using various types of interpolation
        %       queryPt = T(1:3,end)' + 0.5;
        %       gradLinear      = gradient(sdf3D,queryPt,InterpolationMethod="linear");
        %       gradQuadratic   = gradient(sdf3D,queryPt,InterpolationMethod="quadratic");
        %
        %   See also: distance

            arguments
                obj
                pts (:,3) double {mustBeReal, mustBeFinite}
                nv.InterpolationMethod {mustBeMember(nv.InterpolationMethod,["linear","quadratic"])} = "linear"
            end
            switch nv.InterpolationMethod
              case "linear"
                iMethod = 1;
              case "quadratic"
                iMethod = 2;
            end
            grad = obj.VDBBuiltin.gradient(pts,iMethod);
        end

        function varargout = activeVoxels(obj)
        %activeVoxels Return centers, sizes, and values of all active voxels
        %   [CENTERS,SIZES,VALUES] = activeVoxels(OBJ) retrieves voxel
        %   CENTERS (Nx3), SIZES (Nx1), and distance VALUES (Nx1) for all
        %   active voxels in the map.
        %
        %   Example:
        %
        %       % Create empty field with resolution of 10 cells/meter
        %       sdf3D = signedDistanceMap3D(Resolution=10);
        %
        %       % Insert a pointcloud
        %       origin = [0 0 0];
        %       endPoints = rand(10,3)*10;
        %       sdf3D.insertPointCloud(origin,endPoints);
        %
        %       % Extract all voxels in the truncated area
        %       [ctrs,sizes,dists] = activeVoxels(sdf3D);

        % Retrieve voxel information for all active voxels
            output = obj.VDBBuiltin.activeVoxels();
            [varargout{1:nargout}] = deal(output{1:nargout});
        end

        function [h,hBar] = show(obj,nv)
        %show Display the signed distance map
        %
        %   [H,HBAR] = show(OBJ) Displays all active voxels as a scatter
        %   plot, where red->green marker colors corresponding to most
        %   negative to most positive distance values.
        %
        %   [H,HBAR] = show(OBJ,Name=Value) takes in optional name-value
        %   pairs:
        %
        %       Parent          - Axes to plot the map, specified as an
        %                         axes handle.
        %
        %                           Default: gca
        %
        %   Example:
        %
        %       % Create empty field with resolution of 10 cells/meter
        %       sdf3D = signedDistanceMap3D(Resolution=10);
        %
        %       % Insert a pointcloud
        %       origin = [0 0 0];
        %       endPoints = rand(10,3)*10;
        %       sdf3D.insertPointCloud(origin,endPoints);
        %
        %       % Compute distances along one of the rays
        %       rayLen = norm(endPoints(1,:),2);
        %       L = linspace(0, rayLen + sdf.TruncationDistance*2);
        %       signedDistance = sdf3D.distance(endPoints(1,:)/rayLen*L);
        %
        %       % Display the map
        %       ax1 = subplot(1,2,1);
        %       subplot(1,2,2);
        %       sdf3D.show(Parent=ax);  % Display on 1st axes

            arguments
                obj
                nv.Parent (1,1) {mustBeA(nv.Parent,'matlab.graphics.axis.Axes')} = gca
            end

            % Retrieve ax handle
            ax = nv.Parent;

            % Extract voxel information
            [voxCtrs,~,voxVals] = obj.activeVoxels();

            % Create colormap
            greenColorMap = [zeros(1, 132), linspace(0, 1, 124)];
            redColorMap = [linspace(1, 0, 124), zeros(1, 132)];
            cmap = [redColorMap; greenColorMap; zeros(1, 256)]';

            % Convert voxel values to map values
            maxVal = max(abs(voxVals));
            C = .5 + voxVals/maxVal/2;
            cIdx = discretize(C,linspace(0,1,256));

            % Update colormap and display points
            hold(ax,"on");
            h = scatter3(ax,voxCtrs(:,1),voxCtrs(:,2),voxCtrs(:,3),[],cIdx,'.');
            colormap(ax,cmap);
            hBar = colorbar(ax);
            ticks = hBar.Ticks;
            tickLabels = linspace(-maxVal,maxVal,numel(ticks));
            hBar.TickLabels = tickLabels;
        end

        function numVoxel = get.Resolution(obj)
        %get.NumActiveVoxels
            numVoxel = obj.VDBBuiltin.Resolution;
        end

        function numVoxel = get.FullTracing(obj)
        %get.NumActiveVoxels
            numVoxel = obj.VDBBuiltin.FullTracing;
        end

        function numVoxel = get.TruncationDistance(obj)
        %get.NumActiveVoxels
            numVoxel = obj.VDBBuiltin.TruncationDistance;
        end

        function numVoxel = get.NumActiveVoxels(obj)
        %get.NumActiveVoxels
            numVoxel = obj.VDBBuiltin.NumActiveVoxel;
        end

        function bounds = get.MapLimits(obj)
        %get.MapLimits
            bounds = obj.VDBBuiltin.ActiveBoundingBox;
        end

        function cObj = copy(obj)
        % Create empty object
            cObj = signedDistanceMap3D;

            % Copy over serialized information
            data = obj.VDBBuiltin.serialize();
            cObj.VDBBuiltin.deserialize(data);
        end
    end

    methods (Hidden,Static)
        function obj = loadobj(S)
        % Create empty map
            obj = signedDistanceMap3D;

            % Populate with serialized data
            obj.VDBBuiltin.deserialize(S.Data);
        end
    end

    methods (Hidden)
        function S = saveobj(obj)
            S = struct('Data',obj.VDBBuiltin.serialize);
        end
        function [vertices,faces] = mesh(obj,nv)
        %mesh Generate an iso-surface mesh from active voxels
        %
        %   [VERTICES,FACES] = mesh(OBJ) Uses the Marching Cubes
        %   algorithm to generate a mesh from voxels making up the
        %   zero-level set. The mesh is returned as an Vx3 set of xyz
        %   VERTICES and an Fx3 set of FACES, where each row contains 3
        %   vertex-indices making up a triangular face.
        %
        %   [VERTICES,FACES] = mesh(OBJ,Name=Value) takes in
        %   optional name-value pairs:
        %
        %       FillHoles - If true, the algorithm attempts to fill holes
        %                   in the mesh by connecting zero-level set voxels
        %                   within TruncationDistance of each other, if
        %                   false (default), voxel faces will only
        %                   connected immediate neighbors.
        %
        %       MinWeight - Minimum weight-value required for a voxel to
        %                   be considered in the mesh.
        %
        %   Example:
        %
        %       % Create empty field with resolution of 10 cells/meter
        %       sdf3D = signedDistanceMap3D(Resolution=10);
        %
        %       % Insert a pointcloud
        %       origin = [0 0 0];
        %       endPoints = rand(10,3)*10;
        %       sdf3D.insertPointCloud(origin,endPoints);
        %
        %       % Generate mesh along the iso-boundary
        %       [V,F] = mesh(sdf3D);
        %
        %       % Display the result
        %       patch(Vertices=V,Faces=F);
        %
        %   See also: insertPointCloud

            arguments
                obj
                nv.FillHoles (1,1) logical {mustBeMember(nv.FillHoles,[0,1])} = false   % logical
                nv.MinWeight (1,1) double {} = 0
            end
            output = obj.VDBBuiltin.mesh(nv.FillHoles, nv.MinWeight);
            vertices = output{1};
            faces = output{2};
        end
    end

    methods (Hidden,Static)
        function validateTruncationDistance(minEdgeLength,truncDist)
            arguments
                minEdgeLength (1,1) {mustBePositive} %#ok<INUSA>
                truncDist (1,1) {mustBeNumeric, mustBeGreaterThanOrEqual(truncDist,minEdgeLength)} %#ok<INUSA>
            end
        end
    end
end
