-- |
-- Module: Symtegration.Integration.Substitution
-- Description: Integration by substitution.
-- Copyright: Copyright 2025 Yoo Chung
-- License: Apache-2.0
-- Maintainer: dev@chungyc.org
module Symtegration.Integration.Substitution (integrate) where

import Data.Foldable (asum)
import Data.Text (Text)
import Symtegration.Differentiation
import Symtegration.Integration.Factor
import Symtegration.Symbolic

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

-- | Integrates by substitution.
--
-- Specifically, if for
--
-- \[ \int f(g(x)) h(x) \, dx\]
--
-- it is the case that \(\frac{dg(x)}{dx} = h(x)\), then compute \(\int f(v) \, dv\) and substitute with \(v=g(x)\).
--
-- >>> import Symtegration.Integration.Trigonometric qualified as Trigonometric
-- >>> toHaskell <$> simplify <$> integrate [Trigonometric.integrate] "x" (sin ("a" * "x" + 1))
-- Just "(-1) * 1 / a * cos (1 + a * x)"
integrate ::
  -- | Integration algorithms to try after substitution.
  [Text -> Expression -> Maybe Expression] ->
  -- | Symbol for the variable.
  Text ->
  -- | Expression to integrate.
  Expression ->
  -- | Integral, if derived.
  Maybe Expression
integrate :: [Text -> Expression -> Maybe Expression]
-> Text -> Expression -> Maybe Expression
integrate [Text -> Expression -> Maybe Expression]
fs Text
v (Expression
x :*: UnaryApply UnaryFunction
func Expression
y)
  | Number Integer
0 <- Expression
d = Maybe Expression
forall a. Maybe a
Nothing -- Argument is constant.
  | Expression
x' Expression -> Expression -> Bool
forall a. Eq a => a -> a -> Bool
== Expression
y',
    -- Re-use v as the variable, as it is the one symbol guaranteed not to appear outside the argument.
    Just Expression
e <- [Text -> Expression -> Maybe Expression]
-> Text -> Expression -> Maybe Expression
integrateSubstitution [Text -> Expression -> Maybe Expression]
fs Text
v (UnaryFunction -> Expression -> Expression
UnaryApply UnaryFunction
func (Text -> Expression
Symbol Text
v)) =
      Expression -> Maybe Expression
forall a. a -> Maybe a
Just (Expression -> Maybe Expression) -> Expression -> Maybe Expression
forall a b. (a -> b) -> a -> b
$ (Expression
c Expression -> Expression -> Expression
:/: Expression
d) Expression -> Expression -> Expression
:*: Expression -> (Text -> Maybe Expression) -> Expression
substitute Expression
e (\Text
s -> if Text
s Text -> Text -> Bool
forall a. Eq a => a -> a -> Bool
== Text
v then Expression -> Maybe Expression
forall a. a -> Maybe a
Just Expression
y else Maybe Expression
forall a. Maybe a
Nothing)
  | Bool
