{-# 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(..)

    ) where

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

import Foreign.MathLink.Internal

import Data.Ix.Shapable
import Data.Array.Unboxed

-- | 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.
            | MLExpr :@ [MLExpr]   -- ^ Compound expression.
              deriving (Eq, Ord, Show)

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

-- | Sugar for the @MessageName@ /Mathematica/ expression
(-::-) :: String -> String -> MLExpr
sym -::- tag = (Sy "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 -> ML 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 -> ML ()
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 -> ML ()
sendString str = send $ (Sy "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.
    -> ML ()
defineMessage sym tag defn = 
  send (sym-::-tag -=- St defn)

-- | Causes the specified message to be sent to /Mathematica/.
--
--   Making the call (corresponding to 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.
  -> ML ()
printMessage sym tag args = do
  v <- evaluate ((Sy "Message"):@[sym-::-tag,(Sy "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 -> ML ()

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 
    putType $ MLType TypFn
    putArgCount (length exs)
    mapM_ mlPut (hd: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

putArrayWith :: ( Ix ix
                , Shapable ix
                , IArray a e
                ) => ([e] -> [(Int,String)] -> ML ())
                  -> a ix e -> ML ()
putArrayWith fn ary = fn xs dims
  where xs = elems ary
        dims = zip (shape ary) (repeat "List")

instance ( Ix ix
         , Shapable ix
         ) => MLPut (UArray ix Int16) where
  mlPut = putArrayWith putInt16Array

instance ( Ix ix
         , Shapable ix
         ) => MLPut (UArray ix Int32) where
  mlPut = putArrayWith putInt32Array

instance ( Ix ix
         , Shapable ix
         ) => MLPut (UArray ix Int) where
  mlPut = putArrayWith putIntArray

instance ( Ix ix
         , Shapable ix
         ) => MLPut (UArray ix Float) where
  mlPut = putArrayWith putFloatArray

instance ( Ix ix
         , Shapable ix
         ) => MLPut (UArray ix Double) where
  mlPut = putArrayWith putDoubleArray

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 (Maybe a) where
  mlPut (Just v) = do putFunction "Just" 1
                      mlPut v
  mlPut Nothing  = putSymbol "Nothing"

instance (MLPut a, MLPut b) => MLPut (Either a b) where
  mlPut (Left v) = do putFunction "Left" 1
                      mlPut v
  mlPut (Right v) = do putFunction "Right" 1
                       mlPut v

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 :: ML 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 nargs <- getArgCount
                          hd    <- mlGet
                          exs   <- replicateM nargs mlGet
                          return $ hd:@exs
      MLType TypErr -> throwMsg "mlGet/Expression" "type error"
      _             -> 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

getArrayWith :: ( Ix ix
                , Shapable ix
                , IArray a e
                ) => ML ([e], [(Int,String)])
                  -> ML (a ix e)
getArrayWith fn = do
  (xs,dims) <- fn
  return $ listArray (sBounds $ fst $ unzip dims) xs

instance ( Ix ix
         , Shapable ix
         ) => MLGet (UArray ix Int16) where
  mlGet = getArrayWith getInt16Array

instance ( Ix ix
         , Shapable ix
         ) => MLGet (UArray ix Int32) where
  mlGet = getArrayWith getInt32Array

instance ( Ix ix
         , Shapable ix
         ) => MLGet (UArray ix Int) where
  mlGet = getArrayWith getIntArray

instance ( Ix ix
         , Shapable ix
         ) => MLGet (UArray ix Float) where
  mlGet = getArrayWith getFloatArray

instance ( Ix ix
         , Shapable ix
         ) => MLGet (UArray ix Double) where
  mlGet = getArrayWith getDoubleArray

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 (Maybe a) where
  mlGet = do
    typ <- getType
    case typ of
      MLType TypSy -> do s <- getSymbol
                         if s == "Nothing" then 
                             return Nothing
                           else
                             throwMsg 
                               "mlGet/Maybe a" "unexpected symbol"
      MLType TypFn -> do testFunction (== "Just") (== 1)
                         mlGet >>= (return . Just)
      _ -> throwMsg "mlGet/Maybe a" "unexpected type"

instance (MLGet a, MLGet b) => MLGet (Either a b) where
  mlGet = do (hd,_) <- testFunction (const True) (== 1)
             case hd of
               "Left" -> mlGet >>= (return . Left)
               "Right" -> mlGet >>= (return . Right)
               _ -> throwMsg "mlGet/Either a b" "unexpected head"

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)