Expr.hs

{-# LANGUAGE OverloadedStrings, TypeFamilies #-}

module Expr ( Real, Var, Expr, literalFunction, varE, constE ) where

import           GHC.Exts (IsString (..))
import           Prelude hiding (Real)
import qualified Data.Map  as Map
import qualified Data.List as List

import           Basis
import           VectorSpace
import           InnerSpace
import           Differentiation

type Real = Double
type Var  = String

newtype Expr = Expr (Map.Map Prod Real) deriving (Eq,Ord)

newtype Prod = Prod (Map.Map Atom Int) deriving (Eq,Ord)

data Atom = Var Var
          | App String Expr
          deriving (Eq,Ord,Show)

------------------------------
-- Constructors
------------------------------

constE :: Real -> Expr
constE c = Expr $ Map.singleton (Prod $ Map.empty) c

atomE :: Atom -> Expr
atomE a = Expr $ Map.singleton (Prod $ Map.singleton a 1) 1

varE :: String -> Expr
varE = atomE . Var

------------------------------
-- Predicates/selectors for atoms
------------------------------

varA :: Var -> Atom
varA = Var

appA :: Var -> Expr -> Atom
appA = App

isVarA :: Atom -> Bool
isVarA (Var _) = True
isVarA _       = False

getVarA :: Atom -> Var
getVarA (Var v) = v
getVarA _       = error "Not a variable!"

isAppA :: Atom -> Bool
isAppA (App _ _) = True
isAppA _         = False

getAppA :: Atom -> (String, Expr)
getAppA (App f e) = (f, e)
getAppA _         = error "Not a variable!"

modifyA :: (Var -> Var) -> Atom -> Atom
modifyA g (Var v)   = Var (g v)
modifyA g (App f e) = App (g f) e

------------------------------
-- Predicates/selectors for products
------------------------------

isConstP :: Prod -> Bool
isConstP (Prod m) = Map.null m

------------------------------
-- Constructors for expressions
------------------------------

-- |Create a literal symbolic function.
literalFunction :: Expr -> Expr -> Expr
literalFunction f expr = atomE $ App (getVar f) expr

------------------------------
-- Predicates/selectors for expressions
------------------------------

isConst :: Expr -> Bool
isConst (Expr m) = Map.size m == 1 && Map.size p == 0
    where
        (Prod p, _) = Map.findMin m

getConst :: Expr -> Real
getConst e@(Expr m) = if isConst e
    then c
    else error "Not a constant!"
    where
        (Prod p, c) = Map.findMin m

isAtom :: Expr -> Bool
isAtom (Expr m) = Map.size m == 1 && c == 1 && Map.size p == 1 && n == 1
    where
        (Prod p, c) = Map.findMin m
        (atom, n)   = Map.findMin p

getAtom :: Expr -> Atom
getAtom e@(Expr m) = if not (isAtom e)
    then error "Not an atom!"
    else atom
        where
            (Prod p, c) = Map.findMin m
            (atom, _)   = Map.findMin p

isVar :: Expr -> Bool
isVar e = isAtom e && isVarA (getAtom e)

getVar :: Expr -> Var
getVar e = if isVar e
    then getVarA (getAtom e)
    else error "Not a variable!"

------------------------------
-- Convert to string
------------------------------

showAtom :: Atom -> String
showAtom (Var s)   = s
showAtom (App f s) = f ++ "(" ++ show s ++ ")"

showProd :: Prod -> String
showProd (Prod m) = List.intercalate " " $ map shw $ Map.toList m
    where
        shw (e,n) = showAtom e ++ (case n of
                                    1 -> ""
                                    2 -> "²"
                                    3 -> "³"
                                    4 -> "⁴"
                                    5 -> "⁵"
                                    6 -> "⁶"
                                    7 -> "⁷"
                                    8 -> "⁸"
                                    9 -> "⁹"
                                    n -> "^" ++ show n)

showExpr :: Expr -> String
showExpr (Expr m) =
    if Map.null m
        then "0"
        else List.intercalate " + " . map shw $ Map.toList m
    where
        shw (p,n) = pre ++ showProd p
            where
                pre = if isConstP p
                        then show n
                        else case n of
                                1    -> ""
                                (-1) -> "-"
                                _    -> show n

instance Show Expr where
    show = showExpr

instance IsString Expr where
    fromString = varE

------------------------------
-- Numeric instances
------------------------------

instance Num Expr where
    Expr as + Expr bs =
        Expr $ Map.filter (/=0) $ Map.unionWith (+) as bs

    Expr as * Expr bs =
        Expr $ Map.fromListWith (+) [ mul a b | a <- Map.toList as, b <- Map.toList bs ]
        where
            mul (Prod xs, a) (Prod ys, b) = (Prod $ Map.unionWith (+) xs ys, a * b)

    fromInteger = constE . fromInteger

    negate (Expr as) = Expr $ Map.map negate as

    abs    _ = undefined
    signum _ = undefined

instance Fractional Expr where
    fromRational = constE . fromRational

    recip = undefined

instance Floating Expr where
    pi   = constE pi
    exp  = atomE . App "exp"
    log  = atomE . App "log"
    sqrt = atomE . App "sqrt"
    cos  = atomE . App "cos"
    sin  = atomE . App "sin"
    tan  = atomE . App "tan"
    acos = atomE . App "acos"
    asin = atomE . App "asin"
    atan = atomE . App "atan"
    sinh = atomE . App "sinh"
    cosh = atomE . App "cosh"
    tanh = atomE . App "tanh"
    acosh = atomE . App "acosh"
    asinh = atomE . App "asinh"
    atanh = atomE . App "atanh"

instance AdditiveGroup Expr where
    zeroV = 0
    (<+>) = (+)
    (<->) = (-)

instance VectorSpace Expr where
    type Scalar Expr = Expr

    (*>) = (*)

instance InnerSpace Expr where
    dot = (*)

instance HasBasis Expr where
    type Basis Expr = ()

    basisValue _ = 1
    decompose x  = const x

instance Differentiable Expr where
    d f x = fromCoords . fmap diffExpr . toCoords $ f x
        where
            diffExpr = atomE . modifyA ('d':) . getAtom