{-# LANGUAGE RankNTypes
           , FlexibleContexts
           , FlexibleInstances
           , TypeSynonymInstances
           , OverlappingInstances
  #-}
module Foreign.MathLink.Expression 
    ( 
    -- * Generic /Mathematica/ expression
      MLExpr(..)
    , (-=-)
    , (-::-)

    -- * /Mathematica/ evaluation
    , send
    , sendString
    , evaluate
    , defineMessage
    , printMessage

    -- * Data marshaling 
    , MLPut(..)
    , MLGet(..)
    , Packable
    , putArray
    , getArray

    ) where

import Control.Monad
import Data.Word
import Data.Int
import Data.Complex
import GHC.Real

import Foreign.MathLink.Internal
import Foreign.Storable

import Data.Array.IArray
import Data.Ix.Shapable
import Data.Array (Array)
import Data.Array.Unboxed (UArray)

import System.IO

-- | Haskell representation of a generic /Mathematica/ expression.
data MLExpr = I Integer            -- ^ Atomic integer value.
            | R Rational           -- ^ Atomic real value.
            | St String            -- ^ Atomic string value, as a string.
            | Sy String            -- ^ Atomic string value, as a symbol.
            | String :@ [MLExpr]   -- ^ Compound expression.
              deriving (Eq, Ord, Show)

-- | Sugar for the @Set@ /Mathematica/ expression
(-=-) :: MLExpr -> MLExpr -> MLExpr
ex1 -=- ex2 = "Set":@[ex1,ex2]

-- | Sugar for the @MessageName@ /Mathematica/ expression
(-::-) :: String -> String -> MLExpr
sym -::- tag = "MessageName":@[Sy sym, St tag]


------------------- Evaluation ---------------------

-- | Sends a value (most likely an 'MLExpr') to /Mathematica/ for
--   evaluation and gets the result (appropriately marshaled to the
--   desired type).
--
--   Does not call 'endPacket'.
evaluate :: (MLPut a, MLGet b) => a -> IO b
evaluate v = do
  send v
  waitForPacket (== (MLPacket PktReturn))
  res <- mlGet
  newPacket
  return res

-- | Puts a value to the /MathLink/ connection.
--
--   Does not wait for a return packet or call 'endPacket'.
send :: MLPut a => a -> IO ()
send v = do
  putFunction "EvaluatePacket" 1
  mlPut v

-- | Puts a value, as a string, to the /MathLink/ connection.
--
--   Does not wait for a return packet or call 'endPacket'.
sendString :: String -> IO ()
sendString str = send $ "ToExpression":@[St str]

-- | Defines a message.
--
--   The message definition can contain backquoted numbers to
--   represent the positions of where string representations of
--   the corresponding extra values passed in a call to 
--   'printMessage' are to be interpolated. An example of a message 
--   definition that interpolates two values is:
--
--   @
--   'defineMessage' \"Foo\" \"bar\" \"Thingy `1` doesn't go with thingy `2`.\"
--   @
--
--   Does not wait for a return packet or call 'endPacket'.
defineMessage 
    :: String -- ^ The symbol associated with this message.
    -> String -- ^ The tag associated with this message.
    -> String -- ^ The message definition.
    -> IO ()
defineMessage sym tag defn = 
  send (sym-::-tag -=- St defn)

-- | Causes the specified message to be sent to /Mathematica/.
--
--   Making the call (corresponding the the 'defineMessage' example):
--
--   @
--   'printMessage' \"Foo\" \"bar\" ['St' \"fish\", 'St' \"bicycle\"]
--   @
--
--   would yield a message to appear on the /Mathematica/ end looking
--   something like:
--
--   @
--   Foo::bar : Thingy fish doesn't go with thingy bicycle.
--   @
--
--   Does not call 'endPacket'.
printMessage 
  :: String    -- ^ The symbol associated with the message.
  -> String    -- ^ The tag associated with the message.
  -> [MLExpr]  -- ^ A list of expressions whose string values are
               --   to be interpolated in to the message.
  -> IO ()
printMessage sym tag args = do
  v <- evaluate ("Message":@[sym-::-tag,"Sequence":@args])
  (v :: MLExpr) `seq` return ()

--------------------- MLPut -------------------------

-- | The class of types that can be marshaled to /Mathematica/.
class MLPut a where
  -- | Send a value to /Mathematica/.
  mlPut :: a -> IO ()

instance MLPut MLExpr where
  mlPut (I i)    = mlPut i
  mlPut (R r)    = mlPut r
  mlPut (St str) = mlPut str
  mlPut (Sy sym) = putSymbol sym
  mlPut (hd :@ exs) = do 
    putFunction hd (length exs)
    mapM_ mlPut exs

instance MLPut Bool where
  mlPut True = putSymbol "True"
  mlPut False = putSymbol "False"

instance MLPut Int16 where
  mlPut = putInt16

instance MLPut Int32 where
  mlPut = putInt32

instance MLPut Int where
  mlPut = putInt

instance MLPut Integer where
  mlPut i = do putFunction "ToExpression" 1
               putString $ show i

instance MLPut Float where
  mlPut = putFloat

instance MLPut Double where
  mlPut = putDouble

instance MLPut String where
  mlPut = putString

instance MLPut [Int16] where
  mlPut = putInt16List

instance MLPut [Int32] where
  mlPut = putInt32List

instance MLPut [Int] where
  mlPut = putIntList

instance MLPut [Float] where
  mlPut = putFloatList

instance MLPut [Double] where
  mlPut = putDoubleList

instance (RealFloat a, MLPut a) => MLPut (Complex a) where
  mlPut (r :+ i) = do putFunction "Complex" 2
                      mlPut r
                      mlPut i

instance (Integral a, MLPut a) => MLPut (Ratio a) where
  mlPut (n :% d) = do putFunction "Rational" 2
                      mlPut n
                      mlPut d

instance (MLPut a) => MLPut [a] where
  mlPut xs = do putFunction "List" (length xs)
                mapM_ mlPut xs

instance MLPut () where
  mlPut () = mlPut $ Sy "Null"

instance ( MLPut t1
         , MLPut t2
         ) => MLPut (t1,t2) where
  mlPut (v1,v2) = do
      putFunction "List" 2
      mlPut v1
      mlPut v2

instance ( MLPut t1
         , MLPut t2
         , MLPut t3
         ) => MLPut (t1,t2,t3) where
  mlPut (v1,v2,v3) = do
      putFunction "List" 3
      mlPut v1
      mlPut v2
      mlPut v3

instance ( MLPut t1
         , MLPut t2
         , MLPut t3
         , MLPut t4
         ) => MLPut (t1,t2,t3,t4) where
  mlPut (v1,v2,v3,v4) = do
      putFunction "List" 4
      mlPut v1
      mlPut v2
      mlPut v3
      mlPut v4

instance ( MLPut t1
         , MLPut t2
         , MLPut t3
         , MLPut t4
         , MLPut t5
         ) => MLPut (t1,t2,t3,t4,t5) where
  mlPut (v1,v2,v3,v4,v5) = do
      putFunction "List" 5
      mlPut v1
      mlPut v2
      mlPut v3
      mlPut v4
      mlPut v5

instance ( MLPut t1
         , MLPut t2
         , MLPut t3
         , MLPut t4
         , MLPut t5
         , MLPut t6
         ) => MLPut (t1,t2,t3,t4,t5,t6) where
  mlPut (v1,v2,v3,v4,v5,v6) = do
      putFunction "List" 6
      mlPut v1
      mlPut v2
      mlPut v3
      mlPut v4
      mlPut v5
      mlPut v6

instance ( MLPut t1
         , MLPut t2
         , MLPut t3
         , MLPut t4
         , MLPut t5
         , MLPut t6
         , MLPut t7
         ) => MLPut (t1,t2,t3,t4,t5,t6,t7) where
  mlPut (v1,v2,v3,v4,v5,v6,v7) = do
      putFunction "List" 7
      mlPut v1
      mlPut v2
      mlPut v3
      mlPut v4
      mlPut v5
      mlPut v6
      mlPut v7


-------------------------- MLGet ---------------------------

-- | The class of types that can be marshaled from /Mathematica/.
class MLGet a where
  -- | Get a value from /Mathematica/.
  mlGet :: IO a

instance MLGet MLExpr where
  mlGet = do
    typ <- getType
    case typ of
      MLType TypI   -> mlGet     >>= (return . I)
      MLType TypR   -> mlGet     >>= (return . R)
      MLType TypSt  -> mlGet     >>= (return . St)
      MLType TypSy  -> getSymbol >>= (return . Sy)
      MLType TypFn  -> do (hd,nargs) <- getFunction
                          exs <- replicateM nargs mlGet
                          return $ hd:@exs
      MLType TypErr -> throwMLError
      _             -> throwMsg "mlGet/Expression: unexpected type"

instance MLGet Bool where
  mlGet = do
    str <- getSymbol
    case str of
      "True"  -> return True
      "False" -> return False
      _       -> throwMsg "mlGet/Bool: unexpected symbol"

instance MLGet Int16 where
  mlGet = getInt16

instance MLGet Int32 where
  mlGet = getInt32

instance MLGet Int where
  mlGet = getInt

instance MLGet Integer where
  mlGet = getString >>= (return . read)

instance MLGet Float where
  mlGet = getFloat

instance MLGet Double where
  mlGet = getDouble

instance MLGet String where
  mlGet = getString

instance MLGet [Int16] where
  mlGet = getInt16List

instance MLGet [Int32] where
  mlGet = getInt32List

instance MLGet [Int] where
  mlGet = getIntList

instance MLGet [Float] where
  mlGet = getFloatList

instance MLGet [Double] where
  mlGet = getDoubleList

instance (RealFloat a, MLGet a) => MLGet (Complex a) where
  mlGet = do
      typ <- getType
      case typ of
        MLType TypI  -> fromReal
        MLType TypFn -> fromFunc
        _            -> throwMsg "mlGet/Complex: unexpected type"
    where fromReal = do
            r <- mlGet
            return $ r :+ 0

          fromFunc = do
            testFunction (== "Complex") (== 2)
            r <- mlGet
            i <- mlGet
            return (r :+ i)

instance (Integral a, MLGet a) => MLGet (Ratio a) where
  mlGet = do
      typ <- getType
      case typ of
        MLType TypI  -> fromInt
        MLType TypR  -> fromReal
        MLType TypFn -> fromFunc
        _            -> throwMsg "mlGet/Ratio: unexpected type"
    where fromInt = do 
            i <- mlGet
            return $ i :% 1

          fromReal = do
            r <- mlGet
            let (m,n) = decodeFloat (r :: Double)
                b     = floatRadix r
            return $ (fromIntegral m) * (fromIntegral b)^^n
          
          fromFunc = do 
            testFunction (== "Rational") (== 2)
            n <- mlGet
            d <- mlGet
            return (n :% d)

instance (MLGet a) => MLGet [a] where
  mlGet = do (_,nargs) <- testFunction (== "List") (>= 0)
             replicateM nargs mlGet

instance MLGet () where
  mlGet = do
      typ <- getType
      case typ of
        (MLType TypSy) -> do 
          sy <- getSymbol
          if sy == "Null" then 
              return () 
            else
              throwMsg "mlGet/(): unexpected symbol"
        (MLType TypFn) -> do 
          testFunction (== "List") (== 0)
          return ()
        _ -> throwMsg "mlGet/(): unexpected type"

instance ( MLGet t1
         , MLGet t2
         ) => MLGet (t1,t2) where
  mlGet = do testFunction (== "List") (== 2)
             v1 <- mlGet
             v2 <- mlGet
             return (v1,v2)

instance ( MLGet t1
         , MLGet t2
         , MLGet t3
         ) => MLGet (t1,t2,t3) where
  mlGet = do testFunction (== "List") (== 3)
             v1 <- mlGet
             v2 <- mlGet
             v3 <- mlGet
             return (v1,v2,v3)

instance ( MLGet t1
         , MLGet t2
         , MLGet t3
         , MLGet t4
         ) => MLGet (t1,t2,t3,t4) where
  mlGet = do testFunction (== "List") (== 4)
             v1 <- mlGet
             v2 <- mlGet
             v3 <- mlGet
             v4 <- mlGet
             return (v1,v2,v3,v4)

instance ( MLGet t1
         , MLGet t2
         , MLGet t3
         , MLGet t4
         , MLGet t5
         ) => MLGet (t1,t2,t3,t4,t5) where
  mlGet = do testFunction (== "List") (== 5)
             v1 <- mlGet
             v2 <- mlGet
             v3 <- mlGet
             v4 <- mlGet
             v5 <- mlGet
             return (v1,v2,v3,v4,v5)

instance ( MLGet t1
         , MLGet t2
         , MLGet t3
         , MLGet t4
         , MLGet t5
         , MLGet t6
         ) => MLGet (t1,t2,t3,t4,t5,t6) where
  mlGet = do testFunction (== "List") (== 6)
             v1 <- mlGet
             v2 <- mlGet
             v3 <- mlGet
             v4 <- mlGet
             v5 <- mlGet
             v6 <- mlGet
             return (v1,v2,v3,v4,v5,v6)

instance ( MLGet t1
         , MLGet t2
         , MLGet t3
         , MLGet t4
         , MLGet t5
         , MLGet t6
         , MLGet t7
         ) => MLGet (t1,t2,t3,t4,t5,t6,t7) where
  mlGet = do testFunction (== "List") (== 7)
             v1 <- mlGet
             v2 <- mlGet
             v3 <- mlGet
             v4 <- mlGet
             v5 <- mlGet
             v6 <- mlGet
             v7 <- mlGet
             return (v1,v2,v3,v4,v5,v6,v7)


-------------------------- Packable --------------------------

-- | The class of types that can be marshaled to and from
--   /Mathematica/ via packed storage.
class Storable a => Packable a where
    putPackedArray :: [a] -> [(Int,String)] -> IO ()
    getPackedArray :: IO ([a],[(Int,String)])

instance Packable Word8 where
    putPackedArray xs dims = putInt16Array (map fromIntegral xs) dims
    getPackedArray = do (xs,dims) <- getInt16Array
                        return (map fromIntegral xs, dims)

instance Packable Int16 where
    putPackedArray = putInt16Array
    getPackedArray = getInt16Array

instance Packable Int32 where
    putPackedArray = putInt32Array
    getPackedArray = getInt32Array

instance Packable Int where
    putPackedArray = putIntArray
    getPackedArray = getIntArray

instance Packable Float where
    putPackedArray = putFloatArray
    getPackedArray = getFloatArray

instance Packable Double where
    putPackedArray = putDoubleArray
    getPackedArray = getDoubleArray

-- | Send an array to /Mathematica/.
putArray :: ( IArray a e
            , Ix ix 
            , Shapable ix
            , Packable e
            ) => a ix e -> IO ()
putArray arr = putPackedArray xs dims
  where xs = elems arr
        sh = shape arr
        dims = zip sh $ repeat "List"

-- | Get an array from /Mathematica/.
getArray :: ( IArray a e
            , Ix ix
            , Shapable ix
            , Packable e
            ) => IO (a ix e)
getArray = do 
  (xs,dims) <- getPackedArray
  let (shp,_) = unzip dims
      bnds = sBounds shp
  return $ listArray bnds xs

instance (Packable e, Ix ix, Shapable ix) => MLPut (Array ix e) where
  mlPut = putArray

instance (Packable e, Ix ix, Shapable ix) => MLGet (Array ix e) where
  mlGet = getArray

instance ( IArray UArray e
         , Packable e
         , Ix ix
         , Shapable ix 
         ) => MLPut (UArray ix e) where
  mlPut = putArray

instance ( IArray UArray e
         , Packable e
         , Ix ix
         , Shapable ix
         ) => MLGet (UArray ix e) where
  mlGet = getArray