{-# OPTIONS_GHC -O #-}
{- The -O option is necessary to get the rewrite rules which use more
   efficient marshaling for certain structured types to fire. 
-}

module Main where

import Foreign.MathLink
import Foreign.MathLink.Expressible

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 :: (Int,Int) -> IO Int
addTwo (m1,m2) = return (m1+m2)

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

-- Example 2: reverse a list of numbers (coercing them to Doubles).
reverseNumbers :: [Double] -> IO [Double]
reverseNumbers is = return $ reverse is

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

-- Example 3: 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 :: UArray (Int,Int) Int -> IO (UArray (Int,Int) Int)
reverseArray arr = return $ listArray (bounds arr) $ reverse $ elems arr

reverseArrayFunction = 
  mkFunction "ReverseArray[a_?(ArrayQ[#,2,IntegerQ]&)]"
             "a"
             reverseArray


-- Example 4: gets an arbitrary Mathematica expression and returns 
-- a tweaked version of it.
tweakExpression :: Expression -> IO Expression
tweakExpression expr = return $ 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 = 
  mkFunction"TweakExpression[expr_]"
            "expr"
            tweakExpression

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

instance Expressible ExtendedComplex where
    toExpression Infinity = 
        ExSymbol "Infinity"
    toExpression (Finite (r :+ i)) = 
        ExFunction "Complex" [ExReal r, ExReal i]
    fromExpression expr =
      case expr of
        ExFunction "DirectedInfinity" [_] -> 
          Right $ Infinity
        ExFunction "Complex" [ExReal r, ExReal i] -> 
          Right $ Finite (r :+ i)
        _ -> Left $ ExpressibleErrorMsg $ 
          "Unexpected expression: " ++ show expr

addExtendedComplexes :: (ExtendedComplex,ExtendedComplex) -> IO ExtendedComplex
addExtendedComplexes (ec1,ec2) =
  case (ec1,ec2) of
    (Infinity,_) -> return Infinity
    (_,Infinity) -> return Infinity
    (Finite c1,Finite c2) -> return $ Finite (c1 + c2)

addExtendedComplexesFunction = 
  mkFunction "AddExtendedComplexes[\
               \ec1:(Infinity|Complex[_,_]),\
               \ec2:(Infinity|Complex[_,_])]"
             "N[{ec1,ec2}]"
             addExtendedComplexes

-- Example 6: check for abort
abortTest :: Int -> IO Bool
abortTest n = do
  liftIO $ timeout (100000*n) (bottom)
  checkAbort
  where bottom = (bottom :: IO ())

abortTestFunction = 
  mkFunction "AbortTest[i_Integer]"
             "i"
             abortTest

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