module Language.Sifflet.ExprTree
    (
      ExprTree, ExprNode(..), ExprNodeLabel(..)
    , exprNodeIoletCounter -- needs work ****** get rid of it???
    , exprToTree, treeToExpr, exprToReprTree
    , evalTree, unevalTree
    )

where

import Data.Number.Sifflet

import Data.Sifflet.Tree as T
import Data.Sifflet.TreeLayout (IoletCounter)

import Language.Sifflet.Expr

import Text.Sifflet.Repr ()

import Language.Sifflet.Util


-- | 
-- EXPRESSION TREES
-- For pure Sifflet, so not defined for extended expressions.

type ExprTree = Tree ExprNode
data ExprNode = ENode ExprNodeLabel EvalResult
              deriving (Eq, Show)

data ExprNodeLabel = NUndefined 
                   | NSymbol Symbol 
                   |
                     -- formerly NLit Value
                     NBool Bool | NChar Char | NNumber Number 
                   | NString String
                   | NList [Expr] -- ???
              deriving (Eq, Show)

instance Repr ExprNode where
    reprl (ENode label evalRes) =
        case label of
          NUndefined ->
              case evalRes of
                EvalUntried -> ["undefined"]
                EvalError e -> ["undefined", "error: " ++ e]
                EvalOk _ -> 
                    errcats ["reprl of ExprNode: NUndefined with EvalOk",
                             "should not happen!"]
          NSymbol s ->
              case evalRes of
                EvalOk v -> [repr s, repr v]
                EvalError e -> [repr s, "error: " ++ e]
                EvalUntried -> reprl s

          -- NLit l -> reprl l
          NBool b -> reprl b
          NChar c -> reprl c
          NNumber n -> reprl n
          NString s -> [show s]
          NList es -> reprl (EList es) -- check ***

-- | Counts the number of inlets and outlets for an ExprNode
exprNodeIoletCounter :: Env -> [ArgSpec] -> IoletCounter ExprNode
exprNodeIoletCounter env aspecs (ENode nodeLabel _nodeResult) =
    case nodeLabel of
      NUndefined -> (0, 1)
      NSymbol (Symbol "if") -> (3, 1) 
      NSymbol (Symbol s) -> 
          case envLookup env s of
            Nothing ->    -- probably a parameter of the function
                case aspecsLookup s aspecs of
                  Nothing -> (0, 1)
                  Just i -> (i, 1)
            Just value ->
                case value of
                  VFun function -> (functionNArgs function, 1)
                  _ -> (0, 1)   -- symbol bound to non-function value
      _ -> (0, 1)

exprToTree :: Expr -> ExprTree
exprToTree expr =
    let leafnode :: ExprNodeLabel -> T.Tree ExprNode
        leafnode e = node e []
        node :: ExprNodeLabel -> [T.Tree ExprNode] -> T.Tree ExprNode
        node e ts = T.Node (ENode e EvalUntried) ts
        errext = error ("exprToTree: extended expr: " ++ show expr)
    in case expr of
         -- EUndefined, ESymbol, and literals map directly 
         -- to NUndefined, NSymbol, E(literal-type) 
         EUndefined -> leafnode NUndefined
         ESymbol s -> leafnode (NSymbol s)

         -- Literals
         EBool b -> leafnode (NBool b)
         EChar c -> leafnode (NChar c)
         ENumber n -> leafnode (NNumber n)
         EString s -> leafnode (NString s)

         -- EIf maps to symbol "if" at the root, 3 subtrees
         EIf t a b -> node (NSymbol (Symbol "if")) (map exprToTree [t, a, b])

         ELambda _x _body -> error "exprToTree: not implemented for lambda expr"
         EApp f arg -> node (NSymbol (Symbol "@")) (map exprToTree [f, arg])
         -- ECall maps to symbol f (function name) at the root,
         -- each argument forms a subtree
         ECall f args -> node (NSymbol f) (map exprToTree args)
         EList xs -> leafnode (NList xs)
         -- Extended Exprs not supported!
         EGroup _ -> errext
         EOp _ _ _ -> errext

