## Copyright (C) 2024 David Legland
## All rights reserved.
## 
## Redistribution and use in source and binary forms, with or without
## modification, are permitted provided that the following conditions are met:
## 
##     1 Redistributions of source code must retain the above copyright notice,
##       this list of conditions and the following disclaimer.
##     2 Redistributions in binary form must reproduce the above copyright
##       notice, this list of conditions and the following disclaimer in the
##       documentation and/or other materials provided with the distribution.
## 
## 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 AUTHOR 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.
## 
## The views and conclusions contained in the software and documentation are
## those of the authors and should not be interpreted as representing official
## policies, either expressed or implied, of the copyright holders.

function varargout = drawArrow3d(varargin)
%DRAWARROW3D Plot a quiver of 3D arrows.
%
%   drawArrow3d(pos, vec) 
%   Plots 3D arrows given the (pos)ition array [x1 y1 z1; x2 y2 z2; ...] 
%   and the (vec)tor array [dx1 dy1 dz1; dx2 dy2 dz2; ...].
%
%   drawArrow3d(pos, vec, color)
%   Optional positional argument color conforms to 'ColorSpec.'  
%   For example, 'r','red',[1 0 0] will all plot a quiver with all arrows 
%   as red. This can also be in the form of Nx3 where 'N' is the number of 
%   arrows and each column corresponds to the RGB values. Default color is 
%   black.
%
%   drawArrow3d(...,Name,Value) Optional name-value pair arguments:
%   'stemRatio': Ratio of the arrow head (cone) to the arrow stem (cylinder)
%       For example, setting this value to 0.94 will produce arrows with 
%       arrow stems 94% of the length and short, 6% cones as arrow heads.
%       Values above 0 and below 1 are valid. Default is 0.75.
%   'arrowRadius': changes the radius of the arrowstem. Percentage of the
%       lenght of the arrow. Values between 0.01 and 0.3 are valid. 
%       Default is 0.025.
%   Uses the 'patch' function to plot the arrows. 'patch' properties can be  
%   used to control the appearance of the arrows.
%
%   drawArrow3d(AX,...) plots into AX instead of GCA.
%
%   H = drawArrow3d(...) returns the handles of the arrows.
%
% Example:
%    [X,Y] = meshgrid(1:5, -2:2);
%    Z = zeros(size(X));
%    pos = [X(:),Y(:),Z(:)];
%    vec = zeros(size(pos));
%    vec(:,1) = 1;
%    drawArrow3d(pos, vec, 'g', 'stemRatio', 0.6);
%    view(3); lighting('phong'); camlight('head'); axis('equal')
%

% ------
% Authors: Shawn Arseneau, oqilipo
% E-mail: N/A
% Created: 2006-09-14, by Shawn Arseneau
% Copyright 2006-2023

% extract handle of axis to draw on
[ax, varargin] = parseAxisHandle(varargin{:});

% retrieve positions and vectors
pos = varargin{1};
vec = varargin{2};
varargin(1:2) = [];

numArrows = size(pos,1);
if numArrows ~= size(vec,1)
    error(['Number of rows of position and magnitude inputs do not agree. ' ...
        'Type ''help drawArrow3d'' for details']);
end

% Parsing
p = inputParser;
p.KeepUnmatched = true;
isPointArray3d = @(x) validateattributes(x,{'numeric'},...
    {'nonempty','nonnan','real','finite','size',[nan,3]});
addRequired(p,'pos',isPointArray3d)
addRequired(p,'vec',isPointArray3d);
addOptional(p,'color', 'k', @(x) validateColor(x, numArrows));
isStemRatio = @(x) validateattributes(x,{'numeric'},{'vector','>', 0, '<', 1});
addParameter(p,'stemRatio', 0.75, isStemRatio);
isArrowRadius = @(x) validateattributes(x,{'numeric'},{'scalar','>=', 0.01, '<=', 0.3});
addParameter(p,'arrowRadius',0.025, isArrowRadius);

parse(p,pos,vec,varargin{:});
pos = p.Results.pos;
vec = p.Results.vec;
[~, color] = validateColor(p.Results.color, numArrows);
stemRatio = p.Results.stemRatio;
if numel(stemRatio) == 1; stemRatio = repmat(stemRatio,numArrows,1); end
arrowRadius = p.Results.arrowRadius;
if numel(arrowRadius) == 1; arrowRadius = repmat(arrowRadius,numArrows,1); end
drawOptions = p.Unmatched;

