function [dx,dv,dxq,dextrap] = internal_interp1Backward(dvq,Flags,Method,x,v,xq,doExtrap)
% This function is undocumented and will change in a future release.

% Backward method for Interp1Op.
%
% vq = interp1(x, v, xq, method, extrap), compute reverse derivative.
%
% Inputs:
%   dvq      = gradient of loss w.r.t vq
%   Flags    = logical vector indicating which outputs to compute
%   Method   = method description: 'linear' or 'nearest'
%   x        = grid vector
%   v        = values
%   xq       = query points
%   extrap   = boolean, true if input was 'extrap', false if it was a numeric scalar.
%
% Outputs:
%   dx       = gradient of loss w.r.t x
%   dv       = gradient of loss w.r.t v
%   dxq      = gradient of loss w.r.t xq
%   dextrap  = gradient of loss w.r.t extrap

%   Copyright 2020-2022 The MathWorks, Inc.

dxNeeded = Flags(1);
dvNeeded = Flags(2);
dxqNeeded = Flags(3);
dextrapNeeded = Flags(4);

dx = [];
dv = [];
dxq = [];
dextrap = [];

% Get original sizes of the inputs
szX = size(x);
szV = size(v);
szXq = size(xq);

% Extract relevant sizes for this operation
numInterp = numel(x);
numBatch = numel(v) / numInterp;
numQuery = numel(xq);

% Standarize sizes of inputs
x = x(:);                                % numInterp-by-1
v = reshape(v, numInterp, numBatch);     % numInterp-by-numBatch
xq = xq(:);                              % numQuery-by-1
dvq = reshape(dvq, numQuery, numBatch);  % numQuery-by-numBatch

% Note: dx, dv, dxq, dvq have the same size as x, v, xq,
% vq, respectively.

if Method == "nearest"
    mapQueryToGrid = iLookupNearest(x, xq, doExtrap);
    
    if dvNeeded
        dv = iValueNearest(x, xq, dvq, mapQueryToGrid);
        dv = reshape(dv, szV);
    end
    
    if dextrapNeeded
        if doExtrap
            % extrap == 'extrap', can't request dextrap in this case
            dextrap = [];
        else
            dextrap = sum(dvq(mapQueryToGrid == 0, :), 'all');
        end
    end
    
else % op.Method == "linear"
    
    mapQueryToGrid = iLookupLinear(x, xq, doExtrap);
    
    if dxNeeded
        dx = iCoordLinear(x, xq, v, dvq, mapQueryToGrid);
        dx = reshape(dx, szX);
    end
    
    if dxqNeeded
        dxq = iQueryLinear(x, xq, v, dvq, mapQueryToGrid);
        dxq = reshape(dxq, szXq);
    end
    
    if dvNeeded
        dv = iValueLinear(x, xq, dvq, mapQueryToGrid);
        dv = reshape(dv, szV);
    end
    
    if dextrapNeeded
        if doExtrap
            % extrap == 'extrap', can't request dextrap in this case
            dextrap = [];
        else
            dextrap = sum(dvq(mapQueryToGrid == 0, :), 'all');
        end
    end
end
end


function mapQueryToGrid = iLookupNearest(x, xq, doExtrap)
% Use another call to 'nearest', but set the values to become indices into
% the values array.
%
%  If doExtrap is false:
%    j = mapQueryToGrid(i) is a value from 0:length(x)
%    If j==0, xq(i) is not in bounds.
%    If 1<j<length(x), xq(i) is closest to x(j).
%
%  If doExtrap is true:
%    j = mapQueryToGrid(i) is a value from 1:length(x)
%    xq(i) is closest to x(j).

if doExtrap
    % x = 2:5, xq = 1:6 -> mapQueryToGrid = [1 1 2 3 4 4]
    mapQueryToGrid = interp1(x, 1:length(x), xq, 'nearest', 'extrap');
else
    % x = 2:5, xq = 1:6 -> mapQueryToGrid = [0 1 2 3 4 0]
    mapQueryToGrid = interp1(x, 1:length(x), xq, 'nearest', 0);
end

end