-- | Convert an expression tree (back) to an expression
-- It will not give back the *same* expression in the case of an EList.
treeToExpr :: ExprTree -> SuccFail Expr
treeToExpr (T.Node (ENode label _) trees) =
    let lit e = if null trees then Succ e
                    else Fail "literal node with non-empty subtrees"
    in case label of
         NUndefined -> Succ EUndefined
         NBool b -> lit (EBool b)
         NChar c -> lit (EChar c)
         NNumber n -> lit (ENumber n)
         NString s -> lit (EString s)
         NList xs -> lit (EList xs)
         NSymbol (Symbol "@") ->
             case trees of
               [f, arg] -> 
                   do
                     {
                       f' <- treeToExpr f
                     ; arg' <- treeToExpr arg
                     ; Succ $ EApp f' arg'
                     }
               _ -> Fail "'@' node with /= 2 subtrees"
         NSymbol (Symbol "if") ->
             case trees of
               [q, a, b] -> 
                   do
                     {
                       q' <- treeToExpr q
                     ; a' <- treeToExpr a
                     ; b' <- treeToExpr b
                     ; Succ $ EIf q' a' b'
                     }
               _ -> Fail ("An 'if' node has the wrong number of subtrees" ++
                          " (should be 3)")
         NSymbol s ->
             -- VVV Do I really need to distinguish these two cases?
             if null trees 
             then 
                 -- s = terminal symbol
                 Succ $ ESymbol s 
             else -- s = function symbol in function call
                 do
                   {
                     trees' <- mapM treeToExpr trees
                   ; Succ $ ECall s trees'
                   }

-- Convert an expression to a repr tree (of string elements)
-- (Why?)

exprToReprTree :: Expr -> Tree String
exprToReprTree = fmap repr . exprToTree


-- Evaluate an expression tree showing the evaluation at each node.
-- There's a lot of redundancy in this computation, but does it matter?

evalTree :: ExprTree -> Env -> ExprTree
evalTree atree env = evalTreeWithLimit atree env stackSize

evalTreeWithLimit :: ExprTree -> Env -> Int -> ExprTree
evalTreeWithLimit atree env stacksize =

    let T.Node root subtrees = atree
        ss' = pred stacksize
    in case root of
         ENode (NSymbol (Symbol "if")) _ ->
             case subtrees of
               [tt, ta, tb] ->
                   let tt' = evalTreeWithLimit tt env ss'
                       ENode _ testResult = rootLabel tt'
                       subEval subtree =
                           let subtree' = evalTreeWithLimit subtree env ss'
                               ENode _ subresult = rootLabel subtree'
                           in (subresult, subtree')
                       ifNode result = ENode (NSymbol (Symbol "if")) result
                   in case testResult of
                        EvalOk (VBool True) -> 
                            let (taValue, ta') = subEval ta in
                            T.Node (ifNode taValue) [tt', ta', tb]

                        EvalOk (VBool False) -> 
                            let (tbValue, tb') = subEval tb in
                            T.Node (ifNode tbValue) [tt', ta, tb']

                        EvalOk weirdValue ->
                            -- This shouldn't happen with proper type
                            -- checking!
                            let msg = "if: non-boolean condition value: " ++
                                      repr weirdValue
                            in T.Node (ifNode (EvalError msg)) [tt', ta, tb]

                        EvalError msg ->
                            T.Node (ifNode (EvalError msg)) [tt', ta, tb]

                        _ -> errcats ["evalTreeWithLimit (if):",
                                      "unexpected test result"]

               _ -> error "evalTreeWithLimit: if: wrong number of subtrees"

         ENode rootOper _ ->
             let evalResult =
                     case treeToExpr atree of
                       Succ expr -> evalWithLimit expr env ss'
                       Fail msg -> EvalError msg
             in T.Node (ENode rootOper evalResult)
                       [evalTreeWithLimit s env ss' | s <- subtrees]

-- remove the values from the ExprNodes
-- "inverse" of evalTree 
unevalTree :: ExprTree -> ExprTree
unevalTree atree = 
    let unevalNode (ENode oper _) = ENode oper EvalUntried
    in fmap unevalNode atree