########################################################################
##
## Copyright (C) 2022 Kasper H. Filtenborg
##
## 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}, @dots{}, "NumDimensionsA", @var{value})
## Compute the tensor product between tensors @var{A} and @var{B}.
##
## The dimensions of @var{A} and @var{B} that are contracted are defined by
## @var{dimA} and @var{dimB}, respectively. @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 @var{dimA} = @var{dimB} =
## @var{dim}.
##
## Running without an additional argument is equivalent to @var{dimA} =
## @var{dimB} = []. This computes the outer product between @var{A} and
## @var{B}.
##
## Using the "all" option results in the inner product between @var{A} and
## @var{B}. For this, it is required that size(@var{A}) = size(@var{B}).
##
## Use the property-value pair with the property name "NumDimensionsA"
## when @var{A} has trailing singleton dimensions that should be transfered to
## @var{C}. The specified value should be the number of dimensions of @var{A}.
##
## @seealso{kron, dot, mtimes}
## @end deftypefn

function C = tensorprod (A, B, varargin)

  if (nargin < 2)
    print_usage ();
  endif

  ## Check for NumDimensionsA property
  if (nargin > 2)
    if (strcmpi (varargin{end}, "NumDimensionsA"))
      error ("tensorprod: a value for the NumDimensionsA property much be provided");
    endif
  endif

  if (nargin > 3)
    if (strcmpi (varargin{end - 1}, "NumDimensionsA"))
      NumDimensionsA = varargin{end};
    endif
  endif

  existNumDimensionsA = exist("NumDimensionsA");
  ndimargs = nargin - 2 - 2 * existNumDimensionsA;

  ## Set dimA and dimB
  if (ndimargs == 0)
    dimA = [];
    dimB = [];
  elseif (ndimargs == 1)
    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 (ndimargs == 2)
    dimA = varargin{1};
    dimB = varargin{2};
    if (length (dimA) != length (dimB))
      error ("tensorprod: an equal number of dimensions must be matched for A and B");
    endif
  endif

  ## Find size and ndims of A and B
  sizeA = size (A);
  ndimsA = ndims (A);
  sizeB = size (B);
  ndimsB = ndims (B);

  if (existNumDimensionsA)                                                      # Take NumDimensionsA property into account
    if (NumDimensionsA != ndimsA)
      sizeA = [sizeA, ones(1, NumDimensionsA - ndimsA)];
    endif

    if (NumDimensionsA < ndimsA)
      error ("tensorprod: NumDimensionsA cannot be smaller than the number of dimensions of A");
    endif
    ndimsA = NumDimensionsA;                                                    # Overwrite ndimsA with provided number of dimensions for A
  endif

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

  if (sum (ndimsA < dimA) || sum (ndimsB < dimB))
    error ("tensorprod: dimA and dimB cannot exceed the number of dimensions of A and B, respectively");
  endif

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

  ## Interchange the dimension to sum over the end of A and the front of B
  ## Prepare for A
  remainDimA = setdiff (1:ndimsA, dimA);                                        # Dimensions of A to keep
  newDimOrderA =  [remainDimA, dimA];                                           # New order of dimensions (dimensions to keep first, dimensions to contract last)
  newSizeA = [prod(sizeA(remainDimA)), prod(sizeA(dimA))];                      # Size of temporary 2D representation of A
  remainSizeA = sizeA(remainDimA);                                              # Contribution to size of C from remaining dimensions of A

  ## 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 * transpose (newB);

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

endfunction
