From efffda378ae168ce5ba6fdf72435e72e7d914a60 Mon Sep 17 00:00:00 2001
From: ndwarshuis <ndwar@yavin4.ch>
Date: Tue, 13 Jun 2023 20:12:29 -0400
Subject: [PATCH] ADD calculations for running balances in statements

---
 app/Main.hs             | 12 +++----
 lib/Internal/Budget.hs  |  4 ---
 lib/Internal/History.hs | 79 ++++++++++++++++++++++++++---------------
 lib/Internal/Utils.hs   |  4 +++
 4 files changed, 61 insertions(+), 38 deletions(-)

diff --git a/app/Main.hs b/app/Main.hs
index 666c943..597a8c1 100644
--- a/app/Main.hs
+++ b/app/Main.hs
@@ -180,13 +180,13 @@ runSync c = do
 
     -- update the DB
     runSqlQueryT pool $ withTransaction $ flip runReaderT state $ do
-      let hTransRes = mapErrors insertHistTransfer hTs
-      let bgtRes = mapErrors insertBudget $ budget config
+      let runHist = do
+            ts <- catMaybes <$> mapErrors readHistTransfer hTs
+            insertHistory $ bSs ++ ts
+      let runBudget = mapErrors insertBudget $ budget config
       updateDBState updates -- TODO this will only work if foreign keys are deferred
-      res <- runExceptT $ do
-        mapM_ (uncurry insertHistStmt) bSs
-        combineError hTransRes bgtRes $ \_ _ -> ()
-      rerunnableIO $ fromEither res
+      res <- runExceptT $ combineError runHist runBudget $ \_ _ -> ()
+      rerunnableIO $ fromEither res -- TODO why is this here?
   where
     root = takeDirectory c
     err (InsertException es) = do
diff --git a/lib/Internal/Budget.hs b/lib/Internal/Budget.hs
index ec92a72..1e0a3db 100644
--- a/lib/Internal/Budget.hs
+++ b/lib/Internal/Budget.hs
@@ -74,10 +74,6 @@ balanceTransfers = snd . L.mapAccumR go M.empty . reverse . L.sortOn ftWhen
     amtToMove bal BTPercent x = -(x / 100 * bal)
     amtToMove bal BTTarget x = x - bal
 
--- TODO this seems too general for this module
-mapAdd_ :: (Ord k, Num v) => k -> v -> M.Map k v -> M.Map k v
-mapAdd_ k v = M.alter (maybe (Just v) (Just . (+ v))) k
-
 insertBudgetTx
   :: (MonadInsertError m, MonadSqlQuery m, MonadFinance m)
   => BalancedTransfer
diff --git a/lib/Internal/History.hs b/lib/Internal/History.hs
index 4c6d630..3042569 100644
--- a/lib/Internal/History.hs
+++ b/lib/Internal/History.hs
@@ -1,8 +1,8 @@
 module Internal.History
