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