#!/usr/bin/runghc

module RabinMiller where

-- https://programmingpraxis.com/wp-content/uploads/2012/09/primenumbers.pdf
-- https://programmingpraxis.com/programming-with-prime-numbers-source-code-in-haskell/

import Control.Monad (forM_, when)
import Control.Monad.ST
import Data.Array.ST
import Data.Array.Unboxed
import Data.List (sort)

ancientSieve :: Int -> UArray Int Bool
ancientSieve n = runSTUArray $ do
  bits <- newArray (2, n) True
  forM_ [2 .. n] $ \p -> do
    isPrime <- readArray bits p
    when isPrime $ do
      forM_ [2 * p, 3 * p .. n] $ \i -> do
        writeArray bits i False
  return bits

ancientPrimes :: Int -> [Int]
ancientPrimes n =
  [ p
    | (p, True) <-
        assocs $ ancientSieve n
  ]

sieve :: Int -> UArray Int Bool
sieve n = runSTUArray $ do
  let m = (n - 1) `div` 2
      r = floor . sqrt $ fromIntegral n
  bits <- newArray (0, m - 1) True
  forM_ [0 .. r `div` 2 - 1] $ \i -> do
    isPrime <- readArray bits i
    when isPrime $ do
      let a = 2 * i * i + 6 * i + 3
          b = 2 * i * i + 8 * i + 6
      forM_ [a, b .. (m - 1)] $ \j -> do
        writeArray bits j False
  return bits

primes :: Int -> [Int]
primes n = 2 : [2 * i + 3 | (i, True) <- assocs $ sieve n]

tdPrime :: Int -> Bool
tdPrime n = prime (2 : [3, 5 ..])
  where
    prime (d : ds)
      | n < d * d = True
      | n `mod` d == 0 = False
      | otherwise = prime ds

tdFactors :: Int -> [Int]
tdFactors n = facts n (2 : [3, 5 ..])
  where
    facts n (f : fs)
      | n < f * f = [n]
      | n `mod` f == 0 =
          f : facts (n `div` f) (f : fs)
      | otherwise = facts n fs

powmod :: Integer -> Integer -> Integer -> Integer
powmod b e m =
  let times p q = (p * q) `mod` m
      pow b e x
        | e == 0 = x
        | even e =
            pow
              (times b b)
              (e `div` 2)
              x
        | otherwise =
            pow
              (times b b)
              (e `div` 2)
              (times b x)
   in pow b e 1

isSpsp :: Integer -> Integer -> Bool
isSpsp n a =
  let getDandS d s =
        if even d
          then getDandS (d `div` 2) (s + 1)
          else (d, s)
      spsp (d, s) =
        let t = powmod a d n
         in if t == 1 then True else doSpsp t s
      doSpsp t s
        | s == 0 = False
        | t == (n - 1) = True
        | otherwise = doSpsp ((t * t) `mod` n) (s - 1)
   in spsp $ getDandS (n - 1) 0

isPrime :: Integer -> Bool
isPrime n =
  let ps =
        [ 2,
          3,
          5,
          7,
          11,
          13,
          17,
          19,
          23,
          29,
          31,
          37,
          41,
          43,
          47,
          53,
          59,
          61,
          67,
          71,
          73,
          79,
          83,
          89,
          97
        ]
   in n `elem` ps || all (isSpsp n) ps

rhoFactor :: Integer -> Integer -> Integer
rhoFactor n c =
  let f x = (x * x + c) `mod` n
      fact t h
        | d == 1 = fact t' h'
        | d == n = rhoFactor n (c + 1)
        | isPrime d = d
        | otherwise = rhoFactor d (c + 1)
        where
          t' = f t
          h' = f (f h)
          d = gcd (t' - h') n
   in fact 2 2

rhoFactors :: Integer -> [Integer]
rhoFactors n =
  let facts n
        | n == 2 = [2]
        | even n = 2 : facts (n `div` 2)
        | isPrime n = [n]
        | otherwise =
            let f = rhoFactor n 1
             in f : facts (n `div` f)
   in sort $ facts n

-- main = do
--  print $ ancientPrimes 100
--  print $ primes 100
--  print $ length $ primes 1000000
--  print $ tdPrime 716151937
--  print $ tdFactors 8051
--  print $ powmod 437 13 1741
--  print $ isSpsp 2047 2
--  print $ isPrime 600851475143
--  print $ isPrime 2305843009213693951
--  print $ rhoFactors 600851475142