-- |
-- Module: Symtegration.Polynomial.Symbolic
-- Description: Conversion between data structures storing general mathematical expressions and those specialized for storing polynomials.
-- Copyright: Copyright 2024 Yoo Chung
-- License: Apache-2.0
-- Maintainer: dev@chungyc.org
module Symtegration.Polynomial.Symbolic
  ( -- * Converting expression to polynomial
    fromExpression,
    forVariable,
    withSymbolicCoefficients,

    -- * Converting polynomial to expression
    toExpression,
    toRationalCoefficient,
    toSymbolicCoefficient,
  )
where

import Data.Maybe (fromMaybe)
import Data.Monoid (Sum (..))
import Data.Ratio (denominator, numerator)
import Data.Text (Text)
import Symtegration.Polynomial
import Symtegration.Symbolic

-- $setup
-- >>> import Symtegration
-- >>> import Symtegration.Polynomial
-- >>> import Symtegration.Polynomial.Indexed

-- | Converts an 'Expression' into a 'Polynomial'.
-- 'Nothing' will be returned if the conversion is not possible.
--
-- Specify the symbol representing the variable for the polynomial with 'forVariable'.
-- For example,
--
-- >>> fromExpression (forVariable "x") (("x" + 4) ** 3) :: Maybe IndexedPolynomial
-- Just x^3 + 12x^2 + 48x + 64
--
-- By default, symbols other than the variable for the polynomial are not allowed.
-- To use symbols representing constants, use 'withSymbolicCoefficients' as well.
-- Note that the polynomial type the expression is being converted into
-- must be able to handle symbolic mathematical expressions for the coefficients.
-- For example,
--
-- >>> let expr = ("a" + "b") * "x" + "c" :: Expression
-- >>> let (Just p) = fromExpression (withSymbolicCoefficients (forVariable "x")) expr :: Maybe IndexedSymbolicPolynomial
-- >>> toHaskell $ simplify $ coefficient p 1
-- "a + b"
--
-- The expressions which can be converted must only use 'negate', '(+)', '(*)', '(-)',
-- '(/)' with only numbers, coefficients which do not contain the variable,
-- '(**)' with a non-negative integral exponent, and expressions formed thereof.
fromExpression ::
  (Polynomial p e c, Num (p e c), Fractional c) =>
  (Text -> Maybe (p e c), Expression -> Maybe c) ->
  Expression ->
  Maybe (p e c)