function mapQueryToGrid = iLookupLinear(x, xq, doExtrap)
% Returns a vector mapping every query point in xq to an interval in the
% grid vector x (which is strictly monotonically increasing).
%
%  If doExtrap is false:
%    j = mapQueryToGrid(i) is a value from 0:length(x)-1.
%    If j==0, xq(i) is not in bounds          xq(i) > x(end) or xq(i) < x(1)
%    If 0 < j < length(x)-1, xq(i) satisfies  x(j) <= xq(i) < x(j+1)
%    If j==length(x)-1, xq(i) satisfies       x(end-1) <= xq(i) <= x(end).
%
%  If doExtrap is true:
%    j = mapQueryToGrid(i) is a value from 1:length(x)-1
%    If j==1, xq(i) satisfies                 xq(i) < x(2).
%    If 1 < j < length(x)-1, xq(i) satisfies  x(j) <= xq(i) < x(j+1)
%    If j==length(x)-1, xq(i) satisfies       x(end-1) <= xq(i).

% Need to move from gpuArray to CPU, because option 'previous' is not
% supported for gpuArray/interp1.
x = gather(x);
xq = gather(xq);

rm = realmax(class(x));

if doExtrap
    xmod = x;
    xmod(1) = -rm;
    xmod(end) = rm;
    mapQueryToGrid = interp1(xmod, 1:length(x), xq, 'previous');
