{-# LANGUAGE RecordWildCards #-}

module General.Timing(Timing, withTiming, timed, timedOverwrite) where

import Data.List.Extra
import System.Time.Extra
import Data.IORef
import Control.Monad.Extra
import System.IO
import General.Util
import Control.Monad.IO.Class


data Timing = Timing
    {Timing -> IO Seconds
timingOffset :: IO Seconds
    ,Timing -> IORef [(String, Seconds)]
timingStore :: IORef [(String, Seconds)] -- records for writing to a file
    ,Timing -> IORef (Maybe (Seconds, Int))
timingOverwrite :: IORef (Maybe (Seconds, Int)) -- if you are below T you may overwrite N characters
    ,Timing -> Bool
timingTerminal :: Bool -- is this a terminal
    }


withTiming :: Maybe FilePath -> (Timing -> IO a) -> IO a
withTiming :: Maybe String -> (Timing -> IO a) -> IO a
withTiming Maybe String
file Timing -> IO a
f = do
    IO Seconds
timingOffset <- IO (IO Seconds)
offsetTime
    IORef [(String, Seconds)]
timingStore <- [(String, Seconds)] -> IO (IORef [(String, Seconds)])
forall a. a -> IO (IORef a)
newIORef []
    IORef (Maybe (Seconds, Int))
timingOverwrite <- Maybe (Seconds, Int) -> IO (IORef (Maybe (Seconds, Int)))
forall a. a -> IO (IORef a)
newIORef Maybe (Seconds, Int)
forall a. Maybe a
Nothing
    Bool
timingTerminal <- Handle -> IO Bool
hIsTerminalDevice Handle
stdout

    a
res <- Timing -> IO a
f Timing :: IO Seconds
-> IORef [(String, Seconds)]
-> IORef (Maybe (Seconds, Int))
-> Bool
-> Timing
Timing{Bool
IO Seconds
IORef [(String, Seconds)]
IORef (Maybe (Seconds, Int))
timingTerminal :: Bool
timingOverwrite :: IORef (Maybe (Seconds, Int))
timingStore :: IORef [(String, Seconds)]
timingOffset :: IO Seconds
timingTerminal :: Bool
timingOverwrite :: IORef (Maybe (Seconds, Int))
timingStore :: IORef [(String, Seconds)]
timingOffset :: IO Seconds
..}
    Seconds
total <- IO Seconds
timingOffset
    Maybe String -> (String -> IO ()) -> IO ()
forall (m :: * -> *) a.
Applicative m =>
Maybe a -> (a -> m ()) -> m ()
whenJust Maybe String
file ((String -> IO ()) -> IO ()) -> (String -> IO ()) -> IO ()
forall a b. (a -> b) -> a -> b
$ \String
file -> do
        [(String, Seconds)]
xs <- IORef [(String, Seconds)] -> IO [(String, Seconds)]
forall a. IORef a -> IO a
readIORef IORef [(String, Seconds)]
timingStore
        -- Expecting unrecorded of ~2s
        -- Most of that comes from the pipeline - we get occasional 0.01 between items as one flushes
        -- Then at the end there is ~0.5 while the final item flushes
        [(String, Seconds)]
xs <- [(String, Seconds)] -> IO [(String, Seconds)]
forall (f :: * -> *) a. Applicative f => a -> f a
pure ([(String, Seconds)] -> IO [(String, Seconds)])
-> [(String, Seconds)] -> IO [(String, Seconds)]
forall a b. (a -> b) -> a -> b
$ ((String, Seconds) -> Seconds)
-> [(String, Seconds)] -> [(String, Seconds)]
forall b a. Ord b => (a -> b) -> [a] -> [a]
sortOn (Seconds -> Seconds
forall a. Num a => a -> a
negate (Seconds -> Seconds)
-> ((String, Seconds) -> Seconds) -> (String, Seconds) -> Seconds
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (String, Seconds) -> Seconds
forall a b. (a, b) -> b
snd) ([(String, Seconds)] -> [(String, Seconds)])
-> [(String, Seconds)] -> [(String, Seconds)]
forall a b. (a -> b) -> a -> b
$ (String
"Unrecorded", Seconds
total Seconds -> Seconds -> Seconds
forall a. Num a => a -> a -> a
- [Seconds] -> Seconds
forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
sum (((String, Seconds) -> Seconds) -> [(String, Seconds)] -> [Seconds]
forall a b. (a -> b) -> [a] -> [b]
map (String, Seconds) -> Seconds
forall a b. (a, b) -> b
snd [(String, Seconds)]
xs)) (String, Seconds) -> [(String, Seconds)] -> [(String, Seconds)]
forall a. a -> [a] -> [a]
: [(String, Seconds)]
xs
        String -> String -> IO ()
writeFile String
file (String -> IO ()) -> String -> IO ()
forall a b. (a -> b) -> a -> b
$ [String] -> String
unlines ([String] -> String) -> [String] -> String
forall a b. (a -> b) -> a -> b
$ Int -> String -> [(String, Seconds)] -> [String]
prettyTable Int
2 String
"Secs" [(String, Seconds)]
xs
    String -> IO ()
putStrLn (String -> IO ()) -> String -> IO ()
forall a b. (a -> b) -> a -> b
$ String
"Took " String -> String -> String
forall a. [a] -> [a] -> [a]
++ Seconds -> String
showDuration Seconds
total
    a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure a
res


-- skip it if have written out in the last 1s and takes < 0.1


timed :: MonadIO m => Timing -> String -> m a -> m a
timed :: Timing -> String -> m a -> m a
timed = Bool -> Timing -> String -> m a -> m a
forall (m :: * -> *) a.
MonadIO m =>
Bool -> Timing -> String -> m a -> m a
timedEx Bool
False

timedOverwrite :: MonadIO m => Timing -> String -> m a -> m a
timedOverwrite :: Timing -> String -> m a -> m a
timedOverwrite = Bool -> Timing -> String -> m a -> m a
forall (m :: * -> *) a.
MonadIO m =>
Bool -> Timing -> String -> m a -> m a
timedEx Bool
True

timedEx :: MonadIO m => Bool -> Timing -> String -> m a -> m a
timedEx :: Bool -> Timing -> String -> m a -> m a
timedEx Bool
overwrite Timing{Bool
IO Seconds
IORef [(String, Seconds)]
IORef (Maybe (Seconds, Int))
timingTerminal :: Bool
timingOverwrite :: IORef (Maybe (Seconds, Int))
timingStore :: IORef [(String, Seconds)]
timingOffset :: IO Seconds
timingTerminal :: Timing -> Bool
timingOverwrite :: Timing -> IORef (Maybe (Seconds, Int))
timingStore :: Timing -> IORef [(String, Seconds)]
timingOffset :: Timing -> IO Seconds
..} String
msg m a
act = do
    Seconds
start <- IO Seconds -> m Seconds
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO IO Seconds
timingOffset
    IO () -> m ()
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO () -> m ()) -> IO () -> m ()
forall a b. (a -> b) -> a -> b
$ IO (Maybe (Seconds, Int)) -> ((Seconds, Int) -> IO ()) -> IO ()
forall (m :: * -> *) a.
Monad m =>
m (Maybe a) -> (a -> m ()) -> m ()
whenJustM (IORef (Maybe (Seconds, Int)) -> IO (Maybe (Seconds, Int))
forall a. IORef a -> IO a
readIORef IORef (Maybe (Seconds, Int))
timingOverwrite) (((Seconds, Int) -> IO ()) -> IO ())
-> ((Seconds, Int) -> IO ()) -> IO ()
forall a b. (a -> b) -> a -> b
$ \(Seconds
t,Int
n) ->
        if Bool
