-- ParserFunction provides utilities for parsing and evaluating mathematical expressions. The central parsing 
-- function in this package is stringToExpr, which parses a string-expression and returns a maybe expression tree.
--
-- EXAMPLE:
-- > stringToExpr "e^(1-x)*cos(pi*y)"
-- > Just (Mul (Pow (Var "e") (Sub (Num 1.0) (Var "x"))) (Cos (Mul (Var "pi") (Var "y"))))
--
-- This type is suitable for performing symbolic manipulation.
-- 
-- Expressions can then be evaluated using the function evalExpr. 
--
-- EXAMPLE:
-- > evalExpr ((Mul (Pow (Var "e") (Sub (Num 1.0) (Var "x"))) (Cos (Mul (Var "pi") (Var "y"))))) [("x",1),("y",0)]
-- > Just (1.0 :+ 0.0)
--
-- If you wish to evaluate a string-expression without any intermediate operations, simply use the function evalString.
--
-- EXAMPLE:
-- > evalString "e^(1-x)*cos(pi*y)" [("x",1),("y",0)]
-- > Just (1.0 :+ 0.0)
--
-- EXAMPLE:
-- > evalString "e^(-pi*i)+1" []
-- > Just (0.0 :+ (-1.2246467991473532e-16))
--

module Text.ParserCombinators.Parsec.ParserFunction
  (Expr,Variable,evalString,evalExpr,stringToExpr,buildExpr,eval) where

import Control.Monad (liftM,liftM2)
import Text.ParserCombinators.Parsec.Expr 
import Text.ParserCombinators.Parsec
import qualified Data.Map as M
import Data.Maybe (fromMaybe)
import Data.List (isInfixOf)
import Data.Char (toLower)
import Data.Complex

type Variable = String

-- |The Expr data type provides a basis for ordering mathematical operations.
data Expr = 
  Num Double    | Var String    | Sub Expr Expr |
  Div Expr Expr | Pow Expr Expr | Log Expr      |
  Abs Expr      | Sqrt Expr     | Cbrt Expr     |
  ArcSinh Expr  | ArcCosh Expr  | ArcTanh Expr  |
  ArcSin Expr   | ArcCos Expr   | ArcTan Expr   |
  Sinh Expr     | Cosh Expr     | Tanh Expr     |
  Sin Expr      | Cos Expr      | Tan Expr      |
  ArcSech Expr  | ArcCsch Expr  | ArcCoth Expr  |
  ArcSec Expr   | ArcCsc Expr   | ArcCot Expr   |
  Sech Expr     | Csch Expr     | Coth Expr     |
  Sec Expr      | Csc Expr      | Cot Expr      |
  Mul Expr Expr | Add Expr Expr | Exp Expr      deriving (Show, Eq)

-- |@evalExpr@ evaluates an expression tree using a list of variable definitions with values. 
evalExpr :: Expr -> [(Variable,Complex Double)] -> Maybe (Complex Double)
evalExpr e m = eval (M.fromAscList $ caseMap m) (Just e)
  where caseMap x = fmap (\(a,b)->(map toLower a, b)) x

-- |@evalString@ evaluates a string-expression using a list of variable definitions with values. 
evalString :: String -> [(Variable,Complex Double)] -> Maybe (Complex Double)
evalString s m = eval (M.fromAscList $ caseMap m) (stringToExpr s)
  where caseMap x = fmap (\(a,b)->(map toLower a, b)) x

-- |@stringToExpr@ parses a string-expression and returns a maybe expression tree.
stringToExpr :: String -> Maybe Expr
stringToExpr xs =
  if null xs || any (==True) (symbols failingSymbols xs)
  then Nothing
  else either (const Nothing) (Just) (parse buildExpr "" handleString)
  where
    handleString   = "(" ++ (map toLower $ filter (/=' ') xs) ++ ")"
    symbols [] y   = []
    symbols x  y   = [isInfixOf (head x) y] ++ (symbols (drop 1 x) y)
    failingSymbols = [
      "^^","^*","^/","^+","^-","*^","**","*/","*+","*-",
      "/^","/*","//","/+","/-","+^","+*","+/","++","+-",
      "-^","-*","-/","-+","--","()"]

buildExpr :: Parser Expr
buildExpr = buildExpressionParser expressionTable factor

expressionTable :: [[Operator Char st Expr]]
expressionTable =  [
  [pr "arcsinh" ArcSinh, pr "arcsin" ArcSin, pr "sinh" Sinh, pr "sin" Sin],
  [pr "arccosh" ArcCosh, pr "arccos" ArcCos, pr "cosh" Cosh, pr "cos" Cos],
  [pr "arctanh" ArcTanh, pr "arctan" ArcTan, pr "tanh" Tanh, pr "tan" Tan],
  [pr "arcsech" ArcSech, pr "arcsec" ArcSec, pr "sech" Sech, pr "sec" Sec],
  [pr "arccsch" ArcCsch, pr "arccsc" ArcCsc, pr "csch" Csch, pr "csc" Csc],
  [pr "arccoth" ArcCoth, pr "arccot" ArcCot, pr "coth" Coth, pr "cot" Cot],
  [pr "log" Log, pr "abs" Abs,pr "exp" Exp],
  [pr "sqrt" Sqrt, pr "cbrt" Cbrt],
  [op "^" Pow AssocRight],
  [op "*" Mul AssocLeft, op "/" Div AssocLeft],
  [op "+" Add AssocLeft, op "-" Sub AssocLeft]]
  where
    op s f assoc = Infix  (do{ string s; return f}) assoc
    pr s f       = Prefix (try (string s) >> return f)

