add progress bar

4 files changed, 29 insertions, 5 deletions

diff --git a/bin/cluster.hs b/bin/cluster.hs
index 6b0d917..4c182a1 100644
--- a/bin/cluster.hs
+++ b/bin/cluster.hs
@@ -1,4 +1,6 @@
 {-# LANGUAGE ViewPatterns #-}
+{-# LANGUAGE ScopedTypeVariables #-}
+{-# LANGUAGE LambdaCase #-}
 
 module Main where
 
@@ -11,7 +13,12 @@ import GHC.Compact
 import Numeric.Log hiding (sum)
 import Options.Applicative
 import PPL hiding (binom)
-import System.Random (mkStdGen, setStdGen)
+import System.Random (mkStdGen, setStdGen, StdGen, random, split)
+import qualified Streaming as S
+import qualified Streaming.Prelude as S
+import Streaming.Prelude (each, fold, Of, Stream, yield)
+import System.IO ( hSetBuffering, BufferMode(NoBuffering), stdout)
+import System.ProgressBar
 
 cumsum = scanl1 (+)
 
@@ -116,12 +123,23 @@ main = run =<< execParser opts
       [(r, "")] -> if r <= 1 && r > 0 then Right r else Left "mhfrac not a valid probability"
       _ -> Left "mhfrac not a valid probability"
 
+
+takeWithProgress :: S.MonadIO m => Int -> Stream (Of a) m r -> Stream (Of a) m ()
+takeWithProgress n str = do
+  pb <- S.liftIO $ newProgressBar defStyle 10 (Progress 0 n ())
+  S.mapM (update pb) $ S.take n str
+  where
+    update pb x = do
+      S.liftIO $ incProgress pb 1
+      pure x
+
 run opts = do
-  setStdGen . mkStdGen $ seed opts
   (hdr : lines) <- lines <$> readFile (input opts)
   let dbl = round . read :: String -> Int
+      g = mkStdGen $ seed opts
   parsed <- compact $ map (map dbl . tail . words) lines
-  ((ps, cl), _) <- foldl1' (\a c -> if mml a < mml c then a else c) . take (nsamples opts) <$> mh (mhfrac opts) 0.5 (model $ getCompact parsed)
+  hSetBuffering stdout NoBuffering
+  ((ps, cl), _) <- S.fold_ (\l r -> if mml l < mml r then l else r) (([[]],[]), -1/0) id . takeWithProgress (nsamples opts) $ mh g (mhfrac opts) 0.5 (model $ getCompact parsed)
   writeFile (propsPath opts) . unlines $ map (intercalate "," . map show) ps
   writeFile (clusterPath opts) . unlines $ map show cl
   where
@@ -129,3 +147,5 @@ run opts = do
       where
         tab = tabulate cl
         sum' f = sum . map f
+
+
diff --git a/cabal.project b/cabal.project
index 1869ff4..e2b3d7d 100644
--- a/cabal.project
+++ b/cabal.project
@@ -3,5 +3,5 @@ packages: *.cabal
 source-repository-package
   type: git
   location: https://vk3.wtf/cgit/ppl.git
-  tag: 275ed22d5050488e6d40bb5800f1ade9c30d8a76
-  --sha256: 0s3knn77s4pc277j3z0n80qh90kzlayl712b4i7sq43ib3p340ii
+  tag: 7653e357f04aa39c1e96037bf1ea2e4338f8ae76
+  --sha256: 1zdx396937q5pxyp1yca1ns6b19bz0vwp5c39zqkr93968nar25f
diff --git a/package.yaml b/package.yaml
index f72db5a..9e8b5e4 100644
--- a/package.yaml
+++ b/package.yaml
@@ -17,6 +17,8 @@ executables:
       - random
       - optparse-applicative
       - ghc-compact
+      - terminal-progress-bar
+      - streaming
   draw:
     main: draw.hs
     source-dirs: bin
diff --git a/phylogenies.cabal b/phylogenies.cabal
index ae63597..be88d4e 100644
--- a/phylogenies.cabal
+++ b/phylogenies.cabal
@@ -27,6 +27,8 @@ executable cluster
     , optparse-applicative
     , ppl
     , random
+    , streaming
+    , terminal-progress-bar
   default-language: Haskell2010
 
 executable draw