-- Pretty printing of an ART expression with a sub-expression highlit.
-- First convert a value (of class HatRep) into an S-expression,
-- then convert the S-expression to a Doc for pretty printing.
module SExp
  ( SExp(..)
  , hatRep2SExp
  , sExp2Doc
  , arity, child, label
  ) where

import HatTrace           (HatRep(..),HatNodeType(..),HatNode,isValidNode
                          ,HatInfixType(HatInfix,HatInfixL,HatInfixR,HatNoInfix)
                          ,HatSourceRef(..))
-- import HatExpressionTree  (HatLimit,toHatLimit)
import PrettyLibHighlight (Doc,text,(<>),delimiter,fdelimiter,nil,group,parens
                          ,groupNest,pretty,highlight)
import Char               (isAlpha)
import List               (unzip3)


data SExp a
  = SApp a [SExp a] -- n-ary application of at least 2 expressions
  | SId a String SFixity
  | SLiteral a String
  | SLambda a
  | SIf a (SExp a)
  | SCase a (SExp a)
  | SGuard a (SExp a)
  | SInvalid a		-- subexpression that was cut off (to limit depth)
  | SUnevaluated a
  | SBottom a
  | SCycle a String (SExp a)
			-- cyclic expression to be shown as `id where id = ..'

data SFixity = 
  SInfix Int | SInfixL Int | SInfixR Int | SAssoc Int String | SInfixDefault
  -- need own type for some hardcoded operators that are known to be
  -- semantically associative

transFixity :: HatInfixType -> SFixity
transFixity (HatInfix pri) = SInfix pri
transFixity (HatInfixL pri) = SInfixL pri
transFixity (HatInfixR pri) = SInfixR pri
transFixity HatNoInfix = SInfixDefault


-- arity of an S-expression
arity :: SExp a -> Int
arity (SApp _ exps) = length exps
arity (SIf _ _)     = 1
arity (SCase _ _)   = 1
arity (SGuard _ _)  = 1
arity _             = 0

-- get child of an S-expression
-- precondition: i < arity exp
child :: Int -> SExp a -> SExp a
child i (SApp _ exps)  = exps!!i
child 0 (SIf _ exp)    = exp
child 0 (SCase _ exp)  = exp
child 0 (SGuard _ exp) = exp
child i exp            = error ("SExp.child: "++show i++" is too large.")

label :: SExp a -> a
label (SApp l _)      = l
label (SId l _ _)     = l
label (SLiteral l _)  = l
label (SLambda l)     = l
label (SIf l _)       = l
label (SCase l _)     = l
label (SGuard l _)    = l
label (SInvalid l)    = l
label (SUnevaluated l)= l
label (SBottom l)     = l
label (SCycle l _ _)  = l


-- conversion function
-- if first boolean is True, then cycles are located and expressed as SCycles
-- if False, there is the danger of obtaining an infinite expression
-- if second boolean is True, then unevaluated arguments appear in result,
-- otherwise they are represented by SUnevaluated.
hatRep2SExp :: HatRep a =>
               Bool -> Bool -> a -> SExp HatNode
hatRep2SExp cyc uneval expObj =
    case go cyc uneval [] expObj of (e,_,_) -> e  

fst3 (x,_,_) = x
snd3 (_,x,_) = x
thd3 (_,_,x) = x

-- worker for definition above
go :: HatRep a =>
      Bool			-- look for cycles?
      -> Bool			-- transform unevaluated args?
      -> [(HatNode,String)]	-- enclosing nodes w/ variable name for `where'
      -> a 			-- root node of expression
      -> ( SExp HatNode		-- expression 
         , [HatNode]		-- nodes that start cycle
         , [String] )		-- variable names occurring (except for cycles)
