-- |
-- Module:      Data.Poly.Semiring
-- Copyright:   (c) 2019 Andrew Lelechenko
-- Licence:     BSD3
-- Maintainer:  Andrew Lelechenko <andrew.lelechenko@gmail.com>
--
-- Dense polynomials and a 'Semiring'-based interface.
--
-- @since 0.2.0.0

{-# LANGUAGE CPP              #-}
{-# LANGUAGE DataKinds        #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE PatternSynonyms  #-}

module Data.Poly.Semiring
  ( Poly
  , VPoly
  , UPoly
  , unPoly
  , leading
  , toPoly
  , monomial
  , scale
  , pattern X
  , eval
  , subst
  , deriv
  , integral
  , timesRing
#ifdef SupportSparse
  , denseToSparse
  , sparseToDense
#endif
  , dft
  , inverseDft
  , dftMult
  ) where

import Data.Bits
import Data.Euclidean (Field)
import Data.Semiring (Semiring(..))
import qualified Data.Vector.Generic as G

import Data.Poly.Internal.Dense (Poly(..), VPoly, UPoly, leading, timesRing)
import qualified Data.Poly.Internal.Dense as Dense
import Data.Poly.Internal.Dense.Field ()
import Data.Poly.Internal.Dense.DFT
import Data.Poly.Internal.Dense.GcdDomain ()

#ifdef SupportSparse
import qualified Data.Vector.Unboxed.Sized as SU
import qualified Data.Poly.Internal.Multi as Sparse
import qualified Data.Poly.Internal.Convert as Convert
#endif

-- | Make a 'Poly' from a vector of coefficients
-- (first element corresponds to the constant term).
--
-- >>> :set -XOverloadedLists
-- >>> toPoly [1,2,3] :: VPoly Integer
-- 3 * X^2 + 2 * X + 1
-- >>> toPoly [0,0,0] :: UPoly Int
-- 0
--
-- @since 0.2.0.0
toPoly :: (Eq a, Semiring a, G.Vector v a) => v a -> Poly v a
toPoly :: forall a (v :: * -> *).
(Eq a, Semiring a, Vector v a) =>
v a -> Poly v a
toPoly = forall a (v :: * -> *).
(Eq a, Semiring a, Vector v a) =>
v a -> Poly v a
Dense.toPoly'

-- | Create a monomial from a power and a coefficient.
--
-- @since 0.3.0.0
monomial :: (Eq a, Semiring a, G.Vector v a) => Word -> a -> Poly v a
monomial :: forall a (v :: * -> *).
(Eq a, Semiring a, Vector v a) =>
Word -> a -> Poly v a
monomial = forall a (v :: * -> *).
(Eq a, Semiring a, Vector v a) =>
Word -> a -> Poly v a
Dense.monomial'

-- | Multiply a polynomial by a monomial, expressed as a power and a coefficient.
--
-- >>> scale 2 3 (X^2 + 1) :: UPoly Int
-- 3 * X^4 + 0 * X^3 + 3 * X^2 + 0 * X + 0
--
-- @since 0.3.0.0
scale :: (Eq a, Semiring a, G.Vector v a) => Word -> a -> Poly v a -> Poly v a
scale :: forall a (v :: * -> *).
(Eq a, Semiring a, Vector v a) =>
Word -> a -> Poly v a -> Poly v a
scale = forall a (v :: * -> *).
(Eq a, Semiring a, Vector v a) =>
Word -> a -> Poly v a -> Poly v a
Dense.scale'

-- | The polynomial 'X'.
--
-- > X == monomial 1 one
--
-- @since 0.2.0.0
pattern X :: (Eq a, Semiring a, G.Vector v a) => Poly v a
pattern $bX :: forall a (v :: * -> *). (Eq a, Semiring a, Vector v a) => Poly v a
$mX :: forall {r} {a} {v :: * -> *}.
(Eq a, Semiring a, Vector v a) =>
Poly v a -> ((# #) -> r) -> ((# #) -> r) -> r
X = Dense.X'

-- | Evaluate the polynomial at a given point.
--
-- >>> eval (X^2 + 1 :: UPoly Int) 3
-- 10
--
-- @since 0.2.0.0
eval :: (Semiring a, G.Vector v a) => Poly v a -> a -> a
eval :: forall a (v :: * -> *).
(Semiring a, Vector v a) =>
Poly v a -> a -> a
eval = forall a (v :: * -> *).
(Semiring a, Vector v a) =>
Poly v a -> a -> a
Dense.eval'

-- | Substitute another polynomial instead of 'X'.
--
-- >>> subst (X^2 + 1 :: UPoly Int) (X + 1 :: UPoly Int)
-- 1 * X^2 + 2 * X + 2
--
-- @since 0.3.3.0
subst :: (Eq a, Semiring a, G.Vector v a, G.Vector w a) => Poly v a -> Poly w a -> Poly w a
subst :: forall a (v :: * -> *) (w :: * -> *).
(Eq a, Semiring a, Vector v a, Vector w a) =>
Poly v a -> Poly w a -> Poly w a
subst = forall a (v :: * -> *) (w :: * -> *).
(Eq a, Semiring a, Vector v a, Vector w a) =>
Poly v a -> Poly w a -> Poly w a
Dense.subst'

-- | Take the derivative of the polynomial.
--
-- >>> deriv (X^3 + 3 * X) :: UPoly Int
-- 3 * X^2 + 0 * X + 3
--
-- @since 0.2.0.0
deriv :: (Eq a, Semiring a, G.Vector v a) => Poly v a -> Poly v a
deriv :: forall a (v :: * -> *).
(Eq a, Semiring a, Vector v a) =>
Poly v a -> Poly v a
deriv = forall a (v :: * -> *).
(Eq a, Semiring a, Vector v a) =>
Poly v a -> Poly v a
Dense.deriv'

-- | Compute an indefinite integral of the polynomial,
-- setting the constant term to zero.
--
-- >>> integral (3 * X^2 + 3) :: UPoly Double
-- 1.0 * X^3 + 0.0 * X^2 + 3.0 * X + 0.0
--
-- @since 0.3.2.0
integral :: (Eq a, Field a, G.Vector v a) => Poly v a -> Poly v a
integral :: forall a (v :: * -> *).
(Eq a, Field a, Vector v a) =>
Poly v a -> Poly v a
integral = forall a (v :: * -> *).
(Eq a, Field a, Vector v a) =>
Poly v a -> Poly v a
Dense.integral'

-- | Multiplication of polynomials using
-- <https://en.wikipedia.org/wiki/Fast_Fourier_transform discrete Fourier transform>.
-- It could be faster than '(*)' for large polynomials
-- if multiplication of coefficients is particularly expensive.
--
-- @since 0.5.0.0
dftMult
  :: (Eq a, Field a, G.Vector v a)
  => (Int -> a) -- ^ mapping from \( N = 2^n \) to a primitive root \( \sqrt[N]{1} \)
  -> Poly v a
  -> Poly v a
  -> Poly v a
dftMult :: forall a (v :: * -> *).
(Eq a, Field a, Vector v a) =>
(Int -> a) -> Poly v a -> Poly v a -> Poly v a
dftMult Int -> a
getPrimRoot (Poly v a
xs) (Poly v a
ys) =
  forall a (v :: * -> *).
(Eq a, Semiring a, Vector v a) =>
v a -> Poly v a
toPoly forall a b. (a -> b) -> a -> b
$ forall a (v :: * -> *). (Field a, Vector v a) => a -> v a -> v a
inverseDft a
primRoot forall a b. (a -> b) -> a -> b
$ forall (v :: * -> *) a b c.
(Vector v a, Vector v b, Vector v c) =>
(a -> b -> c) -> v a -> v b -> v c
G.zipWith forall a. Semiring a => a -> a -> a
times (forall a (v :: * -> *). (Ring a, Vector v a) => a -> v a -> v a
dft a
primRoot v a
xs') (forall a (v :: * -> *). (Ring a, Vector v a) => a -> v a -> v a
dft a
primRoot v a
ys')
  where
    nextPowerOf2 :: Int -> Int
    nextPowerOf2 :: Int -> Int
nextPowerOf2 Int
0 = Int
1
    nextPowerOf2 Int
1 = Int
1
    nextPowerOf2 Int
x = Int
1 forall a. Bits a => a -> Int -> a
`unsafeShiftL` (forall b. FiniteBits b => b -> Int
finiteBitSize (Int
0 :: Int) forall a. Num a => a -> a -> a
- forall b. FiniteBits b => b -> Int
countLeadingZeros (Int
x forall a. Num a => a -> a -> a
- Int
1))

    padTo :: Int -> v a -> v a
padTo Int
l v a
vs = forall (v :: * -> *) a. Vector v a => Int -> (Int -> a) -> v a
G.generate Int
l (\Int
k -> if Int
k forall a. Ord a => a -> a -> Bool
< forall (v :: * -> *) a. Vector v a => v a -> Int
G.length v a
vs then v a
vs forall (v :: * -> *) a.
(HasCallStack, Vector v a) =>
v a -> Int -> a
G.! Int
k else forall a. Semiring a => a
zero)

    zl :: Int
zl = Int -> Int
nextPowerOf2 (forall (v :: * -> *) a. Vector v a => v a -> Int
G.length v a
xs forall a. Num a => a -> a -> a
+ forall (v :: * -> *) a. Vector v a => v a -> Int
G.length v a
ys)
    xs' :: v a
xs' = forall {v :: * -> *} {a} {v :: * -> *}.
(Vector v a, Vector v a, Semiring a) =>
Int -> v a -> v a
padTo Int
zl v a
xs
    ys' :: v a
ys' = forall {v :: * -> *} {a} {v :: * -> *}.
(Vector v a, Vector v a, Semiring a) =>
Int -> v a -> v a
padTo Int
zl v a
ys
    primRoot :: a
primRoot = Int -> a
getPrimRoot Int
zl
{-# INLINABLE dftMult #-}

#ifdef SupportSparse
-- | Convert from dense to sparse polynomials.
--
-- >>> :set -XFlexibleContexts
-- >>> denseToSparse (1 `Data.Semiring.plus` Data.Poly.X^2) :: Data.Poly.Sparse.UPoly Int
-- 1 * X^2 + 1
--
-- @since 0.5.0.0
denseToSparse :: (Eq a, Semiring a, G.Vector v a, G.Vector v (SU.Vector 1 Word, a)) => Dense.Poly v a -> Sparse.Poly v a
denseToSparse :: forall a (v :: * -> *).
(Eq a, Semiring a, Vector v a, Vector v (Vector 1 Word, a)) =>
Poly v a -> Poly v a
denseToSparse = forall a (v :: * -> *).
(Eq a, Semiring a, Vector v a, Vector v (Vector 1 Word, a)) =>
Poly v a -> Poly v a
Convert.denseToSparse'

-- | Convert from sparse to dense polynomials.
--
-- >>> :set -XFlexibleContexts
-- >>> sparseToDense (1 `Data.Semiring.plus` Data.Poly.Sparse.X^2) :: Data.Poly.UPoly Int
-- 1 * X^2 + 0 * X + 1
--
-- @since 0.5.0.0
sparseToDense :: (Semiring a, G.Vector v a, G.Vector v (SU.Vector 1 Word, a)) => Sparse.Poly v a -> Dense.Poly v a
sparseToDense :: forall a (v :: * -> *).
(Semiring a, Vector v a, Vector v (Vector 1 Word, a)) =>
Poly v a -> Poly v a
sparseToDense = forall a (v :: * -> *).
(Semiring a, Vector v a, Vector v (Vector 1 Word, a)) =>
Poly v a -> Poly v a
Convert.sparseToDense'
#endif