factor :: Parser Expr
factor = do
  char '('
  e <- buildExpr
  char ')'
  return e
  <|> variables

variables :: Parser Expr
variables = do
  ds <- many1 letter
  return $ Var ds
  <|> number

number :: Parser Expr
number = do
    br  <- many digit
    let d :: Double
        d = fromInteger (foldl ((. ch2num) . (+) . (*10)) 0 br)
    option (Num (d)) (try (do
        char '.'
        ar <- many1 digit
        return $ (Num (d + foldr (fd) 0 ar)) ))
        where
            fd a b = (fromInteger (ch2num a) + b) / 10
            fe = toInteger . fromEnum
            ch2num = (subtract $ fe '0') . fe

-- |@eval@ takes a map of variable definitions and values, and a maybe expression tree, to produce maybe a numerical value.
eval :: M.Map Variable (Complex Double) -> Maybe Expr -> Maybe (Complex Double)
eval m expr =
  case expr of
    Just (Num d)      -> Just $ d :+ 0
    Just (Var "pi")   -> Just $ pi
    Just (Var "i")    -> Just $ 0 :+ 1
    Just (Var "e")    -> Just $ exp 1
    Just (Var c)      -> M.lookup c m
    Just (Add e1 e2)  -> liftM2 (+)  (eval m $ Just e1) (eval m $ Just e2)
    Just (Sub e1 e2)  -> liftM2 (-)  (eval m $ Just e1) (eval m $ Just e2)
    Just (Mul e1 e2)  -> liftM2 (*)  (eval m $ Just e1) (eval m $ Just e2)
    Just (Div e1 e2)  -> liftM2 (/)  (eval m $ Just e1) (eval m $ Just e2)
    Just (Pow e1 e2)  -> liftM2 (**) (eval m $ Just e1) (eval m $ Just e2)
    Just (Exp e1)     -> liftM (exp) (eval m $ Just e1)
    Just (Sqrt e1)    -> liftM (\x->x**(0.5)) (eval m $ Just e1)
    Just (Cbrt e1)    -> liftM (\x->x**(1/3)) (eval m $ Just e1)
    Just (Log e1)     -> liftM (log) (eval m $ Just e1)
    Just (Abs e1)     -> liftM (abs) (eval m $ Just e1)
    Just (Sin e1)     -> liftM (sin) (eval m $ Just e1)
    Just (Cos e1)     -> liftM (cos) (eval m $ Just e1)
    Just (Tan e1)     -> liftM (tan) (eval m $ Just e1)
    Just (Sec e1)     -> liftM (\x->1/sin x) (eval m $ Just e1)
    Just (Csc e1)     -> liftM (\x->1/cos x) (eval m $ Just e1)
    Just (Cot e1)     -> liftM (\x->1/tan x) (eval m $ Just e1)
    Just (Sinh e1)    -> liftM (sinh) (eval m $ Just e1)
    Just (Cosh e1)    -> liftM (cosh) (eval m $ Just e1)
    Just (Tanh e1)    -> liftM (tanh) (eval m $ Just e1)
    Just (Sech e1)    -> liftM (\x->1/sinh x) (eval m $ Just e1)
    Just (Csch e1)    -> liftM (\x->1/cosh x) (eval m $ Just e1)
    Just (Coth e1)    -> liftM (\x->1/tanh x) (eval m $ Just e1)
    Just (ArcSin e1)  -> liftM (asin) (eval m $ Just e1)
    Just (ArcCos e1)  -> liftM (acos) (eval m $ Just e1)
    Just (ArcTan e1)  -> liftM (atan) (eval m $ Just e1)
    Just (ArcSec e1)  -> liftM (\x->1/asin x) (eval m $ Just e1)
    Just (ArcCsc e1)  -> liftM (\x->1/acos x) (eval m $ Just e1)
    Just (ArcCot e1)  -> liftM (\x->1/atan x) (eval m $ Just e1)
    Just (ArcSinh e1) -> liftM (asinh) (eval m $ Just e1)
    Just (ArcCosh e1) -> liftM (acosh) (eval m $ Just e1)
    Just (ArcTanh e1) -> liftM (atanh) (eval m $ Just e1)
    Just (ArcSech e1) -> liftM (\x->1/asinh x) (eval m $ Just e1)
    Just (ArcCsch e1) -> liftM (\x->1/acosh x) (eval m $ Just e1)
    Just (ArcCoth e1) -> liftM (\x->1/atanh x) (eval m $ Just e1)
    _                      -> Nothing