-- |
-- Module: Symtegration.Symbolic.Simplify.AlgebraicRingOrder
-- Description: Order terms in a mathematical expression according to a deterministic order.
-- Copyright: Copyright 2024 Yoo Chung
-- License: Apache-2.0
-- Maintainer: dev@chungyc.org
module Symtegration.Symbolic.Simplify.AlgebraicRingOrder (order) where

import Data.List (sortBy)
import Data.Set qualified as Set
import Data.Text (Text)
import Symtegration.Symbolic

-- $setup
-- >>> import Symtegration.Symbolic.Haskell

-- | Order terms in an mathematical expression.
--
-- Terms will be ordered according to a deterministic set of rules.
-- The re-ordering aims to make it easier to identify common factors and terms.
-- Terms with higher integral powers of the variable are sorted later.
-- Addition and multiplication will be re-arranged to associate to the left.
--
-- >>> toHaskell $ order "x" $ "x" + 1
-- "1 + x"
-- >>> toHaskell $ order "x" $ 2 + 3 * "x"**2 + "x"
-- "2 + x + 3 * x ** 2"
order ::
  -- | Symbol representing the variable.
  Text ->
  -- | Expression to be ordered.
  Expression ->
  -- | Expression with terms re-ordered.
  Expression
order :: Text -> Expression -> Expression
order Text
_ e :: Expression
e@(Number Integer
_) = Expression
e
order Text
_ e :: Expression
e@(Symbol Text
_) = Expression
e
order Text
v (UnaryApply UnaryFunction
func Expression
x) = UnaryFunction -> Expression -> Expression
UnaryApply UnaryFunction
func (Expression -> Expression) -> Expression -> Expression
forall a b. (a -> b) -> a -> b
$ Text -> Expression -> Expression
order Text
v Expression
x
order Text
v (Expression
x :/: Expression
y) = Text -> Expression -> Expression
order Text
v Expression
x Expression -> Expression -> Expression
:/: Text -> Expression -> Expression
order Text
v Expression
y
order Text
v (Expression
x :**: Expression
y) = Text -> Expression -> Expression
order Text
v Expression
x Expression -> Expression -> Expression
:**: Text -> Expression -> Expression
order Text
v Expression
y
order Text
v (LogBase' Expression
x Expression
y) = Expression -> Expression -> Expression
LogBase' (Text -> Expression -> Expression
order Text
v Expression
x) (Text -> Expression -> Expression
order Text
v Expression
y)
order Text
v Expression
e = [Expression] -> Expression
fromAddList ([Expression] -> Expression) -> [Expression] -> Expression
forall a b. (a -> b) -> a -> b
$ (Expression -> Expression -> Ordering)
-> [Expression] -> [Expression]
forall a. (a -> a -> Ordering) -> [a] -> [a]
sortBy (Text -> Expression -> Expression -> Ordering
compareExpressions Text
v) [Expression]
orderedAddTerms
  where
    terms :: [[Expression]]
terms = Text -> Expression -> [[Expression]]
toAddMultiplyList Text
v Expression
e
    orderedAddTerms :: [Expression]
orderedAddTerms = ([Expression] -> Expression) -> [[Expression]] -> [Expression]
forall a b. (a -> b) -> [a] -> [b]
map ([Expression] -> Expression
fromMultiplyList ([Expression] -> Expression)
-> ([Expression] -> [Expression]) -> [Expression] -> Expression
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Expression -> Expression -> Ordering)
-> [Expression] -> [Expression]
forall a. (a -> a -> Ordering) -> [a] -> [a]
sortBy (Text -> Expression -> Expression -> Ordering
compareExpressions Text
v)) [[Expression]]
terms

