{-# LANGUAGE BangPatterns #-}
{-# OPTIONS_GHC -fno-warn-name-shadowing #-}
{-# OPTIONS_HADDOCK not-home #-}
-----------------------------------------------------------------------------
-- |
-- Copyright   :  (c) Edward Kmett 2010-2021
-- License     :  BSD3
-- Maintainer  :  ekmett@gmail.com
-- Stability   :  experimental
-- Portability :  GHC only
--
-- common guts for Sparse.Double and Sparse mode
--
-- Handle with care.
-----------------------------------------------------------------------------
module Numeric.AD.Internal.Sparse.Common
  ( Monomial(..)
  , emptyMonomial
  , addToMonomial
  , indices
  , skeleton
  , terms
  ) where

import Data.IntMap (IntMap, toAscList, insertWith)
import qualified Data.IntMap as IntMap
import Data.Traversable

newtype Monomial = Monomial (IntMap Int)

emptyMonomial :: Monomial
emptyMonomial :: Monomial
emptyMonomial = IntMap Int -> Monomial
Monomial forall a. IntMap a
IntMap.empty
{-# INLINE emptyMonomial #-}

addToMonomial :: Int -> Monomial -> Monomial
addToMonomial :: Int -> Monomial -> Monomial
addToMonomial Int
k (Monomial IntMap Int
m) = IntMap Int -> Monomial
Monomial (forall a. (a -> a -> a) -> Int -> a -> IntMap a -> IntMap a
insertWith forall a. Num a => a -> a -> a
(+) Int
k Int
1 IntMap Int
m)
{-# INLINE addToMonomial #-}

indices :: Monomial -> [Int]
indices :: Monomial -> [Int]
indices (Monomial IntMap Int
as) = forall a b c. (a -> b -> c) -> (a, b) -> c
uncurry (forall a b c. (a -> b -> c) -> b -> a -> c
flip forall a. Int -> a -> [a]
replicate) forall (t :: * -> *) a b. Foldable t => (a -> [b]) -> t a -> [b]
`concatMap` forall a. IntMap a -> [(Int, a)]
toAscList IntMap Int
as
{-# INLINE indices #-}

skeleton :: Traversable f => f a -> f Int
skeleton :: forall (f :: * -> *) a. Traversable f => f a -> f Int
skeleton = forall a b. (a, b) -> b
snd forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (t :: * -> *) s a b.
Traversable t =>
(s -> a -> (s, b)) -> s -> t a -> (s, t b)
mapAccumL (\ !Int
n a
_ -> (Int
n forall a. Num a => a -> a -> a
+ Int
1, Int
n)) Int
0
{-# INLINE skeleton #-}

terms :: Monomial -> [(Integer,Monomial,Monomial)]
terms :: Monomial -> [(Integer, Monomial, Monomial)]
terms (Monomial IntMap Int
m) = [(Int, Int)] -> [(Integer, Monomial, Monomial)]
t (forall a. IntMap a -> [(Int, a)]
toAscList IntMap Int
m) where
  t :: [(Int, Int)] -> [(Integer, Monomial, Monomial)]
t [] = [(Integer
1,Monomial
emptyMonomial,Monomial
emptyMonomial)]
  t ((Int
k,Int
a):[(Int, Int)]
ts) = forall (t :: * -> *) a b. Foldable t => (a -> [b]) -> t a -> [b]
concatMap (forall {a}.
Num a =>
[(a, Monomial, Monomial)] -> (a, Int) -> [(a, Monomial, Monomial)]
f ([(Int, Int)] -> [(Integer, Monomial, Monomial)]
t [(Int, Int)]
ts)) (forall a b. [a] -> [b] -> [(a, b)]
zip ([[Integer]]
binsforall a. [a] -> Int -> a
!!Int
a) [Int
0..Int
a]) where
    f :: [(a, Monomial, Monomial)] -> (a, Int) -> [(a, Monomial, Monomial)]
f [(a, Monomial, Monomial)]
ps (a
b,Int
i) = forall a b. (a -> b) -> [a] -> [b]
map (\(a
w,Monomial IntMap Int
mf,Monomial IntMap Int
mg) -> (a
wforall a. Num a => a -> a -> a
*a
b,IntMap Int -> Monomial
Monomial (forall a. Int -> a -> IntMap a -> IntMap a
IntMap.insert Int
k Int
i IntMap Int
mf), IntMap Int -> Monomial
Monomial (forall a. Int -> a -> IntMap a -> IntMap a
IntMap.insert Int
k (Int
aforall a. Num a => a -> a -> a
-Int
i) IntMap Int
mg))) [(a, Monomial, Monomial)]
ps
  bins :: [[Integer]]
bins = forall a. (a -> a) -> a -> [a]
iterate forall {a}. Num a => [a] -> [a]
next [Integer
1]
  next :: [a] -> [a]
next xs :: [a]
xs@(a
_:[a]
ts) = a
1 forall a. a -> [a] -> [a]
: forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith forall a. Num a => a -> a -> a
(+) [a]
xs [a]
ts forall a. [a] -> [a] -> [a]
++ [a
1]
  next [] = forall a. HasCallStack => [Char] -> a
error [Char]
"impossible"