content: module Stats(
      -- mutable
      Stats,
      new,
      tick,
      setPrintStats,
      ticks,
      theStats,
      isEmpty,
      null,
      Stats.print,
      clear,
      combine,
      -- pure
      printStat,
      printLStat,
      Stat,
      Stats.singleton,
      Stats.singleStat,
      prependStat,
      -- monad
      MonadStats(..),
      StatT,
      StatM,
      mtick,
      mtick',
      mticks,
      runStatT,
      runStatIO,
      runStatM,
      -- combined
      tickStat,
      readStat
      ) where
  
  import Control.Monad.Identity
  import Control.Monad.Reader
  import Control.Monad.Writer.Strict
  import Data.IORef
  import Data.Tree
  import Data.List(sort,groupBy)
  import Prelude hiding(null)
  import System.IO.Unsafe
  import qualified Data.Map as Map
  import qualified Prelude(null)
  
  import GenUtil
  import StringTable.Atom
  import qualified Doc.Chars as C
  import qualified Util.IntBag as IB
  
  splitUp :: Int -> String -> [String]
  splitUp n str = filter (not . Prelude.null) (f n str)  where
      f 0 str = []
      f n str = case span (`notElem` "/.{") str  of
          (x,"") -> [x]
          (x,('/':rs)) -> x:f (n - 1) rs
          (x,('.':rs)) -> x:f n rs
          (x,('{':rs)) -> case span (/= '}') rs of
              (a,'}':b) -> x:a:f n b
              (a,"") -> [x,a]
              _ -> error "this can't happen"
          _ -> error "this can't happen"
  
  print greets stats = do
      l <- toList stats
      let fs = createForest 0 $ sort [(splitUp (-1) $ fromAtom x,y) | (x,y) <- l]
      mapM_ putStrLn $ ( draw . fmap p ) (Node (greets,0) fs)  where
          p (x,0) = x
          p (x,n) = x ++ ": " ++ show n
  
  createForest :: a -> [([String],a)] -> Forest (String,a)
  createForest def xs = map f gs where
      f [(xs,ys)] =  Node (intercalate "." xs,ys) []
      f xs@((x:_,_):_) = Node (x,def) (createForest def [ (xs,ys) | (_:xs@(_:_),ys)<- xs])
      f _ = error "createForest: should not happen."
      gs = groupBy (\(x:_,_) (y:_,_) -> x == y) xs
  
  draw :: Tree String -> [String]
  draw (Node x ts0) = x : drawSubTrees ts0
    where drawSubTrees [] = []
          drawSubTrees [t] =
                  {-[vLine] :-} shift lastBranch "  " (draw t)
          drawSubTrees (t:ts) =
                  {-[vLine] :-} shift branch (C.vLine  ++ " ") (draw t) ++ drawSubTrees ts
  
          branch     = C.lTee ++ C.hLine
          lastBranch = C.llCorner ++ C.hLine
  
          shift first other = zipWith (++) (first : repeat other)
          --vLine = chr 0x254F
  
  -- Pure varients
  
  newtype Stat = Stat IB.IntBag
      deriving(Eq,Ord,Monoid)
  
  prependStat :: String -> Stat -> Stat
  prependStat name (Stat m) = Stat $ IB.fromList [ (fromAtom $ mappend (toAtom $ "{" ++ name ++ "}.")  (unsafeIntToAtom x),y) | (x,y) <- IB.toList m ]
  
  printStat greets (Stat s) = do
      let fs = createForest 0 $ sort [(splitUp (-1) $ fromAtom (unsafeIntToAtom x),y) | (x,y) <- IB.toList s]
      mapM_ putStrLn $ ( draw . fmap p ) (Node (greets,0) fs)  where
          p (x,0) = x
          p (x,n) = x ++ ": " ++ show n
  
  printLStat n greets (Stat s) = do
      let fs = createForest 0 $ [ (x,y) | (x,y) <- Map.toList $ Map.fromListWith (+) [( splitUp n (fromAtom (unsafeIntToAtom x)),y) | (x,y) <- IB.toList s]]
      mapM_ putStrLn $ ( draw . fmap p ) (Node (greets,0) fs)  where
          p (x,0) = x
          p (x,n) = x ++ ": " ++ show n
  
  --------------
  -- monad stats
  --------------
  
  class Monad m => MonadStats m where
      mticks' ::  Int -> Atom -> m ()
      mtickStat :: Stat -> m ()
  
  newtype StatT m a = StatT (WriterT Stat m a)
      deriving(MonadIO, Functor, MonadFix, MonadTrans, Monad)
  
  runStatT :: Monad m => StatT m a -> m (a,Stat)
  runStatT (StatT m) =  runWriterT m
  
  data StatM a = StatM a !Stat
  
  instance Functor StatM where
      fmap f (StatM a s) = StatM (f a) s
  
  instance Monad StatM where
      StatM _ s1 >> StatM y s2 = StatM y (s1 `mappend` s2)
      return x = StatM x mempty
      StatM x s1 >>= y = case y x of StatM z s2 -> StatM z (s1 `mappend` s2)
  
  instance Stats.MonadStats StatM where
     mticks' 0 k = StatM () mempty
     mticks' n k = StatM () $ Stats.singleStat n k
     mtickStat s = StatM () s
  
  runStatM ::  StatM a -> (a,Stat)
  runStatM (StatM a s) = (a,s)
  
  -- These are inlined so the 'toAtom' can become a caf and be shared
  {-# INLINE mtick  #-}
  {-# INLINE mticks #-}
  mtick k = mticks 1 k
  mtick' k = mticks' 1 k
  mticks 0 _ = return ()
  mticks n k = let k' = toAtom k in k' `seq` n `seq` mticks' n k'
  
  instance MonadStats Identity where
      mticks' _ _ = return ()
      mtickStat _ = return ()
  
  instance MonadReader r m => MonadReader r (StatT m) where
      ask = lift $ ask
      local f (StatT m) = StatT $ local f m
  
  instance (Monad m, Monad (t m), MonadTrans t, MonadStats m) => MonadStats (t m) where
      mticks' n k = lift $ mticks' n k
      mtickStat s = lift $ mtickStat s
  
  instance Monad m => MonadStats (StatT m) where
      mticks' n k = StatT $ tell (Stat $ IB.msingleton (fromAtom k) n)
      mtickStat s =  StatT $ tell s
  
  singleton n = Stat $ IB.singleton (fromAtom $ toAtom n)
  
  singleStat :: ToAtom a => Int -> a -> Stat
  singleStat 0 _ = mempty
  singleStat n k = Stat $ IB.msingleton (fromAtom $ toAtom k) n
  
  null (Stat r) = IB.null r
  
  instance MonadStats IO where
      mticks' 0 _ = return ()
      mticks' n a = do
          p <- readIORef printStats
          when p (putStrLn $ (show a ++ ": " ++ show n))
          ticks theStats n a
      mtickStat (Stat s) = do
          tickStat theStats (Stat s)
          p <- readIORef printStats
          when p $ forM_ (IB.toList s) $ \ (x,y) -> do
              putStrLn (show (unsafeIntToAtom x) ++ ": " ++ show y)
  
  --------------------
  -- Stateful IO stats
  --------------------
  
  newtype Stats = Stats (IORef Stat)
  
  {-# NOINLINE theStats #-}
  theStats :: Stats
  theStats = unsafePerformIO new
  
  {-# NOINLINE printStats #-}
  printStats :: IORef Bool
  printStats = unsafePerformIO $ newIORef False
  
  setPrintStats :: Bool -> IO ()
  setPrintStats b = writeIORef printStats b
  
  combine :: Stats -> Stats -> IO ()
  combine (Stats s1) (Stats s2) = do
      s <- readIORef s2
      modifyIORef s1 (mappend s)
  
  new = Stats `liftM` newIORef mempty
  
  clear (Stats h) = writeIORef h mempty
  
  toList (Stats r) = do
      Stat s <- readIORef r
      return [(unsafeIntToAtom x,y) | (x,y) <- IB.toList s]
  
  isEmpty (Stats r) = null `liftM` readIORef r
  
  tick stats k = ticks stats 1 k
  
  ticks (Stats r) c k = modifyIORef r (mappend $ singleStat c k)
  
  -----------------
  -- pure + mutable
  -----------------
  
  tickStat ::  Stats -> Stat -> IO ()
  tickStat (Stats r) s = modifyIORef r (mappend s)
  
  runStatIO :: MonadIO m =>  Stats -> StatT m a -> m a
  runStatIO stats action = do
      (a,s) <- runStatT action
      liftIO $ tickStat stats s
      return a
  
  readStat :: Stats -> IO Stat
  readStat (Stats r) = readIORef r
  