-- | Gather additive terms formed out of multiplicative terms.
-- No particular ordering should be expected.
toAddMultiplyList :: Text -> Expression -> [[Expression]]
toAddMultiplyList :: Text -> Expression -> [[Expression]]
toAddMultiplyList Text
v (x :: Expression
x@(Expression
_ :+: Expression
_) :+: y :: Expression
y@(Expression
_ :+: Expression
_)) = Text -> Expression -> [[Expression]]
toAddMultiplyList Text
v Expression
x [[Expression]] -> [[Expression]] -> [[Expression]]
forall a. [a] -> [a] -> [a]
++ Text -> Expression -> [[Expression]]
toAddMultiplyList Text
v Expression
y
toAddMultiplyList Text
v (x :: Expression
x@(Expression
_ :+: Expression
_) :+: Expression
y) = Text -> Expression -> [Expression]
toMultiplyList Text
v Expression
y [Expression] -> [[Expression]] -> [[Expression]]
forall a. a -> [a] -> [a]
: Text -> Expression -> [[Expression]]
toAddMultiplyList Text
v Expression
x
toAddMultiplyList Text
v (Expression
x :+: y :: Expression
y@(Expression
_ :+: Expression
_)) = Text -> Expression -> [Expression]
toMultiplyList Text
v Expression
x [Expression] -> [[Expression]] -> [[Expression]]
forall a. a -> [a] -> [a]
: Text -> Expression -> [[Expression]]
toAddMultiplyList Text
v Expression
y
toAddMultiplyList Text
v (Expression
x :+: Expression
y) = (Expression -> [Expression]) -> [Expression] -> [[Expression]]
forall a b. (a -> b) -> [a] -> [b]
map (Text -> Expression -> [Expression]
toMultiplyList Text
v) [Expression
x, Expression
y]
toAddMultiplyList Text
v (Expression
x :-: Expression
y) = Text -> Expression -> [[Expression]]
toAddMultiplyList Text
v (Expression
x Expression -> Expression -> Expression
:+: (Integer -> Expression
Number (-Integer
1) Expression -> Expression -> Expression
:*: Expression
y))
toAddMultiplyList Text
v Expression
x = [Text -> Expression -> [Expression]
toMultiplyList Text
v Expression
x]

-- | Gather multiplicative terms.
-- No particular ordering should be expected.
toMultiplyList :: Text -> Expression -> [Expression]
toMultiplyList :: Text -> Expression -> [Expression]
toMultiplyList Text
v (x :: Expression
x@(Expression
_ :*: Expression
_) :*: y :: Expression
y@(Expression
_ :*: Expression
_)) = Text -> Expression -> [Expression]
toMultiplyList Text
v Expression
x [Expression] -> [Expression] -> [Expression]
forall a. [a] -> [a] -> [a]
++ Text -> Expression -> [Expression]
toMultiplyList Text
v Expression
y
toMultiplyList Text
v (x :: Expression
x@(Expression
_ :*: Expression
_) :*: Expression
y) = Expression
y Expression -> [Expression] -> [Expression]
forall a. a -> [a] -> [a]
: Text -> Expression -> [Expression]
toMultiplyList Text
v Expression
x
toMultiplyList Text
v (Expression
x :*: y :: Expression
y@(Expression
_ :*: Expression
_)) = Expression
x Expression -> [Expression] -> [Expression]
forall a. a -> [a] -> [a]
: Text -> Expression -> [Expression]
toMultiplyList Text
v Expression
y
toMultiplyList Text
v (Expression
x :*: Expression
y) = [Text -> Expression -> Expression
order Text
v Expression
x, Text -> Expression -> Expression
order Text
v Expression
y]
toMultiplyList Text
_ x :: Expression
x@(Number Integer
_) = [Expression
x]
toMultiplyList Text
_ x :: Expression
x@(Symbol Text
_) = [Expression
x]
toMultiplyList Text
v (Negate' Expression
x) = Integer -> Expression
Number (-Integer
1) Expression -> [Expression] -> [Expression]
forall a. a -> [a] -> [a]
: Text -> Expression -> [Expression]
toMultiplyList Text
v Expression
x
toMultiplyList Text
v (UnaryApply UnaryFunction
func Expression
x) = [UnaryFunction -> Expression -> Expression
UnaryApply UnaryFunction
func (Expression -> Expression) -> Expression -> Expression
forall a b. (a -> b) -> a -> b
$ Text -> Expression -> Expression
order Text
v Expression
x]
toMultiplyList Text
v (BinaryApply BinaryFunction
func Expression
x Expression
y) = [BinaryFunction -> Expression -> Expression -> Expression
BinaryApply BinaryFunction
func (Text -> Expression -> Expression
order Text
v Expression
x) (Text -> Expression -> Expression
order Text
v Expression
y)]

-- | Convert a list of sub-expressions for a multiplicative term into a single expression.
fromMultiplyList :: [Expression] -> Expression
fromMultiplyList :: [Expression] -> Expression
fromMultiplyList [] = Integer -> Expression
Number Integer
1
fromMultiplyList [Expression
x] = Expression
x
fromMultiplyList (Expression
x : [Expression]
xs) = Expression
x Expression -> Expression -> Expression
:*: [Expression] -> Expression
fromMultiplyList [Expression]
xs

