{-# LANGUAGE Unsafe, ScopedTypeVariables #-}
-- | Progress estimates.
module Data.Progress (progress, progressWithFile, progressWithCalls, progress', progressWithCalls') where

import System.IO.Unsafe
import System.IO
import System.Console.ANSI
import Data.Data
import Control.Monad.Identity
import Control.Monad
import Control.Concurrent
import Control.Exception
import Data.List

newtype Size t = Size { unSize :: Integer }

size :: (Data t) => t -> Integer
size x = unSize (gfoldl (\(Size n) y -> Size (n + size y)) (const (Size 1)) x)

{-# NOINLINE active #-}
active :: MVar Integer
active = unsafePerformIO (newMVar $ -1)

putBar :: Integer -> IO ()
putBar percent = modifyMVar_ active $ \n ->
	if n < percent then do
		setCursorColumn 0
		let s = show $ min 100 percent
		putStr $ replicate (3 - length s) ' ' ++ s
		return percent
	else
		return n

-- Prevent simultaneous execution of multiple progress estimates.
withActive m = do
	b <- modifyMVar active (\n -> return $ if n == -1 then (0, False) else (n, True))
	if b then
		m
		else
		finally m (modifyMVar_ active (const $ return $ -1))

fork' m = do
	caps <- getNumCapabilities
	if caps >= 2 then
			void $ forkIO m
		else
			void m

-- | Estimate progress based on thunks forced.
{-# NOINLINE progress #-}
progress f dat = withActive $ do
	putStr "\n  0%"
	sz <- newMVar $ -1
	count <- newMVar 0

	-- Compute the size in a separate thread, so as not to be embarrassed
	-- by the parallelism.
	fork' $ do
		let s = size dat
		evaluate s
		modifyMVar_ sz $ const $ return s

	-- The 'rec' function will make a copy of the input data
	-- structure, with I/O effects added that print a progress bar
	-- as the data structure is forced.
	let
		rec :: (Data t) => t -> t
		rec dat = runIdentity $ gfoldl
			(\(Identity f) x -> unsafePerformIO $ do
				s <- readMVar sz
				modifyMVar_ count $ \n ->
					if n /= -1 then do
						let n' = n + 1
						unless (s == -1) $ putBar (n' * 100 `quot` s)
						return n'
					else
						return n
				return $ Identity $ f $ rec x)
			Identity
			dat

	-- Run the function on the copy.
	finally
		(do
		res <- f $ rec dat
		evaluate res
		return res)
		(do
		-- Record that the function is done so no more bars are printed.
		modifyMVar_ count $ const $ return $ -1
		putBar 100
		putStrLn "")

try' :: IO t -> IO (Either SomeException t)
try' = try

-- | ...based on amount of file consumed.
{-# NOINLINE progressWithFile #-}
progressWithFile f hdl = withActive $ do
	putStr "\n  0%"

	-- Check the position of the handle periodically and print
	-- a progress bar.
	thd <- try' $ do
		sz <- liftM fromInteger $ hFileSize hdl
		forkIO $ foldM_ (\_ _ -> do
			n <- liftM fromInteger $ hTell hdl
			putBar (n * 100 `quot` sz)
			threadDelay 500000
			return ())
			()
			(repeat ())

	finally
	-- Run the function.
		(f hdl)
		(do
		-- Again, prevent the progress bar from being printed once
		-- the function is done.
		try' $ either (\_ -> return ()) killThread thd
		putBar 100
		putStrLn "")

-- | ...based on number of recursive calls.
--
--   It returns a result equivalent to that of /fix f x/.
{-# NOINLINE progressWithCalls #-}
progressWithCalls f x = withActive $ do
	putStr "\n  0%"

	-- As the function runs, the procedure will estimate the
	-- geometric sequence giving the recursion costs.
	parms <- newMVar (-1, 20000, 1, 0)
	let rec depth count x = do
		-- Do a recursive call. The call gets a fresh recursion counter.
		count' <- newMVar (1, 1)
		res <- f (rec (depth + 1) count') x
		evaluate res

		(x, y) <- readMVar count'
		modifyMVar_ count $ \(_, z) -> return (y, z + y)

		modifyMVar_ parms $ \tup@(dep, rPrev, yPrev, total) -> do
			-- Calculate the new parameters.
			let tup'@(dep', ratio', y', total') = if total < 0 then
					tup
				else if dep == -1 || depth <= dep then
					(depth, (4 * rPrev + y * 65536 `quot` x) `quot` 5, y, total + 1)
				else
					(dep, rPrev, yPrev, total + 1)

			-- Print a progress bar with the new estimate.
			when (y' >= 100) $ putBar (total' * 100 * 65536 ^ dep' `quot` (y' * ratio' ^ dep'))

			return tup'

		return res
	count <- newMVar (1, 1)
	finally
		(do
		res <- rec 0 count x
		return res)
		(do
		modifyMVar_ parms $ const $ return (0, 0, 0, -1)
		putBar 100
		putStrLn "")

-- | Adapters for pure functions.
{-# INLINE progress' #-}
progress' f = progress (return . f)

{-# INLINE progressWithCalls' #-}
progressWithCalls' f = progressWithCalls (\g -> return . f (unsafePerformIO . g))

{-# RULES
"progress" forall f (g :: forall t. t -> f t) x. progress (f . progress g) x = g x >>= progress (f . return)
"progressWithCalls" forall f g. progressWithCalls (f (progressWithCalls g)) = progressWithCalls (\h -> either (g (h . Left)) (f (h . Left) (h . Right))) . Right
  #-}

{-quicksort _ [] = []
quicksort f (x:xs) = f tk ++ x : f dr where
	(tk, dr) = partition (<x) xs

ex :: IO ()
ex = void $ do
	rs :: [Int] <- liftM (take 1000000) getRandoms
	progress' (fix $ \f -> (`using` evalList rseq) . quicksort f) rs-}