{-# LANGUAGE TypeFamilies #-}

-- | Monoid action \(f: x \rightarrow ax + b\).
--
-- - Use @Mat2x2@ if inverse operations are required, or if it's necessary to store the monoid
-- length in the acted monoid (@V2@).
--
-- @since 1.0.0.0
module AtCoder.Extra.Monoid.Affine1
  ( -- * Affine1
    Affine1 (..),
    Affine1Repr,

    -- * Constructors
    new,
    unAffine1,
    ident,
    zero,

    -- * Actions
    act,
  )
where

import AtCoder.Extra.Math qualified as ACEM
import AtCoder.LazySegTree (SegAct (..))
import Data.Coerce (coerce)
import Data.Foldable (foldl')
import Data.List.NonEmpty (NonEmpty (..))
import Data.Semigroup (Dual (..), Semigroup (..), Sum (..))
import Data.Vector.Generic qualified as VG
import Data.Vector.Generic.Mutable qualified as VGM
import Data.Vector.Unboxed qualified as VU
import Data.Vector.Unboxed.Mutable qualified as VUM

-- | Monoid action \(f: x \rightarrow ax + b\).
--
-- - Use @Mat2x2@ if inverse operations are required, or if it's necessary to store the monoid
-- length in the acted monoid (@V2@).
--
-- ==== Composition and dual
-- The affine transformation acts as a left monoid action: \(f_2 (f_1 v) = (f_2 \circ f_1) v\). To
-- apply the leftmost transformation first in a segment tree, wrap `Affine1` in @Data.Monoid.Dual@.
--
-- ==== __Example__
-- >>> import AtCoder.Extra.Monoid (SegAct(..), Affine1(..))
-- >>> import AtCoder.LazySegTree qualified as LST
-- >>> seg <- LST.build @_ @(Affine1 Int) @(Sum Int) $ VU.generate 3 Sum -- [0, 1, 2]
-- >>> LST.applyIn seg 0 3 $ Affine1 (2, 1) -- [1, 3, 5]
-- >>> getSum <$> LST.allProd seg
-- 9
--
-- @since 1.0.0.0
newtype Affine1 a = Affine1 (Affine1Repr a)
  deriving newtype
    ( -- | @since 1.0.0.0
      Affine1 a -> Affine1 a -> Bool
(Affine1 a -> Affine1 a -> Bool)
-> (Affine1 a -> Affine1 a -> Bool) -> Eq (Affine1 a)
forall a. Eq a => Affine1 a -> Affine1 a -> Bool
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
$c== :: forall a. Eq a => Affine1 a -> Affine1 a -> Bool
== :: Affine1 a -> Affine1 a -> Bool
$c/= :: forall a. Eq a => Affine1 a -> Affine1 a -> Bool
/= :: Affine1 a -> Affine1 a -> Bool
Eq,
      -- | @since 1.0.0.0
      Eq (Affine1 a)
Eq (Affine1 a) =>
(Affine1 a -> Affine1 a -> Ordering)
-> (Affine1 a -> Affine1 a -> Bool)
-> (Affine1 a -> Affine1 a -> Bool)
-> (Affine1 a -> Affine1 a -> Bool)
-> (Affine1 a -> Affine1 a -> Bool)
-> (Affine1 a -> Affine1 a -> Affine1 a)
-> (Affine1 a -> Affine1 a -> Affine1 a)
-> Ord (Affine1 a)
Affine1 a -> Affine1 a -> Bool
Affine1 a -> Affine1 a -> Ordering
Affine1 a -> Affine1 a -> Affine1 a
forall a.
Eq a =>
(a -> a -> Ordering)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> a)
-> (a -> a -> a)
-> Ord a
forall a. Ord a => Eq (Affine1 a)
forall a. Ord a => Affine1 a -> Affine1 a -> Bool
forall a. Ord a => Affine1 a -> Affine1 a -> Ordering
forall a. Ord a => Affine1 a -> Affine1 a -> Affine1 a
$ccompare :: forall a. Ord a => Affine1 a -> Affine1 a -> Ordering
compare :: Affine1 a -> Affine1 a -> Ordering
$c< :: forall a. Ord a => Affine1 a -> Affine1 a -> Bool
< :: Affine1 a -> Affine1 a -> Bool
$c<= :: forall a. Ord a => Affine1 a -> Affine1 a -> Bool
<= :: Affine1 a -> Affine1 a -> Bool
$c> :: forall a. Ord a => Affine1 a -> Affine1 a -> Bool
> :: Affine1 a -> Affine1 a -> Bool
$c>= :: forall a. Ord a => Affine1 a -> Affine1 a -> Bool
>= :: Affine1 a -> Affine1 a -> Bool
$cmax :: forall a. Ord a => Affine1 a -> Affine1 a -> Affine1 a
max :: Affine1 a -> Affine1 a -> Affine1 a
$cmin :: forall a. Ord a => Affine1 a -> Affine1 a -> Affine1 a
min :: Affine1 a -> Affine1 a -> Affine1 a
Ord,
      -- | @since 1.0.0.0
      Int -> Affine1 a -> ShowS
[Affine1 a] -> ShowS
Affine1 a -> String
(Int -> Affine1 a -> ShowS)
-> (Affine1 a -> String)
-> ([Affine1 a] -> ShowS)
-> Show (Affine1 a)
forall a. Show a => Int -> Affine1 a -> ShowS
forall a. Show a => [Affine1 a] -> ShowS
forall a. Show a => Affine1 a -> String
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
$cshowsPrec :: forall a. Show a => Int -> Affine1 a -> ShowS
showsPrec :: Int -> Affine1 a -> ShowS
$cshow :: forall a. Show a => Affine1 a -> String
show :: Affine1 a -> String
$cshowList :: forall a. Show a => [Affine1 a] -> ShowS
showList :: [Affine1 a] -> ShowS
Show
    )

-- | `Affine1` internal representation. Tuples are not the fastest representation, but it's easier
-- to implement `Data.Vector.Unboxed.Unbox`.
--
-- @since 1.0.0.0
type Affine1Repr a = (a, a)

-- | \(O(1)\) Creates a one-dimensional affine transformation: \(f: x \rightarrow a \times x + b\).
--
-- @since 1.0.0.0
{-# INLINE new #-}
new :: a -> a -> Affine1 a
new :: forall a. a -> a -> Affine1 a
new !a
a !a
b = Affine1Repr a -> Affine1 a
forall a. Affine1Repr a -> Affine1 a
Affine1 (a
a, a
b)

-- | \(O(1)\) Retrieves the two components of `Affine1`.
--
-- @since 1.1.0.0
{-# INLINE unAffine1 #-}
unAffine1 :: Affine1 a -> Affine1Repr a
unAffine1 :: forall a. Affine1 a -> Affine1Repr a
unAffine1 (Affine1 Affine1Repr a
a) = Affine1Repr a
a

-- | \(O(1)\) Identity transformation.
--
-- @since 1.1.0.0
{-# INLINE ident #-}
ident :: (Num a) => Affine1 a
ident :: forall a. Num a => Affine1 a
ident = Affine1Repr a -> Affine1 a
forall a. Affine1Repr a -> Affine1 a
Affine1 (a
1, a
0)

-- | \(O(1)\) Transformation to zero.
--
-- @since 1.1.0.0
{-# INLINE zero #-}
zero :: (Num a) => Affine1 a
zero :: forall a. Num a => Affine1 a
zero = Affine1Repr a -> Affine1 a
forall a. Affine1Repr a -> Affine1 a
Affine1 (a
0, a
0)

-- | \(O(1)\) Applies the one-dimensional affine transformation \(f: x \rightarrow a \times x + b\).
--
-- @since 1.0.0.0
{-# INLINE act #-}
act :: (Num a) => Affine1 a -> a -> a
act :: forall a. Num a => Affine1 a -> a -> a
act (Affine1 (!a
a, !a
b)) a
x = a
a a -> a -> a
forall a. Num a => a -> a -> a
* a
x a -> a -> a
forall a. Num a => a -> a -> a
+ a
b

-- | \(O(1)\) Acts on @a@ with length in terms of `SegAct`. Works for `Sum a` only.
--
-- @since 1.0.0.0
{-# INLINE actWithLength #-}
actWithLength :: (Num a) => Int -> Affine1 a -> a -> a
actWithLength :: forall a. Num a => Int -> Affine1 a -> a -> a
actWithLength Int
len (Affine1 (!a
a, !a
b)) !a
x = a
a a -> a -> a
forall a. Num a => a -> a -> a
* a
x a -> a -> a
forall a. Num a => a -> a -> a
+ a
b a -> a -> a
forall a. Num a => a -> a -> a
* Int -> a
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
len

-- | @since 1.0.0.0
instance (Num a) => Semigroup (Affine1 a) where
  {-# INLINE (<>) #-}
  (Affine1 (!a
a1, !a
b1)) <> :: Affine1 a -> Affine1 a -> Affine1 a
<> (Affine1 (!a
a2, !a
b2)) = (a, a) -> Affine1 a
forall a. Affine1Repr a -> Affine1 a
Affine1 (a
a', a
b')
    where
      !a' :: a
a' = a
a1 a -> a -> a
forall a. Num a => a -> a -> a
* a
a2
      !b' :: a
b' = a
a1 a -> a -> a
forall a. Num a => a -> a -> a
* a
b2 a -> a -> a
forall a. Num a => a -> a -> a
+ a
b1
  {-# INLINE sconcat #-}
  sconcat :: NonEmpty (Affine1 a) -> Affine1 a
sconcat (Affine1 a
x :| [Affine1 a]
xs) = (Affine1 a -> Affine1 a -> Affine1 a)
-> Affine1 a -> [Affine1 a] -> Affine1 a
forall b a. (b -> a -> b) -> b -> [a] -> b
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl' Affine1 a -> Affine1 a -> Affine1 a
forall a. Semigroup a => a -> a -> a
(<>) Affine1 a
x [Affine1 a]
xs
  {-# INLINE stimes #-}
  stimes :: forall b. Integral b => b -> Affine1 a -> Affine1 a
stimes b
b = (Affine1 a -> Affine1 a -> Affine1 a)
-> Int -> Affine1 a -> Affine1 a
forall a. (a -> a -> a) -> Int -> a -> a
ACEM.power Affine1 a -> Affine1 a -> Affine1 a
forall a. Semigroup a => a -> a -> a
(<>) (b -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral b
b)

-- | @since 1.0.0.0
instance (Num a) => Monoid (Affine1 a) where
  {-# INLINE mempty #-}
  mempty :: Affine1 a
mempty = Affine1Repr a -> Affine1 a
forall a. Affine1Repr a -> Affine1 a
Affine1 (a
1, a
0)
  {-# INLINE mconcat #-}
  mconcat :: [Affine1 a] -> Affine1 a
mconcat [] = Affine1 a
forall a. Monoid a => a
mempty
  mconcat (Affine1 a
x : [Affine1 a]
xs) = (Affine1 a -> Affine1 a -> Affine1 a)
-> Affine1 a -> [Affine1 a] -> Affine1 a
forall b a. (b -> a -> b) -> b -> [a] -> b
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl' Affine1 a -> Affine1 a -> Affine1 a
forall a. Semigroup a => a -> a -> a
(<>) Affine1 a
x [Affine1 a]
xs

-- | @since 1.0.0.0
instance (Num a) => SegAct (Affine1 a) (Sum a) where
  {-# INLINE segActWithLength #-}
  segActWithLength :: Int -> Affine1 a -> Sum a -> Sum a
segActWithLength !Int
len Affine1 a
f (Sum !a
x) = a -> Sum a
forall a. a -> Sum a
Sum (a -> Sum a) -> a -> Sum a
forall a b. (a -> b) -> a -> b
$! Int -> Affine1 a -> a -> a
forall a. Num a => Int -> Affine1 a -> a -> a
actWithLength Int
len Affine1 a
f a
x

-- | @since 1.0.0.0
instance (Num a) => SegAct (Affine1 (Sum a)) (Sum a) where
  {-# INLINE segActWithLength #-}
  segActWithLength :: Int -> Affine1 (Sum a) -> Sum a -> Sum a
segActWithLength = Int -> Affine1 (Sum a) -> Sum a -> Sum a
forall a. Num a => Int -> Affine1 a -> a -> a
actWithLength

-- | @since 1.0.0.0
instance (Num a) => SegAct (Dual (Affine1 a)) (Sum a) where
  {-# INLINE segActWithLength #-}
  segActWithLength :: Int -> Dual (Affine1 a) -> Sum a -> Sum a
segActWithLength !Int
len (Dual Affine1 a
f) (Sum !a
x) = a -> Sum a
forall a. a -> Sum a
Sum (a -> Sum a) -> a -> Sum a
forall a b. (a -> b) -> a -> b
$! Int -> Affine1 a -> a -> a
forall a. Num a => Int -> Affine1 a -> a -> a
actWithLength Int
len Affine1 a
f a
x

-- | @since 1.0.0.0
instance (Num a) => SegAct (Dual (Affine1 (Sum a))) (Sum a) where
  {-# INLINE segActWithLength #-}
  segActWithLength :: Int -> Dual (Affine1 (Sum a)) -> Sum a -> Sum a
segActWithLength !Int
len (Dual Affine1 (Sum a)
f) (Sum !a
x) = a -> Sum a
forall a. a -> Sum a
Sum (a -> Sum a) -> a -> Sum a
forall a b. (a -> b) -> a -> b
$! Int -> Affine1 a -> a -> a
forall a. Num a => Int -> Affine1 a -> a -> a
actWithLength Int
len (Affine1 (Sum a) -> Affine1 a
forall a b. Coercible a b => a -> b
coerce Affine1 (Sum a)
f) a
x

-- not works as SegAct for Product, Min, and Max.

-- | @since 1.0.0.0
newtype instance VU.MVector s (Affine1 a) = MV_Affine1 (VU.MVector s (Affine1Repr a))

-- | @since 1.0.0.0
newtype instance VU.Vector (Affine1 a) = V_Affine1 (VU.Vector (Affine1Repr a))

-- | @since 1.0.0.0
deriving instance (VU.Unbox a) => VGM.MVector VUM.MVector (Affine1 a)

-- | @since 1.0.0.0
deriving instance (VU.Unbox a) => VG.Vector VU.Vector (Affine1 a)

-- | @since 1.0.0.0
instance (VU.Unbox a) => VU.Unbox (Affine1 a)