overwrite Bool -> Bool -> Bool
&& Seconds
start Seconds -> Seconds -> Bool
forall a. Ord a => a -> a -> Bool
< Seconds
t then
            String -> IO ()
putStr (String -> IO ()) -> String -> IO ()
forall a b. (a -> b) -> a -> b
$ Int -> Char -> String
forall a. Int -> a -> [a]
replicate Int
n Char
'\b' String -> String -> String
forall a. [a] -> [a] -> [a]
++ Int -> Char -> String
forall a. Int -> a -> [a]
replicate Int
n Char
' ' String -> String -> String
forall a. [a] -> [a] -> [a]
++ Int -> Char -> String
forall a. Int -> a -> [a]
replicate Int
n Char
'\b'
        else
            String -> IO ()
putStrLn String
""

    let out :: String -> m Int
out String
msg = IO Int -> m Int
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO Int -> m Int) -> IO Int -> m Int
forall a b. (a -> b) -> a -> b
$ String -> IO ()
putStr String
msg IO () -> IO Int -> IO Int
forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> Int -> IO Int
forall (f :: * -> *) a. Applicative f => a -> f a
pure (String -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length String
msg)
    Int
undo1 <- String -> m Int
forall (m :: * -> *). MonadIO m => String -> m Int
out (String -> m Int) -> String -> m Int
forall a b. (a -> b) -> a -> b
$ String
msg String -> String -> String
forall a. [a] -> [a] -> [a]
++ String
"... "
    IO () -> m ()
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO () -> m ()) -> IO () -> m ()
forall a b. (a -> b) -> a -> b
$ Handle -> IO ()
hFlush Handle
stdout

    a
