-- |
-- Module:      Data.Poly.Internal.Convert
-- Copyright:   (c) 2020 Andrew Lelechenko
-- Licence:     BSD3
-- Maintainer:  Andrew Lelechenko <andrew.lelechenko@gmail.com>
--
-- Conversions between polynomials.
--

{-# LANGUAGE DataKinds        #-}
{-# LANGUAGE FlexibleContexts #-}

module Data.Poly.Internal.Convert
  ( denseToSparse
  , denseToSparse'
  , sparseToDense
  , sparseToDense'
  ) where

import Control.Monad.ST
import Data.Semiring (Semiring(..))
import qualified Data.Vector.Generic as G
import qualified Data.Vector.Generic.Mutable as MG
import qualified Data.Vector.Unboxed.Sized as SU

import qualified Data.Poly.Internal.Dense as Dense
import qualified Data.Poly.Internal.Multi as Sparse

-- | Convert from dense to sparse polynomials.
--
-- >>> :set -XFlexibleContexts
-- >>> denseToSparse (1 + Data.Poly.X^2) :: Data.Poly.Sparse.UPoly Int
-- 1 * X^2 + 1
--
-- @since 0.5.0.0
denseToSparse
  :: (Eq a, Num 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, Num a, Vector v a, Vector v (Vector 1 Word, a)) =>
Poly v a -> Poly v a
denseToSparse = forall a (v :: * -> *).
(Eq a, Vector v a, Vector v (Vector 1 Word, a)) =>
a -> Poly v a -> Poly v a
denseToSparseInternal a
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, Vector v a, Vector v (Vector 1 Word, a)) =>
a -> Poly v a -> Poly v a
denseToSparseInternal forall a. Semiring a => a
zero

denseToSparseInternal
  :: (Eq a, G.Vector v a, G.Vector v (SU.Vector 1 Word, a))
  => a
  -> Dense.Poly v a
  -> Sparse.Poly v a
denseToSparseInternal :: forall a (v :: * -> *).
(Eq a, Vector v a, Vector v (Vector 1 Word, a)) =>
a -> Poly v a -> Poly v a
denseToSparseInternal a
z = forall (v :: * -> *) (n :: Nat) a.
v (Vector n Word, a) -> MultiPoly v n a
Sparse.MultiPoly forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (v :: * -> *) a b.
(Vector v a, Vector v b) =>
(Int -> a -> Maybe b) -> v a -> v b
G.imapMaybe (\Int
i a
c -> if a
c forall a. Eq a => a -> a -> Bool
== a
z then forall a. Maybe a
Nothing else forall a. a -> Maybe a
Just (forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
i, a
c)) forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (v :: * -> *) a. Poly v a -> v a
Dense.unPoly

-- | Convert from sparse to dense polynomials.
--
-- >>> :set -XFlexibleContexts
-- >>> sparseToDense (1 + Data.Poly.Sparse.X^2) :: Data.Poly.UPoly Int
-- 1 * X^2 + 0 * X + 1
--
-- @since 0.5.0.0
sparseToDense
  :: (Num 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 :: * -> *).
(Num a, Vector v a, Vector v (Vector 1 Word, a)) =>
Poly v a -> Poly v a
sparseToDense = forall (v :: * -> *) a.
(Vector v a, Vector v (Vector 1 Word, a)) =>
a -> Poly v a -> Poly v a
sparseToDenseInternal a
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 (v :: * -> *) a.
(Vector v a, Vector v (Vector 1 Word, a)) =>
a -> Poly v a -> Poly v a
sparseToDenseInternal forall a. Semiring a => a
zero

sparseToDenseInternal
  :: (G.Vector v a, G.Vector v (SU.Vector 1 Word, a))
  => a
  -> Sparse.Poly v a
  -> Dense.Poly v a
sparseToDenseInternal :: forall (v :: * -> *) a.
(Vector v a, Vector v (Vector 1 Word, a)) =>
a -> Poly v a -> Poly v a
sparseToDenseInternal a
z (Sparse.MultiPoly v (Vector 1 Word, a)
xs)
  | forall (v :: * -> *) a. Vector v a => v a -> Bool
G.null v (Vector 1 Word, a)
xs = forall (v :: * -> *) a. v a -> Poly v a
Dense.Poly forall (v :: * -> *) a. Vector v a => v a
G.empty
  | Bool
otherwise = forall a. (forall s. ST s a) -> a
runST forall a b. (a -> b) -> a -> b
$ do
  let len :: Int
len = forall a b. (Integral a, Num b) => a -> b
fromIntegral (forall (n :: Nat) a. Unbox a => Vector (1 + n) a -> a
SU.head (forall a b. (a, b) -> a
fst (forall (v :: * -> *) a. Vector v a => v a -> a
G.unsafeLast v (Vector 1 Word, a)
xs)) forall a. Num a => a -> a -> a
+ Word
1)
  Mutable v (PrimState (ST s)) a
ys <- forall (m :: * -> *) (v :: * -> * -> *) a.
(PrimMonad m, MVector v a) =>
Int -> m (v (PrimState m) a)
MG.unsafeNew Int
len
  forall (m :: * -> *) (v :: * -> * -> *) a.
(PrimMonad m, MVector v a) =>
v (PrimState m) a -> a -> m ()
MG.set Mutable v (PrimState (ST s)) a
ys a
z
  let go :: Int -> Int -> ST s ()
go Int
xi Int
yi
        | Int
xi forall a. Ord a => a -> a -> Bool
>= forall (v :: * -> *) a. Vector v a => v a -> Int
G.length v (Vector 1 Word, a)
xs = forall (f :: * -> *) a. Applicative f => a -> f a
pure ()
        | (Vector 1 Word
yi', a
c) <- forall (v :: * -> *) a. Vector v a => v a -> Int -> a
G.unsafeIndex v (Vector 1 Word, a)
xs Int
xi
        , Int
yi forall a. Eq a => a -> a -> Bool
== forall a b. (Integral a, Num b) => a -> b
fromIntegral (forall (n :: Nat) a. Unbox a => Vector (1 + n) a -> a
SU.head Vector 1 Word
yi')
        = forall (m :: * -> *) (v :: * -> * -> *) a.
(PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> a -> m ()
MG.unsafeWrite Mutable v (PrimState (ST s)) a
ys Int
yi a
c forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> Int -> Int -> ST s ()
go (Int
xi forall a. Num a => a -> a -> a
+ Int
1) (Int
yi forall a. Num a => a -> a -> a
+ Int
1)
        | Bool
otherwise = Int -> Int -> ST s ()
go Int
xi (Int
yi forall a. Num a => a -> a -> a
+ Int
1)
  Int -> Int -> ST s ()
go Int
0 Int
0
  forall (v :: * -> *) a. v a -> Poly v a
Dense.Poly forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall (m :: * -> *) (v :: * -> *) a.
(PrimMonad m, Vector v a) =>
Mutable v (PrimState m) a -> m (v a)
G.unsafeFreeze Mutable v (PrimState (ST s)) a
ys