% save hold state
holdState = ishold(ax);
hold(ax, 'on');


%% Loop through all arrows and plot in 3D
qHandle = gobjects(numArrows,1);
for i = 1:numArrows
    qHandle(i) = drawSingleVector3d(ax, pos(i,:), vec(i,:), color(i,:), ...
        stemRatio(i), arrowRadius(i), drawOptions);
end

% restore hold state
if ~holdState
    hold(ax, 'off');
end

if nargout > 0
    varargout = {qHandle};
end

end

function [valid, color] = validateColor(color, numArrows)
valid = true;
[arrowRow, arrowCol] = size(color);
if arrowRow == 1
    if ischar(color) %in ShortName or LongName color format
        color = repmat(color,numArrows,1);
    else
        if arrowCol ~= 3
            error('color in RGBvalue must be of the form 1x3.');
        end
        color = repmat(color, numArrows, 1);
    end
elseif arrowRow ~= numArrows
    error('color in RGBvalue must be of the form Nx3.');
end

end

function arrowHandle = drawSingleVector3d(hAx, pos, vec, color, stemRatio, arrowRadius, drawOptions)
%ARROW3D Plot a single 3D arrow with a cylindrical stem and cone arrowhead
%
% See header of drawArrow3d

X = pos(1); Y = pos(2); Z = pos(3);

[~, ~, srho] = cart2sph(vec(1), vec(2), vec(3));

%% CYLINDER == STEM
cylinderRadius = srho*arrowRadius;
cylinderLength = srho*stemRatio;
[CX,CY,CZ] = cylinder(cylinderRadius);
CZ = CZ.*cylinderLength; % lengthen

% Rotate Cylinder
[row, col] = size(CX); % initial rotation to coincide with x-axis

newCyl = transformPoint3d([CX(:), CY(:), CZ(:)], createRotationVector3d([1 0 0],[0 0 -1]));
CX = reshape(newCyl(:,1), row, col);
CY = reshape(newCyl(:,2), row, col);
CZ = reshape(newCyl(:,3), row, col);

[row, col] = size(CX);
newCyl = transformPoint3d([CX(:), CY(:), CZ(:)], createRotationVector3d([1 0 0],vec));
stemX = reshape(newCyl(:,1), row, col);
stemY = reshape(newCyl(:,2), row, col);
stemZ = reshape(newCyl(:,3), row, col);

% Translate cylinder
stemX = stemX + X;
stemY = stemY + Y;
stemZ = stemZ + Z;

%% CONE == ARROWHEAD
RADIUS_RATIO = 1.5;
coneLength = srho*(1-stemRatio);
coneRadius = cylinderRadius*RADIUS_RATIO;
incr = 4;  % Steps of cone increments
coneincr = coneRadius/incr;
[coneX, coneY, coneZ] = cylinder(cylinderRadius*2:-coneincr:0); % Cone
coneZ = coneZ.*coneLength;

% Rotate cone
[row, col] = size(coneX);
newCyl = transformPoint3d([coneX(:), coneY(:), coneZ(:)], createRotationVector3d([1 0 0], [0 0 -1]));
coneX = reshape(newCyl(:,1), row, col);
coneY = reshape(newCyl(:,2), row, col);
coneZ = reshape(newCyl(:,3), row, col);

newCyl = transformPoint3d([coneX(:), coneY(:), coneZ(:)], createRotationVector3d([1 0 0], vec));
headX = reshape(newCyl(:,1), row, col);
headY = reshape(newCyl(:,2), row, col);
headZ = reshape(newCyl(:,3), row, col);

% Translate cone
% centerline for cylinder: the multiplier is to set the cone 'on the rim' of the cylinder
V = [0, 0, srho*stemRatio];
Vp = transformPoint3d(V, createRotationVector3d([1 0 0], [0 0 -1]));
Vp = transformPoint3d(Vp, createRotationVector3d([1 0 0], vec));
headX = headX + Vp(1) + X;
headY = headY + Vp(2) + Y;
headZ = headZ + Vp(3) + Z;

% Draw cylinder & cone
hStem = patch(hAx, surf2patch(stemX, stemY, stemZ), 'FaceColor', color, 'EdgeColor', 'none', drawOptions);
hHead = patch(hAx, surf2patch(headX, headY, headZ), 'FaceColor', color, 'EdgeColor', 'none', drawOptions);
arrowHandle = hggroup(hAx);
set([hStem, hHead], 'Parent', arrowHandle);

end