-- | Convert a list of sub-expressions for an additive term into a single expression.
fromAddList :: [Expression] -> Expression
fromAddList :: [Expression] -> Expression
fromAddList [] = Integer -> Expression
Number Integer
0
fromAddList [Expression
x] = Expression
x
fromAddList (Expression
x : [Expression]
xs) = Expression
x Expression -> Expression -> Expression
:+: [Expression] -> Expression
fromAddList [Expression]
xs

-- | Defines a total order among expressions.
-- In particular, higher integral powers of the variable are ordered later.
compareExpressions :: Text -> Expression -> Expression -> Ordering
compareExpressions :: Text -> Expression -> Expression -> Ordering
compareExpressions Text
v Expression
x Expression
y
  | (Just Ordering
LT) <- Maybe Ordering
compareDegree = Ordering
LT
  | (Just Ordering
GT) <- Maybe Ordering
compareDegree = Ordering
GT
  | Ordering
LT <- Ordering
comparePseudoDegree = Ordering
LT
  | Ordering
GT <- Ordering
comparePseudoDegree = Ordering
GT
  | Ordering
LT <- Ordering
compareSymbolCount = Ordering
LT
  | Ordering
GT <- Ordering
compareSymbolCount = Ordering
GT
  | Ordering
LT <- Ordering
compareOp = Ordering
LT
  | Ordering
GT <- Ordering
compareOp = Ordering
GT
  | Number Integer
n <- Expression
x, Number Integer
m <- Expression
y = Integer -> Integer -> Ordering
forall a. Ord a => a -> a -> Ordering
compare Integer
n Integer
m
  | Symbol Text
s <- Expression
x, Symbol Text
r <- Expression
y = Text -> Text -> Ordering
forall a. Ord a => a -> a -> Ordering
compare Text
s Text
r
  | UnaryApply UnaryFunction
_ Expression
x' <- Expression
x, UnaryApply UnaryFunction
_ Expression
y' <- Expression
y = Text -> Expression -> Expression -> Ordering
compareExpressions Text
v Expression
x' Expression
y'
  | BinaryApply BinaryFunction
_ Expression
x' Expression
x'' <- Expression
x,
    BinaryApply BinaryFunction
_ Expression
y' Expression
y'' <- Expression
y =
      case Text -> Expression -> Expression -> Ordering
compareExpressions Text
v Expression
x' Expression
y' of
        Ordering
EQ -> Text -> Expression -> Expression -> Ordering
compareExpressions Text
v Expression
x'' Expression
y''
        Ordering
c -> Ordering
c
  | Bool
otherwise = Ordering
EQ
  where
    compareDegree :: Maybe Ordering
compareDegree = do
      xd <- Text -> Expression -> Maybe Integer
degree Text
v Expression
x
      yd <- degree v y
      case (xd, yd) of
        (Integer
0, Integer
0) -> Ordering -> Maybe Ordering
forall a. a -> Maybe a
forall (m :: * -> *) a. Monad m => a -> m a
return Ordering
EQ
        (Integer
0, Integer
_) -> Ordering -> Maybe Ordering
forall a. a -> Maybe a
forall (m :: * -> *) a. Monad m => a -> m a
return Ordering
LT
        (Integer
_, Integer
0) -> Ordering -> Maybe Ordering
forall a. a -> Maybe a
forall (m :: * -> *) a. Monad m => a -> m a
return Ordering
GT
        (Integer, Integer)
_ -> Ordering -> Maybe Ordering
forall a. a -> Maybe a
forall (m :: * -> *) a. Monad m => a -> m a
return (Ordering -> Maybe Ordering) -> Ordering -> Maybe Ordering
forall a b. (a -> b) -> a -> b
$ Integer -> Integer -> Ordering
forall a. Ord a => a -> a -> Ordering
compare Integer
xd Integer
yd
    comparePseudoDegree :: Ordering
comparePseudoDegree = Integer -> Integer -> Ordering
forall a. Ord a => a -> a -> Ordering
compare (Text -> Expression -> Integer
pseudoDegree Text
v Expression
x) (Text -> Expression -> Integer
pseudoDegree Text
v Expression
y)
    compareSymbolCount :: Ordering
compareSymbolCount = Int -> Int -> Ordering
forall a. Ord a => a -> a -> Ordering
compare (Expression -> Int
symbolCount Expression
x) (Expression -> Int
symbolCount Expression
y)
    compareOp :: Ordering
