{-# LANGUAGE TypeSynonymInstances #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE MultiParamTypeClasses #-}

module Math.Algebra.Hspray
  ( Spray
  , lone
  , unitSpray
  , constantSpray
  , fromList
  , toList
  , sprayTerms
  , (*^)
  , (.^)
  , (^+^)
  , (^-^)
  , (^*^)
  , (^**^)
  , evalSpray
  , prettySpray
  , composeSpray
  , bombieriSpray
  ) where
import qualified Algebra.Additive              as AlgAdd
import qualified Algebra.Module                as AlgMod
import qualified Algebra.Ring                  as AlgRing
import qualified Data.Foldable                 as DF
import           Data.Function                  ( on )
import           Data.HashMap.Strict            ( HashMap )
import qualified Data.HashMap.Strict           as HM
import           Data.Hashable
import           Data.List                      ( sortBy )
import qualified Data.Sequence                 as S
import           Data.Sequence                  ( (><)
                                                , Seq
                                                , dropWhileR
                                                , (|>)
                                                )
import           Data.Text                      ( Text
                                                , append
                                                , cons
                                                , intercalate
                                                , pack
                                                , snoc
                                                , unpack
                                                )


infixr 7 *^, .^

infixl 6 ^+^, ^-^

infixl 7 ^*^

infixr 8 ^**^


data Powers = Powers
  { Powers -> Seq Int
exponents  :: Seq Int
  , Powers -> Int
nvariables :: Int
  }
  deriving Int -> Powers -> ShowS
[Powers] -> ShowS
Powers -> String
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
showList :: [Powers] -> ShowS
$cshowList :: [Powers] -> ShowS
show :: Powers -> String
$cshow :: Powers -> String
showsPrec :: Int -> Powers -> ShowS
$cshowsPrec :: Int -> Powers -> ShowS
Show

growSequence :: Seq Int -> Int -> Int -> Seq Int
growSequence :: Seq Int -> Int -> Int -> Seq Int
growSequence Seq Int
s Int
m Int
n = Seq Int
s forall a. Seq a -> Seq a -> Seq a
>< Seq Int
t where t :: Seq Int
t = forall a. Int -> a -> Seq a
S.replicate (Int
n forall a. Num a => a -> a -> a
- Int
m) Int
0

harmonize :: (Powers, Powers) -> (Powers, Powers)
harmonize :: (Powers, Powers) -> (Powers, Powers)
harmonize (Powers
pows1, Powers
pows2) = (Seq Int -> Int -> Powers
Powers Seq Int
e1' Int
n, Seq Int -> Int -> Powers
Powers Seq Int
e2' Int
n)
 where
  e1 :: Seq Int
e1            = Powers -> Seq Int
exponents Powers
pows1
  e2 :: Seq Int
e2            = Powers -> Seq Int
exponents Powers
pows2
  n1 :: Int
n1            = Powers -> Int
nvariables Powers
pows1
  n2 :: Int
n2            = Powers -> Int
nvariables Powers
pows2
  (Seq Int
e1', Seq Int
e2', Int
n) = if Int
n1 forall a. Ord a => a -> a -> Bool
< Int
n2
    then (Seq Int -> Int -> Int -> Seq Int
growSequence Seq Int
e1 Int
n1 Int
n2, Seq Int
e2, Int
n2)
    else (Seq Int
e1, Seq Int -> Int -> Int -> Seq Int
growSequence Seq Int
e2 Int
n2 Int
n1, Int
n1)

instance Eq Powers where
  Powers
pows1 == :: Powers -> Powers -> Bool
== Powers
pows2 = (Powers -> Seq Int
exponents Powers
pows1') forall a. Eq a => a -> a -> Bool
== (Powers -> Seq Int
exponents Powers
pows2')
    where (Powers
pows1', Powers
pows2') = (Powers, Powers) -> (Powers, Powers)
harmonize (Powers
pows1, Powers
pows2)

instance Hashable Powers where
  hashWithSalt :: Int -> Powers -> Int
hashWithSalt Int
k Powers
pows = forall a. Hashable a => Int -> a -> Int
hashWithSalt Int
k (Powers -> Seq Int
exponents Powers
pows, Powers -> Int
nvariables Powers
pows)

type Spray a = HashMap Powers a

instance (AlgAdd.C a, Eq a) => AlgAdd.C (Spray a) where
  Spray a
p + :: Spray a -> Spray a -> Spray a
+ Spray a
q = forall a. (C a, Eq a) => Spray a -> Spray a -> Spray a
addSprays Spray a
p Spray a
q
  zero :: Spray a
zero   = forall k v. HashMap k v
HM.empty
  negate :: Spray a -> Spray a
negate = forall a. C a => Spray a -> Spray a
negateSpray

instance (AlgRing.C a, Eq a) => AlgMod.C a (Spray a) where
  a
lambda *> :: a -> Spray a -> Spray a
*> Spray a
p = forall a. (C a, Eq a) => a -> Spray a -> Spray a
scaleSpray a
lambda Spray a
p

instance (AlgRing.C a, Eq a) => AlgRing.C (Spray a) where
  Spray a
p * :: Spray a -> Spray a -> Spray a
* Spray a
q = forall a. (C a, Eq a) => Spray a -> Spray a -> Spray a
multSprays Spray a
p Spray a
q
  one :: Spray a
one = forall a. C a => Int -> Spray a
lone Int
0

-- | Addition of two sprays
(^+^) :: (AlgAdd.C a, Eq a) => Spray a -> Spray a -> Spray a
^+^ :: forall a. (C a, Eq a) => Spray a -> Spray a -> Spray a
(^+^) Spray a
p Spray a
q = Spray a
p forall a. C a => a -> a -> a
AlgAdd.+ Spray a
q

-- | Substraction of two sprays
(^-^) :: (AlgAdd.C a, Eq a) => Spray a -> Spray a -> Spray a
^-^ :: forall a. (C a, Eq a) => Spray a -> Spray a -> Spray a
(^-^) Spray a
p Spray a
q = Spray a
p forall a. C a => a -> a -> a
AlgAdd.- Spray a
q

-- | Multiply two sprays
(^*^) :: (AlgRing.C a, Eq a) => Spray a -> Spray a -> Spray a
^*^ :: forall a. (C a, Eq a) => Spray a -> Spray a -> Spray a
(^*^) Spray a
p Spray a
q = Spray a
p forall a. C a => a -> a -> a
AlgRing.* Spray a
q

-- | Power of a spray
(^**^) :: (AlgRing.C a, Eq a) => Spray a -> Int -> Spray a
^**^ :: forall a. (C a, Eq a) => Spray a -> Int -> Spray a
(^**^) Spray a
p Int
n = forall a. C a => [a] -> a
AlgRing.product (forall a. Int -> a -> [a]
replicate Int
n Spray a
p)

-- | Scale spray by a scalar
(*^) :: (AlgRing.C a, Eq a) => a -> Spray a -> Spray a
*^ :: forall a. (C a, Eq a) => a -> Spray a -> Spray a
(*^) a
lambda Spray a
pol = a
lambda forall a v. C a v => a -> v -> v
AlgMod.*> Spray a
pol

-- | Scale spray by an integer
(.^) :: (AlgAdd.C a, Eq a) => Int -> Spray a -> Spray a
.^ :: forall a. (C a, Eq a) => Int -> Spray a -> Spray a
(.^) Int
k Spray a
pol = if Int
k forall a. Ord a => a -> a -> Bool
>= Int
0 
  then forall a. C a => [a] -> a
AlgAdd.sum (forall a. Int -> a -> [a]
replicate Int
k Spray a
pol)
  else forall a. C a => a -> a
AlgAdd.negate forall a b. (a -> b) -> a -> b
$ forall a. C a => [a] -> a
AlgAdd.sum (forall a. Int -> a -> [a]
replicate (-Int
k) Spray a
pol)

simplifyPowers :: Powers -> Powers
simplifyPowers :: Powers -> Powers
simplifyPowers Powers
pows = Seq Int -> Int -> Powers
Powers Seq Int
s (forall a. Seq a -> Int
S.length Seq Int
s)
  where s :: Seq Int
s = forall a. (a -> Bool) -> Seq a -> Seq a
dropWhileR (forall a. Eq a => a -> a -> Bool
== Int
0) (Powers -> Seq Int
exponents Powers
pows)

simplifySpray :: Spray a -> Spray a
simplifySpray :: forall a. Spray a -> Spray a
simplifySpray Spray a
p = forall k2 k1 v.
(Eq k2, Hashable k2) =>
(k1 -> k2) -> HashMap k1 v -> HashMap k2 v
HM.mapKeys Powers -> Powers
simplifyPowers Spray a
p

cleanSpray :: (AlgAdd.C a, Eq a) => Spray a -> Spray a
cleanSpray :: forall a. (C a, Eq a) => Spray a -> Spray a
cleanSpray Spray a
p = forall v k. (v -> Bool) -> HashMap k v -> HashMap k v
HM.filter (forall a. Eq a => a -> a -> Bool
/= forall a. C a => a
AlgAdd.zero) (forall a. Spray a -> Spray a
simplifySpray Spray a
p)

addSprays :: (AlgAdd.C a, Eq a) => Spray a -> Spray a -> Spray a
addSprays :: forall a. (C a, Eq a) => Spray a -> Spray a -> Spray a
addSprays Spray a
p Spray a
q = forall a. (C a, Eq a) => Spray a -> Spray a
cleanSpray forall a b. (a -> b) -> a -> b
$ forall a k v. (a -> k -> v -> a) -> a -> HashMap k v -> a
HM.foldlWithKey' forall {k} {v}.
(Hashable k, C v) =>
HashMap k v -> k -> v -> HashMap k v
f Spray a
p Spray a
q
  where f :: HashMap k v -> k -> v -> HashMap k v
f HashMap k v
s k
powers v
coef = forall k v.
(Eq k, Hashable k) =>
(v -> v -> v) -> k -> v -> HashMap k v -> HashMap k v
HM.insertWith forall a. C a => a -> a -> a
(AlgAdd.+) k
powers v
coef HashMap k v
s

negateSpray :: AlgAdd.C a => Spray a -> Spray a
negateSpray :: forall a. C a => Spray a -> Spray a
negateSpray = forall v1 v2 k. (v1 -> v2) -> HashMap k v1 -> HashMap k v2
HM.map forall a. C a => a -> a
AlgAdd.negate

scaleSpray :: (AlgRing.C a, Eq a) => a -> Spray a -> Spray a
scaleSpray :: forall a. (C a, Eq a) => a -> Spray a -> Spray a
scaleSpray a
lambda Spray a
p = forall a. (C a, Eq a) => Spray a -> Spray a
cleanSpray forall a b. (a -> b) -> a -> b
$ forall v1 v2 k. (v1 -> v2) -> HashMap k v1 -> HashMap k v2
HM.map (a
lambda forall a. C a => a -> a -> a
AlgRing.*) Spray a
p

multMonomial :: AlgRing.C a => (Powers, a) -> (Powers, a) -> (Powers, a)
multMonomial :: forall a. C a => (Powers, a) -> (Powers, a) -> (Powers, a)
multMonomial (Powers
pows1, a
coef1) (Powers
pows2, a
coef2) = (Powers
pows, a
coef1 forall a. C a => a -> a -> a
AlgRing.* a
coef2)
 where
  (Powers
pows1', Powers
pows2') = (Powers, Powers) -> (Powers, Powers)
harmonize (Powers
pows1, Powers
pows2)
  expts :: Seq Int
expts            = forall a b c. (a -> b -> c) -> Seq a -> Seq b -> Seq c
S.zipWith forall a. Num a => a -> a -> a
(+) (Powers -> Seq Int
exponents Powers
pows1') (Powers -> Seq Int
exponents Powers
pows2')
  pows :: Powers
pows             = Seq Int -> Int -> Powers
Powers Seq Int
expts (Powers -> Int
nvariables Powers
pows1')

multSprays :: (AlgRing.C a, Eq a) => Spray a -> Spray a -> Spray a
multSprays :: forall a. (C a, Eq a) => Spray a -> Spray a -> Spray a
multSprays Spray a
p Spray a
q = forall a. (C a, Eq a) => Spray a -> Spray a
cleanSpray forall a b. (a -> b) -> a -> b
$ forall k v.
(Eq k, Hashable k) =>
(v -> v -> v) -> [(k, v)] -> HashMap k v
HM.fromListWith forall a. C a => a -> a -> a
(AlgAdd.+) [(Powers, a)]
prods
 where
  p' :: [(Powers, a)]
p'    = forall k v. HashMap k v -> [(k, v)]
HM.toList Spray a
p
  q' :: [(Powers, a)]
q'    = forall k v. HashMap k v -> [(k, v)]
HM.toList Spray a
q
  prods :: [(Powers, a)]
prods = [ forall a. C a => (Powers, a) -> (Powers, a) -> (Powers, a)
multMonomial (Powers, a)
mp (Powers, a)
mq | (Powers, a)
mp <- [(Powers, a)]
p', (Powers, a)
mq <- [(Powers, a)]
q' ]

-- | Spray corresponding to polynomial x_n
lone :: AlgRing.C a => Int -> Spray a
lone :: forall a. C a => Int -> Spray a
lone Int
n = forall k v. Hashable k => k -> v -> HashMap k v
HM.singleton Powers
pows forall a. C a => a
AlgRing.one
 where
  pows :: Powers
pows = if Int
n forall a. Eq a => a -> a -> Bool
== Int
0
    then Seq Int -> Int -> Powers
Powers forall a. Seq a
S.empty Int
0
    else Seq Int -> Int -> Powers
Powers (forall a. Int -> a -> Seq a
S.replicate (Int
n forall a. Num a => a -> a -> a
- Int
1) forall a. C a => a
AlgAdd.zero forall a. Seq a -> a -> Seq a
|> forall a. C a => a
AlgRing.one) Int
n

-- | Unit spray
unitSpray :: AlgRing.C a => Spray a
unitSpray :: forall a. C a => Spray a
unitSpray = forall a. C a => Int -> Spray a
lone Int
0

-- | Constant spray
constantSpray :: (AlgRing.C a, Eq a) => a -> Spray a
constantSpray :: forall a. (C a, Eq a) => a -> Spray a
constantSpray a
c = a
c forall a. (C a, Eq a) => a -> Spray a -> Spray a
*^ (forall a. C a => Int -> Spray a
lone Int
0)

evalMonomial :: AlgRing.C a => [a] -> (Powers, a) -> a
evalMonomial :: forall a. C a => [a] -> (Powers, a) -> a
evalMonomial [a]
xyz (Powers
powers, a
coeff) = a
coeff
  forall a. C a => a -> a -> a
AlgRing.* forall a. C a => [a] -> a
AlgRing.product (forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith forall a. C a => a -> Integer -> a
(AlgRing.^) [a]
xyz [Integer]
pows)
  where pows :: [Integer]
pows = forall (t :: * -> *) a. Foldable t => t a -> [a]
DF.toList (forall a b. (Integral a, Num b) => a -> b
fromIntegral forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Powers -> Seq Int
exponents Powers
powers)

-- | Evaluate a spray
evalSpray :: AlgRing.C a => Spray a -> [a] -> a
evalSpray :: forall a. C a => Spray a -> [a] -> a
evalSpray Spray a
p [a]
xyz = forall a. C a => [a] -> a
AlgAdd.sum forall a b. (a -> b) -> a -> b
$ forall a b. (a -> b) -> [a] -> [b]
map (forall a. C a => [a] -> (Powers, a) -> a
evalMonomial [a]
xyz) (forall k v. HashMap k v -> [(k, v)]
HM.toList Spray a
p)

identify :: (AlgRing.C a, Eq a) => Spray a -> Spray (Spray a)
identify :: forall a. (C a, Eq a) => Spray a -> Spray (Spray a)
identify Spray a
p = forall v1 v2 k. (v1 -> v2) -> HashMap k v1 -> HashMap k v2
HM.map forall a. (C a, Eq a) => a -> Spray a
constantSpray Spray a
p

-- | Compose a spray with a change of variables
composeSpray :: (AlgRing.C a, Eq a) => Spray a -> [Spray a] -> Spray a
composeSpray :: forall a. (C a, Eq a) => Spray a -> [Spray a] -> Spray a
composeSpray Spray a
p [Spray a]
newvars = forall a. C a => Spray a -> [a] -> a
evalSpray (forall a. (C a, Eq a) => Spray a -> Spray (Spray a)
identify Spray a
p) [Spray a]
newvars

-- | Create a spray from list of terms
fromList :: (AlgRing.C a, Eq a) => [([Int], a)] -> Spray a
fromList :: forall a. (C a, Eq a) => [([Int], a)] -> Spray a
fromList [([Int], a)]
x = forall a. (C a, Eq a) => Spray a -> Spray a
cleanSpray forall a b. (a -> b) -> a -> b
$ forall k v. (Eq k, Hashable k) => [(k, v)] -> HashMap k v
HM.fromList forall a b. (a -> b) -> a -> b
$ forall a b. (a -> b) -> [a] -> [b]
map
  (\([Int]
expts, a
coef) -> (Seq Int -> Int -> Powers
Powers (forall a. [a] -> Seq a
S.fromList [Int]
expts) (forall (t :: * -> *) a. Foldable t => t a -> Int
length [Int]
expts), a
coef)) [([Int], a)]
x

prettyPowers :: String -> [Int] -> Text
prettyPowers :: String -> [Int] -> Text
prettyPowers String
var [Int]
pows = Text -> Text -> Text
append (String -> Text
pack String
x) (Char -> Text -> Text
cons Char
'(' forall a b. (a -> b) -> a -> b
$ Text -> Char -> Text
snoc Text
string Char
')')
 where
  x :: String
x      = String
" " forall a. [a] -> [a] -> [a]
++ String
var forall a. [a] -> [a] -> [a]
++ String
"^"
  string :: Text
string = Text -> [Text] -> Text
intercalate (String -> Text
pack String
", ") (forall a b. (a -> b) -> [a] -> [b]
map (String -> Text
pack forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a. Show a => a -> String
show) [Int]
pows)

-- | Pretty form of a spray
prettySpray :: (a -> String) -> String -> Spray a -> String
prettySpray :: forall a. (a -> String) -> String -> Spray a -> String
prettySpray a -> String
prettyCoef String
var Spray a
p = Text -> String
unpack forall a b. (a -> b) -> a -> b
$ Text -> [Text] -> Text
intercalate (String -> Text
pack String
" + ") [Text]
stringTerms
 where
  stringTerms :: [Text]
stringTerms = forall a b. (a -> b) -> [a] -> [b]
map (Powers, a) -> Text
stringTerm (forall a. (a -> a -> Ordering) -> [a] -> [a]
sortBy (forall a. Ord a => a -> a -> Ordering
compare forall b c a. (b -> b -> c) -> (a -> b) -> a -> a -> c
`on` forall {b}. (Powers, b) -> Seq Int
fexpts) (forall k v. HashMap k v -> [(k, v)]
HM.toList Spray a
p))
  fexpts :: (Powers, b) -> Seq Int
fexpts (Powers, b)
term = Powers -> Seq Int
exponents forall a b. (a -> b) -> a -> b
$ forall a b. (a, b) -> a
fst (Powers, b)
term
  stringTerm :: (Powers, a) -> Text
stringTerm (Powers, a)
term = Text -> Text -> Text
append
    (Text -> Char -> Text
snoc (Text -> Char -> Text
snoc (Char -> Text -> Text
cons Char
'(' forall a b. (a -> b) -> a -> b
$ Text -> Char -> Text
snoc Text
stringCoef Char
')') Char
' ') Char
'*')
    (String -> [Int] -> Text
prettyPowers String
var [Int]
pows)
   where
    pows :: [Int]
pows       = forall (t :: * -> *) a. Foldable t => t a -> [a]
DF.toList forall a b. (a -> b) -> a -> b
$ Powers -> Seq Int
exponents (forall a b. (a, b) -> a
fst (Powers, a)
term)
    stringCoef :: Text
stringCoef = String -> Text
pack forall a b. (a -> b) -> a -> b
$ a -> String
prettyCoef (forall a b. (a, b) -> b
snd (Powers, a)
term)

-- | Terms of a spray
sprayTerms :: Spray a -> HashMap (Seq Int) a
sprayTerms :: forall a. Spray a -> HashMap (Seq Int) a
sprayTerms Spray a
p = forall k2 k1 v.
(Eq k2, Hashable k2) =>
(k1 -> k2) -> HashMap k1 v -> HashMap k2 v
HM.mapKeys Powers -> Seq Int
exponents Spray a
p

-- | Spray as list
toList :: Spray a -> [([Int], a)]
toList :: forall a. Spray a -> [([Int], a)]
toList Spray a
p = forall k v. HashMap k v -> [(k, v)]
HM.toList forall a b. (a -> b) -> a -> b
$ forall k2 k1 v.
(Eq k2, Hashable k2) =>
(k1 -> k2) -> HashMap k1 v -> HashMap k2 v
HM.mapKeys (forall (t :: * -> *) a. Foldable t => t a -> [a]
DF.toList forall b c a. (b -> c) -> (a -> b) -> a -> c
. Powers -> Seq Int
exponents) Spray a
p

-- | Bombieri spray
bombieriSpray :: AlgAdd.C a => Spray a -> Spray a
bombieriSpray :: forall a. C a => Spray a -> Spray a
bombieriSpray Spray a
p = forall k v1 v2. (k -> v1 -> v2) -> HashMap k v1 -> HashMap k v2
HM.mapWithKey forall {a}. C a => Powers -> a -> a
f Spray a
p 
 where
  f :: Powers -> a -> a
f Powers
pows a
coef     = forall {a}. C a => Int -> a -> a
times (forall {a}. (Num a, Enum a, Eq a) => Seq a -> a
pfactorial forall a b. (a -> b) -> a -> b
$ Powers -> Seq Int
exponents Powers
pows) a
coef
  pfactorial :: Seq a -> a
pfactorial Seq a
pows = forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
product forall a b. (a -> b) -> a -> b
$ forall (t :: * -> *) a. Foldable t => t a -> [a]
DF.toList forall a b. (a -> b) -> a -> b
$ forall {a}. (Num a, Enum a) => a -> a
factorial forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (forall a. (a -> Bool) -> Seq a -> Seq a
S.filter (forall a. Eq a => a -> a -> Bool
/= a
0) Seq a
pows)
  factorial :: a -> a
factorial a
n     = forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
product [a
1 .. a
n]
  times :: Int -> a -> a
times Int
k a
x       = forall a. C a => [a] -> a
AlgAdd.sum (forall a. Int -> a -> [a]
replicate Int
k a
x)