########################################################################
##
## Copyright (C) 2000-2021 The Octave Project Developers
##
## See the file COPYRIGHT.md in the top-level directory of this
## distribution or <https://octave.org/copyright/>.
##
## This file is part of Octave.
##
## Octave 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.
##
## Octave 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 Octave; see the file COPYING.  If not, see
## <https://www.gnu.org/licenses/>.
##
########################################################################

## -*- texinfo -*-
## @deftypefn {} {} isprime2 (@var{x})
## Return a logical array which is true where the elements of @var{x} are prime
## numbers and false where they are not.
##
## A prime number is conventionally defined as a positive integer greater than
## 1 (e.g., 2, 3, @dots{}) which is divisible only by itself and 1.  Octave
## extends this definition to include both negative integers and complex
## values.  A negative integer is prime if its positive counterpart is prime.
## This is equivalent to @code{isprime2 (abs (x))}.
##
## If @code{class (@var{x})} is complex, then primality is tested in the domain
## of Gaussian integers (@url{https://en.wikipedia.org/wiki/Gaussian_integer}).
## Some non-complex integers are prime in the ordinary sense, but not in the
## domain of Gaussian integers.  For example, @math{5 = (1+2i)*(1-2i)} shows
## that 5 is not prime because it has a factor other than itself and 1.
## Exercise caution when testing complex and real values together in the same
## matrix.
##
## Examples:
##
## @example
## @group
## isprime2 (1:6)
##   @result{}  0  1  1  0  1  0
## @end group
## @end example
##
## @example
## @group
## isprime2 ([i, 2, 3, 5])
##   @result{}  0  0  1  0
## @end group
## @end example
##
## Programming Note: @code{isprime2} uses direct division when the maximum
## value in @var{x} is within 1e14. For larger values up to 2^64-1
## it uses the deterministic Miller-Rabin test with all primes <= 37.
##
## Compatibility Note: @sc{matlab} does not extend the definition of prime
## numbers and will produce an error if given negative or complex inputs.
## @seealso{primes, factor, gcd, lcm}
## @end deftypefn

function t = isprime2 (x)

  if (nargin < 1)
    print_usage ();
  elseif (any (fix (x) != x))
    error ("isprime2: X contains non-integer entries");
  endif

  if (isempty (x))
    t = x;
    return;
  endif

  if (iscomplex (x))
    t = isgaussianprime (x);
    return;
  endif

  ## Code strategy is to build a table with the list of possible primes
  ## and then quickly compare entries in x with the table of primes using
  ## lookup().  The table size is limited to save memory and computation
  ## time during its creation.  All entries larger than the maximum in the
  ## table are checked by straightforward division or by Miller-Rabin.

  THRESHOLD = 1e14;
  # For numbers <= THRESHOLD, will use direct division.
  # For numbers > THRESHOLD, will use Miller-Rabin test.
  # THRESHOLD should be somewhere between 1e14 and flintmax == 2^53-1.

  x = abs (x);              # handle negative entries
  maxn = max (x(:));
  maxp = floor (sqrt (min (maxn, THRESHOLD)));
  pr = primes (maxp);   ## generate prime table of suitable length.
  t = lookup (pr, x, "b");  # quick search for table matches.

  ## process large entries that are still suitable for direct division
  m = x(x > maxp & x <= THRESHOLD);
  if (! isempty (m))
    if (maxn <= intmax ("uint32"))
      m = uint32 (m);
    elseif (maxn <= intmax ("uint64"))
      m = uint64 (m);
    else
      warning ("isprime2: X contains integers too large to be tested");
    endif

    ## Start by dividing through by the small primes until the remaining
    ## list of entries is small (and most likely prime themselves).
    pr = cast (pr(pr <= sqrt (maxn)), class (m));
    for p = pr
      m = m(rem (m, p) != 0);
      if (numel (m) < numel (pr) / 10)
        break;
      endif
    endfor

    ## Check the remaining list of possible primes against the
    ## remaining prime factors which were not tested in the for loop.
    ## This is just an optimization to use arrayfun over for loo
    pr = pr(pr > p);
    mm = arrayfun (@(x) all (rem (x, pr)), m);
    m = m(mm);

    ## Add any remaining entries, which are truly prime, to the results.
    if (! isempty (m))
      m = cast (sort (m), class (x));
      t |= lookup (m, x, "b");
    endif
  endif
  
  ## Process any large values in x > THRESHOLD serially using Miller-Rabin
  for i = x(x > THRESHOLD)
    t(x == i) = isprime2large (i);
  end

endfunction

function t = isgaussianprime (z)

  ## Assume prime unless proven otherwise
  t = true (size (z));

  x = real (z);
  y = imag (z);

  ## If purely real or purely imaginary, ordinary prime test for
  ## that complex part if that part is 3 mod 4.
  xidx = y==0 & mod (x, 4) == 3;
  yidx = x==0 & mod (y, 4) == 3;

  t(xidx) &= isprime2 (x(xidx));
  t(yidx) &= isprime2 (y(yidx));

  ## Otherwise, prime if x^2 + y^2 is prime
  zidx = ! (xidx | yidx);          # Skip entries that were already evaluated
  zabs = x(zidx).^2 + y(zidx).^2;
  t(zidx) &= isprime2 (zabs);

endfunction

