{-# OPTIONS_GHC -O #-}
{- The -O option is necessary to get the rewrite rules which use more
   efficient marshaling for certain types of lists to fire. According
   to the GHC documentation, it should be sufficient to set
   -fenable-rewrite-rules, but this didn't work in the test cases I
   examined, while -O (which implies -fenable-rewrite-rules) did. 

   In the case of 'String,' there is not just an efficiency difference,
   but a semantic one, as well. Without the rule rewrite, the string is
   marshaled out as a list of characters, which yields a list of single-
   character strings on the Mathematica side (and probably breaks when
   trying a get). In the absence of the rewrite rules for strings,
   there are two alternatives: (1) marshal strings as a generic
   Expression (or part of an expression), or, for just a single string
   in isolation you can (2) use the putString/getString exported from
   the Expressible module.
-}

module Main where

import Foreign.MathLink
import Data.Array.Unboxed
import Data.Complex
import Control.Monad.Error (throwError)

-- Example 1: get a pair of integers and return their sum.
addTwo :: ML ()
addTwo = do
  (m1,m2) <- get
  put ((m1 + m2) :: Int)

-- callPattern <=> mprep's :Pattern: directive
-- argumentPattern <=> mprep's :Arguments: directive
-- Note that, in argumentPattern, the desired 2-tuple of integers
-- is represented as a Mathematica list of length 2.
addTwoFunction = 
  Function { callPattern = "AddTwo[i_Integer,j_Integer]"
           , argumentPattern = "{i,j}"
           , function = addTwo
           }

-- Example 2: reverse a list of numbers (coercing them to Doubles).
reverseNumbers :: ML ()
reverseNumbers = do
  is <- get
  put $ reverse (is :: [Double])

-- Like tuples, lists are represented on the Mathematica side as
-- lists, but of arbitrary length. Here, the Mathematica pattern
-- would accept a list of zero or more numeric values.
-- Notice that argumentPattern provides some preprocessing on the
-- arguments, taking its real part and coercing the values to
-- machine precision. This trick allows for a callPattern that
-- matches more expressions, while still ensuring that on the
-- Haskell side, the arguments have the desired, more specific,
-- form when marshaled.
reverseNumbersFunction = 
  Function { callPattern = "ReverseNumbers[is___?NumericQ]"
           , argumentPattern = "N[Re[{is}]]"
           , function = reverseNumbers
           }

-- Example 3: put a string.
-- It's a good test to see if the rewrite rules are firing: if the
-- rule fired, you should see a single string in Mathematica,
-- otherwise, a list of single-character strings.
greetWorld :: ML ()
greetWorld = do
  put "Hello, world!"

greetWorldFunction = 
  Function { callPattern = "GreetWorld[]"
           , argumentPattern = ""
           , function = greetWorld
           }

-- Example 4: get and put a 2D array, reversing its contents.
-- NB: When marshaling in an array, make sure that you test for the
-- array rank in callPattern to match that expected on the Haskell
-- side, as below. Otherwise, you're likely to have the underlying 
-- call to fromDimensions raise an error that halts the program.
reverseArray :: ML ()
reverseArray = do
  arr <- get
  let xs = elems (arr :: UArray (Int,Int) Int)
  put $ ((listArray (bounds arr) $ reverse xs) :: UArray (Int,Int) Int)

reverseArrayFunction = 
  Function { callPattern = "ReverseArray[a_?(ArrayQ[#,2,IntegerQ]&)]"
           , argumentPattern = "a"
           , function = reverseArray
           }

-- Example 5: gets an arbitrary Mathematica expression and returns 
-- a tweaked version of it.
tweakExpression :: ML ()
tweakExpression = do
  expr <- get
  put $ tweak expr
  where tweak ex = 
          case ex of
            ExInt i            -> ExInt (-i)
            ExReal r           -> ExReal (-r)
            ExString s         -> ExString $ reverse s
            ExSymbol s         -> ExSymbol $ reverse s
            ExFunction hd args -> ExFunction (reverse hd) $ map tweak args

tweakExpressionFunction = 
  Function { callPattern = "TweakExpression[expr_]"
           , argumentPattern = "expr"
           , function = tweakExpression
           }

-- Example 6: use your own instance of Expressible
data ExtendedComplex = Finite (Complex Double) 
                     | Infinity
                       deriving (Eq,Show)

instance Expressible ExtendedComplex where
    put Infinity          = put $ ExSymbol "Infinity"
    put (Finite (r :+ i)) = put $ ExFunction "Complex" [ExReal r,ExReal i]
    get = do
      expr <- get
      case (expr :: Expression) of
        ExFunction "DirectedInfinity" [_]         -> 
          return Infinity
        ExFunction "Complex" [ExReal r, ExReal i] -> 
          return $ Finite (r :+ i)
        _                                         -> 
          throwError $ "Unexpected expression: " ++ show expr

addExtendedComplexes :: ML ()
addExtendedComplexes = do
  (ec1,ec2) <- get
  case (ec1,ec2) of
    (Infinity,_) -> put Infinity
    (_,Infinity) -> put Infinity
    (Finite c1,Finite c2) -> put $ Finite (c1 + c2)

addExtendedComplexesFunction = 
  Function { callPattern = "AddExtendedComplexes[ec1:(Infinity|Complex[_,_]),ec2:(Infinity|Complex[_,_])]"
           , argumentPattern = "N[{ec1,ec2}]"
           , function = addExtendedComplexes
           }

-- run mathlink exposing the given list of actions
main = runMathLink [ addTwoFunction
                   , reverseNumbersFunction
                   , greetWorldFunction
                   , reverseArrayFunction
                   , tweakExpressionFunction
                   , addExtendedComplexesFunction
                   ]