{- Primality testing

Copyright 2016 Ken Takusagawa

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 <http://www.gnu.org/licenses/>.
-}
{-# LANGUAGE ScopedTypeVariables, LambdaCase, MagicHash, FlexibleContexts #-}
module Primality2 where {

-- these are the most useful functions (primes,isPrime_slow,isPrime_x, isPrime_verified, isPrime_small, isPrime_big, fast_primes, modular_exponentiation, isPrime, randomsInteger, isPrime_baillie)
{-
isPrime_x isPrime_small isPrime_big are based on only the bpsw test, with no miller-rabin verification.  We tempt fate because finding a bpsw false positive is a very notable event.
-}

import Data.Function((&));
import Control.Category((>>>));
import Prelude hiding((.),(>>));
import Data.List;
import Data.Maybe;
import Control.Monad.Except;
import GHC.Integer.Logarithms(integerLog2#);
import GHC.Exts(Int(I#),quotInt#);
import GHC.Integer.GMP.Internals(bitInteger);
import System.Random(StdGen,randoms,RandomGen);
import qualified System.Random as Random;
import qualified Control.Monad.State as State;
import Control.Monad.State(State,runState);

-- most conservative approach, not very fast.
isPrime :: Integer -> Bool;
-- xxx use the optimal number of trial divisions for the given number size
isPrime = isPrime_verified large_divnum;

-- in general, two knobs:
-- 1. whether to verify a prime with 20 rounds of miller-rabin.
-- 2. how much trial division to do beforehand.

type Matrix = ((Integer,Integer),(Integer,Integer));

matrix_sqr :: Matrix -> Matrix;
-- strassen style, 5 is min multiplications
matrix_sqr ((a,b),(c,d)) = ((sqr a+bc, b*(a+d)), (c*(a+d), sqr d+bc)) where {bc = b*c};

-- multiply the matrix [ p -q ; 1 0 ]
advance :: Integer -> Integer -> Matrix -> Matrix;
advance p q ((a,b),(c,d)) = ((a*p-c*q, b*p-d*q),(a,b));

matrix_pow :: Integer -> Integer -> Integer -> Matrix;
matrix_pow _ _ 0 = ((1,0),(0,1));
-- matrix_pow p q 1 = ((p,negate q),(1,0));
matrix_pow p q nn = case divMod nn 2 of
{ (n,0) -> matrix_pow p q n & matrix_sqr
; _ -> pred nn & matrix_pow p q & advance p q
};

get_u :: Matrix -> Integer;
get_u = snd >>> fst;

get_v :: Integer -> Matrix -> Integer;
get_v p (_,(c,d)) = p*c + 2*d;

do_matrix_mod :: Integer -> Matrix -> Matrix;
do_matrix_mod modulus ((a,b),(c,d)) = ((mod a modulus, mod b modulus), (mod c modulus, mod d modulus));

jacobi_symbol :: Integer -> Integer -> Integer;
-- jacobi_symbol a n | gcd a n > 1 = 0; -- apparently this one is not needed because the algorithm does gcd implicitly
jacobi_symbol _ n | n < 0 = error "jacobi_symbol lower negative";
jacobi_symbol _ n | mod n 2 == 0 = error "jacobi_symbol lower even";
jacobi_symbol a n | a >= n = jacobi_symbol (mod a n) n;
jacobi_symbol 2 n
| mod8 == 1 || mod8 == 7 = 1
| mod8 == 3 || mod8 == 5 = negate 1
| True = error "jacobi_symbol mod8 even"
where { mod8 = mod n 8 };
jacobi_symbol 0 1 = 1; -- check this
jacobi_symbol 0 _ = 0; -- check this, this is probably where the gcd>1 case goes
jacobi_symbol a n | a < 0 && mod4 == 1 = jacobi_symbol (negate a) n
| a < 0 && mod4 == 3 = negate $ jacobi_symbol (negate a) n
| a < 0 = error "jacobi_symbol negative, even mod4"
where { mod4 = mod n 4};
jacobi_symbol a n | mod2 == 0 = jacobi_symbol 2 n * jacobi_symbol div2 n
where { (div2,mod2) = divMod a 2 };
jacobi_symbol a n | a < n && 3 == mod a 4 && 3 == mod n 4 = negate $ jacobi_symbol n a;
jacobi_symbol a n | a < n && (1 == mod a 4 || 1 == mod n 4) = jacobi_symbol n a;
jacobi_symbol a n | a < n = error "jacobi_symbol even reciprocity";
jacobi_symbol a n = error $ "jacobi_symbol fail " ++ show a ++ " " ++ show n;  -- if the compiler were smarter, it could tell that this case is not reached.

selfridge_method_a :: [Integer];
selfridge_method_a = enumFromThen 5 7 & alternating_signs;

-- yay haskell idiom, i.e., unfoldr and point-free
alternating_signs :: forall a . Num a => [a] -> [a];
alternating_signs = ((,) True) >>> unfoldr (\(positive :: Bool, l :: [a]) ->
if null l then Nothing
else Just (head l & if positive then id else negate, (not positive, tail l)));

data Composite_special = Zero_Jacobi | Perfect_square | Even_n deriving (Show);

-- FlexibleContexts
choose_d :: MonadError Composite_special m => Integer -> m Integer;
choose_d n = if is_perfect_square n
then throwError Perfect_square
else case (head $ do {
d <- selfridge_method_a;
-- it is possible that this will take a long time or never terminate
let {j = jacobi_symbol d n};
guard $ j < 1;
guard $ abs d /= n;  -- for 5 and 11
return (d,j)
}) of {
(_,0) -> throwError Zero_Jacobi ; -- d divides n
(d,_) -> return d;
};

get_pq :: MonadError Composite_special m => Integer -> m(Integer,Integer);
get_pq n = do { d <- choose_d n;
return (1, div_cleanly (1-d) 4);
};

div_cleanly :: Integer -> Integer -> Integer;
div_cleanly x y = case divMod x y of
{ (q,0) -> q
; _ -> error "div_cleanly"
};

data Strong_prime_annotation = Positive_unity | Negative_unity | Two deriving Show;

{- Incidentally, both annotation get used roughly equally frequently:
[3,5..] & map (miller_rabin_raw 2 >>> fst >>> get_strong_annotation) & catMaybes & map strong_annotation_to_char & putStr
-}
{-
strong_annotation_to_char :: Strong_prime_annotation -> Char;
strong_annotation_to_char = \case
{ Two -> '2'
; Positive_unity -> '+'
; Negative_unity -> '.'
};
-}

miller_rabin_structure :: Integral i => Operations a -> i -> State (Maybe Strong_prime_annotation) a;
miller_rabin_structure operation expo = if expo == 0
then error "miller_rabin_structure_state: 0 power"
else case divMod expo 2 of {
(j,0) -> do {
  half <- miller_rabin_structure operation j;
  when (half & is_negative_unity operation) $ state_write Negative_unity;
  half & double operation & return;
};
_ -> do {
  let {answer = boring_exponentiate operation expo};
  when (answer & is_unity operation) $ state_write Positive_unity;
  return answer;
}};

state_write :: Strong_prime_annotation -> State (Maybe Strong_prime_annotation) ();
state_write new = do {
  old <- State.get;
  if isNothing old then State.put $ Just new
  else error "state already set";
};

-- type ESimple = Either Strong_prime_annotation Matrix;

boring_exponentiate :: Integral i => Operations a -> i -> a;
boring_exponentiate operation expo = if expo==0
then unity operation
else case divMod expo 2 of {
(j,0) -> boring_exponentiate operation j & double operation;
_ -> pred expo & boring_exponentiate operation & increment operation;
};

-- portfolio of operations, essentially a dictionary for a typeclass
-- names for operations, one step down the hyperoperation hierarchy, borrowed from elliptic curves
data Operations a = Operations
{ unity :: a
, double :: a -> a
, increment :: a -> a
, is_unity :: a -> Bool
, is_negative_unity :: a -> Bool
};

arithmetic_mod_base :: Integer -> Integer -> Operations Integer;
arithmetic_mod_base modulus base = Operations
{ unity = 1
, double = \x -> mod (x * x) modulus
, increment = \x -> mod (x * base) modulus
, is_unity = \x -> x == 1
, is_negative_unity = \x -> x == pred modulus
};

miller_rabin_raw :: Integer -> Integer -> (Test_result Strong_prime_annotation (Maybe Composite_special), Maybe Integer);
miller_rabin_raw a n = if n == 2 then (Probable_prime Two, Nothing)
else if mod n 2 == 0
then (Composite $ Just Even_n, Nothing) -- even numbers are composite
-- the test mostly just works for even numbers, but we special case them out of tradition, so that the lists of strong pseudoprimes line up with published results.
else case run_miller_rabin_structure (arithmetic_mod_base n a) $ pred n of {
(answer,annot:: (Maybe Strong_prime_annotation)) -> (case annot of {
Nothing -> Composite Nothing;
Just se -> Probable_prime se;}, Just answer)};

miller_rabin_bool :: (Test_result Strong_prime_annotation x, y) -> Bool;
miller_rabin_bool = fst >>> translate_test_result;

translate_test_result :: Test_result a b -> Bool;
translate_test_result = \case {
Composite _ -> False;
Probable_prime _ -> True;
};

{-
get_strong_annotation :: Test_result a b -> Maybe a;
get_strong_annotation (Probable_prime x) = Just x;
get_strong_annotation _ = Nothing;
-}

miller_rabin_2 :: Integer -> Bool;
miller_rabin_2 = miller_rabin_many [2];

miller_rabin_many :: [Integer] -> Integer -> Bool;
miller_rabin_many bases n = if n<2 then False
else (map (flip miller_rabin_raw n >>> miller_rabin_bool) bases & and);

lucas_operations :: Integer -> (Integer,Integer) -> Operations Matrix;
lucas_operations modulus (p,q) = Operations
{ unity = matrix_pow undefined undefined 0  -- identity
, double = matrix_sqr >>> do_matrix_mod modulus
, increment = advance p q >>> do_matrix_mod modulus
, is_unity = get_u >>> flip mod modulus >>> (== 0)
, is_negative_unity = get_v p >>> flip mod modulus >>> (== 0)
};

data Test_result a b = Probable_prime a | Composite b deriving Show; -- isomorphic with Either

run_miller_rabin_structure :: Integral i => Operations a -> i -> (a,Maybe Strong_prime_annotation);
run_miller_rabin_structure = (uncurry miller_rabin_structure >>> flip runState Nothing) & curry;
-- awkward to have to call uncurry and curry to make point-free

strong_lucas_raw :: Integer -> (Test_result Strong_prime_annotation (Maybe Composite_special), Maybe (Matrix, (Integer, Integer)));
strong_lucas_raw n = if n == 2 then (Probable_prime Two,Nothing)
else if 0 == mod n 2
then (Composite $ Just Even_n, Nothing)
else case get_pq n of {
Left d -> (Composite $ Just d, Nothing);
Right pq -> case run_miller_rabin_structure (lucas_operations n pq) $ succ n of {
(matrix,annot) -> (case annot of {
Nothing -> Composite Nothing;
Just se -> Probable_prime se;
}, Just (matrix,pq))
}};

-- V(n+1) = 2Q (mod n), from Wikipedia
-- similar to (perhaps identical) to frobenius pseudoprime
v2q_test :: Integer -> (Integer,Integer) -> Matrix -> Bool;
-- v2q_test n (p,q) m = mod (m & get_v p) n == mod (2*q) n;
v2q_test n (p,q) m = 0 == mod (get_v p m - 2 * q) n;

strong_lucas_with_frobenius :: Integer -> Maybe Bool;
strong_lucas_with_frobenius n = case strong_lucas_raw n of {
(Composite _,_) -> Nothing;
(Probable_prime Two,Nothing) -> Just True;
(Probable_prime _,Just (m,pq)) -> Just $ v2q_test n pq m;
(Probable_prime _, Nothing) -> error "strong_lucas: should not happen"
};

is_perfect_square :: Integer -> Bool;
is_perfect_square n = if n < 0 then False else n == (sqrt_integer n & sqr);

sqrt_integer :: Integer -> Integer;
sqrt_integer n = case compare n 0 of {
LT -> error "sqrt_integer negative";
EQ -> 0 ; -- newton fails at 0
GT -> sqrt_integer_seq n & find_converged;
};

-- GHC.Integer.Type does not yet use mpn_sqr;
sqr :: Num a => a -> a;
sqr x = x*x;

sqrt_integer_seq :: Integer -> [Integer];
sqrt_integer_seq n = approx_sqrt n & iterate_newton n;

find_converged :: (Eq a) => [a] -> a;
find_converged [_] = error "no equality found on finite list";
find_converged (x:rest) = if (x == head rest) then x else find_converged rest;
find_converged [] = error "no equality found on empty list";

approx_sqrt_via_double :: Integer -> Integer;
approx_sqrt_via_double n = ((fromInteger n)::Double) & sqrt & round;

approx_sqrt_via_log :: Integer -> Integer;
approx_sqrt_via_log n = bitInteger (quotInt# (integerLog2# n) 2#);

approx_sqrt :: Integer -> Integer;
approx_sqrt n = if integerLog2 n < 1023 then approx_sqrt_via_double n
else approx_sqrt_via_log n;

integerLog2 :: Integer -> Int;
integerLog2 i = I# (integerLog2# i);

newton :: Integer -> Integer -> Integer;
newton a x = div (x + (div a x)) 2;

newton2 :: Integer -> Integer -> Integer;
newton2 a x = newton a x & newton a;

-- doing 2 iterations at a time avoids problems with oscillatory behavior
iterate_newton :: Integer -> Integer -> [Integer];
iterate_newton a x = iterate (newton2 a) x;

strong_lucas :: Integer -> Bool;
-- strong_lucas = strong_lucas_raw >>> isLeft;
strong_lucas = strong_lucas_with_frobenius >>> fromMaybe False;


-- the optimal amount of trial division might depend of the size of the number.
-- 1000 is better than 500 for  fast-test 1000000 (sum first million primes)
-- at 2^20480 minus, 3000 and 5000 give approximately the same performance, better than 1500
-- at 2^16384 minus, 10000 (19m23s) gives slightly better performance than 5000 (20m33) and 3000 (21m43).  This is because more composites searched before finding one at -13797.
-- at 5 billion frobeniusly-strong-bench 10000 36.9, 5000 31.9, 2000 21.7, 1000 18.2, 500 16.1 200 14.7 100 14.4 , double work size 200 29.7, 100 28.7, 50 29.0 (possibly outlier), 40 28.4, 20 28.6, 10 28.6, 2 31.4.  a 1 billion 10 33.8, 100 34.6
-- again 5 billion larger work size 4000000, 100 56.9, 20 56.27 10 56.84

-- to cache results
divisibility_test_primes :: [Integer];
-- cannot use fast_primes here or else infinite loop, fortunately, primes() is pretty fast up to the first million primes
--divisibility_test_primes = primes ();
divisibility_test_primes = small_primes;

baillie_psw :: Integer -> Bool;
baillie_psw n = (is_almost_square n & not) && miller_rabin_2 n && strong_lucas n;

isPrime_baillie :: Integer -> Integer -> Bool;
isPrime_baillie = isPrime_framework baillie_psw;

isPrime_verified :: Integer -> Integer -> Bool;
isPrime_verified= isPrime_framework baillie_verified;

-- the additional miller_rabin tests might not be a big deal if searching among many composites; there will have been many failed miller_rabin_2 tests.
-- at 8192 bits, 2.2 seconds vs 16.3 seconds to verify
baillie_verified :: Integer -> Bool;
baillie_verified n = if baillie_psw n
then if miller_rabin_by_random 19 (mkStdGen2 $ Random_seed 0) n -- fixme xxx , unsafePerformIO
then True
else error $ "baillie_verified false positive " ++ show n
else False;

isPrime_framework :: (Integer -> Bool) -> Integer -> Integer -> Bool;

isPrime_framework final_test divisibility_test_num n = msum [obvious_tests n, test_divisibility_num n divisibility_test_num, final_test n & Just] & fromJust;

-- space leak, runs out of 5GB memory before 100million primes.
test_divisibility_num :: MonadPlus m => Integer -> Integer -> m Bool;
test_divisibility_num n divnum = test_divisibility n (genericTake divnum divisibility_test_primes);

-- lucas does poorly on this form, according to Wikipedia
-- though miller_rabin_2 does seem to get them
is_almost_square :: Integer -> Bool;
is_almost_square 3 = False;
is_almost_square n = is_perfect_square $ n+1;

fast_primes :: [Integer];
fast_primes = 2:3:5:7:([9,11..] & filter (isPrime_baillie small_divnum));
-- using baillie_psw instead of isPrime_baillie is about 2.5x slower at 500,000 primes.
-- primes() runs much faster when compiled with -O.  primes() is competitive with isPrime_baillie even around 1 million; just slightly slower 40 vs 33 seconds.  baillie_psw is 68 seconds.  isPrime_verified also 33 seconds
-- 2 million: isPrime_baillie 88 seconds, primes() 111 seconds, isPrime_verified 88 seconds.
-- the 1 millionth prime is about 15 million.

-------

lengthen_integer :: Integer -> Bool -> Integer;
lengthen_integer i False = i*2;
lengthen_integer i True = i*2+1;

newtype Random_seed = Random_seed { unRandom_seed :: Int };

mkStdGen2 :: Random_seed -> StdGen;
mkStdGen2 = unRandom_seed >>> Random.mkStdGen >>> Random.next >>> snd;   --discard the first sample because bad

rand_and_range :: [Bool] -> [(Integer,Integer)];
rand_and_range = scanl' lengthen_integer 0 >>> zip two_pow_range;

two_pow_range :: [Integer];
two_pow_range = 1 : map (*2) two_pow_range; -- start with because first element of scanl uses zero elements of the list

-- XXX integerLog2 might make this faster
big_enough :: Integer -> [(Integer,Integer)] -> Integer;
big_enough i = find (fst >>> (\x -> x>=i)) >>> fromJust >>> snd;

good_big :: Integer -> [Bool] -> Maybe Integer;
good_big i b = let { t = rand_and_range b & big_enough i } in
if t < i then Just t else Nothing;

many_random_gen :: RandomGen g => g-> [g];
--many_random_gen g = let {s = Random.split g} in fst s:(many_random_gen $ snd s);
many_random_gen = unfoldr (Random.split >>> Just);

good_big_r :: RandomGen g => Integer -> g -> Maybe Integer;
good_big_r i g = randoms g & good_big i;

jgood_big :: RandomGen g => Integer -> g -> Integer;
jgood_big i = many_random_gen >>> map (good_big_r i) >>> find isJust >>> fromJust >>> fromJust;

many_jgood :: RandomGen g => Integer -> g -> [Integer];
many_jgood i g = many_random_gen g & map (jgood_big i);

randomsInteger :: RandomGen g => Integer -> g -> [Integer];
randomsInteger onePlusMax = many_jgood onePlusMax;

miller_rabin_bases :: RandomGen g => g -> Integer -> [Integer];
miller_rabin_bases g n = if n < 4 then repeat 1 else many_jgood (n-3) g & map (+2);
-- 0^(n-1) is boring, 1^(n-1) is boring, (n-1)^(n-1) == 1 for all odd n

-- xxx try to do distinct bases
miller_rabin_by_random :: RandomGen g => Integer -> g -> Integer -> Bool;
miller_rabin_by_random num_iterations g n = miller_rabin_many (miller_rabin_bases g n & genericTake num_iterations) n;

-- isPrime_baillie is approximately 3 iterations of miller_rabin_by_random at 8192 bits
-- about 5 or 5.5 iterations for deterministic bases

-- Arnault examples (and more) at
-- https://github.com/danaj/Math-Prime-Util-GMP/blob/master/t/17-pseudoprime.t

test_divisibility :: MonadPlus m => Integer -> [Integer] -> m Bool;
test_divisibility _ [] = mzero;
test_divisibility x (h:t) = if h*h > x then return True
else case mod x h of {
0 -> return False;
_ -> test_divisibility x t
};

obvious_tests :: MonadPlus m => Integer -> m Bool;
obvious_tests n = case compare n 2 of {
LT -> return False;
EQ -> return True;
GT -> mzero;
};

-- number of divisibility tests to run when testing large numbers
large_divnum :: Integer;
large_divnum = 2000000;

-- number of divisibility tests to run when testing small numbers
small_divnum :: Integer;
small_divnum = 100;

{-
:set +s
*Main> isPrime_baillie'  (2^(8192::Integer)-2439)
(2.22 secs, 386,778,080 bytes)

*Main> isPrime_verified'   (2^(8192::Integer)-2439)
(16.28 secs, 2,332,141,528 bytes)

*Main> isPrime_verified'   (2^(8192::Integer)-2439)
(16.31 secs, 2,375,636,648 bytes)

*Main> isPrime_baillie' (2^(19937::Integer)-1)
(15.74 secs, 1,504,986,712 bytes)

*Main> isPrime_baillie' (2^(23209::Integer)-1)
(22.83 secs, 1,980,001,568 bytes)

*Main> isPrime_baillie' (2^(44497::Integer)-1)
(102.37 secs, 6,975,033,264 bytes)

*Main> isPrime_baillie' (2^(86243::Integer)-1)
(517.47 secs, 25,568,248,968 bytes)

*Main> isPrime_baillie' (2^(110503::Integer)-1)
(954.82 secs, 41,664,425,448 bytes)

*Main> isPrime_baillie' (2^(132049::Integer)-1)
(1219.95 secs, 59,366,052,576 bytes)

*Main> isPrime_baillie' (2^(216091::Integer)-1)
(3696.46 secs, 157,866,346,896 bytes)
-}

isPrime_slow :: Integer -> Bool;
isPrime_slow n = mplus (obvious_tests n) (test_divisibility n small_primes) & fromJust;

-- need only sqrt memory
small_primes :: [Integer];
small_primes = 2:(filter isPrime_slow [3,5..]);

primes :: () -> [Integer];
primes _ = filter isPrime_slow $ enumFrom 0;

isPrime_x :: Integer -> Integer -> Bool;
isPrime_x = isPrime_baillie;

isPrime_small :: Integer -> Bool;
isPrime_small = isPrime_x small_divnum;

isPrime_big :: Integer -> Bool;
isPrime_big = isPrime_x large_divnum;

-- computes base^exponent mod modulus.
-- exponent is curried.
modular_exponentiation :: Integer -> Integer -> Integer -> Integer;
modular_exponentiation modulus base = boring_exponentiate $ arithmetic_mod_base modulus base;

} --end