fromExpression :: forall (p :: * -> * -> *) e c.
(Polynomial p e c, Num (p e c), Fractional c) =>
(Text -> Maybe (p e c), Expression -> Maybe c)
-> Expression -> Maybe (p e c)
fromExpression (Text -> Maybe (p e c), Expression -> Maybe c)
_ (Number Integer
n) = p e c -> Maybe (p e c)
forall a. a -> Maybe a
Just (p e c -> Maybe (p e c)) -> p e c -> Maybe (p e c)
forall a b. (a -> b) -> a -> b
$ Integer -> p e c
forall a. Num a => Integer -> a
fromInteger Integer
n
fromExpression (Text -> Maybe (p e c)
cf, Expression -> Maybe c
_) (Symbol Text
x) = Text -> Maybe (p e c)
cf Text
x
fromExpression (Text -> Maybe (p e c), Expression -> Maybe c)
t (Negate' Expression
x) = p e c -> p e c
forall a. Num a => a -> a
negate (p e c -> p e c) -> Maybe (p e c) -> Maybe (p e c)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (Text -> Maybe (p e c), Expression -> Maybe c)
-> Expression -> Maybe (p e c)
forall (p :: * -> * -> *) e c.
(Polynomial p e c, Num (p e c), Fractional c) =>
(Text -> Maybe (p e c), Expression -> Maybe c)
-> Expression -> Maybe (p e c)
fromExpression (Text -> Maybe (p e c), Expression -> Maybe c)
t Expression
x
fromExpression (Text -> Maybe (p e c), Expression -> Maybe c)
t (Expression
x :+: Expression
y) = p e c -> p e c -> p e c
forall a. Num a => a -> a -> a
(+) (p e c -> p e c -> p e c)
-> Maybe (p e c) -> Maybe (p e c -> p e c)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (Text -> Maybe (p e c), Expression -> Maybe c)
-> Expression -> Maybe (p e c)
forall (p :: * -> * -> *) e c.
(Polynomial p e c, Num (p e c), Fractional c) =>
(Text -> Maybe (p e c), Expression -> Maybe c)
-> Expression -> Maybe (p e c)
fromExpression (Text -> Maybe (p e c), Expression -> Maybe c)
t Expression
x Maybe (p e c -> p e c) -> Maybe (p e c) -> Maybe (p e c)
forall a b. Maybe (a -> b) -> Maybe a -> Maybe b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> (Text -> Maybe (p e c), Expression -> Maybe c)
-> Expression -> Maybe (p e c)
forall (p :: * -> * -> *) e c.
(Polynomial p e c, Num (p e c), Fractional c) =>
(Text -> Maybe (p e c), Expression -> Maybe c)
-> Expression -> Maybe (p e c)
fromExpression (Text -> Maybe (p e c), Expression -> Maybe c)
t Expression
y
fromExpression (Text -> Maybe (p e c), Expression -> Maybe c)
t (Expression
x :*: Expression
y) = p e c -> p e c -> p e c
forall a. Num a => a -> a -> a
(*) (p e c -> p e c -> p e c)
-> Maybe (p e c) -> Maybe (p e c -> p e c)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (Text -> Maybe (p e c), Expression -> Maybe c)
-> Expression -> Maybe (p e c)
forall (p :: * -> * -> *) e c.
(Polynomial p e c, Num (p e c), Fractional c) =>
(Text -> Maybe (p e c), Expression -> Maybe c)
-> Expression -> Maybe (p e c)
fromExpression (Text -> Maybe (p e c), Expression -> Maybe c)
t Expression
x Maybe (p e c -> p e c) -> Maybe (p e c) -> Maybe (p e c)
forall a b. Maybe (a -> b) -> Maybe a -> Maybe b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> (Text -> Maybe (p e c), Expression -> Maybe c)
-> Expression -> Maybe (p e c)
forall (p :: * -> * -> *) e c.
(Polynomial p e c, Num (p e c), Fractional c) =>
(Text -> Maybe (p e c), Expression -> Maybe c)
-> Expression -> Maybe (p e c)
fromExpression (Text -> Maybe (p e c), Expression -> Maybe c)
t Expression
y
fromExpression (Text -> Maybe (p e c), Expression -> Maybe c)
t (Expression
x :-: Expression
y) = (-) (p e c -> p e c -> p e c)
-> Maybe (p e c) -> Maybe (p e c -> p e c)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (Text -> Maybe (p e c), Expression -> Maybe c)
-> Expression -> Maybe (p e c)
forall (p :: * -> * -> *) e c.
(Polynomial p e c, Num (p e c), Fractional c) =>
(Text -> Maybe (p e c), Expression -> Maybe c)
-> Expression -> Maybe (p e c)
fromExpression (Text -> Maybe (p e c), Expression -> Maybe c)
t Expression
x Maybe (p e c -> p e c) -> Maybe (p e c) -> Maybe (p e c)
forall a b. Maybe (a -> b) -> Maybe a -> Maybe b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> (Text -> Maybe (p e c), Expression -> Maybe c)
-> Expression -> Maybe (p e c)
forall (p :: * -> * -> *) e c.
(Polynomial p e c, Num (p e c), Fractional c) =>
(Text -> Maybe (p e c), Expression -> Maybe c)
-> Expression -> Maybe (p e c)
fromExpression (Text -> Maybe (p e c), Expression -> Maybe c)
t Expression
y
fromExpression (Text -> Maybe (p e c), Expression -> Maybe c)
t (Expression
x :**: (Number Integer
n))
  | Integer
n Integer -> Integer -> Bool
forall a. Ord a => a -> a -> Bool
>= Integer
0 = (p e c -> Integer -> p e c
forall a b. (Num a, Integral b) => a -> b -> a
^ Integer
n) (p e c -> p e c) -> Maybe (p e c) -> Maybe (p e c)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (Text -> Maybe (p e c), Expression -> Maybe c)
-> Expression -> Maybe (p e c)
forall (p :: * -> * -> *) e c.
(Polynomial p e c, Num (p e c), Fractional c) =>
(Text -> Maybe (p e c), Expression -> Maybe c)
-> Expression -> Maybe (p e c)
fromExpression (Text -> Maybe (p e c), Expression -> Maybe c)
t Expression
x
  | Bool
