content: {-# LANGUAGE MagicHash #-}
  {-# LANGUAGE RecordWildCards #-}
  
  -- | A simple HxW matrix backed by a vector, mainly for binary exponentiation.
  --
  -- The matrix is a left semigroup action: \(m_2 (m_1 v) = (m_2 \circ m_1) v\).
  --
  -- @since 1.1.0.0
  module AtCoder.Extra.Semigroup.Matrix
    ( -- * Matrix
      Matrix (..),
  
      -- * Constructors
      new,
      square,
      zero,
      ident,
      diag,
  
      -- * Mapping
      map,
  
      -- * Multiplications
      mulToCol,
      mul,
      mulMod,
      mulMint,
  
      -- * Powers
      pow,
      powMod,
      powMint,
  
      -- * Rank
      rank,
  
      -- * Inverse
      inv,
      invRaw,
  
      -- * Determinant
      detMod,
      detMint,
    )
  where
  
  import AtCoder.Extra.Math qualified as ACEM
  import AtCoder.Internal.Assert qualified as ACIA
  import AtCoder.Internal.Barrett qualified as BT
  import AtCoder.ModInt qualified as M
  import Control.Monad (when)
  import Control.Monad.Primitive (PrimMonad, PrimState)
  import Control.Monad.ST (runST)
  import Data.Foldable (for_)
  import Data.Semigroup (Semigroup (..))
  import Data.Vector qualified as V
  import Data.Vector.Generic qualified as VG
  import Data.Vector.Generic.Mutable qualified as VGM
  import Data.Vector.Mutable qualified as VM
  import Data.Vector.Unboxed qualified as VU
  import Data.Vector.Unboxed.Mutable qualified as VUM
  import Data.Word (Word64)
  import GHC.Exts (proxy#)
  import GHC.Stack (HasCallStack)
  import GHC.TypeNats (KnownNat, natVal')
  import Prelude hiding (map)
  
  -- | A simple HxW matrix backed by a vector, mainly for binary exponentiation.
  --
  -- The matrix is a left semigroup action: \(m_2 (m_1 v) = (m_2 \circ m_1) v\).
  --
  --
  -- @since 1.1.0.0
  data Matrix a = Matrix
    { -- | @since 1.1.0.0
      hM :: {-# UNPACK #-} !Int,
      -- | @since 1.1.0.0
      wM :: {-# UNPACK #-} !Int,
      -- | @since 1.1.0.0
      vecM :: !(VU.Vector a)
    }
    deriving
      ( -- | @since 1.1.0.0
        Show,
        -- | @since 1.1.0.0
        Eq
      )
  
  -- | Type alias of a column vector.
  --
  -- @since 1.1.0.0
  type Col a = VU.Vector a
  
  -- | \(O(hw)\) Creates an HxW matrix.
  --
  -- @since 1.1.0.0
  {-# INLINE new #-}
  new :: (HasCallStack, VU.Unbox a) => Int -> Int -> VU.Vector a -> Matrix a
  new h w vec
    | VU.length vec /= h * w = error "AtCoder.Extra.Matrix: size mismatch"
    | otherwise = Matrix h w vec
  
  -- | \(O(n^2)\) Creates an NxN square matrix.
  --
  -- @since 1.1.1.0
  {-# INLINE square #-}
  square :: (HasCallStack, VU.Unbox a) => Int -> VU.Vector a -> Matrix a
  square n = new n n
  
  -- | \(O(n^2)\) Creates an NxN zero matrix.
  --
  -- @since 1.1.0.0
  {-# INLINE zero #-}
  zero :: (VU.Unbox a, Num a) => Int -> Matrix a
  zero n = Matrix n n $ VU.replicate (n * n) 0
  
  -- | \(O(n^2)\) Creates an NxN identity matrix.
  --
  -- @since 1.1.0.0
  {-# INLINE ident #-}
  ident :: (VU.Unbox a, Num a) => Int -> Matrix a
  ident n = Matrix n n $ VU.create $ do
    vec <- VUM.replicate (n * n) 0
    for_ [0 .. n - 1] $ \i -> do
      VGM.write vec (i + n * i) 1
    pure vec
  
  -- | \(O(n^2)\) Creates an NxN diagonal matrix.
  --
  -- @since 1.1.0.0
  {-# INLINE diag #-}
  diag :: (VU.Unbox a, Num a) => VU.Vector a -> Matrix a
  diag xs = Matrix n n $ VU.create $ do
    vec <- VUM.replicate (n * n) 0
    VU.iforM_ xs $ \i x -> do
      VGM.write vec (i + n * i) x
    pure vec
    where
      n = VU.length xs
  
  -- | \(O(n^2)\) Maps the `Matrix` elements.
  --
  -- @since 1.1.0.0
  {-# INLINE map #-}
  map :: (VU.Unbox a, VU.Unbox b) => (a -> b) -> Matrix a -> Matrix b
  map f Matrix {..} = Matrix hM wM $ VU.map f vecM
  
  -- | \(O(hw)\) Multiplies HxW matrix to a Hx1 column vector.
  --
  -- @since 1.1.0.0
  {-# INLINE mulToCol #-}
  mulToCol :: (Num a, VU.Unbox a) => Matrix a -> Col a -> Col a
  mulToCol Matrix {..} !col = VU.convert $ V.map (VU.sum . VU.zipWith (*) col) rows
    where
      !n = VU.length col
      !_ = ACIA.runtimeAssert (n == wM) "AtCoder.Extra.Matrix.mulToCol: size mismatch"
      rows = V.unfoldrExactN hM (VU.splitAt wM) vecM
  
  -- | \(O(h_1 K w_2)\) Multiplies H1xK matrix to a KxW2 matrix.
  --
  -- @since 1.1.0.0
  {-# INLINE mul #-}
  mul :: forall e. (Num e, VU.Unbox e) => Matrix e -> Matrix e -> Matrix e
  mul (Matrix h w vecA) (Matrix h' w' vecB) = Matrix h w' $ VU.create $ do
    c <- VUM.replicate (h * w') (0 :: e)
    for_ [0 .. h - 1] $ \i -> do
      for_ [0 .. w - 1] $ \k -> do
        for_ [0 .. w' - 1] $ \j -> do
          let !aik = VG.unsafeIndex vecA (i * w + k)
          let !bkj = VG.unsafeIndex vecB (k * w' + j)
          VGM.unsafeModify c (+ (aik * bkj)) (i * w' + j)
    pure c
    where
      !_ = ACIA.runtimeAssert (w == h') "AtCoder.Extra.Matrix.mul: matrix size mismatch"
  
  -- | \(O(h_1 w_2 K)\) Multiplies H1xK matrix to a KxW2 matrix, taking the mod.
  --
  -- @since 1.1.0.0
  {-# INLINE mulMod #-}
  mulMod :: Int -> Matrix Int -> Matrix Int -> Matrix Int
  mulMod !m (Matrix h w vecA) (Matrix h' w' vecB) =
    Matrix h w' $
      VU.unfoldrExactN
        (h * w')
        ( \(!row, !col) ->
            let !x = f row col
             in if col + 1 >= w'
                  then (x, (row + 1, 0))
                  else (x, (row, col + 1))
        )
        (0, 0)
    where
      !bt = BT.new32 $ fromIntegral m
      -- NOTE: this is unsafe if the matrix is too large
      f row col =
        fromIntegral
          . (`rem` fromIntegral m)
          . VU.sum
          $ VU.imap
            (\iRow x -> BT.mulMod bt (fromIntegral x) (fromIntegral (VG.unsafeIndex vecB (col + (iRow * w')))))
            (VU.unsafeSlice (w * row) w vecA)
      !_ = ACIA.runtimeAssert (w == h') "AtCoder.Extra.Matrix.mulMod: matrix size mismatch"
  
  -- | \(O(h_1 w_2 K)\) `mul` specialized to `M.ModInt`.
  --
  -- @since 1.1.0.0
  {-# INLINE mulMint #-}
  mulMint :: forall a. (KnownNat a) => Matrix (M.ModInt a) -> Matrix (M.ModInt a) -> Matrix (M.ModInt a)
  mulMint = mulMintImpl bt
    where
      !bt = BT.new32 $ fromIntegral (natVal' (proxy# @a))
  
  {-# INLINE mulMintImpl #-}
  mulMintImpl :: forall a. (KnownNat a) => BT.Barrett -> Matrix (M.ModInt a) -> Matrix (M.ModInt a) -> Matrix (M.ModInt a)
  mulMintImpl !bt (Matrix h w vecA) (Matrix h' w' vecB) =
    Matrix h w' $
      VU.unfoldrExactN
        (h * w')
        ( \(!row, !col) ->
            let !x = f row col
             in if col + 1 >= w'
                  then (x, (row + 1, 0))
                  else (x, (row, col + 1))
        )
        (0, 0)
    where
      -- NOTE: this is unsafe if the matrix is too large
      f :: Int -> Int -> M.ModInt a
      f row col =
        M.new64
          . VU.sum
          $ VU.imap
            ( \iRow x ->
                BT.mulMod
                  bt
                  (fromIntegral (M.unModInt x))
                  (fromIntegral (M.unModInt (VG.unsafeIndex vecB (col + (iRow * w')))))
            )
            (VU.unsafeSlice (w * row) w vecA)
      !_ = ACIA.runtimeAssert (w == h') "AtCoder.Extra.Matrix.mulMintImpl: matrix size mismatch"
  
  -- | \(O(w n^3)\) Returns \(M^k\).
  --
  -- @since 1.1.0.0
  {-# INLINE pow #-}
  pow :: Int -> Matrix Int -> Matrix Int
  pow k mat
    | k < 0 = error "AtCoder.Extra.Matrix.pow: the exponential must be non-negative"
    | k == 0 = ident $ hM mat
    | otherwise = ACEM.power mul k mat
    where
      !_ = ACIA.runtimeAssert (hM mat == wM mat) "AtCoder.Extra.Matrix.pow: matrix size mismatch"
  
  -- | \(O(w n^3)\) Returns \(M^k\), taking the mod.
  --
  -- @since 1.1.0.0
  {-# INLINE powMod #-}
  powMod :: Int -> Int -> Matrix Int -> Matrix Int
  powMod m k mat
    | k < 0 = error "AtCoder.Extra.Matrix.powMod: the exponential must be non-negative"
    | k == 0 = ident $ hM mat
    | otherwise = ACEM.power (mulMod m) k mat
    where
      !_ = ACIA.runtimeAssert (hM mat == wM mat) "AtCoder.Extra.Matrix.powMod: matrix size mismatch"
  
  -- | \(O(w n^3)\) Returns \(M^k\), specialized to `M.ModInt`.
  --
  -- @since 1.1.0.0
  {-# INLINE powMint #-}
  powMint :: forall m. (KnownNat m) => Int -> Matrix (M.ModInt m) -> Matrix (M.ModInt m)
  powMint k mat
    | k < 0 = error "AtCoder.Extra.Matrix.powMint: the exponential must be non-negative"
    | k == 0 = ident $ hM mat
    | otherwise = ACEM.power (mulMintImpl bt) k mat
    where
      !_ = ACIA.runtimeAssert (hM mat == wM mat) "AtCoder.Extra.Matrix.powMint: matrix size mismatch"
      !bt = BT.new32 $ fromIntegral (natVal' (proxy# @m))
  
  -- | (Internal)
  {-# INLINE read2d #-}
  read2d ::
    (PrimMonad m, VU.Unbox a) =>
    VM.MVector (PrimState m) (VUM.MVector (PrimState m) a) ->
    Int ->
    Int ->
    m a
  read2d view i j = do
    row <- VGM.unsafeRead view i
    VGM.unsafeRead row j
  
  -- | \(O(hw \min(h, w))\) Returns the rank of the matrix.
  --
  -- @since 1.1.1.0
  {-# INLINE rank #-}
  rank :: (Fractional a, Eq a, VU.Unbox a) => Matrix a -> Int
  rank (Matrix h w vec) = runST $ do
    vm <- VU.thaw vec
    view <- V.thaw $ V.unfoldrExactN h (VUM.splitAt w) vm
    let inner rk j
          | rk == h || j == w = pure rk
          | otherwise = do
              xrj <- read2d view rk j
              when (xrj == 0) $ do
                let runSwap i
                      | i == h = pure ()
                      | otherwise = do
                          xij <- read2d view i j
                          if xij /= 0
                            then VGM.unsafeSwap view rk i
                            else runSwap (i + 1)
                runSwap (rk + 1)
              xrj' <- read2d view rk j
              if xrj' == 0
                then inner rk (j + 1)
                else do
                  let c = 1 / xrj'
                  rowRk <- VGM.read view rk
                  -- for_ [j .. w - 1] $ \k -> do
                  VGM.iforM_ (VGM.unsafeDrop j rowRk) $ \k_ x -> do
                    VGM.unsafeWrite rowRk (k_ + j) $! c * x
                  for_ [rk + 1 .. h - 1] $ \i -> do
                    c_ <- read2d view i j
                    rowI <- VGM.read view i
                    -- for_ [j .. w - 1] $ \k -> do
                    VGM.iforM_ (VGM.unsafeDrop j rowRk) $ \k_ ark -> do
                      VGM.unsafeModify rowI (subtract (ark * c_)) (k_ + j)
                  inner (rk + 1) (j + 1)
    inner 0 0
  
  -- TODO: add HasCallStack and compare their speeds
  
  -- | \(O(n^3)\) Returns @(det, invMatrix)@ or `Nothing` if the matrix does not have inverse (the
  -- determinant is zero).
  --
  -- ==== Constraints
  -- - The input must be a square matrix.
  --
  -- @since 1.1.1.0
  {-# INLINE inv #-}
  inv :: forall a. (Fractional a, Eq a, VU.Unbox a) => Matrix a -> Maybe (a, Matrix a)
  inv mat@(Matrix n _ _) = do
    (!det, !invMat) <- invRaw mat
    let !invMat' = VU.concat $ V.toList invMat
    pure (det, Matrix n n invMat')
  
  -- | \(O(n^3)\) Returns @(det, invMatrix)@ or `Nothing` if the matrix does not have inverse (the
  -- determinant is zero).
  --
  -- ==== Constraints
  -- - The input must be a square matrix.
  --
  -- @since 1.1.1.0
  {-# INLINE invRaw #-}
  invRaw :: forall a. (Fractional a, Eq a, VU.Unbox a) => Matrix a -> Maybe (a, V.Vector (VU.Vector a))
  invRaw (Matrix h w vec) = runST $ do
    vecA <- VU.thaw vec
    viewA <- V.thaw $ V.unfoldrExactN n (VUM.splitAt n) vecA
  
    vecB <- VUM.replicate (n * n) (0 :: a)
    for_ [0 .. n - 1] $ \i -> do
      VGM.unsafeWrite vecB (n * i + i) (1 :: a)
    viewB <- V.thaw $ V.unfoldrExactN n (VUM.splitAt n) vecB
  
    let inner i !det
          | i >= n = do
              viewB' <- V.mapM VU.unsafeFreeze =<< V.unsafeFreeze viewB
              pure $ Just (det, viewB')
          | otherwise = do
              let swapLoop k !det_
                    | k >= n = pure det_
                    | otherwise = do
                        aki <- read2d viewA k i
                        if aki /= 0
                          then do
                            if k /= i
                              then do
                                VGM.unsafeSwap viewA i k
                                VGM.unsafeSwap viewB i k
                                pure (-det_)
                              else pure det_
                          else do
                            swapLoop (k + 1) det_
              det' <- swapLoop i det
              aii <- read2d viewA i i
              if aii == 0
                then pure Nothing
                else do
                  let !c = (1 :: a) / aii
                  let !det'' = det' * aii
                  rowAI <- VGM.unsafeRead viewA i
                  rowBI <- VGM.unsafeRead viewB i
                  VUM.iforM_ (VUM.unsafeDrop i rowAI) $ \j_ x -> do
                    VGM.unsafeWrite rowAI (j_ + i) $! x * c
                  VUM.iforM_ rowBI $ \j x -> do
                    VGM.unsafeWrite rowBI j $! x * c
                  for_ [0 .. n - 1] $ \k -> do
                    when (i /= k) $ do
                      c_ <- read2d viewA k i
                      rowAK <- VGM.unsafeRead viewA k
                      rowBK <- VGM.unsafeRead viewB k
                      VGM.iforM_ (VGM.unsafeDrop i rowAI) $ \j_ aij -> do
                        VGM.unsafeModify rowAK (subtract (aij * c_)) (j_ + i)
                      VGM.iforM_ rowBI $ \j bij -> do
                        VGM.unsafeModify rowBK (subtract (bij * c_)) j
                  inner (i + 1) det''
  
    inner 0 (1 :: a)
    where
      !_ = ACIA.runtimeAssert (h == w) $ "AtCoder.Extra.Semigroup.Matrix.inv: given non-square matrix of size " ++ show (h, w)
      n = h
  
  -- | \(O(hw \min(h, w))\) Returns the rank of the matrix.
  --
  -- @since 1.1.1.0
  {-# INLINE detMod #-}
  detMod :: Int -> Matrix Int -> Int
  detMod m (Matrix h w vecA) = runST $ do
    vm <- VU.thaw vecA
    view <- V.thaw $ V.unfoldrExactN n (VUM.splitAt n) vm
  
    let inner i (!det :: Int)
          | i >= n = pure det
          | otherwise = do
              let swapLoop j !det_
                    | j >= n = pure det_
                    | otherwise = do
                        aji <- read2d view j i
                        if aji == 0
                          then swapLoop (j + 1) det_
                          else do
                            if i /= j
                              then do
                                VGM.unsafeSwap view i j
                                pure $! m - det_
                              else pure det_
              det' <- swapLoop i det
              det'' <-
                VU.foldM'
                  ( \ !acc j -> do
                      let visitDiag !det_ = do
                            aii <- read2d view i i
                            if aii == 0
                              then pure det_
                              else do
                                aji <- read2d view j i
                                let !c = m - aji `div` aii
                                rowI <- VGM.unsafeRead view i
                                rowJ <- VGM.unsafeRead view j
                                -- NOTE: it's a reverse loop!
                                VGM.ifoldrM'
                                  ( \k_ aik () -> do
                                      VGM.unsafeModify rowJ ((`mod` m) . (+ aik * c)) (k_ + i)
                                  )
                                  ()
                                  (VGM.unsafeDrop i rowI)
                                VGM.unsafeSwap view i j
                                visitDiag (m - det_)
                      acc' <- visitDiag acc
                      VGM.unsafeSwap view i j
                      pure $! m - acc'
                  )
                  det'
                  (VU.generate (n - (i + 1)) (+ (i + 1)))
  
              inner (i + 1) det''
  
    det <- inner 0 (1 :: Int)
    fromIntegral
      <$> VU.foldM'
        ( \(!acc :: Word64) i -> do
            aii <- read2d view i i
            pure $! BT.mulMod bt acc $! fromIntegral aii
        )
        (fromIntegral det)
        (VU.generate n id)
    where
      !_ = ACIA.runtimeAssert (h == w) $ "AtCoder.Extra.Semigroup.Matrix.detMod: given non-square matrix of size " ++ show (h, w)
      !n = h
      !bt = BT.new32 $ fromIntegral m
  
  -- | \(O(hw \min(h, w))\) Returns the rank of the matrix.
  --
  -- @since 1.1.1.0
  {-# INLINE detMint #-}
  detMint :: forall a. (KnownNat a) => Matrix (M.ModInt a) -> M.ModInt a
  detMint matA = M.new . detMod m $ map M.val matA
    where
      !m = fromIntegral (natVal' (proxy# @a))
  
  -- | @since 1.1.0.0
  instance (Num a, VU.Unbox a) => Semigroup (Matrix a) where
    {-# INLINE (<>) #-}
    (<>) = mul
  
    -- Prefer `powMod` or `powMint` as specialized, much efficient variant.
    {-# INLINE stimes #-}
    stimes = ACEM.power (<>) . fromIntegral
  