otherwise = Maybe Expression
forall a. Maybe a
Nothing
  where
    (Expression
c, Expression
x') = Text -> Expression -> (Expression, Expression)
factor Text
v Expression
x
    (Expression
d, Expression
y') = Text -> Expression -> (Expression, Expression)
factor Text
v (Expression -> (Expression, Expression))
-> Expression -> (Expression, Expression)
forall a b. (a -> b) -> a -> b
$ Text -> Expression -> Expression
differentiate Text
v Expression
y
integrate [Text -> Expression -> Maybe Expression]
fs Text
v (e :: Expression
e@(UnaryApply UnaryFunction
_ Expression
_) :*: Expression
x) = [Text -> Expression -> Maybe Expression]
-> Text -> Expression -> Maybe Expression
integrate [Text -> Expression -> Maybe Expression]
fs Text
v (Expression -> Maybe Expression) -> Expression -> Maybe Expression
forall a b. (a -> b) -> a -> b
$ Expression
x Expression -> Expression -> Expression
:*: Expression
e
integrate [Text -> Expression -> Maybe Expression]
fs Text
v e :: Expression
e@(UnaryApply UnaryFunction
_ Expression
_) = [Text -> Expression -> Maybe Expression]
-> Text -> Expression -> Maybe Expression
integrate [Text -> Expression -> Maybe Expression]
fs Text
v (Expression -> Maybe Expression) -> Expression -> Maybe Expression
forall a b. (a -> b) -> a -> b
$ Integer -> Expression
Number Integer
1 Expression -> Expression -> Expression
:*: Expression
e
integrate [Text -> Expression -> Maybe Expression]
fs Text
v (Expression
x :*: BinaryApply BinaryFunction
func Expression
y Expression
z)
  -- Re-use v as the variable, as it is the one symbol guaranteed not to appear outside the argument.
  | Expression
c Expression -> Expression -> Bool
forall a. Eq a => a -> a -> Bool
/= Integer -> Expression
Number Integer
0,
    Expression
x' Expression -> Expression -> Bool
forall a. Eq a => a -> a -> Bool
== Expression
y',
    Text -> Expression -> Bool
isConstant Text
v Expression
z,
    Just Expression
e <- [Text -> Expression -> Maybe Expression]
-> Text -> Expression -> Maybe Expression
integrateSubstitution [Text -> Expression -> Maybe Expression]
fs Text
v (BinaryFunction -> Expression -> Expression -> Expression
BinaryApply BinaryFunction
func (Text -> Expression
Symbol Text
v) Expression
z) =
      Expression -> Maybe Expression
forall a. a -> Maybe a
Just (Expression -> Maybe Expression) -> Expression -> Maybe Expression
forall a b. (a -> b) -> a -> b
$ (Expression
b Expression -> Expression -> Expression
:/: Expression
c) Expression -> Expression -> Expression
:*: Expression -> (Text -> Maybe Expression) -> Expression
substitute Expression
e (\Text
s -> if Text
s Text -> Text -> Bool
forall a. Eq a => a -> a -> Bool
== Text
v then Expression -> Maybe Expression
forall a. a -> Maybe a
Just Expression
y else Maybe Expression
forall a. Maybe a
Nothing)
  | Expression
d Expression -> Expression -> Bool
forall a. Eq a => a -> a -> Bool
/= Integer -> Expression
Number Integer
0,
    Expression
x' Expression -> Expression -> Bool
forall a. Eq a => a -> a -> Bool
== Expression
z',
    Text -> Expression -> Bool
isConstant Text
v Expression
y,
    Just Expression
e <- [Text -> Expression -> Maybe Expression]
-> Text -> Expression -> Maybe Expression
integrateSubstitution [Text -> Expression -> Maybe Expression]
fs Text
v (BinaryFunction -> Expression -> Expression -> Expression
BinaryApply BinaryFunction
func Expression
y (Text -> Expression
Symbol Text
v)) =
      Expression -> Maybe Expression
forall a. a -> Maybe a
Just (Expression -> Maybe Expression) -> Expression -> Maybe Expression
forall a b. (a -> b) -> a -> b
$ (Expression
b Expression -> Expression -> Expression
:/: Expression
d) Expression -> Expression -> Expression
:*: Expression -> (Text -> Maybe Expression) -> Expression
substitute Expression
e (\Text
s -> if Text
s Text -> Text -> Bool
forall a. Eq a => a -> a -> Bool
== Text
v then Expression -> Maybe Expression
forall a. a -> Maybe a
Just Expression
z else Maybe Expression
forall a. Maybe a
Nothing)
  | Bool
otherwise = Maybe Expression
forall a. Maybe a
Nothing
  where
    (Expression
b, Expression
x') = Text -> Expression -> (Expression, Expression)
factor Text
v Expression
x
    (Expression
c, Expression
y') = Text -> Expression -> (Expression, Expression)
factor Text
v (Expression -> (Expression, Expression))
-> Expression -> (Expression, Expression)
forall a b. (a -> b) -> a -> b
$ Text -> Expression -> Expression
differentiate Text
v Expression
y
    (Expression
d, Expression
z') = Text -> Expression -> (Expression, Expression)
factor Text
v (Expression -> (Expression, Expression))
-> Expression -> (Expression, Expression)
forall a b. (a -> b) -> a -> b
$ Text -> Expression -> Expression
differentiate Text
v Expression
z
integrate [Text -> Expression -> Maybe Expression]
fs Text
v (e :: Expression
e@(BinaryApply BinaryFunction
_ Expression
_ Expression
_) :*: Expression
x) = [Text -> Expression -> Maybe Expression]
-> Text -> Expression -> Maybe Expression
integrate [Text -> Expression -> Maybe Expression]
fs Text
v (Expression -> Maybe Expression) -> Expression -> Maybe Expression
forall a b. (a -> b) -> a -> b
$ Expression
x Expression -> Expression -> Expression
:*: Expression
e
integrate [Text -> Expression -> Maybe Expression]
fs Text
v e :: Expression
e@(BinaryApply BinaryFunction
func Expression
_ Expression
_)
  | BinaryFunction
func BinaryFunction -> BinaryFunction -> Bool
forall a. Eq a => a -> a -> Bool
/= BinaryFunction
Multiply = [Text -> Expression -> Maybe Expression]
-> Text -> Expression -> Maybe Expression
integrate [Text -> Expression -> Maybe Expression]
fs Text
v (Expression -> Maybe Expression) -> Expression -> Maybe Expression
forall a b. (a -> b) -> a -> b
$ Integer -> Expression
Number Integer
1 Expression -> Expression -> Expression
:*: Expression
e
  | Bool
otherwise = Maybe Expression
forall a. Maybe a
Nothing
integrate [Text -> Expression -> Maybe Expression]
_ Text
_ Expression
_ = Maybe Expression
forall a. Maybe a
Nothing

-- | Use the given functions to integrate the given expression.
integrateSubstitution :: [Text -> Expression -> Maybe Expression] -> Text -> Expression -> Maybe Expression
integrateSubstitution :: [Text -> Expression -> Maybe Expression]
-> Text -> Expression -> Maybe Expression
integrateSubstitution [Text -> Expression -> Maybe Expression]
fs Text
v Expression
e = [Maybe Expression] -> Maybe Expression
forall (t :: * -> *) (f :: * -> *) a.
(Foldable t, Alternative f) =>
t (f a) -> f a
asum ([Maybe Expression] -> Maybe Expression)
-> [Maybe Expression] -> Maybe Expression
forall a b. (a -> b) -> a -> b
$ ((Text -> Expression -> Maybe Expression) -> Maybe Expression)
-> [Text -> Expression -> Maybe Expression] -> [Maybe Expression]
forall a b. (a -> b) -> [a] -> [b]
map (\Text -> Expression -> Maybe Expression
f -> Text -> Expression -> Maybe Expression
f Text
v Expression
e) [Text -> Expression -> Maybe Expression]
fs