-- |
-- Module:      Math.NumberTheory.Utils
-- Copyright:   (c) 2011 Daniel Fischer
-- Licence:     MIT
-- Maintainer:  Daniel Fischer <daniel.is.fischer@googlemail.com>
--
-- Some utilities, mostly for bit twiddling.
--

{-# LANGUAGE BangPatterns   #-}
{-# LANGUAGE MagicHash      #-}
{-# LANGUAGE UnboxedTuples  #-}
{-# LANGUAGE RankNTypes     #-}
{-# LANGUAGE KindSignatures #-}
{-# LANGUAGE DataKinds      #-}
{-# LANGUAGE GADTs          #-}

module Math.NumberTheory.Utils
    ( SomeKnown(..)
    , shiftToOddCount
    , shiftToOdd
    , shiftToOdd#
    , shiftToOddCount#
    , shiftToOddCountBigNat
    , splitOff
    , splitOff#

    , mergeBy

    , recipMod

    , toWheel30
    , fromWheel30
    , withSomeKnown
    , intVal
    ) where

import Prelude hiding (mod, quotRem)
import qualified Prelude as P

import Data.Bits
import Data.Euclidean
import Data.List.Infinite (Infinite(..))
import Data.Semiring (Semiring(..), isZero)
import GHC.Base
import GHC.Num.BigNat
import GHC.Num.Integer
import GHC.Num.Natural
import qualified Math.NumberTheory.Utils.FromIntegral as UT
import GHC.Natural
import GHC.TypeNats
import Math.NumberTheory.Utils.FromIntegral (intToWord)

-- | Remove factors of @2@ and count them. If
--   @n = 2^k*m@ with @m@ odd, the result is @(k, m)@.
--   Precondition: argument not @0@ (not checked).
{-# RULES
"shiftToOddCount/Int"       shiftToOddCount = shiftOCInt
"shiftToOddCount/Word"      shiftToOddCount = shiftOCWord
"shiftToOddCount/Integer"   shiftToOddCount = shiftOCInteger
"shiftToOddCount/Natural"   shiftToOddCount = shiftOCNatural
  #-}
{-# INLINE [1] shiftToOddCount #-}
shiftToOddCount :: Integral a => a -> (Word, a)
shiftToOddCount n = case shiftOCInteger (toInteger n) of
                      (z, o) -> (z, fromInteger o)

-- | Specialised version for @'Word'@.
--   Precondition: argument strictly positive (not checked).
shiftOCWord :: Word -> (Word, Word)
shiftOCWord (W# w#) = case shiftToOddCount# w# of
                        (# z# , u# #) -> (W# z#, W# u#)

-- | Specialised version for @'Int'@.
--   Precondition: argument nonzero (not checked).
shiftOCInt :: Int -> (Word, Int)
shiftOCInt (I# i#) = case shiftToOddCount# (int2Word# i#) of
                        (# z#, u# #) -> (W# z#, I# (word2Int# u#))

-- | Specialised version for @'Integer'@.
--   Precondition: argument nonzero (not checked).
shiftOCInteger :: Integer -> (Word, Integer)
shiftOCInteger n@(IS i#) =
    case shiftToOddCount# (int2Word# i#) of
      (# 0##, _ #) -> (0, n)
      (# z#, w# #) -> (W# z#, integerFromWord# w#)
shiftOCInteger n@(IP bn#) = case bigNatZeroCount bn# of
                                 0## -> (0, n)
                                 z#  -> (W# z#, integerFromBigNat# (bn# `bigNatShiftR#` z#))
shiftOCInteger n@(IN bn#) = case bigNatZeroCount bn# of
                                 0## -> (0, n)
                                 z#  -> (W# z#, integerFromBigNatNeg# (bn# `bigNatShiftR#` z#))

-- | Specialised version for @'Natural'@.
--   Precondition: argument nonzero (not checked).
shiftOCNatural :: Natural -> (Word, Natural)
shiftOCNatural n@(NatS# i#) =
    case shiftToOddCount# i# of
      (# 0##, _ #) -> (0, n)
      (# z#, w# #) -> (W# z#, NatS# w#)
shiftOCNatural n@(NatJ# (BN# bn#)) = case bigNatZeroCount bn# of
                                 0## -> (0, n)
                                 z#  -> (W# z#, naturalFromBigNat# (bn# `bigNatShiftR#` z#))

shiftToOddCountBigNat :: BigNat# -> (# Word, BigNat# #)
shiftToOddCountBigNat bn# = case bigNatZeroCount bn# of
  0## -> (# 0, bn# #)
  z#  -> (# W# z#, bn# `bigNatShiftR#` z# #)

-- | Count trailing zeros in a @'BigNat'@.
--   Precondition: argument nonzero (not checked, Integer invariant).
bigNatZeroCount :: BigNat# -> Word#
bigNatZeroCount bn# = count 0## 0#
  where
    !(W# bitSize#) = intToWord (finiteBitSize (0 :: Word))
    count a# i# =
          case bigNatIndex# bn# i# of
            0## -> count (a# `plusWord#` bitSize#) (i# +# 1#)
            w#  -> a# `plusWord#` ctz# w#

-- | Remove factors of @2@. If @n = 2^k*m@ with @m@ odd, the result is @m@.
--   Precondition: argument not @0@ (not checked).
{-# RULES
"shiftToOdd/Int"       shiftToOdd = shiftOInt
"shiftToOdd/Word"      shiftToOdd = shiftOWord
"shiftToOdd/Integer"   shiftToOdd = shiftOInteger
  #-}
{-# INLINE [1] shiftToOdd #-}
shiftToOdd :: Integral a => a -> a
shiftToOdd n = fromInteger (shiftOInteger (toInteger n))

-- | Specialised version for @'Int'@.
--   Precondition: argument nonzero (not checked).
shiftOInt :: Int -> Int
shiftOInt (I# i#) = I# (word2Int# (shiftToOdd# (int2Word# i#)))

-- | Specialised version for @'Word'@.
--   Precondition: argument nonzero (not checked).
shiftOWord :: Word -> Word
shiftOWord (W# w#) = W# (shiftToOdd# w#)

-- | Specialised version for @'Int'@.
--   Precondition: argument nonzero (not checked).
shiftOInteger :: Integer -> Integer
shiftOInteger (IS i#) = integerFromWord# (shiftToOdd# (int2Word# i#))
shiftOInteger n@(IP bn#) = case bigNatZeroCount bn# of
                                 0## -> n
                                 z#  -> integerFromBigNat# (bn# `bigNatShiftR#` z#)
shiftOInteger n@(IN bn#) = case bigNatZeroCount bn# of
                                 0## -> n
                                 z#  -> integerFromBigNatNeg# (bn# `bigNatShiftR#` z#)

-- | Shift argument right until the result is odd.
--   Precondition: argument not @0@, not checked.
shiftToOdd# :: Word# -> Word#
shiftToOdd# w# = uncheckedShiftRL# w# (word2Int# (ctz# w#))

-- | Like @'shiftToOdd#'@, but count the number of places to shift too.
shiftToOddCount# :: Word# -> (# Word#, Word# #)
shiftToOddCount# w# = case ctz# w# of
                        k# -> (# k#, uncheckedShiftRL# w# (word2Int# k#) #)

splitOff :: (Eq a, GcdDomain a) => a -> a -> (Word, a)
splitOff p n
  | isZero n  = (0, zero) -- prevent infinite loop
  | otherwise = go 0 n
  where
    go !k m = case m `divide` p of
      Just q -> go (k + 1) q
      _      -> (k, m)
{-# INLINABLE splitOff #-}

-- | It is difficult to convince GHC to unbox output of 'splitOff' and 'splitOff.go',
-- so we fallback to a specialized unboxed version to minimize allocations.
splitOff# :: Word# -> Word# -> (# Word#, Word# #)
splitOff# _ 0## = (# 0##, 0## #)
splitOff# p n = go 0## n
  where
    go k m = case m `quotRemWord#` p of
      (# q, 0## #) -> go (k `plusWord#` 1##) q
      _            -> (# k, m #)
{-# INLINABLE splitOff# #-}

-- | Merges two ordered lists into an ordered list. Checks for neither its
-- precondition or postcondition.
mergeBy :: (a -> a -> Ordering) -> Infinite a -> Infinite a -> Infinite a
mergeBy cmp = loop
  where
    loop ( x:< xs) (y :< ys)
      = case cmp x y of
         GT -> y :< loop (x :< xs) ys
         _  -> x :< loop xs (y :< ys)

-- | Work around https://ghc.haskell.org/trac/ghc/ticket/14085
recipMod :: Integer -> Integer -> Maybe Integer
recipMod x m = case integerRecipMod# (x `P.mod` m) (fromInteger m) of
  (# | _ #) -> Nothing
  (# y | #) -> Just (toInteger y)

-------------------------------------------------------------------------------
-- Helpers for mapping to rough numbers and back.
-- Copypasted from Data.BitStream.WheelMapping

toWheel30 :: (Integral a, Bits a) => a -> a
toWheel30 i = q `shiftL` 3 + (r + r `shiftR` 4) `shiftR` 2
  where
    (q, r) = i `P.quotRem` 30

fromWheel30 :: (Num a, Bits a) => a -> a
fromWheel30 i = ((i `shiftL` 2 - i `shiftR` 2) .|. 1)
              + ((i `shiftL` 1 - i `shiftR` 1) .&. 2)

-------------------------------------------------------------------------------
-- Helpers for dealing with data types parametrised by natural numbers.

data SomeKnown (f :: Nat -> Type) where
  SomeKnown :: KnownNat k => f k -> SomeKnown f

withSomeKnown :: (forall k. KnownNat k => f k -> a) -> SomeKnown f -> a
withSomeKnown f (SomeKnown x) = f x

intVal :: KnownNat k => a k -> Int
intVal = UT.naturalToInt . natVal