module Crux.ProgressBar where

import System.IO
import System.Console.ANSI
import Control.Monad(zipWithM)

-- | Callback called with index when a goal is started
type ProverMilestoneStartGoal = Integer -> IO ()

-- | Callback called with index when a goal has ended
type ProverMilestoneEndGoal = Integer -> IO ()

-- | Callback called when all goals have ended
type ProverMilestoneFinish = IO ()

-- | Set of three callbacks called by the prover to let loggers indicate
-- progress of the proving process.
type ProverMilestoneCallbacks =
  ( ProverMilestoneStartGoal
  , ProverMilestoneEndGoal
  , ProverMilestoneFinish
  )

silentProverMilestoneCallbacks :: ProverMilestoneCallbacks
silentProverMilestoneCallbacks :: ProverMilestoneCallbacks
silentProverMilestoneCallbacks =
  (IO () -> Integer -> IO ()
forall a b. a -> b -> a
const (() -> IO ()
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return ()), IO () -> Integer -> IO ()
forall a b. a -> b -> a
const (() -> IO ()
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return ()), () -> IO ()
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return ())

prepStatus :: String -> Int -> IO ProverMilestoneCallbacks
prepStatus :: String -> Int -> IO ProverMilestoneCallbacks
prepStatus String
pref Int
tot =
   do Bool
ansi <- Handle -> IO Bool
hSupportsANSI Handle
stdout
      if Bool
ansi then
        ProverMilestoneCallbacks -> IO ProverMilestoneCallbacks
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return (Integer -> IO ()
start,Integer -> IO ()
forall {p}. p -> IO ()
end,IO ()
finish)
      else
        ProverMilestoneCallbacks -> IO ProverMilestoneCallbacks
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return ProverMilestoneCallbacks
silentProverMilestoneCallbacks

  where
  start :: Integer -> IO ()
start Integer
n = do Handle -> IO ()
hSaveCursor Handle
stdout
               Handle -> String -> IO ()
hPutStr Handle
stdout (Integer -> String
msg Integer
n)
               Handle -> IO ()
hFlush Handle
stdout
  end :: p -> IO ()
end p
_n  = do Handle -> IO ()
hRestoreCursor Handle
stdout
               Handle -> IO ()
hFlush Handle
stdout

  finish :: IO ()
finish = do Handle -> IO ()
hClearLine Handle
stdout
              Handle -> IO ()
hFlush Handle
stdout

  totS :: String
totS  = Int -> String
forall a. Show a => a -> String
show Int
tot
  totW :: Int
totW  = String -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length String
totS
  sh :: p -> String
sh p
x  = let y :: String
y = p -> String
forall a. Show a => a -> String
show p
x
          in Int -> Char -> String
forall a. Int -> a -> [a]
replicate (Int
totW Int -> Int -> Int
forall a. Num a => a -> a -> a
- String -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length String
y) Char
' ' String -> String -> String
forall a. [a] -> [a] -> [a]
++ String
y

  msg :: Integer -> String
msg Integer
n = String
pref String -> String -> String
forall a. [a] -> [a] -> [a]
++ Integer -> String
forall a. Show a => a -> String
sh (Integer
n::Integer) String -> String -> String
forall a. [a] -> [a] -> [a]
++ String
" / " String -> String -> String
forall a. [a] -> [a] -> [a]
++ String
totS



withProgressBar' :: String -> [a] -> (a -> IO b) -> IO [b]
withProgressBar' :: forall a b. String -> [a] -> (a -> IO b) -> IO [b]
withProgressBar' String
pref [a]
xs a -> IO b
f =
    do (Integer -> IO ()
start,Integer -> IO ()
end,IO ()
finish) <- String -> Int -> IO ProverMilestoneCallbacks
prepStatus String
pref ([a] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [a]
xs)
       let one :: Integer -> a -> IO b
one Integer
n a
a =
            do Integer -> IO ()
start Integer
n
               b
b <- a -> IO b
f a
a
               Integer -> IO ()
end Integer
n
               b -> IO b
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return b
b
       (Integer -> a -> IO b) -> [Integer] -> [a] -> IO [b]
forall (m :: * -> *) a b c.
Applicative m =>
(a -> b -> m c) -> [a] -> [b] -> m [c]
zipWithM Integer -> a -> IO b
one [ Integer
1 .. ] [a]
xs IO [b] -> IO () -> IO [b]
forall a b. IO a -> IO b -> IO a
forall (f :: * -> *) a b. Applicative f => f a -> f b -> f a
<* IO ()
finish

withProgressBar :: Int -> [a] -> (a -> IO b) -> IO [b]
withProgressBar :: forall a b. Int -> [a] -> (a -> IO b) -> IO [b]
withProgressBar Int
w [a]
xs a -> IO b
f =
  do String -> IO ()
prt String
"|"
     Int -> Float -> [a] -> [b] -> IO [b]
go Int
0 Float
0 [a]
xs []
  where
  prt :: String -> IO ()
prt String
x = do String -> IO ()
putStr String
x
             Handle -> IO ()
hFlush Handle
stdout

  step :: Float
step = Int -> Float
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
w Float -> Float -> Float
forall a. Fractional a => a -> a -> a
/ Int -> Float
forall a b. (Integral a, Num b) => a -> b
fromIntegral ([a] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [a]
xs) :: Float
  go :: Int -> Float -> [a] -> [b] -> IO [b]
go Int
shown Float
loc [a]
todo [b]
done =
    do let new :: Int
new = Float -> Int
forall b. Integral b => Float -> b
forall a b. (RealFrac a, Integral b) => a -> b
floor Float
loc Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
shown
       String -> IO ()
prt (Int -> Char -> String
forall a. Int -> a -> [a]
replicate (Float -> Int
forall b. Integral b => Float -> b
forall a b. (RealFrac a, Integral b) => a -> b
floor Float
loc Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
shown) Char
'=')
       case [a]
todo of
         [] -> String -> IO ()
prt String
"|\n" IO () -> IO [b] -> IO [b]
forall a b. IO a -> IO b -> IO b
forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> [b] -> IO [b]
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return ([b] -> [b]
forall a. [a] -> [a]
reverse [b]
done)
         a
a : [a]
as ->
           do b
b <- a -> IO b
f a
a
              Int -> Float -> [a] -> [b] -> IO [b]
go (Int
shown Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
new) (Float
loc Float -> Float -> Float
forall a. Num a => a -> a -> a
+ Float
step) [a]
as (b
b b -> [b] -> [b]
forall a. a -> [a] -> [a]
: [b]
done)