else
    % Compute index vector such that j = indxqInx(i) implies
    % that xq(i) is between x(j) and x(j+1), x(j) <= xq(i) < x(j+1).
    % If j == 0, this means that xq(i) is outside of [x(1) x(end)[.
    mapQueryToGrid = interp1([-rm; x(:); rm], [0:length(x)-1 0 0], xq, 'previous');
    
    % Consistent treatment of the case where xq(ii) is equal to x(end): Treat
    % this as part of the linear interpolation, not as outside of the interval.
    xend = x(end);
    mapQueryToGrid(xq == xend) = length(x) - 1;
end

end


function dv = iValueNearest(x, xq, dvq, mapQueryToGrid)

szV = size(dvq);
szV(1) = length(x);

numInterp = length(x);
numQuery = length(xq);

% Vectorized version:
iivec = find(mapQueryToGrid(:));
jjvec = double(mapQueryToGrid(iivec));

AM = sparse(jjvec, iivec, 1, numInterp, numQuery);

AM = iEnsureSameBackend(AM, dvq);
dv = AM * reshape(double(dvq), numQuery, []);

dv = reshape(dv, szV);
dv = cast(dv, 'like', dvq);

%{
% Direct indexing variant
dv = zeros(szV, 'like', dvq);
for ii=1:numQuery
    jj = indxqInx(ii);
    if jj == 0, continue; end
    dv(jj, :) = dv(jj, :) + dvq(ii, :);
end
%}

end


function dx = iCoordLinear(x, xq, v, dvq, mapQueryToGrid)

numInterp = length(x);

% Vectorized version:
iivec = find(mapQueryToGrid(:));
jjvec = mapQueryToGrid(iivec);
iivec = iivec(:);
jjvec = jjvec(:);

commonFactor = sum(dvq(iivec, :) .* (v(jjvec+1, :) - v(jjvec, :)), 2) ./ (x(jjvec+1) - x(jjvec)).^2;

% Need to cast dx to real, but it's cheaper to go one step earlier and cast
% commonFact to real already, since the following formula will not mix up
% the real and imaginary parts of commonFactor.
commonFactor = real(commonFactor);

dx = -accumarray(jjvec+1, commonFactor .* (xq(iivec) - x(jjvec)), [numInterp 1]) ...
    -accumarray(jjvec, commonFactor .* (x(jjvec+1) - xq(iivec)), [numInterp 1]);

%{
% Direct indexing variant
numQuery = length(xq);
dx = zeros(numInterp, 1, 'like', dvq);
for ii=1:numQuery
    jj = indxqInx(ii);
    if jj == 0, continue; end
    commonFactor = sum(dvq(ii, :) .* (v(jj+1, :) - v(jj, :))) / (x(jj+1) - x(jj)).^2;
    dx(jj+1) = dx(jj+1) - commonFactor * (xq(ii)  - x(jj));
    dx(jj)   = dx(jj)   - commonFactor * (x(jj+1) - xq(ii));
end
    %}
end


function dxq = iQueryLinear(x, xq, v, dvq, mapQueryToGrid)

% Vectorized version:
iivec = find(mapQueryToGrid(:));
jjvec = mapQueryToGrid(iivec);
iivec = iivec(:);
jjvec = jjvec(:);

commonFactor = sum(dvq(iivec, :) .* (v(jjvec+1, :) - v(jjvec, :)), 2);

% Need to cast dx to real, but it's cheaper to go one step earlier and cast
% commonFact to real already, since the following formula will not mix up
% the real and imaginary parts of commonFactor.
commonFactor = real(commonFactor);

dxq = iRealZerosLike([length(xq), 1], dvq);
dxq(iivec) = commonFactor ./ (x(jjvec+1) - x(jjvec));

%{
% Direct indexing variant
dxq = zeros(length(xq), 1, 'like', dvq);
for ii=1:length(xq)
    jj = indxqInx(ii);
    if jj == 0, continue; end
    dxq(ii) = sum(dvq(ii, :) .* (v(jj+1, :) - v(jj, :)) / (x(jj+1) - x(jj)));
end
%}
end

function X = iRealZerosLike(szX, dZ)
% Return real array of all zeros with given size and datatype.
X = zeros(szX, 'like', real(dZ([])));
end

function dv = iValueLinear(x, xq, dvq, mapToQueryGrid)

szV = size(dvq);
szV(1) = length(x);

numInterp = length(x);
numQuery = length(xq);

iivec = find(mapToQueryGrid(:));
jjvec = double(mapToQueryGrid(iivec));
iivec = iivec(:);
jjvec = jjvec(:);


xdiff = x(jjvec+1) - x(jjvec);
v1 = (xq(iivec) - x(jjvec)) ./ xdiff;
v2 = (x(jjvec+1) - xq(iivec)) ./ xdiff;

AM1 = sparse(jjvec+1, iivec, double(v1), numInterp, numQuery);
AM2 = sparse(jjvec, iivec, double(v2), numInterp, numQuery);
AM = AM1 + AM2;

AM = iEnsureSameBackend(AM, dvq);
dvq = reshape(dvq, size(dvq, 1), []);
dv = AM*double(dvq);
dv = reshape(dv, szV);
dv = cast(dv, 'like', dvq);


%{
% Direct indexing variant:
dv = zeros(szV, 'like', dvq);
for ii=1:numQuery
    jj = indxqInx(ii);
    if jj == 0, continue; end
    commonFactor = dvq(ii, :) / (x(jj+1) - x(jj));
    dv(jj+1, :) = dv(jj+1, :) + commonFactor .* (xq(ii) - x(jj));
    dv(jj, :) = dv(jj, :) + commonFactor .* (x(jj+1) - xq(ii));
end
%}

end


% Performance considerations: Other variants for computing iValueLinear
%{
numBatch = numel(dvq) / numQuery;

% Variant using accumarray and loop:
iivec = find(indxqInx(:));
jjvec = indxqInx(iivec);
iivec = iivec(:);
jjvec = jjvec(:);
x = x(:);
xq = xq(:);

dv = zeros(szV, 'like', dvq);

dvqLeft = (xq(iivec) - x(jjvec)) ./ (x(jjvec+1) - x(jjvec)) .* dvq(iivec, :);
dvqRight = (x(jjvec+1) - xq(iivec)) ./ (x(jjvec+1) - x(jjvec)) .* dvq(iivec, :);

dvqBoth = [dvqLeft; dvqRight];
jjBoth = [jjvec+1; jjvec];

for ll=1:numBatch
    dv(:, ll) = accumarray(jjBoth, dvqBoth(:, ll), [numInterp 1]);
end
%}
%{
% Another variant
%AM = sparse(jjvec+1, iivec, (xq(iivec) - x(jjvec)) ./ (x(jjvec+1) - x(jjvec)), numInterp, numQuery) + ...
%  sparse(jjvec, iivec, (x(jjvec+1) - xq(iivec)) ./ (x(jjvec+1) - x(jjvec)), numInterp, numQuery);
%}
%{
% Another idea to consider: Construct one large linear index to pass into
accumarray? Seems a bad idea, because gpuArray/accumarray starts out by
sorting the complete index vector that's passed in.
%}

function AM = iEnsureSameBackend(AM, dvq)
% Convert AM to a gpuArray before matrix multiply to allocate the memory
% now
if isgpuarray(dvq)
    AM = gpuArray(AM);
end
end