res <- m a
act
    Seconds
end <- IO Seconds -> m Seconds
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO IO Seconds
timingOffset
    let time :: Seconds
time = Seconds
end Seconds -> Seconds -> Seconds
forall a. Num a => a -> a -> a
- Seconds
start
    IO () -> m ()
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO () -> m ()) -> IO () -> m ()
forall a b. (a -> b) -> a -> b
$ IORef [(String, Seconds)]
-> ([(String, Seconds)] -> [(String, Seconds)]) -> IO ()
forall a. IORef a -> (a -> a) -> IO ()
modifyIORef IORef [(String, Seconds)]
timingStore ((String
msg,Seconds
time)(String, Seconds) -> [(String, Seconds)] -> [(String, Seconds)]
forall a. a -> [a] -> [a]
:)

    String
s <- String -> (String -> String) -> Maybe String -> String
forall b a. b -> (a -> b) -> Maybe a -> b
maybe String
"" (\String
x -> String
" (" String -> String -> String
forall a. [a] -> [a] -> [a]
++ String
x String -> String -> String
forall a. [a] -> [a] -> [a]
++ String
")") (Maybe String -> String) -> m (Maybe String) -> m String
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> IO (Maybe String) -> m (Maybe String)
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO IO (Maybe String)
getStatsPeakAllocBytes
    Int
undo2 <- String -> m Int
forall (m :: * -> *). MonadIO m => String -> m Int
out (String -> m Int) -> String -> m Int
forall a b. (a -> b) -> a -> b
$ Seconds -> String
showDuration Seconds
time String -> String -> String
forall a. [a] -> [a] -> [a]
++ String
s

    Maybe (Seconds, Int)
old <- IO (Maybe (Seconds, Int)) -> m (Maybe (Seconds, Int))
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO (Maybe (Seconds, Int)) -> m (Maybe (Seconds, Int)))
-> IO (Maybe (Seconds, Int)) -> m (Maybe (Seconds, Int))
forall a b. (a -> b) -> a -> b
$ IORef (Maybe (Seconds, Int)) -> IO (Maybe (Seconds, Int))
forall a. IORef a -> IO a
readIORef IORef (Maybe (Seconds, Int))
timingOverwrite
    let next :: Seconds
next = Seconds
-> ((Seconds, Int) -> Seconds) -> Maybe (Seconds, Int) -> Seconds
forall b a. b -> (a -> b) -> Maybe a -> b
maybe (Seconds
start Seconds -> Seconds -> Seconds
forall a. Num a => a -> a -> a
+ Seconds
1.0) (Seconds, Int) -> Seconds
forall a b. (a, b) -> a
fst Maybe (Seconds, Int)
old
    IO () -> m ()
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO () -> m ()) -> IO () -> m ()
forall a b. (a -> b) -> a -> b
$ if Bool
timingTerminal Bool -> Bool -> Bool
&& Bool
overwrite Bool -> Bool -> Bool
&& Seconds
end Seconds -> Seconds -> Bool
forall a. Ord a => a -> a -> Bool
< Seconds
next then
        IORef (Maybe (Seconds, Int)) -> Maybe (Seconds, Int) -> IO ()
forall a. IORef a -> a -> IO ()
writeIORef IORef (Maybe (Seconds, Int))
timingOverwrite (Maybe (Seconds, Int) -> IO ()) -> Maybe (Seconds, Int) -> IO ()
forall a b. (a -> b) -> a -> b
$ (Seconds, Int) -> Maybe (Seconds, Int)
forall a. a -> Maybe a
Just (Seconds
next, Int
undo1 Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
undo2)
     else do
        IORef (Maybe (Seconds, Int)) -> Maybe (Seconds, Int) -> IO ()
forall a. IORef a -> a -> IO ()
writeIORef IORef (Maybe (Seconds, Int))
timingOverwrite Maybe (Seconds, Int)
forall a. Maybe a
Nothing
        String -> IO ()
putStrLn String
""
    a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure a
res