otherwise = Maybe (p e c)
forall a. Maybe a
Nothing
fromExpression (Text -> Maybe (p e c), Expression -> Maybe c)
_ (Expression
_ :**: Expression
_) = Maybe (p e c)
forall a. Maybe a
Nothing
fromExpression (Text -> Maybe (p e c), Expression -> Maybe c)
_ (Expression
_ :/: Number Integer
0) = Maybe (p e c)
forall a. Maybe a
Nothing
fromExpression (Text -> Maybe (p e c), Expression -> Maybe c)
_ (Number Integer
n :/: Number Integer
m) = p e c -> Maybe (p e c)
forall a. a -> Maybe a
Just (p e c -> Maybe (p e c)) -> p e c -> Maybe (p e c)
forall a b. (a -> b) -> a -> b
$ c -> p e c -> p e c
forall (p :: * -> * -> *) e c.
Polynomial p e c =>
c -> p e c -> p e c
scale c
r p e c
1
  where
    r :: c
r = Integer -> c
forall a. Num a => Integer -> a
fromInteger Integer
n c -> c -> c
forall a. Fractional a => a -> a -> a
/ Integer -> c
forall a. Num a => Integer -> a
fromInteger Integer
m
fromExpression (Text -> Maybe (p e c)
_, Expression -> Maybe c
eval) Expression
e
  | Just c
e' <- Expression -> Maybe c
eval Expression
e = p e c -> Maybe (p e c)
forall a. a -> Maybe a
Just (p e c -> Maybe (p e c)) -> p e c -> Maybe (p e c)
forall a b. (a -> b) -> a -> b
$ c -> p e c -> p e c
forall (p :: * -> * -> *) e c.
Polynomial p e c =>
c -> p e c -> p e c
scale c
e' p e c
1
  | Bool
otherwise = Maybe (p e c)
forall a. Maybe a
Nothing

-- | Specifies the symbol representing the variable for 'fromExpression'.
forVariable ::
  (Polynomial p e c, Num (p e c), Fractional c) =>
  Text ->
  (Text -> Maybe (p e c), Expression -> Maybe c)
forVariable :: forall (p :: * -> * -> *) e c.
(Polynomial p e c, Num (p e c), Fractional c) =>
Text -> (Text -> Maybe (p e c), Expression -> Maybe c)
forVariable Text
v = (Text -> Maybe (p e c)
forall {p :: * -> * -> *} {e} {c}.
Polynomial p e c =>
Text -> Maybe (p e c)
fromSymbol, Expression -> Maybe c
forall {a}. Fractional a => Expression -> Maybe a
toCoefficient)
  where
    fromSymbol :: Text -> Maybe (p e c)
fromSymbol Text
s
      | Text
v Text -> Text -> Bool
forall a. Eq a => a -> a -> Bool
== Text
s = p e c -> Maybe (p e c)
forall a. a -> Maybe a
Just (p e c -> Maybe (p e c)) -> p e c -> Maybe (p e c)
forall a b. (a -> b) -> a -> b
$ e -> p e c
forall (p :: * -> * -> *) e c. Polynomial p e c => e -> p e c
power e
1
      | Bool
otherwise = Maybe (p e c)
forall a. Maybe a
Nothing

    toCoefficient :: Expression -> Maybe a
toCoefficient (Symbol Text
_) = Maybe a
forall a. Maybe a
Nothing
    toCoefficient (Number Integer
n) = a -> Maybe a
forall a. a -> Maybe a
Just (a -> Maybe a) -> a -> Maybe a
forall a b. (a -> b) -> a -> b
$ Integer -> a
forall a. Num a => Integer -> a
fromInteger Integer
n
    toCoefficient (Negate' Expression
x) = a -> a
forall a. Num a => a -> a
negate (a -> a) -> Maybe a -> Maybe a
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Expression -> Maybe a
toCoefficient Expression
x
    toCoefficient (Abs' Expression
x) = a -> a
forall a. Num a => a -> a
abs (a -> a) -> Maybe a -> Maybe a
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Expression -> Maybe a
toCoefficient Expression
x
    toCoefficient (Signum' Expression
x) = a -> a
forall a. Num a => a -> a
signum (a -> a) -> Maybe a -> Maybe a
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Expression -> Maybe a
toCoefficient Expression
x
    toCoefficient (Expression
x :+: Expression
y) = a -> a -> a
forall a. Num a => a -> a -> a
(+) (a -> a -> a) -> Maybe a -> Maybe (a -> a)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Expression -> Maybe a
toCoefficient Expression
x Maybe (a -> a) -> Maybe a -> Maybe a
forall a b. Maybe (a -> b) -> Maybe a -> Maybe b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Expression -> Maybe a
toCoefficient Expression
y
    toCoefficient (Expression
x :*: Expression
y) = a -> a -> a
forall a. Num a => a -> a -> a
(*) (a -> a -> a) -> Maybe a -> Maybe (a -> a)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Expression -> Maybe a
toCoefficient Expression
x Maybe (a -> a) -> Maybe a -> Maybe a
forall a b. Maybe (a -> b) -> Maybe a -> Maybe b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Expression -> Maybe a
toCoefficient Expression
y
    toCoefficient (Expression
x :-: Expression
y) = (-) (a -> a -> a) -> Maybe a -> Maybe (a -> a)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Expression -> Maybe a
toCoefficient Expression
x Maybe (a -> a) -> Maybe a -> Maybe a
forall a b. Maybe (a -> b) -> Maybe a -> Maybe b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Expression -> Maybe a
toCoefficient Expression
y
    toCoefficient (Expression
_ :/: Expression
0) = Maybe a
forall a. Maybe a
Nothing
    toCoefficient (Expression
x :/: Expression
y) = a -> a -> a
forall a. Fractional a => a -> a -> a
(/) (a -> a -> a) -> Maybe a -> Maybe (a -> a)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Expression -> Maybe a
toCoefficient Expression
x Maybe (a -> a) -> Maybe a -> Maybe a
forall a b. Maybe (a -> b) -> Maybe a -> Maybe b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Expression -> Maybe a
toCoefficient Expression
y
    toCoefficient (Expression
x :**: (Number Integer
n)) = (a -> Integer -> a
forall a b. (Fractional a, Integral b) => a -> b -> a
^^ Integer
n) (a -> a) -> Maybe a -> Maybe a
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Expression -> Maybe a
toCoefficient Expression
x
    toCoefficient Expression
