{-# LANGUAGE NoRebindableSyntax #-}

-- | Polynomial normal forms for 'Ring' terms.
--
-- The free ring on generators @a@ is the /noncommutative/ polynomial
-- ring ℤ⟨a⟩: integer combinations of words in the generators.
-- 'toWord' is its normal form — keys are words, order preserved.
--
-- The free /commutative/ ring is ℤ[a], reached from ℤ⟨a⟩ by
-- 'abelianize' (sort each word into a monomial bag).  'toPolynomial'
-- is the composite:
--
-- > toPolynomial = abelianize . toWord
--
-- The factoring keeps the bookkeeping honest: 'Ring' the class does
-- not require commutativity, so the sorted-key 'Polynomial' alone is
-- a /sound but incomplete/ decision procedure for free-'Ring'
-- equality (it identifies @xy@ with @yx@).  'toWord' is sound and
-- complete for ℤ⟨a⟩; 'toPolynomial' is sound and complete for ℤ[a].
module NumHask.Free.Polynomial
  ( -- * Commutative normal form — ℤ[a]
    Polynomial (..),
    toPolynomial,
    fromPolynomial,
    evalPolynomial,

    -- * Noncommutative normal form — ℤ⟨a⟩
    NCPolynomial (..),
    toWord,
    abelianize,
    evalNCPolynomial,
  )
where

import Data.List (sort)
import Data.Map (Map)
import Data.Map qualified as M
import NumHask.Algebra.Additive qualified as NHA
import NumHask.Algebra.Multiplicative qualified as NHM
import NumHask.Algebra.Ring qualified as NHR
import NumHask.Data.Integral qualified as NHI
import NumHask.Free.Ring
import Prelude (Eq, Integer, Ord, Show, foldr, (*), (+), (.), (/=))
import Prelude qualified as P

-- $setup
-- >>> import NumHask.Free.Ring
-- >>> import NumHask.Free.Polynomial
-- >>> import Prelude (fromInteger)
-- >>> import Prelude hiding (negate)

-- | Multivariate polynomial in commutative normal form — ℤ[a].
--
-- Keys are sorted monomials (bags of generators); the empty list is
-- the constant term.  Values are integer coefficients.
--
-- >>> let p = plus (times (embed "x") (embed "y")) (negate (embed "z"))
-- >>> toPolynomial p
-- Polynomial {unPolynomial = fromList [(["x","y"],1),(["z"],-1)]}
newtype Polynomial a = Polynomial {forall a. Polynomial a -> Map [a] Integer
unPolynomial :: Map [a] Integer}
  deriving (Polynomial a -> Polynomial a -> Bool
(Polynomial a -> Polynomial a -> Bool)
-> (Polynomial a -> Polynomial a -> Bool) -> Eq (Polynomial a)
forall a. Eq a => Polynomial a -> Polynomial a -> Bool
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
$c== :: forall a. Eq a => Polynomial a -> Polynomial a -> Bool
== :: Polynomial a -> Polynomial a -> Bool
$c/= :: forall a. Eq a => Polynomial a -> Polynomial a -> Bool
/= :: Polynomial a -> Polynomial a -> Bool
Eq, Int -> Polynomial a -> ShowS
[Polynomial a] -> ShowS
Polynomial a -> String
(Int -> Polynomial a -> ShowS)
-> (Polynomial a -> String)
-> ([Polynomial a] -> ShowS)
-> Show (Polynomial a)
forall a. Show a => Int -> Polynomial a -> ShowS
forall a. Show a => [Polynomial a] -> ShowS
forall a. Show a => Polynomial a -> String
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
$cshowsPrec :: forall a. Show a => Int -> Polynomial a -> ShowS
showsPrec :: Int -> Polynomial a -> ShowS
$cshow :: forall a. Show a => Polynomial a -> String
show :: Polynomial a -> String
$cshowList :: forall a. Show a => [Polynomial a] -> ShowS
showList :: [Polynomial a] -> ShowS
Show)

-- | Noncommutative polynomial in normal form — ℤ⟨a⟩, the monoid ring
-- of the free monoid on the generators.
--
-- Keys are /words/: generator order is preserved, so @xy@ and @yx@
-- are distinct keys.  This is the faithful normal form for the free
-- 'Ring'.  ('Ord' is for the 'Map', not for commutativity.)
newtype NCPolynomial a = NCPolynomial {forall a. NCPolynomial a -> Map [a] Integer
unNCPolynomial :: Map [a] Integer}
  deriving (NCPolynomial a -> NCPolynomial a -> Bool
(NCPolynomial a -> NCPolynomial a -> Bool)
-> (NCPolynomial a -> NCPolynomial a -> Bool)
-> Eq (NCPolynomial a)
forall a. Eq a => NCPolynomial a -> NCPolynomial a -> Bool
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
$c== :: forall a. Eq a => NCPolynomial a -> NCPolynomial a -> Bool
== :: NCPolynomial a -> NCPolynomial a -> Bool
$c/= :: forall a. Eq a => NCPolynomial a -> NCPolynomial a -> Bool
/= :: NCPolynomial a -> NCPolynomial a -> Bool
Eq, Int -> NCPolynomial a -> ShowS
[NCPolynomial a] -> ShowS
NCPolynomial a -> String
(Int -> NCPolynomial a -> ShowS)
-> (NCPolynomial a -> String)
-> ([NCPolynomial a] -> ShowS)
-> Show (NCPolynomial a)
forall a. Show a => Int -> NCPolynomial a -> ShowS
forall a. Show a => [NCPolynomial a] -> ShowS
forall a. Show a => NCPolynomial a -> String
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
$cshowsPrec :: forall a. Show a => Int -> NCPolynomial a -> ShowS
showsPrec :: Int -> NCPolynomial a -> ShowS
$cshow :: forall a. Show a => NCPolynomial a -> String
show :: NCPolynomial a -> String
$cshowList :: forall a. Show a => [NCPolynomial a] -> ShowS
showList :: [NCPolynomial a] -> ShowS
Show)

-- | Normalize a 'Ring' term into ℤ⟨a⟩.
--
-- Distributes 'Times' over 'Plus', pushes 'Negate' to coefficients,
-- collects like /words/ — no sorting, so the commutator survives:
--
-- >>> let comm = minus (times (embed "x") (embed "y")) (times (embed "y") (embed "x"))
-- >>> toWord comm
-- NCPolynomial {unNCPolynomial = fromList [(["x","y"],1),(["y","x"],-1)]}
toWord :: (Ord a) => Ring a -> NCPolynomial a
toWord :: forall a. Ord a => Ring a -> NCPolynomial a
toWord = Map [a] Integer -> NCPolynomial a
forall a. Map [a] Integer -> NCPolynomial a
NCPolynomial (Map [a] Integer -> NCPolynomial a)
-> (Ring a -> Map [a] Integer) -> Ring a -> NCPolynomial a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Integer -> Bool) -> Map [a] Integer -> Map [a] Integer
forall a k. (a -> Bool) -> Map k a -> Map k a
M.filter (Integer -> Integer -> Bool
forall a. Eq a => a -> a -> Bool
/= (Integer
0 :: Integer)) (Map [a] Integer -> Map [a] Integer)
-> (Ring a -> Map [a] Integer) -> Ring a -> Map [a] Integer
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Ring a -> Map [a] Integer
forall {a} {a}. (Num a, Ord a) => Ring a -> Map [a] a
go
  where
    go :: Ring a -> Map [a] a