go cyc uneval nodesAbove expObj =
  let expNode = toHatNode expObj
  in
  case hatNodeType expObj of
  HatApplNode -> 
    let partCycles :: [HatNode]
        partCycles = funCycles ++ concat argsCycles
        partVars :: [String]
        partVars = funVars ++ concat argsVars
        isCycle = cyc && expNode `elem` partCycles
        var :: String
        var = head . filter (not . (`elem` partVars)) .
              map (("cyc"++) . show) $ [1..] 
        newNodesAbove :: [(HatNode,String)]
        newNodesAbove = (expNode,var) : nodesAbove
        -- (fun,funCycles,funVars) = ... not accepted by nhc98
        fun = fst3 z
        funCycles = snd3 z
        funVars = thd3 z
        z = go cyc uneval newNodesAbove (hatApplFun expObj) 
        args = fst3 zs
        argsCycles = snd3 zs
        argsVars = thd3 zs
        -- (args,argsCycles,argsVars) = ... not accepted by nhc98
        zs = unzip3 $ map (go cyc uneval newNodesAbove) (hatApplArgs expObj)
        sexp = case fun of
          SIf n (SLambda n') -> 
            if length args == 2 
              then SIf n (head args)
              else error "hatRep2Doc: invalid If node"
          SCase n (SLambda n') -> 
            if length args == 2 
              then SCase n (head args)
              else error "hatRep2Doc: invalid Case node"
          SGuard n (SLambda n') ->
            if length args == 2 
              then SGuard n (head args)
              else error "hatRep2Doc: invalid Guard node"
          SApp n args1 -> SApp n (args1++args)  -- combine applications
          _ -> SApp expNode (fun:args)
    in case lookup expNode nodesAbove of
         Just var -> (SId expNode var SInfixDefault,[expNode],[]) 
                     -- `lower' end of cycle
         Nothing ->
           case sexp of
             SApp n [SId n' "IO" fixity,_] -> 
               -- get hidden argument of IO, otherwise only get IO <hidden>
               let arg = head (hatApplArgs expObj)
                   par = hatParent arg
                   newArg = if hatNodeType arg == HatHiddenNode && 
                              isValidNode par
                              then par else arg
                   z = go cyc uneval newNodesAbove newArg
               in (SApp n [SId n' "IO" fixity,fst3 z],snd3 z,thd3 z)
             _ -> -- normal case 
               (if isCycle then SCycle expNode var sexp else sexp
                ,partCycles,partVars)
  HatConstantNode -> go cyc uneval nodesAbove (hatApplFun expObj)
  HatSAT_ANode -> if uneval 
                    then go cyc uneval nodesAbove (hatProjValue expObj)
                    else simple $ SUnevaluated expNode
  HatSAT_BNode -> simple $ SBottom expNode
  HatSAT_CNode -> error "hatRep2Doc: hatSAT_CNode" -- should never appear
  HatHiddenNode -> simple $ SLiteral expNode "<hidden>"
  HatDummyNode -> simple $ SLiteral expNode "(<dummy>)"
  HatProjNode -> go cyc uneval nodesAbove (hatProjValue expObj)
  HatConstrNode -> let var = hatName expObj 
                   in (SId expNode var (transFixity (hatInfix expObj)),[],[])
  HatIdentNode -> let var = hatName expObj 
                  in (SId expNode var
                       (case var of
                         "." | moduleName (hatSourceRef expObj) == "Prelude" 
                           -> SAssoc 9 var
                         "++" | moduleName (hatSourceRef expObj) == "Prelude" 
                           -> SAssoc 5 var
                         "&&" | moduleName (hatSourceRef expObj) == "Prelude" 
                           -> SAssoc 3 var
                         "||" | moduleName (hatSourceRef expObj) == "Prelude" 
                           -> SAssoc 2 var
                         "*" | moduleName (hatSourceRef expObj) == "Prelude" 
                           -> SAssoc 7 var
                         "+" | moduleName (hatSourceRef expObj) == "Prelude" 
                           -> SAssoc 6 var
                         ">>" | moduleName (hatSourceRef expObj) == "Prelude" 
                           -> SAssoc 1 var
                         ">>=" | moduleName (hatSourceRef expObj) == "Prelude" 
                           -> SAssoc 1 var
                         _ -> transFixity (hatInfix expObj))
                     ,[],[var])
  HatCaseNode -> simple $ SCase expNode (SLambda expNode)
					 -- hack for recognition as app fun
  HatIfNode -> simple $ SIf expNode (SLambda expNode) -- ""
  HatGuardNode -> simple $ SGuard expNode (SLambda expNode) -- ""
  HatLambdaNode -> simple $ SLambda expNode
  HatIntNode -> simpleLitShow expNode . hatValueInt $ expObj
  HatCharNode -> simpleLitShow expNode . hatValueChar $ expObj
  HatIntegerNode -> simpleLitShow expNode . hatValueInteger $ expObj
  HatRationalNode -> simpleLitShow expNode . hatValueRational $ expObj
  HatFloatNode -> simpleLitShow expNode . hatValueFloat $ expObj
  HatDoubleNode -> simpleLitShow expNode . hatValueDouble $ expObj
  HatCStringNode -> simpleLitShow expNode . hatValueString $ expObj
  HatContainerNode -> simple $ SLiteral expNode "(?)"
  HatInvalidNode -> simple $ SInvalid expNode
  HatModuleNode -> error "hatRep2Doc: HatModuleNode"
  where
  simple e = (e,[],[])
  simpleLitShow n = simple . SLiteral n . show

indentation :: Int
indentation = 2


-- useful document combinators:
-- non-breaking space
(<->) :: Doc -> Doc -> Doc
d1 <-> d2 = d1 <> delimiter " " <> d2 

-- breakable space
(<+>) :: Doc -> Doc -> Doc
d1 <+> d2 = d1 <> fdelimiter " " <> d2

-- breakable non-space
(<|>) :: Doc -> Doc -> Doc
d1 <|> d2 = d1 <> fdelimiter "" <> d2

-- breakable before a comma
(<*>) :: Doc -> Doc -> Doc
d1 <*> d2 = d1 <|> text "," <> d2

-- breakable before or after a cons
(<:>) :: Doc -> Doc -> Doc
d1 <:> d2 = d1 <|> text ":" <|> d2


isOpSym :: String -> Bool
isOpSym sym = let c = head sym in not (isAlpha c || c == '[') 

funDoc :: String -> Doc
funDoc var = (if isOpSym var then parens else id) $ text var 

opDoc :: String -> Doc
opDoc var =  text (if isAlpha (head var) then ('`' :var++ "`") else var)


data ArgPos = ALeft | ARight
isRight ARight = True
isRight ALeft = False
isLeft = not . isRight

-- surround by parentheses if necessary
-- first fixity of surrounding expression, then if left or right argument,
-- then fixity of expression itself
optParens :: SFixity -> ArgPos -> SFixity -> Doc -> Doc
optParens surFixity aPos ownFixity =
  case (priority surFixity) `compare` (priority ownFixity) of
    LT -> if priority surFixity == (-1) then groupNest indentation else id
    GT -> groupNest indentation . parens
    EQ -> if (isInfixR surFixity && isInfixR ownFixity && isRight aPos)
            || (isInfixL surFixity && isInfixL ownFixity && isLeft aPos)
            || sameAssoc surFixity ownFixity 
            then id
            else groupNest indentation . parens

sameAssoc :: SFixity -> SFixity -> Bool
sameAssoc (SAssoc _ var1) (SAssoc _ var2) = True
sameAssoc _ _ = False

foldr0 :: (a -> a -> a) -> a -> [a] -> a
foldr0 f c [] = c
foldr0 f c xs = foldr1 f xs

listDoc :: Eq a => a -> SFixity -> ArgPos -> SExp a -> Doc
listDoc w surFixity aPos e =
  case maybeRest of
    Nothing -> if all isCharLiteral elems 
                 then groupNest 1 $ 
                        text "\"" <> 
                        (foldr0 (<|>) nil (map getChar elems)) <> 
                        text "\""
                 else group $
                        text "["
                        <> walk (text ",") (text "]") e
                    --  <> (foldr0 (<*>) nil (map (sExp2Doc w) elems))
                    --  <> text "]"
    Just eRest -> groupNest indentation .
                    optParens surFixity aPos (SInfixR 5) . 
                    walk (text ":") (text "[]") $ e
                 -- foldr (<:>) (sExp2Doc w eRest) . 
                 -- map (sExp2Doc w) $ elems 
  where
  (elems,maybeRest) = getListElems e
  --d1 <*> d2 = d1 <|> text "," <> d2
  --d1 <:> d2 = d1 <|> text ":" <|> d2
  getListElems :: SExp a -> ([SExp a],Maybe (SExp a))
  getListElems (SApp _ [SId _ ":" _,ee,er]) = (ee:ees,mr)
                                            where (ees,mr) = getListElems er
  getListElems (SId _ "[]" _) = ([],Nothing)
  getListElems e = ([],Just e)
  isCharLiteral :: SExp a -> Bool
  isCharLiteral (SLiteral _ ('\'':_)) = True
  isCharLiteral _ = False
  getChar :: SExp a -> Doc
  getChar (SLiteral v ('\'':cs)) = (if v==w then highlight else id) $
                                   text (init cs)
  walk sep term (SApp v [SId c ":" _,hd,tl]) = high v (sExp2Doc w hd
                                                      <|> high c sep
                                                      <|> walk sep term tl)
  walk sep term (SId v "[]" _) = high v term
  walk sep term e = sExp2Doc w e
  high v = if v==w then highlight else id

priority (SInfix p) = p
priority (SInfixL p) = p
priority (SInfixR p) = p
priority (SAssoc p _) = p
priority SInfixDefault = 9

isInfixL (SInfixL _) = True
isInfixL SInfixDefault = True
isInfixL _ = False

isInfixR (SInfixR _) = True
isInfixR _ = False

isNotInfixDefault SInfixDefault = False
isNotInfixDefault _ = False

considerAsOperator :: String -> SFixity -> Bool
considerAsOperator var fixity =  isOpSym var || isNotInfixDefault fixity 

-- a central function
sExp2Doc :: Eq a => a -> SExp a -> Doc
sExp2Doc w = goDoc (SInfix (-1)) ARight 
  where
  high :: Eq a => a -> Doc -> Doc
  high v | v==w = highlight
         | otherwise = id

  -- fixity of surrounding expression and which sort of argument
  goDoc :: SFixity -> ArgPos -> SExp a -> Doc
  goDoc surFixity aPos (SApp va ((SId vt (',':xs) _):args)) =
    if length xs + 2 == length args 
    then high va $
         group (text "(" <> foldr1 (<*>) (map (sExp2Doc w) args)) <> text ")"
           -- print tuple properly
    else high va $
         optParens surFixity aPos ownFixity
         . (high vt (text ("(,"++xs++")")) <+>)
         . foldr1 (<+>) . map (goDoc ownFixity ARight) $ args
           -- partial application of tuple constructor
    where
    ownFixity = SInfix 10

  goDoc surFixity aPos (e@(SApp v [SId _ ":" _,_,_])) =
    high v $ listDoc w surFixity aPos e
  goDoc surFixity aPos (SApp va [SId vf var ownFixity,e1,e2]) 
    | considerAsOperator var ownFixity = 
      high va $
      optParens surFixity aPos ownFixity
        (goDoc ownFixity ALeft e1
        <|> high vf (opDoc var)
        <|> goDoc ownFixity ARight e2)
  goDoc surFixity aPos (SApp va [SId vf var ownFixity,e]) 
    | considerAsOperator var ownFixity =
    -- show infix operator with single argument as section
    groupNest indentation . high va . parens $
      goDoc ownFixity ALeft e <-> high vf (opDoc var)
  goDoc surFixity aPos (SApp va (fun:args)) =
    high va $
    optParens surFixity aPos ownFixity
    . (goDoc ownFixity ALeft fun <+>)
    .  foldr1 (<+>) . map (goDoc ownFixity ARight) $ args
    where
    ownFixity = SInfix 10
  goDoc _ _ (SId v var fixity) = high v $ funDoc var
  goDoc _ _ (SLiteral v lit) = high v $ text lit
  goDoc _ _ (SLambda v) = high v $ text "(\\..)"
  goDoc _ _ (SIf v exp) = 
    (groupNest indentation . high v . parens) (text "if " <> sExp2Doc w exp)
  goDoc _ _ (SCase v exp) = 
    (groupNest indentation . high v . parens) (text "case " <> sExp2Doc w exp)
  goDoc _ _ (SGuard v exp) = 
    (groupNest indentation . high v . parens)  (text "| " <> sExp2Doc w exp)
  goDoc _ _ (SInvalid v) = high v $ text "<cut>"
  goDoc _ _ (SUnevaluated v) = high v $ text "_"
  goDoc _ _ (SBottom v) = high v $ text "_|_"
  goDoc _ _ (SCycle v var exp) = 
    groupNest indentation . high v . parens $
      text var <+> group (text "where" <+> sExp2Doc w exp)


{-
-- main function
hatRep2Doc :: HatRep a => Bool -> Bool -> Int -> a -> Doc
hatRep2Doc explicitCycle uneval depth v = 
--sExp2Doc . hatRep2SExp explicitCycle uneval -- . toHatLimit depth
  sExp2Doc (hatRep2SExp explicitCycle uneval v)

showExpression :: HatNode -> String -> String
showExpression node initial =
  pretty 80 
    (text initial <> groupNest (length initial) (hatRep2Doc True False 10 node))

showReduction :: Bool -> Int -> HatNode -> String -> String -> String
showReduction verboseMode precision node initial final =
  pretty 80 
    (text initial
    <> groupNest (length initial) 
         (hatRep2Doc True verboseMode precision node
         <> (if isValidNode result
             then nil <-> text "=" <->
                  hatRep2Doc True verboseMode precision result
             else nil)
         <> (if null final then nil else delimiter "  " <> text final)))
  where
  result = hatResult node
-}

{-
-- only for testing:
test1 =
  SApp (mkHatNode 0)
    [ SId (mkHatNode 1) "fun" SInfixDefault
    , SLiteral (mkHatNode 2) "24"
    , SLiteral (mkHatNode 3) "True"
    , SApp (mkHatNode 4)
        [ SId (mkHatNode 5) "+" (SInfixL 5)
        , SLiteral (mkHatNode 6) "3"
        , SLiteral (mkHatNode 7) "4"
        ]
    ]

test2 =
  SApp (mkHatNode 10)
    [ SId (mkHatNode 11) "*" (SInfixL 6)
    , SApp (mkHatNode 12)
        [ SId (mkHatNode 13) "+" (SInfixL 5)
        , SApp (mkHatNode 14)
            [ SId (mkHatNode 15) "-" (SInfixL 5)
            , SLiteral (mkHatNode 16) "3"
            , SLiteral (mkHatNode 17) "6"
            ]
        , SApp (mkHatNode 18)
            [ SId (mkHatNode 19) "-" (SInfixL 5)
            , SLiteral (mkHatNode 20) "3"
            , SLiteral (mkHatNode 21) "6"
            ]
        ]
    , test1
    ]
-}