_ = Maybe a
forall a. Maybe a
Nothing

-- | Specifies that non-variable symbols are allowed for 'fromExpression'.
-- The coefficients will be represented by 'Expression' values.
withSymbolicCoefficients ::
  (Polynomial p e Expression, Num (p e Expression), Integral e) =>
  (Text -> Maybe (p e Expression), Expression -> Maybe Expression) ->
  (Text -> Maybe (p e Expression), Expression -> Maybe Expression)
withSymbolicCoefficients :: forall (p :: * -> * -> *) e.
(Polynomial p e Expression, Num (p e Expression), Integral e) =>
(Text -> Maybe (p e Expression), Expression -> Maybe Expression)
-> (Text -> Maybe (p e Expression), Expression -> Maybe Expression)
withSymbolicCoefficients (Text -> Maybe (p e Expression)
fromSymbol, Expression -> Maybe Expression
_) = (Text -> Maybe (p e Expression)
fromSymbol', Expression -> Maybe Expression
toCoefficient)
  where
    fromSymbol' :: Text -> Maybe (p e Expression)
fromSymbol' Text
s = p e Expression -> Maybe (p e Expression)
forall a. a -> Maybe a
Just (p e Expression -> Maybe (p e Expression))
-> p e Expression -> Maybe (p e Expression)
forall a b. (a -> b) -> a -> b
$ p e Expression -> Maybe (p e Expression) -> p e Expression
forall a. a -> Maybe a -> a
fromMaybe (Expression -> p e Expression -> p e Expression
forall (p :: * -> * -> *) e c.
Polynomial p e c =>
c -> p e c -> p e c
scale (Text -> Expression
Symbol Text
s) p e Expression
1) (Text -> Maybe (p e Expression)
fromSymbol Text
s)

    toCoefficient :: Expression -> Maybe Expression
toCoefficient e :: Expression
e@(Symbol Text
s)
      | Maybe (p e Expression)
Nothing <- Text -> Maybe (p e Expression)
fromSymbol Text
s = Expression -> Maybe Expression
forall a. a -> Maybe a
Just Expression
e
      | Bool
otherwise = Maybe Expression
forall a. Maybe a
Nothing
    toCoefficient e :: Expression
e@(Number Integer
_) = Expression -> Maybe Expression
forall a. a -> Maybe a
Just Expression
e
    toCoefficient (UnaryApply UnaryFunction
func Expression
x) = UnaryFunction -> Expression -> Expression
UnaryApply UnaryFunction
func (Expression -> Expression) -> Maybe Expression -> Maybe Expression
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Expression -> Maybe Expression
toCoefficient Expression
x
    toCoefficient (BinaryApply BinaryFunction
func Expression
x Expression
y) = BinaryFunction -> Expression -> Expression -> Expression
BinaryApply BinaryFunction
func (Expression -> Expression -> Expression)
-> Maybe Expression -> Maybe (Expression -> Expression)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Maybe Expression
x' Maybe (Expression -> Expression)
-> Maybe Expression -> Maybe Expression
forall a b. Maybe (a -> b) -> Maybe a -> Maybe b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Maybe Expression
y'
      where
        x' :: Maybe Expression
