content: {-# LANGUAGE GeneralizedNewtypeDeriving #-}
  {-# LANGUAGE ScopedTypeVariables #-}
  {-# LANGUAGE ViewPatterns #-}
  {-# LANGUAGE TypeFamilies #-}
  {-# LANGUAGE BangPatterns #-}
  {-# LANGUAGE ImplicitParams    #-}
  {-# LANGUAGE LambdaCase #-}
  {-# LANGUAGE OverloadedLists   #-}
  {-# LANGUAGE OverloadedStrings #-}
  {-# LANGUAGE FlexibleContexts #-}
  {-# OPTIONS_GHC -fno-warn-name-shadowing #-}
  -- This file is part of the Haskell debugger Hoed.
  --
  -- Copyright (c) Maarten Faddegon, 2014
  {-# LANGUAGE DeriveGeneric     #-}
  
  module Debug.Hoed.Render
  (CompStmt(..)
  ,StmtDetails(..)
  ,stmtRes
  ,renderCompStmts
  ,CDS
  ,eventsToCDS
  ,noNewlines
  ,sortOn
  ) where
  import           Control.DeepSeq
  import           Control.Exception        (assert)
  import           Control.Monad.Primitive
  import           Control.Monad.ST
  import           Control.Monad.Trans.Reader
  import           Control.Monad.Trans.State.Strict
  import           Data.Array               as Array
  import           Data.Char                (isAlpha)
  import           Data.Coerce
  import           Data.Hashable
  import           Data.List                (nub, sort, unfoldr)
  import           Data.IntMap.Strict (IntMap)
  import qualified Data.IntMap.Strict as IntMap
  import           Data.Map.Strict (Map)
  import qualified Data.Map.Strict as Map
  import           Data.Primitive.MutVar
  import           Data.Text (Text, pack, unpack)
  import qualified Data.Text as T
  import qualified Data.Vector as V
  import qualified Data.Vector.Generic as VG
  import           Data.Word
  import           Debug.Hoed.Compat
  import           Debug.Hoed.Observe
  import           GHC.Exts(IsList(..))
  import           GHC.Generics
  import           Text.PrettyPrint.FPretty hiding (sep, (<$>), text)
  import qualified Text.PrettyPrint.FPretty as FPretty
  import           Text.Read ()
  
  
  ------------------------------------------------------------------------
  -- The CompStmt type
  
  -- MF TODO: naming here is a bit of a mess. Needs refactoring.
  -- Indentifier refers to an identifier users can explicitely give
  -- to observe'. But UID is the unique number assigned to each event.
  -- The field equIdentifier is not an Identifier, but the UID of the
  -- event that starts the observation. And stmtUIDs is the list of
  -- UIDs of all events that form the statement.
  
  data CompStmt = CompStmt { stmtLabel      :: !Text
                           , stmtIdentifier :: !UID
                           , stmtDetails    :: !StmtDetails
                           }
                  deriving (Generic)
  
  instance NFData CompStmt
  
  instance Eq CompStmt where c1 == c2 = stmtIdentifier c1 == stmtIdentifier c2
  instance Ord CompStmt where
    compare c1 c2 = compare (stmtIdentifier c1) (stmtIdentifier c2)
  
  instance Hashable CompStmt where
    hashWithSalt s cs = hashWithSalt s (stmtIdentifier cs)
  
  data StmtDetails
    = StmtCon { stmtCon :: Hashed Text
             ,  stmtPretty :: Hashed Text}
    | StmtLam { stmtLamArgs :: [Hashed Text]
             ,  stmtLamRes :: Hashed Text
             ,  stmtPretty :: Hashed Text}
    deriving (Generic)
  
  instance NFData StmtDetails
  
  stmtRes :: CompStmt -> Text
  stmtRes = unhashed . stmtPretty . stmtDetails
  
  instance Show CompStmt where
    show = unpack . stmtRes
    showList eqs eq = unlines (map show eqs) ++ eq
  
  noNewlines :: String -> String
  noNewlines = noNewlines' False
  noNewlines' :: Bool -> String -> String
  noNewlines' _ [] = []
  noNewlines' w (s:ss)
   | w       && (s == ' ' || s == '\n') =       noNewlines' True ss
   | not w && (s == ' ' || s == '\n') = ' ' : noNewlines' True ss
   | otherwise                          = s   : noNewlines' False ss
  
  ------------------------------------------------------------------------
  -- Render equations from CDS set
  
  renderCompStmts :: (?statementWidth::Int) => CDSSet -> [CompStmt]
  renderCompStmts cdss = runMemoM $ concat <$> mapM renderCompStmt cdss
  
  -- renderCompStmt: an observed function can be applied multiple times, each application
  -- is rendered to a computation statement
  
  renderCompStmt :: (?statementWidth::Int) => CDS -> MemoM [CompStmt]
  renderCompStmt (CDSNamed name uid set) = do
          let output = cdssToOutput set
          concat <$> mapM (renderNamedTop name uid) output
  
  renderCompStmt other = error $ show other
  
  prettySet cds = prettySet_noid(coerce cds)
  
  prettySet_noid :: (?statementWidth::Int) => [CDSsansUID] -> MemoM(Hashed Text)
  prettySet_noid = MemoM . memo (prettyW . renderSet . coerce)
  
  renderNamedTop :: (?statementWidth::Int) => Text -> UID -> Output -> MemoM [CompStmt]
  renderNamedTop name observeUid (OutData cds) = mapM f pairs
    where
      f (args, res, Just i) =
        CompStmt name i <$>
        (StmtLam <$> mapM prettySet args <*>
         prettySet res <*>
         pure (prettyW $ renderNamedFn name (args, res)))
      f (_, cons, Nothing) =
        CompStmt name observeUid <$>
        (StmtCon <$> prettySet cons <*>
         pure (prettyW $ renderNamedCons name cons))
      pairs = (nubSorted . sortOn argAndRes) pairs'
      pairs' = findFn [cds]
      argAndRes (arg, res, _) = (arg, res)
  renderNamedTop name _ other = error $ show other
  
  -- local nub for sorted lists
  nubSorted :: Eq a => [a] -> [a]
  nubSorted []        = []
  nubSorted (a:a':as) | a == a' = nubSorted (a' : as)
  nubSorted (a:as)    = a : nubSorted as
  
  -- %************************************************************************
  -- %*                                                                   *
  -- \subsection{The CDS and converting functions}
  -- %*                                                                   *
  -- %************************************************************************
  
  data CDS = CDSNamed      !Text !UID    !CDSSet
           | CDSCons       !UID  !Text   ![CDSSet]
           | CDSFun        !UID  !CDSSet !CDSSet
           | CDSEntered    !UID
           | CDSTerminated !UID
           | CDSChar       !Char   -- only used internally in eventsToCDS
           | CDSString     !String -- only used internally in eventsToCDS
          deriving (Show,Eq,Ord,Generic)
  
  instance NFData CDS
  
  normalizeCDS :: CDS -> CDS
  normalizeCDS (CDSString s) = CDSCons 0 (pack $ show s) []
  normalizeCDS (CDSChar   s) = CDSCons 0 (pack $ show s) []
  normalizeCDS other = other
  type CDSSet = [CDS]
  
  -- Monomorphized [Parent] for compactness
  data ParentList = ParentCons !Int !Word8 ParentList | ParentNil
  instance IsList ParentList where
    type Item ParentList = Parent
    toList = unfoldr (\case ParentNil -> Nothing ; ParentCons pp pc t -> Just (Parent pp pc,t))
    fromList = foldr (\(Parent pp pc) t -> ParentCons pp pc t) ParentNil
  
  eventsToCDS :: Trace -> CDSSet
  eventsToCDS pairs = getChild (-1) 0
     where
  
       -- res i = out_arr VG.! i
       res i = getNode'' i (change (pairs VG.! i))
  
       mid_arr :: V.Vector ParentList
       mid_arr = VG.unsafeAccumulate
                    (\i (Parent pp pc) -> ParentCons pp pc i)
                    (V.replicate (VG.length pairs) ParentNil)
                    ( VG.map (\(node, Event (Parent pnode pport) _) ->
                                (pnode+1, Parent node pport))
                    $ VG.filter (\(_,e) -> change e /= Enter)
                    $ VG.convert
                    $ VG.indexed pairs)
  
       getNode'' ::  Int -> Change -> CDS
       getNode'' node change =
         case change of
          Observe str         -> let chd = normalizeCDS <$> getChild node 0
                                 in CDSNamed str (getId chd node) chd
          Enter               -> CDSEntered node
          Fun                 -> CDSFun node (normalizeCDS <$> getChild node 0)
                                             (normalizeCDS <$> getChild node 1)
          ConsChar char       -> CDSChar char
          Cons portc cons
                              -> simplifyCons node cons
                                   [ getChild node (fromIntegral n)
                                   | n <- [0::Int .. fromIntegral portc - 1]]
  
       getId []                 i  = i
       getId (CDSFun i _ _:_) _    = i
       getId (_:cs)             i  = getId cs i
  
       getChild :: Int -> Word8 -> CDSSet
       getChild pnode pport =
         [ res content
         | Parent content pport' <- toList $ mid_arr VG.! succ pnode
         , pport == pport'
         ]
  
  simplifyCons :: UID -> Text -> [CDSSet] -> CDS
  simplifyCons _ "throw" [[CDSCons _ "ErrorCall" set]]
    = CDSCons 0 "error" set
  simplifyCons _ ":" [[CDSChar !ch], [CDSCons _ "[]" []]]
    = CDSString [ch]
  simplifyCons _ ":" [[CDSChar !ch], [CDSString s]]
    = CDSString (ch:s)
  simplifyCons uid con xx = CDSCons uid con (map (map normalizeCDS) xx)
  
  render :: Int -> Bool -> CDS -> Doc
  render prec par (CDSCons _ ":" [cds1,cds2]) =
          if par && not needParen
          then doc -- dont use paren (..) because we dont want a grp here!
          else paren needParen doc
     where
          doc = grp (renderSet' 5 False cds1 <> text " : ") <>
                renderSet' 4 True cds2
          needParen = prec > 4
  render prec par (CDSCons _ "," cdss) | length cdss > 0 =
          nest 2 (text "(" <> foldl1 (\ a b -> a <> text ", " <> b)
                              (map renderSet cdss) <>
                  text ")")
  render prec _par (CDSCons _ name cdss)
    | not (T.null name)
    , (not . isAlpha . T.head) name && length cdss > 1 = -- render as infix
          paren (prec /= 0)
                    (grp
                      (renderSet' 10 False (head cdss)
                       <> sep <> text name
                       <> nest 2 (foldr (<>) nil
                                   [ if null cds then nil else sep <> renderSet' 10 False cds
                                   | cds <- tail cdss
                                   ]
                                )
                      )
                    )
    | otherwise = -- render as prefix
          paren (not (null cdss) && prec /= 0)
                   ( grp
                     (text name <> nest 2 (foldr (<>) nil
                                            [ sep <> renderSet' 10 False cds
                                            | cds <- cdss
                                            ]
                                         )
                     )
                   )
  
  {- renderSet handles the various styles of CDSSet.
   -}
  
  renderSet :: CDSSet -> Doc
  renderSet = renderSet' 0 False
  
  renderSet' :: Int -> Bool -> CDSSet -> Doc
  renderSet' _ _      [] = text "_"
  renderSet' prec par [cons@(CDSCons {})]    = render prec par cons
  renderSet' prec par cdss                   =
          nest 0 (text "{ " <> foldl1 (\ a b -> a <> line <>
                                      text ", " <> b)
                                      (map renderFn pairs) <>
                  line <> text "}")
  
     where
          findFn_noUIDs :: CDSSet -> [([CDSSet],CDSSet)]
          findFn_noUIDs c = map (\(a,r,_) -> (a,r)) (findFn c)
          pairs = nub (sort (findFn_noUIDs cdss))
          -- local nub for sorted lists
          nub []        = []
          nub (a:a':as) | a == a' = nub (a' : as)
          nub (a:as)    = a : nub as
  
  renderFn :: ([CDSSet],CDSSet) -> Doc
  renderFn (args, res)
          = grp  (nest 3
                  (text "\\ " <>
                   foldr (\ a b -> nest 0 (renderSet' 10 False a) <> sp <> b)
                         nil
                         args <> softline <>
                   text "-> " <> renderSet' 0 False res
                  )
                 )
  
  renderNamedCons :: Text -> CDSSet -> Doc
  renderNamedCons name cons
    = text name <> nest 2
       ( sep <> grp (text "= " <> renderSet cons)
       )
  
  renderNamedFn :: Text -> ([CDSSet],CDSSet) -> Doc
  renderNamedFn name (args,res)
    = text name <> nest 2
       ( sep <> foldr (\ a b -> grp (renderSet' 10 False a) <> sep <> b) nil args
         <> sep <> grp ("= " <> align(renderSet res))
       )
  
  -- | Reconstructs functional values from a CDSSet.
  --   Returns a triple containing:
  --    1. The arguments, if any, or an empty list for non function values
  --    2. The result
  --    3. The id of the CDSFun, if a functional value.
  findFn :: CDSSet -> [([CDSSet],CDSSet, Maybe UID)]
  findFn = foldr findFn' []
  
  findFn' :: CDS -> [([CDSSet], CDSSet, Maybe UID)] -> [([CDSSet], CDSSet, Maybe UID)]
  findFn' (CDSFun i arg res) rest =
      case findFn res of
         [(args',res',_)] -> (arg : args', res', Just i) : rest
         _                -> ([arg], res, Just i) : rest
  findFn' other rest = ([],[other], Nothing) : rest
  
  paren :: Bool -> Doc -> Doc
  paren False doc = grp (nest 0 doc)
  paren True  doc = grp (text "(" <> doc <> text ")")
  
  data Output = OutLabel Text CDSSet [Output]
              | OutData  CDS
                deriving (Eq,Ord,Show)
  
  cdssToOutput :: CDSSet -> [Output]
  cdssToOutput =  map cdsToOutput
  
  cdsToOutput :: CDS -> Output
  cdsToOutput (CDSNamed name _ cdsset)
              = OutLabel name res1 res2
    where
        res1 = [ cdss | (OutData cdss) <- res ]
        res2 = [ out  | out@OutLabel {} <- res ]
        res  = cdssToOutput cdsset
  cdsToOutput cons@CDSCons {} = OutData cons
  cdsToOutput    fn@CDSFun {} = OutData fn
  
  nil :: Doc
  nil = Text.PrettyPrint.FPretty.empty
  grp :: Doc -> Doc
  grp = Text.PrettyPrint.FPretty.group
  sep :: Doc
  sep = softline  -- A space, if the following still fits on the current line, otherwise newline.
  sp :: Doc
  sp = text " "   -- A space, always.
  
  -- TODO fork FPretty to build on Text instead of Strings
  text = FPretty.text . unpack
  
  prettyW :: (?statementWidth::Int) => Doc -> (Hashed Text)
  prettyW doc = hashed $ pack $ pretty ?statementWidth doc
  
  -- %************************************************************************
  -- %*                                                                   *
  -- \subsection{Custom Eq and Ord instances for CDS that gloss over UIDs}
  -- %*                                                                   *
  -- %************************************************************************
  
  newtype CDSsansUID = CDSsansUID CDS
  
  instance Eq CDSsansUID where
    CDSsansUID(CDSNamed t _ xx) == CDSsansUID(CDSNamed t' _ yy) =
      t == t' && coerce xx == (coerce yy :: [CDSsansUID])
    CDSsansUID (CDSCons _ t xx) == CDSsansUID(CDSCons _ t' yy)  =
      t == t'  && coerce xx == (coerce yy :: [[CDSsansUID]])
    CDSsansUID (CDSFun _ res args) == CDSsansUID (CDSFun _ res' args') =
      (coerce res :: [CDSsansUID]) == coerce res' && coerce args == (coerce args' :: [CDSsansUID])
    CDSsansUID x == CDSsansUID y = x == y
  
  instance Ord CDSsansUID where
    CDSsansUID (CDSNamed t _ xx) `compare` CDSsansUID (CDSNamed t' _ yy) =
      (t, coerce xx :: [CDSsansUID]) `compare` (t', coerce yy)
    CDSsansUID (CDSCons _ t xx) `compare` CDSsansUID (CDSCons _ t' yy) =
      (t, coerce xx :: [[CDSsansUID]]) `compare` (t', coerce yy)
    CDSsansUID (CDSFun _ args res) `compare` CDSsansUID (CDSFun _ args' res') =
      (coerce args :: [CDSsansUID], coerce res :: [CDSsansUID]) `compare` (coerce args', coerce res')
    CDSsansUID x `compare` CDSsansUID y = x `compare` y
  
  instance Hashable CDSsansUID where
    s `hashWithSalt` CDSsansUID (CDSNamed t _ xx) = s `hashWithSalt` t `hashWithSalt` (coerce xx :: [CDSsansUID])
    s `hashWithSalt` CDSsansUID (CDSCons _  t xx) = s `hashWithSalt` t `hashWithSalt` (coerce xx :: [[CDSsansUID]])
    s `hashWithSalt` CDSsansUID (CDSFun _ args res) = s `hashWithSalt` (coerce args :: [CDSsansUID]) `hashWithSalt` (coerce res :: [CDSsansUID])
  
  
  -- %************************************************************************
  -- %*                                                                   *
  -- \subsection{Memoization of pretty calls}
  -- %*                                                                   *
  -- %************************************************************************
  
  newtype MemoM a = MemoM (State (Map [CDSsansUID] (Hashed Text)) a) deriving (Applicative, Functor, Monad)
  
  runMemoM :: MemoM a -> a
  runMemoM (MemoM comp) = evalState comp mempty
  
  memo :: Ord a => (a->b) -> a -> State (Map a b) b
  memo f a = do
    table <- get
    case Map.lookup a table of
      Just b -> return b
      Nothing -> do
        let b = f a
        modify (Map.insert a b)
        return b
  