compareOp = Int -> Int -> Ordering
forall a. Ord a => a -> a -> Ordering
compare (Text -> Expression -> Int
expressionOrder Text
v Expression
x) (Text -> Expression -> Int
expressionOrder Text
v Expression
y)

-- | The integral power of the variable for a particular expression.
degree :: Text -> Expression -> Maybe Integer
degree :: Text -> Expression -> Maybe Integer
degree Text
_ (Number Integer
_) = Integer -> Maybe Integer
forall a. a -> Maybe a
Just Integer
0
degree Text
v (Symbol Text
s) | Text
v Text -> Text -> Bool
forall a. Eq a => a -> a -> Bool
== Text
s = Integer -> Maybe Integer
forall a. a -> Maybe a
Just Integer
1 | Bool
otherwise = Integer -> Maybe Integer
forall a. a -> Maybe a
Just Integer
0
degree Text
v (Negate' Expression
x) = Text -> Expression -> Maybe Integer
degree Text
v Expression
x
degree Text
v (Expression
x :+: Expression
y) = Integer -> Integer -> Integer
forall a. Ord a => a -> a -> a
max (Integer -> Integer -> Integer)
-> Maybe Integer -> Maybe (Integer -> Integer)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Text -> Expression -> Maybe Integer
degree Text
v Expression
x Maybe (Integer -> Integer) -> Maybe Integer -> Maybe Integer
forall a b. Maybe (a -> b) -> Maybe a -> Maybe b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Text -> Expression -> Maybe Integer
degree Text
v Expression
y
degree Text
v (Expression
x :-: Expression
y) = Integer -> Integer -> Integer
forall a. Ord a => a -> a -> a
max (Integer -> Integer -> Integer)
-> Maybe Integer -> Maybe (Integer -> Integer)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Text -> Expression -> Maybe Integer
degree Text
v Expression
x Maybe (Integer -> Integer) -> Maybe Integer -> Maybe Integer
forall a b. Maybe (a -> b) -> Maybe a -> Maybe b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Text -> Expression -> Maybe Integer
degree Text
v Expression
y
degree Text
v (Expression
x :*: Expression
y) = Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
(+) (Integer -> Integer -> Integer)
-> Maybe Integer -> Maybe (Integer -> Integer)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Text -> Expression -> Maybe Integer
degree Text
v Expression
x Maybe (Integer -> Integer) -> Maybe Integer -> Maybe Integer
forall a b. Maybe (a -> b) -> Maybe a -> Maybe b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Text -> Expression -> Maybe Integer
degree Text
v Expression
y
degree Text
v (Expression
x :/: Expression
y) = (-) (Integer -> Integer -> Integer)
-> Maybe Integer -> Maybe (Integer -> Integer)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Text -> Expression -> Maybe Integer
degree Text
v Expression
x Maybe (Integer -> Integer) -> Maybe Integer -> Maybe Integer
forall a b. Maybe (a -> b) -> Maybe a -> Maybe b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Text -> Expression -> Maybe Integer
degree Text
v Expression
y
degree Text
v (Expression
x :**: (Number Integer
n)) = (Integer
n Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
*) (Integer -> Integer) -> Maybe Integer -> Maybe Integer
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Text -> Expression -> Maybe Integer
degree Text
v Expression
x
degree Text
v (Expression
x :**: Negate' Expression
y) = Text -> Expression -> Maybe Integer
degree Text
v (Expression -> Maybe Integer) -> Expression -> Maybe Integer
forall a b. (a -> b) -> a -> b
$ Expression
x Expression -> Expression -> Expression
:**: Expression
y
degree Text
_ Expression
_ = Maybe Integer
forall a. Maybe a
Nothing

-- | The number of times the variable appears in an expression.
-- Used as part of a somewhat arbitrary ordering.
pseudoDegree :: Text -> Expression -> Integer
pseudoDegree :: Text -> Expression -> Integer
pseudoDegree Text
_ (Number Integer
_) = Integer
0
pseudoDegree Text
v (Symbol Text
s) | Text
v Text -> Text -> Bool
forall a. Eq a => a -> a -> Bool
== Text
s = Integer
1 | Bool
otherwise = Integer
0
pseudoDegree Text
v (Negate' Expression
x) = Text -> Expression -> Integer
pseudoDegree Text
v Expression
x
pseudoDegree Text
v (UnaryApply UnaryFunction
_ Expression
x) = Text -> Expression -> Integer
pseudoDegree Text
v Expression
x
pseudoDegree Text
v (BinaryApply BinaryFunction
_ Expression
x Expression
y) = Text -> Expression -> Integer
pseudoDegree Text
v Expression
x Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
+ Text -> Expression -> Integer
pseudoDegree Text
v Expression
y