x' = Expression -> Maybe Expression
toCoefficient Expression
x
        y' :: Maybe Expression
y' = Expression -> Maybe Expression
toCoefficient Expression
y

-- | Converts a 'Polynomial' into an 'Expression'.
-- The symbol which will represent the variable is the first argument.
--
-- How the coefficients are converted must also be specified.
-- To evaluate the coefficients to an exact rational number,
-- use 'toRationalCoefficient'.  For example,
--
-- >>> let (Just p) = fromExpression (forVariable "x") (3 * "x"**4 + 1) :: Maybe IndexedPolynomial
-- >>> toHaskell $ simplify $ toExpression "x" toRationalCoefficient p
-- "1 + 3 * x ** 4"
--
-- To evaluate the coefficients symbolically, use 'toSymbolicCoefficient'.
--
-- >>> let (Just p) = fromExpression (withSymbolicCoefficients (forVariable "x")) (("a"+"b") * "x"**4 + 1) :: Maybe IndexedSymbolicPolynomial
-- >>> toHaskell $ simplify $ toExpression "x" toSymbolicCoefficient p
-- "1 + x ** 4 * (a + b)"
toExpression :: (Polynomial p e c) => Text -> (c -> Expression) -> p e c -> Expression
toExpression :: forall (p :: * -> * -> *) e c.
Polynomial p e c =>
Text -> (c -> Expression) -> p e c -> Expression
toExpression Text
x c -> Expression
cf p e c
p = Sum Expression -> Expression
forall a. Sum a -> a
getSum (Sum Expression -> Expression) -> Sum Expression -> Expression
forall a b. (a -> b) -> a -> b
$ (e -> c -> Sum Expression) -> p e c -> Sum Expression
forall m. Monoid m => (e -> c -> m) -> p e c -> m
forall (p :: * -> * -> *) e c m.
(Polynomial p e c, Monoid m) =>
(e -> c -> m) -> p e c -> m
foldTerms e -> c -> Sum Expression
forall {a}. Integral a => a -> c -> Sum Expression
convert p e c
p
  where
    convert :: a -> c -> Sum Expression
convert a
0 c
c = Expression -> Sum Expression
forall a. a -> Sum a
Sum (Expression -> Sum Expression) -> Expression -> Sum Expression
forall a b. (a -> b) -> a -> b
$ c -> Expression
cf c
c
    convert a
e c
c = Expression -> Sum Expression
forall a. a -> Sum a
Sum (Expression -> Sum Expression) -> Expression -> Sum Expression
forall a b. (a -> b) -> a -> b
$ c -> Expression
cf c
c Expression -> Expression -> Expression
forall a. Num a => a -> a -> a
* Expression
xp
      where
        xp :: Expression
xp = Text -> Expression
Symbol Text
x Expression -> Expression -> Expression
forall a. Floating a => a -> a -> a
** Integer -> Expression
Number (a -> Integer
forall a b. (Integral a, Num b) => a -> b
fromIntegral a
e)

-- | Specifies that coefficients are numbers for 'toExpression'.
toRationalCoefficient :: (Real c) => c -> Expression
toRationalCoefficient :: forall c. Real c => c -> Expression
toRationalCoefficient c
c
  | Integer
d Integer -> Integer -> Bool
forall a. Eq a => a -> a -> Bool
== Integer
1 = Integer -> Expression
Number Integer
n
  | Bool
otherwise = Integer -> Expression
Number Integer
n Expression -> Expression -> Expression
:/: Integer -> Expression
Number Integer
d
  where
    r :: Rational
r = c -> Rational
forall a. Real a => a -> Rational
toRational c
c
    n :: Integer
n = Integer -> Integer
forall a. Num a => Integer -> a
fromInteger (Integer -> Integer) -> Integer -> Integer
forall a b. (a -> b) -> a -> b
$ Rational -> Integer
forall a. Ratio a -> a
numerator Rational
r
    d :: Integer
d = Integer -> Integer
forall a. Num a => Integer -> a
fromInteger (Integer -> Integer) -> Integer -> Integer
forall a b. (a -> b) -> a -> b
$ Rational -> Integer
forall a. Ratio a -> a
denominator Rational
r

-- | Specifies that coefficients are symbolic for 'toExpression'.
toSymbolicCoefficient :: Expression -> Expression
toSymbolicCoefficient :: Expression -> Expression
toSymbolicCoefficient = Expression -> Expression
forall a. a -> a
id