content: {-# LANGUAGE TypeFamilies #-}
  
  -- | Math module. It contains number-theoretic algorithms.
  --
  -- @since 1.0.0.0
  module AtCoder.Math
    ( -- * Modulus operations
  
      -- These functions are internally used for `AtCoder.ModInt`.
      powMod,
      invMod,
  
      -- * Chinese Remainder Theorem
      crt,
  
      -- * Floor sum
      floorSum,
    )
  where
  
  import AtCoder.Internal.Assert qualified as ACIA
  import AtCoder.Internal.Math (powMod)
  import AtCoder.Internal.Math qualified as ACIM
  import Data.Bits (bit)
  import Data.Vector.Unboxed qualified as VU
  import GHC.Stack (HasCallStack)
  
  -- `powMod` is re-exported from the internal math module.
  
  -- | Returns an integer \(y\) such that \(0 \le y < m\) and  \(xy \equiv 1 \pmod m\).
  --
  -- ==== Constraints
  -- - \(\gcd(x, m) = 1\)
  -- - \(1 \leq m\)
  --
  -- ==== Complexity
  -- - \(O(\log m)\)
  --
  -- ==== Example
  -- >>> let m = 998244353
  -- >>> (invMod 2 m) * 2 `mod` m -- (2^(-1) mod m) * 2 mod m
  -- 1
  --
  -- @since 1.0.0.0
  {-# INLINE invMod #-}
  invMod ::
    (HasCallStack) =>
    -- | \(x\)
    Int ->
    -- | \(m\)
    Int ->
    -- | \(x^{-1} \bmod m\)
    Int
  invMod x m =
    let !_ = ACIA.runtimeAssert (1 <= m) $ "AtCoder.Math.invMod: given invalid `m` less than 1: " ++ show m
        (!z1, !z2) = ACIM.invGcd (fromIntegral x) (fromIntegral m)
        !_ = ACIA.runtimeAssert (z1 == 1) "AtCoder.Math.invMod: `x^(-1) mod m` cannot be calculated when `gcd x m /= 1`"
     in z2
  
  -- | Given two arrays \(r,m\) with length \(n\), it solves the modular equation system
  --
  -- \[
  -- x \equiv r[i] \pmod{m[i]}, \forall i \in \lbrace 0,1,\cdots, n - 1 \rbrace.
  -- \]
  --
  -- If there is no solution, it returns \((0, 0)\). Otherwise, all the solutions can be written as the form \(x \equiv y \pmod z\), using integers
  -- \(y, z\) \((0 \leq y < z = \mathrm{lcm}(m[i]))\). It returns this \((y, z)\) as a pair. If \(n=0\), it returns \((0, 1)\).
  --
  -- ==== Constraints
  -- - \(|r| = |m|\)
  -- - \(1 \le m[i]\)
  -- - \(\mathrm{lcm}(m[i])\) is in `Int` bounds.
  --
  -- ==== Complexity
  -- - \(O(n \log{\mathrm{lcm}(m[i])})\)
  --
  -- ==== __Example__
  -- `crt` calculates \(y\) such that \(y \equiv r_i \pmod m_i, \forall i \in \lbrace 0,1,\cdots, n - 1 \rbrace\):
  --
  -- >>> import Data.Vector.Unboxed qualified as VU
  -- >>> let rs = VU.fromList @Int [6, 7, 8, 9]
  -- >>> let ms = VU.fromList @Int [2, 3, 4, 5]
  -- >>> crt rs ms
  -- (4,60)
  --
  -- The property can be checked as follows:
  --
  -- >>> let (y, {- lcm ms -} _) = crt rs ms
  -- >>> VU.zipWith mod rs ms
  -- [0,1,0,4]
  --
  -- >>> VU.zipWith mod rs ms == VU.map (y `mod`) ms
  -- True
  --
  -- @since 1.0.0.0
  {-# INLINE crt #-}
  crt :: (HasCallStack) => VU.Vector Int -> VU.Vector Int -> (Int, Int)
  crt r m = loop 0 1 [0 .. VU.length r - 1]
    where
      !_ = ACIA.runtimeAssert (VU.length r == VU.length m) "AtCoder.Math.crt: given `r` and `m` with different lengths"
      loop !r0 !m0 [] = (r0, m0)
      loop !r0 !m0 (i : rest)
        | m0' `mod` m1' == 0 =
            if r0' `mod` m1' /= r1'
              then (0, 0)
              else loop r0' m0' rest
        | otherwise =
            let (!g, !im) = ACIM.invGcd m0' m1'
                u1 = m1' `div` g
             in if ((r1' - r0') `mod` g) /= 0
                  then (0, 0)
                  else
                    let !x = (r1' - r0') `div` g `mod` u1 * im `mod` u1
                        !r0'' = r0' + x * m0'
                        !m0'' = m0' * u1
                     in if r0'' < 0
                          then loop (r0'' + m0'') m0'' rest
                          else loop r0'' m0'' rest
        where
          !_ = ACIA.runtimeAssert (1 <= mi) $ "AtCoder.Math.crt: `m[i]` is not positive: " ++ show mi
          !mi = VU.unsafeIndex m i
          !ri = VU.unsafeIndex r i
          !r1 = ri `mod` mi
          !m1 = mi
          (!m0', !m1', !r0', !r1')
            | m0 < m1 = (m1, m0, r1, r0)
            | otherwise = (m0, m1, r0, r1)
  
  -- | Returns \(\sum\limits_{i = 0}^{n - 1} \left\lfloor \frac{a \times i + b}{m} \right\rfloor\).
  --
  -- ==== Constraints
  -- - \(0 \le n\)
  -- - \(1 \le m\)
  --
  -- ==== Complexity
  -- - \(O(\log m)\)
  --
  -- ==== __Example__
  -- `floorSum` calculates the number of points surrounded by a line
  -- \(y = \frac {a \times x + b} {m} \) and \(x, y\) axes in \(O(\log m)\) time:
  --
  -- >>> floorSum 5 1 1 1 -- floorSum n {- line information -} m a b
  -- 15
  --
  -- @
  --   y
  --   ^
  -- 6 |
  -- 5 |           o           line: y = x + 1
  -- 4 |        o  o           The number of \`o\` is 15
  -- 3 |     o  o  o
  -- 2 |  o  o  o  o
  -- 1 |  o  o  o  o
  -- --+-----------------> x
  --   0  1  2  3  4  5
  --                  n = 5
  -- @
  --
  -- @since 1.0.0.0
  {-# INLINE floorSum #-}
  floorSum ::
    (HasCallStack) =>
    -- | \(n\)
    Int ->
    -- | \(m\)
    Int ->
    -- | \(a\)
    Int ->
    -- | \(b\)
    Int ->
    -- | \(\sum\limits_{i = 0}^{n - 1} \left\lfloor \frac{a \times i + b}{m} \right\rfloor\)
    Int
  floorSum n m a b = ACIM.floorSumUnsigned n m a' b' - da - db
    where
      !_ = ACIA.runtimeAssert (0 <= n && n < bit 32) $ "AtCoder.Math.floorSum: given invalid `n` (`" ++ show n ++ "`)"
      !_ = ACIA.runtimeAssert (1 <= m && m < bit 32) $ "AtCoder.Math.floorSum: given invalid `m` (`" ++ show m ++ "`)"
      a'
        | a < 0 = a `mod` m
        | otherwise = a
      da
        | a < 0 = n * (n - 1) `div` 2 * (((a `mod` m) - a) `div` m)
        | otherwise = 0
      b'
        | b < 0 = b `mod` m
        | otherwise = b
      db
        | b < 0 = n * (((b `mod` m) - b) `div` m)
        | otherwise = 0
  