{-# OPTIONS_GHC -O #-}
{- The -O option is necessary to get the rewrite rules which use more
   efficient marshaling for certain structured types 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. 
-}

module Main where

import Foreign.MathLink
import qualified Foreign.MathLink.Expressible as Ex

import Data.Array.Unboxed
import Data.Complex
import Control.Monad.Error (throwError)
import Control.Monad.Trans (liftIO)
import System.Timeout

-- Example 1: get a pair of integers and return their sum.
addTwo :: IO ()
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 :: IO ()
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.
-- The first gives a list of single-character strings.
-- The second, an atomic string.
greetWorld :: IO ()
greetWorld = do
  put ("Hello, world!", Ex.ExString "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 :: IO ()
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 :: IO ()
tweakExpression = do
  expr <- get
  put $ tweak expr
  where tweak ex = 
          case ex of
            Ex.ExInt i            -> Ex.ExInt (-i)
            Ex.ExReal r           -> Ex.ExReal (-r)
            Ex.ExString s         -> Ex.ExString $ reverse s
            Ex.ExSymbol s         -> Ex.ExSymbol $ reverse s
            Ex.ExFunction hd args -> Ex.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 Ex.Expressible ExtendedComplex where
    toExpression Infinity = 
        Ex.ExSymbol "Infinity"
    toExpression (Finite (r :+ i)) = 
        Ex.ExFunction "Complex" [Ex.ExReal r,Ex.ExReal i]
    fromExpression expr =
      case expr of
        Ex.ExFunction "DirectedInfinity" [_] -> 
          Right $ Infinity
        Ex.ExFunction "Complex" [Ex.ExReal r, Ex.ExReal i] -> 
          Right $ Finite (r :+ i)
        _ -> Left $ Ex.ExpressibleErrorMsg $ 
          "Unexpected expression: " ++ show expr

addExtendedComplexes :: IO ()
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
           }

-- Example 7: check for abort
abortTest :: IO ()
abortTest = do
  n <- (get :: IO Int)
  liftIO $ timeout (100000*n) (bottom)
  checkAbort >>= put
  where bottom = (bottom :: IO ())

abortTestFunction = 
  Function { callPattern = "AbortTest[i_Integer]"
           , argumentPattern = "i"
           , function = abortTest
           }

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