{-# LANGUAGE BangPatterns #-}

import Control.Arrow(first, second, (>>>))
import Control.Concurrent
import Control.Monad.Trans.State.Strict
import Data.Foldable (toList)
import qualified Data.HashSet as HS
import qualified Data.IntSet as IS
import Data.Ratio
import qualified Data.Sequence as Seq
import Data.Sequence(Seq, (|>))
import Data.List
import Test.QuickCheck

{-H3.1.1-}
decomposition :: Int -> [Int]
decomposition = go 2
  where
    go _ 1 = []
    go p n
        | n `mod` p == 0 = p : go p (n `div` p)
        | otherwise = go (succ p) n

decomposition2 :: Int -> [(Int, Int)]
decomposition2 = map (\l -> (head l, length l)) . group . decomposition

{-H3.1.2-}
isPrime :: Int -> Bool
isPrime = (==1) . length . decomposition

primes :: Int -> [Int]
primes = filter isPrime . enumFromTo 2

-- Algorithm as published by David Harvey 2008
-- Takes a few seconds to calculate B_6000, when used with -O2 (and -foject_code in GHCi). Unfortunately, you cannot use the OPTIONS pragma to specify those, as it prevents loading the module into GHCi, unless the GHCi invocation itself already has -fobject-code.
-- I also might have gone slightly overboard with strictness; pretty sure GHC would also be able to figure out that stuff by itself ...

strictPrimes :: Integral a => a -> Seq a
strictPrimes !n = extendPrimes (Seq.singleton 2) 3 n

extendPrimes :: Integral a => Seq a -> a -> a -> Seq a
extendPrimes !a !x !n
        | x > n = a
        | any (\p -> x `mod` p == 0) a = extendPrimes a (x+2) n
        | otherwise = extendPrimes (a |> x) (x+2) n

-- Just to conveniently have it available when reloading GHCi
manyPrimes :: [Integer]
manyPrimes = toList (strictPrimes 10000)

prop_strictPrimes :: Integer -> Property
prop_strictPrimes n = n >= 3 ==> toList (strictPrimes n) == map fromIntegral (primes (fromIntegral n))

intSqrt :: Integral a => a -> a
intSqrt !n = go n n
  where
    go !x !n =
        let !y = (x + n `div` x) `div` 2
        in if y < x then go y n else x

prop_intSqrt :: Integer -> Property
prop_intSqrt n = n > 0 ==> let r = intSqrt n in r * r <= n && (r+1) * (r+1) > n

powMod :: Integral a => a -> a -> a -> a
powMod _ 0 _ = 1
powMod !x 1 !m = x `mod` m
powMod !x !e !m
    | even e = (powMod x (e `div` 2) m ^ 2) `mod` m
    | otherwise = (x * powMod x (e - 1) m) `mod` m

prop_powMod :: Integer -> Integer -> Integer -> Property
prop_powMod x e m = e > 0 ==> m > 0 ==> x^e `mod` m == powMod x e m

nmod :: Integral a => a -> a -> a
(!m) `nmod` (!n)
    | m < 0 = (m + n) `mod` n
    | otherwise = m `mod` n
infixl 7 `nmod`

-- Can't be bothered to quickcheck this (only primes, and primeList needs to be large enough), but for my quick attempts I always got the correct results.
-- Algorithm is the one given on Wikipedia, I tried reading through the paper cited by Harvey but quickly noped out of that.
findRoot :: [Integer] -> Integer -> Integer
findRoot primeList !n =
    head $ filter (\ !m -> all (\ !p -> powMod m (pred n `div` p) n /= 1) orderFactors) [2..pred n]
  where
    orderFactors = filter (\ !p -> pred n `mod` p == 0) $ takeWhile (<n) primeList

-- No idea how modulo properly works for rational numbers, but this at least matches sagemath.
-- There also has to be something more efficient than bruteforce for it
ratioMod :: Integral a => a -> a -> a -> a
ratioMod !a !b !m =
    head (filter (\n -> n `mod` b == 0) [a, a+m..]) `div` b `nmod` m

bernoulliMod :: [Integer] -> Integer -> Integer -> Integer
bernoulliMod primeList !k !p =
    let !g = findRoot primeList p
        !r = powMod g (pred k) p
        !u = ratioMod (pred g) 2 p
        (!sResult, _, _) = foldl'
            (\(!s, !x, !y) _ ->
                let !q = g * x `div` p
                    !s' = (s + (u - q) * y) `nmod` p
                    !x' = g * x `mod` p
                    !y' = r * y `mod` p
                in (s', x', y'))
            (0, 1, r) [1 .. pred p `div` 2]
    in ratioMod (2 * k * sResult) (1 - powMod g k p) p

bezout :: Integral a => a -> a -> (a, a)
bezout !a !b = go 0 1 1 0 b a
  where
    go _ !s' _ !t' 0 _ = (s', t')
    go !s !s' !t !t' !r !r' =
        let !q = r' `div` r
        in go (s' - q * s) s (t' - q * t) t (r' - q * r) r

pairReduce :: (a -> a -> a) -> [a] -> a
pairReduce f xs = case go xs of
    [!x] -> x
    xs' -> pairReduce f xs'
  where
    go [] = []
    go [x] = [x]
    go (!x1:(!x2):xs) = let !r = f x1 x2 in r : go xs

chineseRemainder :: Integral a => [(a, a)] -> a
chineseRemainder l =
    fst $ pairReduce
        (\(!a1, !n1) (!a2, !n2) ->
            let (!m1, !m2) = bezout n1 n2
                !x = a1 * m2 * n2 + a2 * m1 * n1
            in (x `nmod` (n1 * n2), n1 * n2))
        l

-- This is awful, but should do the job
prop_chineseRemainder :: [(Integer, Integer)] -> Property
prop_chineseRemainder l =
    not (null l)
    ==> all (\(_, n) -> n `elem` map fromIntegral (primes $ maximum $ map (fromInteger . snd) l)) l
    ==> length (nub $ map snd l) == length l
    ==> let r = chineseRemainder l
        in all (\(a, n) -> r `mod` n == a `mod` n) l

factorise :: Integral a => [a] -> a -> [a]
factorise _ 1 = []
factorise (!p:ps) !n 
    | n `mod` p == 0 = p : factorise (p:ps) (n `div` p)
    | otherwise = factorise ps n

bernoulli :: Integer -> Rational
bernoulli !k
    -- Lookup table cause the algorithm doesn't work for k <= 11 (beta ends up being negative, so there will be a negative exponent in computeX)
    | k <= 10 = [1, -1%2, 1%6, 0, -1%30, 0, 1%42, 0, -1%30, 0, 5%66] !! fromInteger k
    | odd k = 0
    | otherwise =
    let !kDouble = fromInteger k :: Double
        !y = max 37 (ceiling ((kDouble + 0.5) * logBase 2 kDouble))
        primesToY = toList (strictPrimes y)
        (primes64bit, primesTooBig) = partition (<2^63-1) primesToY
        !primeSet64bit = IS.fromDistinctAscList $ map fromInteger primes64bit
        !primeSetTooBig = HS.fromList primesTooBig
        !primeCheck =
            if HS.null primeSetTooBig
                then \ !n -> IS.member (fromInteger n) primeSet64bit
                else \ !n -> if n < 2^63-1
                    then IS.member (fromInteger n) primeSet64bit
                    else HS.member n primeSetTooBig
        kFactors = 1 : filter (\ !d -> k `mod` d == 0) [2 .. intSqrt k]
        kFactors' = kFactors ++ map (div k) (reverse kFactors)
        !dk = product $ filter primeCheck $ map succ kFactors'
        !beta = ceiling $ (kDouble + 0.5) * logBase 2 kDouble - 4.094*kDouble + 2.470 + logBase 2 (fromInteger dk)
        computeX (!p : (!p') : ps) !m'
            | m' >= 2^(succ beta) = p
            | otherwise =
                let !m'' = if k `mod` pred p /= 0 then p' * m' else m'
                in computeX (p':ps) m''
        !x = computeX (tail primesToY) 1
        -- Could parallelise if some nice abstraction package or IO was available
        !rs = map (\ !p -> (bernoulliMod primesToY k p, p)) $ filter (\ !p -> k `mod` pred p /= 0) $ takeWhile (<= x) primesToY
        !m = product $ map snd rs
        !r = chineseRemainder rs
        !n' = dk * r `mod` m
        !nk = if k `mod` 4 == 2 then n' else n' - m
    in nk % dk