{-|
Module      : $Header$
Description : Operations on the AST of a mathematical equation.
Copyright   : (c) Joshua Gutow, 2016
License     : GNU AGPL-3

Maintainer  : jbgutow@gmail.com
Stability   : experimental
Portability : portable

This module provides a data structure to represent that AST of an equation and
it provides tools to manipulate that structure.
-}

module Expression where
import Atomic
import Operations
import Floating
import Data.List (intercalate)
import Data.Ratio

-- | Node of an expression tree.
data Node = Leaf      {atom :: Atomic}
          | BinaryOp  {bop :: BOp, left :: Node, right :: Node}
          | UnaryOp   {uop :: UOp, val :: Node}
          | Operation {op :: Op, children :: [Node]}
          deriving (Eq)

instance Show Node where
    show (Leaf v)             = show v
    show (BinaryOp o l r)     = "(" ++ show l ++ show o ++ show r ++ ")"
    show (UnaryOp o (Leaf a)) | o == Factorial = "(" ++ show a ++ show o ++ ")"
                              | otherwise      = show o ++ "(" ++ show a ++ ")"
    show (UnaryOp o l)        | o == Factorial = "(" ++ show l ++ show o ++ ")"
                              | otherwise      = show o ++ show l
    show (Operation o c)      = "(" ++ intercalate (show o) (map show c) ++ ")"


-- | Simplifies all identities given to it.
simplifyIds :: [(Node -> Node)] -> Node -> Node
simplifyIds [] node  = node
simplifyIds ids node = simplifyIds (tail ids) $ simplify (head ids) node

-- | Simplifies an expression tree as best as possible.
simplifyAll :: Node -> Node
simplifyAll = simplifyIds [constants, addId, multId]

-- | Given a function that simplifies a single node, this walks the entire
-- expression tree.
-- This is designed to walk through every simplification at every step.
simplify :: (Node -> Node) -> Node -> Node
simplify fn (UnaryOp o v)    = fn $ UnaryOp   o (simplifyAll v)
simplify fn (BinaryOp o l r) = fn $ BinaryOp  o (simplifyAll l) (simplifyAll r)
simplify fn (Operation o c)  = fn $ Operation o (map simplifyAll c)
simplify _  x@(Leaf _)       = x


-- | Simplifies basic addition identities.
addId :: Node -> Node
addId (Operation Addition [x@(Leaf _)]) = x
addId (Operation Addition xs) | length output == 1 = head output
                              | otherwise          = Operation Addition output
  where
    output = (filter (Leaf (Num (0%1)) /=) xs)
addId x = x

-- | Simplifies basic multiplication identities.
multId :: Node -> Node
multId (Operation Multiplication [x@(Leaf _)]) = x
multId (Operation Multiplication xs) | length output == 1 = head output
                                     | otherwise = Operation Multiplication output
  where
    output = (filter (Leaf (Num (1%1)) /=) xs)
multId x = x

-- | This simplifies a node where the chilren are constants.
constants :: Node -> Node
constants (UnaryOp  o (Leaf (Num n)))                = Leaf $ Num (uopFn o n)
constants (BinaryOp o (Leaf (Num l)) (Leaf (Num r))) = Leaf $ Num (bopFn o l r)
constants (Operation o c) = Operation o (nums:vars)
  where
    val  = foldl (opFn o) (baseId o) (map (num . atom) $ filter isNum c)
    nums = Leaf $ Num val
    vars = filter (not . isNum) c
constants x = x

isNum :: Node -> Bool
isNum (Leaf (Num _)) = True
isNum _              = False


uopFn :: UOp -> (Rational -> Rational)
uopFn Negation = negate
-- | I need a formula that works on rational numbers.
-- uopFn Factorial = factorial
uopFn Ln        = log
uopFn (Log r)   = logBase r
uopFn Tan       = tan
uopFn Sin       = sin
uopFn Cos       = cos


bopFn :: BOp -> (Rational -> Rational -> Rational)
bopFn Division     = (/)
bopFn Subtraction  = (-)
bopFn Exp          = (**)

opFn :: Op -> (Rational -> Rational -> Rational)
opFn Multiplication = (*)
opFn Addition       = (+)

baseId :: Num a => Op -> a
baseId Multiplication = 1
baseId Addition       = 0