go Ring a
Zero = Map [a] a
forall k a. Map k a
M.empty
    go Ring a
One = [a] -> a -> Map [a] a
forall k a. k -> a -> Map k a
M.singleton [] a
1
    go (Embed a
a) = [a] -> a -> Map [a] a
forall k a. k -> a -> Map k a
M.singleton [a
a] a
1
    go (Negate Ring a
r) = (a -> a) -> Map [a] a -> Map [a] a
forall a b k. (a -> b) -> Map k a -> Map k b
M.map a -> a
forall a. Num a => a -> a
P.negate (Ring a -> Map [a] a
go Ring a
r)
    go (Plus Ring a
r Ring a
s) = (a -> a -> a) -> Map [a] a -> Map [a] a -> Map [a] a
forall k a. Ord k => (a -> a -> a) -> Map k a -> Map k a -> Map k a
M.unionWith a -> a -> a
forall a. Num a => a -> a -> a
(+) (Ring a -> Map [a] a
go Ring a
r) (Ring a -> Map [a] a
go Ring a
s)
    go (Times Ring a
r Ring a
s) =
      (a -> a -> a) -> [([a], a)] -> Map [a] a
forall k a. Ord k => (a -> a -> a) -> [(k, a)] -> Map k a
M.fromListWith
        a -> a -> a
forall a. Num a => a -> a -> a
(+)
        [ ([a]
m1 [a] -> [a] -> [a]
forall a. [a] -> [a] -> [a]
P.++ [a]
m2, a
c1 a -> a -> a
forall a. Num a => a -> a -> a
* a
c2)
        | ([a]
m1, a
c1) <- Map [a] a -> [([a], a)]
forall k a. Map k a -> [(k, a)]
M.toList (Ring a -> Map [a] a
go Ring a
r),
          ([a]
m2, a
c2) <- Map [a] a -> [([a], a)]
forall k a. Map k a -> [(k, a)]
M.toList (Ring a -> Map [a] a
go Ring a
s)
        ]