-  ( splitHistory
+  ( readHistStmt
   , readHistTransfer
-  , readHistStmt
   , insertHistory
+  , splitHistory
   )
 where
 
@@ -22,26 +22,16 @@ import qualified RIO.Text as T
 import RIO.Time
 import qualified RIO.Vector as V
 
-splitHistory :: [History] -> ([HistTransfer], [Statement])
-splitHistory = partitionEithers . fmap go
-  where
-    go (HistTransfer x) = Left x
-    go (HistStatement x) = Right x
-
-insertHistory
-  :: (MonadInsertError m, MonadSqlQuery m, MonadFinance m)
-  => [(CommitR, [RawTx])]
-  -> m ()
-insertHistory hs = do
-  bs <- balanceTxs $ concatMap (\(c, xs) -> fmap (c,) xs) hs
-  forM_ (groupKey (\(CommitR h _) -> h) bs) $ \(c, ts) -> do
-    ck <- insert c
-    mapM_ (insertTx ck) ts
-
-groupKey :: Ord c => (a -> c) -> [(a, b)] -> [(a, [b])]
-groupKey f = fmap go . NE.groupAllWith (f . fst)
-  where
-    go xs@((c, _) :| _) = (c, NE.toList $ fmap snd xs)
+-- readHistory
+--   :: (MonadInsertError m, MonadFinance m, MonadUnliftIO m)
+--   => FilePath
+--   -> [History]
+--   -> m [(CommitR, [RawTx])]
+-- readHistory root hs = do
+--   let (ts, ss) = splitHistory hs
+--   ts' <- catMaybes <$> mapErrorsIO readHistTransfer ts
+--   ss' <- catMaybes <$> mapErrorsIO (readHistStmt root) ss
+--   return $ ts' ++ ss'
 
 readHistTransfer
   :: (MonadInsertError m, MonadFinance m)
@@ -64,6 +54,11 @@ readHistTransfer
             return $ fmap tx days
       concat <$> mapErrors go amts
 
+groupKey :: Ord c => (a -> c) -> [(a, b)] -> [(a, [b])]
+groupKey f = fmap go . NE.groupAllWith (f . fst)
+  where
+    go xs@((c, _) :| _) = (c, NE.toList $ fmap snd xs)
+
 readHistStmt
   :: (MonadUnliftIO m, MonadFinance m)
   => FilePath
@@ -74,6 +69,22 @@ readHistStmt root i = whenHash_ CTImport i $ do
   bounds <- askDBState kmStatementInterval
   return $ filter (inDaySpan bounds . txDate) bs
 
+splitHistory :: [History] -> ([HistTransfer], [Statement])
+splitHistory = partitionEithers . fmap go
+  where
+    go (HistTransfer x) = Left x
+    go (HistStatement x) = Right x
+
+insertHistory
+  :: (MonadInsertError m, MonadSqlQuery m, MonadFinance m)
+  => [(CommitR, [RawTx])]
+  -> m ()
+insertHistory hs = do
+  bs <- balanceTxs $ concatMap (\(c, xs) -> fmap (c,) xs) hs
+  forM_ (groupKey (\(CommitR h _) -> h) bs) $ \(c, ts) -> do
+    ck <- insert c
+    mapM_ (insertTx ck) ts
+
 --------------------------------------------------------------------------------
 -- low-level transaction stuff
 
@@ -325,16 +336,28 @@ balanceTxs ts = do
     (cs, ts') = L.unzip $ L.sortOn (txDate . snd) ts
 
 balanceTxTargets
-  :: M.Map a Rational
+  :: (Ord a, Ord c)
+  => M.Map (a, c) Rational
   -> Tx (Entry a (Deferred Rational) c t)
-  -> (M.Map a Rational, Tx (Entry a (Maybe Rational) c t))
-balanceTxTargets = undefined
+  -> (M.Map (a, c) Rational, Tx (Entry a (Maybe Rational) c t))
+balanceTxTargets bals t@Tx {txEntries} = (bals', t {txEntries = es})
+  where
+    (bals', es) = L.mapAccumR balanceEntryTargets bals txEntries
 
 balanceEntryTargets
-  :: M.Map a Rational
+  :: (Ord a, Ord c)
+  => M.Map (a, c) Rational
   -> Entry a (Deferred Rational) c t
-  -> (M.Map a Rational, Entry a (Maybe Rational) c t)
-balanceEntryTargets = undefined
+  -> (M.Map (a, c) Rational, Entry a (Maybe Rational) c t)
+balanceEntryTargets bals e@Entry {eValue, eAcnt, eCurrency} = (bals', e {eValue = v})
+  where
+    key = (eAcnt, eCurrency)
+    curBal = M.findWithDefault 0 key bals
+    v = case eValue of
+      ConstD x -> Just x
+      Target x -> Just $ x - curBal
+      Derive -> Nothing
+    bals' = maybe bals (\y -> mapAdd_ key y bals) v
 
 balanceTx
   :: (MonadInsertError m, MonadFinance m)
diff --git a/lib/Internal/Utils.hs b/lib/Internal/Utils.hs
index e69d609..72ce337 100644
--- a/lib/Internal/Utils.hs
+++ b/lib/Internal/Utils.hs
@@ -55,6 +55,7 @@ module Internal.Utils
   , lookupCurrencyKey
   , lookupCurrencyPrec
   , lookupTag
+  , mapAdd_
   )
 where
 
@@ -855,6 +856,9 @@ unlessLefts_ e f = unlessLefts e (\x -> void (f x) >> return mzero)
 -- thrice :: (a -> d) -> (b -> e) -> (c -> f) -> (a, b, c) -> (d, e, f)
 -- thrice f1 f2 f3 (x, y, z) = (f1 x, f2 y, f3 z)
 
+mapAdd_ :: (Ord k, Num v) => k -> v -> M.Map k v -> M.Map k v
+mapAdd_ k v = M.alter (maybe (Just v) (Just . (+ v))) k
+
 uncurry3 :: (a -> b -> c -> d) -> (a, b, c) -> d
 uncurry3 f (a, b, c) = f a b c