########################################################################
##
## Copyright (C) 2022 Kasper H. Filtenborg
##
## Inspired by "tensordot" from NumPy.
## "tensordot" originally implemented by Travis E. Oliphant in 2006
## (https://github.com/numpy/numpy/blob/main/numpy/core/numeric.py)
##
## Referred to Matlab documentation for expected behaviour.
## (https://se.mathworks.com/help/matlab/ref/tensorprod.html)
##
## This program is free software: you can redistribute it and/or modify
## it under the terms of the GNU General Public License as published by
## the Free Software Foundation, either version 3 of the License, or
## (at your option) any later version.
##
## This program is distributed in the hope that it will be useful,
## but WITHOUT ANY WARRANTY; without even the implied warranty of
## MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
## GNU General Public License for more details.
##
## You should have received a copy of the GNU General Public License
## along with this program.  If not, see <https://www.gnu.org/licenses/>.
##
########################################################################

## -*- texinfo -*-
## @deftypefn  {} {@var{C} =} tensorprod (@var{A}, @var{B}, @var{dimA}, @var{dimB})
## @deftypefnx {} {@var{C} =} tensorprod (@var{A}, @var{B}, @var{dim})
## @deftypefnx {} {@var{C} =} tensorprod (@var{A}, @var{B})
## @deftypefnx {} {@var{C} =} tensorprod (@var{A}, @var{B}, "all")
## @deftypefnx {} {@var{C} =} tensorprod (@var{A}, @var{B}, ..., NumDimensionsA=ndimsA))
##
## Compute the tensor product between tensors @var{A} and @var{B}.
##
## @var{dimA} and @var{dimB} define the dimensions of @var{A} and @var{B},
## respectively, which are contracted. @var{dimA} and @var{dimB} are vectors
## with equal length and define the dimensions to match up. The matched
## dimensions of @var{A} and @var{B} must have equal length.
##
## When @var{dim} is used, it is equivalent to use @var{dimA} = @var{dimB} =
## @var{dim}.
##
## Running without an additional argument is equivalent to @var{dimA} =
## @var{dimB} = []. This results in the outer product between @var{A} and
## @var{B}.
##
## Using the "all" option results in the inner product between @var{A} and
## @var{B}. It is assumed that size(@var{A}) = size(@var{B}).
##
## Use NumDimensionsA=ndimsA when @var{A} has trailing singleton dimensions
## that should be transfered to @var{C}.
##
## @seealso{kron, dot, mtimes}
## @end deftypefn

function C = tensorprod (A, B, varargin)

  if (nargin < 2)
    print_usage ();
  endif

  if (nargin > 2)
    if (strcmp ( strtok ( inputname (nargin, false)), 'NumDimensionsA'))
      ndimsA = varargin{end};
      if (ndimsA < ndims (A))
        error ("tensorprod: NumDimensionsA cannot be smaller than ndims (A)");
      endif
    endif
  endif

  if (nargin - exist('ndimsA') == 2)
    dimA = [];
    dimB = [];
  elseif (nargin - exist('ndimsA') == 3)
    if (strcmp (varargin{1}, "all"))
      if (! isequal (size (A), size (B)))
        error ("tensorprod: size of A and B must be identical when using the \"all\" option");
      endif
      dimA = 1:ndims(A);
    else
      dimA = varargin{1};
    endif
    dimB = dimA;
  elseif (nargin - exist('ndimsA') == 4)
    dimA = varargin{1};
    dimB = varargin{2};
  endif

  if (length (dimA) != length (dimB))
    error ("tensorprod: an equal number of dimensions must be matched for A and B");
  endif

  sizeA = size (A);
  if (! exist('ndimsA'))
    ndimsA = ndims (A);
  elseif (ndimsA != ndims (A))
    sizeA = [sizeA, ones(1, ndimsA - ndims (A))];
  endif
  sizeB = size (B);
  ndimsB = ndims (B);

  if ((sum (0 >= dimA) == 1) || (sum (0 >= dimB) == 1))
    error ("tensorprod: dimension(s) must be positive integers");
  endif

  if ((sum (ndimsA < dimA) == 1) || (sum (ndimsB < dimB) == 1))
    error ("tensorprod: dimension(s) cannot exceed the size of A and B");
  endif

  if (! isequal (size (A, dimA), size (B, dimB)))
    error ("tensorprod: matched dimensions of A and B must have same length");
  endif

  ## Interchange the dimension to sum over the end of A and the front of B
  ## Prepare for A
  remainDimA = setdiff (1:ndimsA, dimA);
  newDimOrderA =  [remainDimA, dimA];
  newSizeA = [prod(sizeA(remainDimA)), prod(sizeA(dimA))];
  remainSizeA = sizeA(remainDimA);

  ## Prepare for B
  remainDimB = setdiff (1:ndimsB, dimB);
  newDimOrderB =  [remainDimB, dimB];
  newSizeB = [prod(sizeB(remainDimB)), prod(sizeB(dimB))];
  remainSizeB = sizeB(remainDimB);

  ## Do reshaping into 2D array
  newA = reshape (permute (A, newDimOrderA), newSizeA);
  newB = reshape (permute (B, newDimOrderB), newSizeB);

  ## Compute
  C = newA * newB.';

  ## If not an inner product, reshape back to tensor
  if (! isscalar (C))
    C = reshape (C, [remainSizeA, remainSizeB]);
  endif

endfunction