-- | Quotient ℤ⟨a⟩ → ℤ[a]: sort each word into a monomial bag and
-- collect.  Coefficients that cancel under the identification are
-- dropped.
abelianize :: (Ord a) => NCPolynomial a -> Polynomial a
abelianize :: forall a. Ord a => NCPolynomial a -> Polynomial a
abelianize (NCPolynomial Map [a] Integer
m) =
  Map [a] Integer -> Polynomial a
forall a. Map [a] Integer -> Polynomial a
Polynomial
    ( (Integer -> Bool) -> Map [a] Integer -> Map [a] Integer
forall a k. (a -> Bool) -> Map k a -> Map k a
M.filter
        (Integer -> Integer -> Bool
forall a. Eq a => a -> a -> Bool
/= (Integer
0 :: Integer))
        ((Integer -> Integer -> Integer)
-> [([a], Integer)] -> Map [a] Integer
forall k a. Ord k => (a -> a -> a) -> [(k, a)] -> Map k a
M.fromListWith Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
(+) [([a] -> [a]
forall a. Ord a => [a] -> [a]
sort [a]
w, Integer
c) | ([a]
w, Integer
c) <- Map [a] Integer -> [([a], Integer)]
forall k a. Map k a -> [(k, a)]
M.toList Map [a] Integer
m])
    )

-- | Normalize a 'Ring' term to commutative polynomial form.
--
-- The composite of the faithful normal form and the abelianization.
-- Sound for any 'Ring' target; complete only up to commutativity —
-- the commutator vanishes here:
--
-- >>> let comm = minus (times (embed "x") (embed "y")) (times (embed "y") (embed "x"))
-- >>> toPolynomial comm
-- Polynomial {unPolynomial = fromList []}
toPolynomial :: (Ord a) => Ring a -> Polynomial a
toPolynomial :: forall a. Ord a => Ring a -> Polynomial a
toPolynomial = NCPolynomial a -> Polynomial a
forall a. Ord a => NCPolynomial a -> Polynomial a
abelianize (NCPolynomial a -> Polynomial a)
-> (Ring a -> NCPolynomial a) -> Ring a -> Polynomial a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Ring a -> NCPolynomial a
forall a. Ord a => Ring a -> NCPolynomial a
toWord

-- | Reconstruct a 'Ring' term from polynomial normal form.
fromPolynomial :: Polynomial a -> Ring a
fromPolynomial :: forall a. Polynomial a -> Ring a
fromPolynomial (Polynomial Map [a] Integer
m) =
  (Ring a -> Ring a -> Ring a) -> Ring a -> [Ring a] -> Ring a
forall a b. (a -> b -> b) -> b -> [a] -> b
forall (t :: * -> *) a b.
Foldable t =>
(a -> b -> b) -> b -> t a -> b
foldr Ring a -> Ring a -> Ring a
forall a. Ring a -> Ring a -> Ring a
plus Ring a
forall a. Ring a
zero [Integer -> [a] -> Ring a
forall {t} {a}. Integral t => t -> [a] -> Ring a
term Integer
c [a]
ms | ([a]
ms, Integer
c) <- Map [a] Integer -> [([a], Integer)]
forall k a. Map k a -> [(k, a)]
M.toList Map [a] Integer
m, Integer
c Integer -> Integer -> Bool
forall a. Eq a => a -> a -> Bool
P./= (Integer
0 :: Integer)]
  where
    term :: t -> [a] -> Ring a
term t
c [a]
ms = Ring a -> Ring a -> Ring a
forall a. Ring a -> Ring a -> Ring a
times (t -> Ring a
forall {t} {a}. Integral t => t -> Ring a
coeffToRing t
c) ((Ring a -> Ring a -> Ring a) -> Ring a -> [Ring a] -> Ring a
forall a b. (a -> b -> b) -> b -> [a] -> b
forall (t :: * -> *) a b.
Foldable t =>
(a -> b -> b) -> b -> t a -> b
foldr Ring a -> Ring a -> Ring a
forall a. Ring a -> Ring a -> Ring a
times Ring a
forall a. Ring a
one [a -> Ring a
forall a. a -> Ring a
embed a
x | a
x <- [a]
ms])
    coeffToRing :: t -> Ring a
coeffToRing t
n
      | t