symbolCount :: Expression -> Int
symbolCount :: Expression -> Int
symbolCount Expression
x = Set Text -> Int
forall a. Set a -> Int
Set.size (Set Text -> Int) -> Set Text -> Int
forall a b. (a -> b) -> a -> b
$ Expression -> Set Text
collect Expression
x
  where
    collect :: Expression -> Set Text
collect (Number Integer
_) = Set Text
forall a. Set a
Set.empty
    collect (Symbol Text
s) = Text -> Set Text
forall a. a -> Set a
Set.singleton Text
s
    collect (UnaryApply UnaryFunction
_ Expression
u) = Expression -> Set Text
collect Expression
u
    collect (BinaryApply BinaryFunction
_ Expression
u Expression
v) = Set Text -> Set Text -> Set Text
forall a. Ord a => Set a -> Set a -> Set a
Set.union (Expression -> Set Text
collect Expression
u) (Expression -> Set Text
collect Expression
v)

-- | A fixed order between functions and operators.
-- Ignores the actual values the functins and operators are given.
expressionOrder :: Text -> Expression -> Int
expressionOrder :: Text -> Expression -> Int
expressionOrder Text
_ (Number Integer
_) = Int
0
-- constant symbol has expressionOrder 1
expressionOrder Text
_ (UnaryApply UnaryFunction
Negate Expression
_) = Int
2
expressionOrder Text
_ (UnaryApply UnaryFunction
Signum Expression
_) = Int
3
expressionOrder Text
_ (UnaryApply UnaryFunction
Abs Expression
_) = Int
4
expressionOrder Text
_ (BinaryApply BinaryFunction
Add Expression
_ Expression
_) = Int
5
expressionOrder Text
_ (BinaryApply BinaryFunction
Subtract Expression
_ Expression
_) = Int
6
expressionOrder Text
_ (BinaryApply BinaryFunction
Multiply Expression
_ Expression
_) = Int
7
expressionOrder Text
_ (BinaryApply BinaryFunction
Divide Expression
_ Expression
_) = Int
8
expressionOrder Text
_ (BinaryApply BinaryFunction
Power Expression
_ Expression
_) = Int
9
expressionOrder Text
_ (UnaryApply UnaryFunction
Sqrt Expression
_) = Int
10
expressionOrder Text
_ (UnaryApply UnaryFunction
Exp Expression
_) = Int
11
expressionOrder Text
_ (UnaryApply UnaryFunction
Log Expression
_) = Int
12
expressionOrder Text
_ (BinaryApply BinaryFunction
LogBase Expression
_ Expression
_) = Int
13
expressionOrder Text
_ (UnaryApply UnaryFunction
Sin Expression
_) = Int
14
expressionOrder Text
_ (UnaryApply UnaryFunction
Cos Expression
_) = Int
15
expressionOrder Text
_ (UnaryApply UnaryFunction
Tan Expression
_) = Int
16
expressionOrder Text
_ (UnaryApply UnaryFunction
Asin Expression
_) = Int
17
expressionOrder Text
_ (UnaryApply UnaryFunction
Acos Expression
_) = Int
18
expressionOrder Text
_ (UnaryApply UnaryFunction
Atan Expression
_) = Int
19
expressionOrder Text
_ (UnaryApply UnaryFunction
Sinh Expression
_) = Int
20
expressionOrder Text
_ (UnaryApply UnaryFunction
Cosh Expression
_) = Int
21
expressionOrder Text
_ (UnaryApply UnaryFunction
Tanh Expression
_) = Int
22
expressionOrder Text
_ (UnaryApply UnaryFunction
Asinh Expression
_) = Int
23
expressionOrder Text
_ (UnaryApply UnaryFunction
Acosh Expression
_) = Int
24
expressionOrder Text
_ (UnaryApply UnaryFunction
Atanh Expression
_) = Int
25
expressionOrder Text
v (Symbol Text
s)
  | Text
v Text -> Text -> Bool
forall a. Eq a => a -> a -> Bool
== Text
s = Int
26
  | Bool
otherwise = Int
1