content: module BenchLib.PowMod
    ( powModBarrettWideWord,
      powModBarrett64,
      powModMontgomery,
      powModMod,
      powModRem,
      powModPowerMod,
      powModPowerRem,
    )
  where
  
  import AtCoder.Extra.Math qualified as ACEM
  import AtCoder.Internal.Assert qualified as ACIA
  import BenchLib.MulMod.Barrett64 qualified as Barrett64
  import BenchLib.MulMod.BarrettWideWord qualified as BarrettWideWord
  import BenchLib.MulMod.Montgomery qualified as Montgomery
  import Data.Bits ((!>>.))
  import Data.Word (Word64)
  import GHC.Stack (HasCallStack)
  
  -- safeMod :: Int -> Int -> Int
  -- safeMod = mod
  
  -- TODO: run unit test. probablly wrong.
  
  -- | Wide word Barrett-based
  powModBarrettWideWord :: (HasCallStack) => Int -> Int -> Int -> Int
  powModBarrettWideWord x n0 m0
    | m0 == 1 = 0
    | otherwise = fromIntegral $ inner n0 1 $ fromIntegral (x `mod` m0)
    where
      !_ = ACIA.runtimeAssert (0 <= n0 && 1 <= m0) $ "BenchLib.PowMod.powModBarrettWideWord: given invalid `n` or `m`: " ++ show (n0, m0)
      bt = BarrettWideWord.new64 $ fromIntegral m0
      inner :: Int -> Word64 -> Word64 -> Word64
      inner !n !r !y
        | n == 0 = r
        | otherwise =
            let r' = if odd n then BarrettWideWord.mulMod bt r y else r
                y' = BarrettWideWord.mulMod bt y y
             in inner (n !>>. 1) r' y'
  
  -- | Barrett64-based
  powModBarrett64 :: (HasCallStack) => Int -> Int -> Int -> Int
  powModBarrett64 x n0 m0
    | m0 == 1 = 0
    | otherwise = fromIntegral $ inner n0 1 $ fromIntegral (x `mod` m0)
    where
      !_ = ACIA.runtimeAssert (0 <= n0 && 1 <= m0) $ "BenchLib.PowMod.powModBarrett64: given invalid `n` or `m`: " ++ show (n0, m0)
      bt = Barrett64.new $ fromIntegral m0
      inner :: Int -> Word64 -> Word64 -> Word64
      inner !n !r !y
        | n == 0 = r
        | otherwise =
            let r' = if odd n then Barrett64.mulMod bt r y else r
                y' = Barrett64.mulMod bt y y
             in inner (n !>>. 1) r' y'
  
  -- | Montgomery-based
  powModMontgomery :: (HasCallStack) => Int -> Int -> Int -> Int
  powModMontgomery x n0 m0
    | m0 == 1 = 0
    | otherwise = fromIntegral . Montgomery.reduce mont $ inner n0 (Montgomery.generate mont 1) . Montgomery.generate mont $ fromIntegral (x `mod` m0)
    where
      !_ = ACIA.runtimeAssert (0 <= n0 && 1 <= m0) $ "BenchLib.PowMod.powModMontgomery: given invalid `n` or `m`: " ++ show (n0, m0)
      mont = Montgomery.new $ fromIntegral m0
      inner :: Int -> Word64 -> Word64 -> Word64
      inner !n !r !y
        | n == 0 = r
        | otherwise =
            let r' = if odd n then Montgomery.mulModGenerated mont r y else r
                y' = Montgomery.mulModGenerated mont y y
             in inner (n !>>. 1) r' y'
  
  -- | mod-based
  powModMod :: (HasCallStack) => Int -> Int -> Int -> Int
  powModMod x n0 m0
    | m0 == 1 = 0
    | otherwise = inner n0 1 $ fromIntegral (x `mod` m0)
    where
      !_ = ACIA.runtimeAssert (0 <= n0 && 1 <= m0) $ "BenchLib.PowMod.powModMod: given invalid `n` or `m`: " ++ show (n0, m0)
      inner :: Int -> Int -> Int -> Int
      inner !n !r !y
        | n == 0 = r
        | otherwise =
            let r' = if odd n then r * y `mod` m0 else r
                y' = y * y `mod` m0
             in inner (n !>>. 1) r' y'
  
  -- | rem-based. Be warned
  powModRem :: (HasCallStack) => Int -> Int -> Int -> Int
  powModRem x n0 m0
    | m0 == 1 = 0
    | otherwise = inner n0 1 $ fromIntegral (x `mod` m0)
    where
      !_ = ACIA.runtimeAssert (0 <= n0 && 1 <= m0) $ "BenchLib.PowMod.powModRem: given invalid `n` or `m`: " ++ show (n0, m0)
      inner :: Int -> Int -> Int -> Int
      inner !n !r !y
        | n == 0 = r
        | otherwise =
            let r' = if odd n then r * y `rem` m0 else r
                y' = y * y `rem` m0
             in inner (n !>>. 1) r' y'
  
  -- | power-based
  powModPowerMod :: (HasCallStack) => Int -> Int -> Int -> Int
  powModPowerMod x n0 m0
    | m0 == 1 = 0
    | otherwise = ACEM.power (\ !a !b -> a * b `mod` m0) n0 x
  
  -- | power-based
  powModPowerRem :: (HasCallStack) => Int -> Int -> Int -> Int
  powModPowerRem x n0 m0
    | m0 == 1 = 0
    | otherwise = ACEM.power (\ !a !b -> a * b `rem` m0) n0 x
  