n t -> t -> Bool
forall a. Ord a => a -> a -> Bool
P.> t
0 = (Ring a -> Ring a -> Ring a) -> Ring a -> [Ring a] -> Ring a
forall a b. (a -> b -> b) -> b -> [a] -> b
forall (t :: * -> *) a b.
Foldable t =>
(a -> b -> b) -> b -> t a -> b
foldr Ring a -> Ring a -> Ring a
forall a. Ring a -> Ring a -> Ring a
plus Ring a
forall a. Ring a
one (Int -> Ring a -> [Ring a]
forall a. Int -> a -> [a]
P.replicate (t -> Int
forall a b. (Integral a, Num b) => a -> b
P.fromIntegral (t
n t -> t -> t
forall a. Num a => a -> a -> a
P.- t
1)) Ring a
forall a. Ring a
one)
      | t
n t -> t -> Bool
forall a. Ord a => a -> a -> Bool
P.< t
0 = Ring a -> Ring a
forall a. Ring a -> Ring a
negate (t -> Ring a
coeffToRing (t -> t
forall a. Num a => a -> a
P.negate t
n))
      | Bool
P.otherwise = Ring a
forall a. Ring a
zero

-- | Evaluate a polynomial via a generator assignment.
--
-- >>> let p = toPolynomial (plus (times (embed "x") (embed "y")) (negate (embed "z")))
-- >>> evalPolynomial (\g -> case g of "x" -> 2; "y" -> 3; _ -> 5) p :: Int
-- 1
evalPolynomial ::
  (NHR.Ring b, NHI.FromInteger b) =>
  (a -> b) ->
  Polynomial a ->
  b
evalPolynomial :: forall b a.
(Ring b, FromInteger b) =>
(a -> b) -> Polynomial a -> b
evalPolynomial a -> b
f (Polynomial Map [a] Integer
m) =
  (b -> b -> b) -> b -> [b] -> b
forall a b. (a -> b -> b) -> b -> [a] -> b
forall (t :: * -> *) a b.
Foldable t =>
(a -> b -> b) -> b -> t a -> b
foldr
    b -> b -> b
forall a. Additive a => a -> a -> a
(NHA.+)
    b
forall a. Additive a => a
NHA.zero
    [ Integer -> b
forall a. FromInteger a => Integer -> a
NHI.fromInteger Integer
c b -> b -> b
forall a. Multiplicative a => a -> a -> a
NHM.* (b -> b -> b) -> b -> [b] -> b
forall a b. (a -> b -> b) -> b -> [a] -> b
forall (t :: * -> *) a b.
Foldable t =>
(a -> b -> b) -> b -> t a -> b
foldr b -> b -> b
forall a. Multiplicative a => a -> a -> a
(NHM.*) b
forall a. Multiplicative a => a
NHM.one [a -> b
f a
x | a
x <- [a]
ms]
    | ([a]
ms, Integer
c) <- Map [a] Integer -> [([a], Integer)]
forall k a. Map k a -> [(k, a)]
M.toList Map [a] Integer
m
    ]

-- | Evaluate a noncommutative polynomial via a generator assignment.
--
-- Word order is respected, so this is the unique 'Ring' homomorphism
-- ℤ⟨a⟩ → b extending the assignment — lawful for noncommutative
-- targets where 'evalPolynomial' is not.
evalNCPolynomial ::
  (NHR.Ring b, NHI.FromInteger b) =>
  (a -> b) ->
  NCPolynomial a ->
  b
evalNCPolynomial :: forall b a.
(Ring b, FromInteger b) =>
(a -> b) -> NCPolynomial a -> b
evalNCPolynomial a -> b
f (NCPolynomial Map [a] Integer
m) =
  (b -> b -> b) -> b -> [b] -> b
forall a b. (a -> b -> b) -> b -> [a] -> b
forall (t :: * -> *) a b.
Foldable t =>
(a -> b -> b) -> b -> t a -> b
foldr
    b -> b -> b
forall a. Additive a => a -> a -> a
(NHA.+)
    b
forall a. Additive a => a
NHA.zero
    [ Integer -> b
forall a. FromInteger a => Integer -> a
NHI.fromInteger Integer
c b -> b -> b
forall a. Multiplicative a => a -> a -> a
NHM.* (b -> b -> b) -> b -> [b] -> b
forall a b. (a -> b -> b) -> b -> [a] -> b
forall (t :: * -> *) a b.
Foldable t =>
(a -> b -> b) -> b -> t a -> b
foldr b -> b -> b
forall a. Multiplicative a => a -> a -> a
(NHM.*) b
forall a. Multiplicative a => a
NHM.one [a -> b
f a
x | a
x <- [a]
ms]
    | ([a]
ms, Integer
c) <- Map [a] Integer -> [([a], Integer)]
forall k a. Map k a -> [(k, a)]
M.toList Map [a] Integer
m
    ]