function retval = isprime2large (n)
  # Returns true if n is prime, false otherwise.
  # For use with large scalar n only.
  # Uses Miller-Rabin technique, deterministic variant.
  # FIXME: Can this code be vectorized or is it faster to test each input serially?
  # 
  # Background reading for theory:
  # https://en.wikipedia.org/wiki/Miller%E2%80%93Rabin_primality_test
  # 
  # The basic technique is to do the Miller-Rabin test on n with several divisors.
  # If any of them are false, then n is composite so we can return false immediately.
  # If ALL of them are true, then n really is prime.
  # If only some are true and others are false, such a composite number
  # is called pseudoprime, so don't return true early without testing for all divisors.
  #   
  # The mathematically tricky part is to choose enough divisors that suffice
  # for all 64-bit numbers we may care to test.
  # That part has been proved by mathematicians,
  # and one needs no more than 12 divisors, and it can be as few as 7.
  # In particular, it can be the first 12 primes.
  #
  # Sources:
  # (1) https://oeis.org/A014233
  #     Comment: a(12) > 2^64.  Hence the primality of numbers < 2^64 can be
  #     determined by asserting strong pseudoprimality to all prime bases <= 37
  #     (=prime(12)). Testing to prime bases <=31 does not suffice,
  #     as a(11) < 2^64 and a(11) is a strong pseudoprime
  #     to all prime bases <= 31 (=prime(11)). - Joerg Arndt, Jul 04 2012
  # (2) https://arxiv.org/abs/1509.00864
  #     Strong Pseudoprimes to Twelve Prime Bases
  #     Jonathan P. Sorenson, Jonathan Webster
  #
  # In addition, a source listed here: https://miller-rabin.appspot.com/
  # reports that all 64-bit numbers can be covered with only 7 divisors,
  # namely 2, 325, 9375, 28178, 450775, 9780504, and 1795265022.
  # I couldn't find a peer-reviewed article to back it up.
  # I am using the 12 primes <= 37 in this code.
  #
  # Further practical primality discussion here for 64-bit numbers:
  # https://math.stackexchange.com/questions/2481148/primality-testing-for-64-bit-numbers

  # For this technique, we need n to be odd. If n is even, straightaway return composite.
  if (rem(n,2) == 0 && n > 2)
    retval = false; return
  end

  n = uint64 (n);

  # The divisors we choose, per comments above.
  # Making it persistent speeds up subsequent calls by a small amount, up to 6%.
  persistent divisors = uint64 ([2 3 5 7 11 13 17 19 23 29 31 37]);
  if (any (rem (n, divisors) == 0))  # cheap division by small primes
    retval = false; return
  end
  
  # express n as 2^r * d + 1 for some odd number d, where r >= 0
  d = n-1;
  r = uint64(0);
  while (rem(d,2) == 0)
    d /= 2;
    r += 1;
  end
 
  for i = divisors
    if ( ! millerrabinsingle (i, d, r, n))
      # composite, so return false straightaway
      retval = false; return
    end
    # otherwise it is either pseudoprime or prime, so keep looking.
  end
  
  # if here, then pseudoprime over ALL divisors ==> really is prime
  retval = true;
endfunction

function retval = millerrabinsingle (divisor, d, r, n)
  # Returns false if n is composite,
  #         true when n is pseudoprime or prime.
  x = safepower (divisor, d, n);
  if (x == 1 || x == n-1)
    retval = true; return
  end

  for j = 1:r-1
    x = safemultiply (x, x, n);
    if (x == n-1)
      retval = true; return
    end
  end

  retval = false;
endfunction

function retval = safepower (a, b, modulus)
  # Calculate rem (a ^ b, modulus) without overflow
  retval = uint64 (1);
  while (b)
    if (rem (b,2) == 1)
      retval = safemultiply (retval, a, modulus);
      b = (b-1) / 2;
    else
      b /= 2;
    end
    a = safemultiply (a, a, modulus);
  end
endfunction

function retval = safemultiply (a, b, modulus)
  # Calculate rem (a * b, modulus) without overflow
  retval = uint64 (0);
  while (b > 0)
    if (rem (b,2) == 1)
      retval = safeadd (retval, a, modulus);
      b = (b-1) / 2;
    else
      b /= 2;
    end
    a = safeadd (a, a, modulus);
  end
endfunction

function retval = safeadd (a, b, modulus)
  # Calculate rem (a + b, modulus) without overflow
  persistent IMAX = intmax ("uint64");
  if (b < IMAX - a)
    retval = a+b;
  else
    retval = (IMAX - modulus) + b - (IMAX - a);
    # functionally this is equivalent to (a+b-modulus)
    # but done without overflow or underflow,
    # especially with Octave's saturation semantics.
  end
  if (retval >= modulus)
    retval -= modulus;   # faster than calling rem or mod
  end
endfunction

%!assert (isprime2 (3), true)
%!assert (isprime2 (4), false)
%!assert (isprime2 (5i), false)
%!assert (isprime2 (7i), true)
%!assert (isprime2 ([1+2i, (2+3i)*(-1+2i)]), [true, false])
%!assert (isprime2 (-2), true)
%!assert (isprime2 (complex (-2)), false)
%!assert (isprime2 (2i), false)
%!assert (isprime2 ([i, 2, 3, 5]), [false, false, true, false])
%!assert (isprime2 (0), false)
%!assert (isprime2 (magic (3)), logical ([0, 0, 0; 1, 1, 1; 0, 0, 1]))

## Test input validation
%!error <Invalid call> isprime2 ()
%!error <X contains non-integer entries> isprime2 (0.5i)
%!error <X contains non-integer entries> isprime2 (0.5)
