pwncash/lib/Internal/Budget.hs

module Internal.Budget (insertBudget) where

import Control.Monad.Except
import Data.Foldable
import Database.Persist.Monad
import Internal.Database
import Internal.History
import Internal.Types.Main
import Internal.Utils
import RIO hiding (to)
import qualified RIO.List as L
import qualified RIO.Map as M
import qualified RIO.NonEmpty as NE
import qualified RIO.Text as T
import RIO.Time

-- each budget (designated at the top level by a 'name') is processed in the
-- following steps
-- 1. expand all transactions given the desired date range and date patterns for
--    each directive in the budget
-- 2. sort all transactions by date
-- 3. propagate all balances forward, and while doing so assign values to each
--    transaction (some of which depend on the 'current' balance of the
--    target account)
-- 4. assign shadow transactions
-- 5. insert all transactions

insertBudget
  :: (MonadInsertError m, MonadSqlQuery m, MonadFinance m)
  => Budget
  -> m ()
insertBudget
  b@Budget
    { bgtLabel
    , bgtIncomes
    , bgtTransfers
    , bgtShadowTransfers
    , bgtPretax
    , bgtTax
    , bgtPosttax
    , bgtInterval
    } =
    whenHash CTBudget b () $ \key -> do
      (intAllos, _) <- combineError intAlloRes acntRes (,)
      let res1 = mapErrors (insertIncome key bgtLabel intAllos bgtInterval) bgtIncomes
      let res2 = expandTransfers key bgtLabel bgtInterval bgtTransfers
      txs <- combineError (concat <$> res1) res2 (++)
      shadow <- addShadowTransfers bgtShadowTransfers txs
      (_, toIns) <- balanceTxs $ fmap ToInsert $ txs ++ shadow
      void $ insertBudgetTx toIns
    where
      acntRes = mapErrors isNotIncomeAcnt alloAcnts
      intAlloRes = combineError3 pre_ tax_ post_ (,,)
      pre_ = sortAllos bgtPretax
      tax_ = sortAllos bgtTax
      post_ = sortAllos bgtPosttax
      sortAllos = liftExcept . mapErrors sortAllo
      alloAcnts =
        (alloAcnt <$> bgtPretax)
          ++ (alloAcnt <$> bgtTax)
          ++ (alloAcnt <$> bgtPosttax)

-- TODO need to systematically make this function match the history version,
-- which will allow me to use the same balancing algorithm for both
-- balanceTransfers :: [Tx BudgetMeta] -> [KeyEntry]
-- balanceTransfers = undefined

-- balanceTransfers = snd . L.mapAccumR go M.empty . reverse . L.sortOn ftWhen
--   where
--     go bals f@FlatTransfer {ftFrom, ftTo, ftValue = UnbalancedValue {cvValue, cvType}} =
--       let balTo = M.findWithDefault 0 ftTo bals
--           x = amtToMove balTo cvType cvValue
--           bals' = mapAdd_ ftTo x $ mapAdd_ ftFrom (-x) bals
--        in (bals', f {ftValue = x})
--     -- TODO might need to query signs to make this intuitive; as it is this will
--     -- probably work, but for credit accounts I might need to supply a negative
--     -- target value
--     amtToMove _ TFixed x = x
--     amtToMove bal TPercent x = -(x / 100 * bal)
--     amtToMove bal TBalance x = x - bal

insertBudgetTx
  :: (MonadInsertError m, MonadSqlQuery m, MonadFinance m)
  => [InsertTx BudgetMeta]
  -> m ()
insertBudgetTx toInsert = do
  forM_ (groupKey (commitRHash . bmCommit) $ (\x -> (itxCommit x, x)) <$> toInsert) $
    \(c, ts) -> do
      ck <- insert $ bmCommit c
      mapM_ (insertTx ck) ts
  where
    insertTx c InsertTx {itxDate = d, itxDescr = e, itxEntries = ss, itxCommit = BudgetMeta {bmName}} = do
      let anyDeferred = any (isJust . feDeferred) ss
      k <- insert $ TransactionR c d e anyDeferred
      mapM_ (insertBudgetLabel bmName k) ss
    insertBudgetLabel n k entry = do
      sk <- insertEntry k entry
      insert_ $ BudgetLabelR sk n

-- insertBudgetTx
--   :: (MonadInsertError m, MonadSqlQuery m, MonadFinance m)
--   => BalancedTransfer
--   -> m ()
-- insertBudgetTx FlatTransfer {ftFrom, ftTo, ftMeta, ftCur, ftValue, ftDesc, ftWhen} = do
-- ((sFrom, sTo), exchange) <- entryPair ftFrom ftTo ftCur ftValue
-- insertPair sFrom sTo
-- forM_ exchange $ uncurry insertPair
-- where
--   insertPair from to = do
--     k <- insert $ TransactionR (bmCommit ftMeta) ftWhen ftDesc
--     insertBudgetLabel k from
--     insertBudgetLabel k to
--   insertBudgetLabel k entry = do
--     sk <- insertEntry k entry
--     insert_ $ BudgetLabelR sk $ bmName ftMeta

entryPair
  :: (MonadInsertError m, MonadFinance m)
  => TaggedAcnt
  -> TaggedAcnt
  -> CurID
  -> T.Text
  -> Double
  -> m (EntrySet AcntID CurrencyPrec TagID Rational (EntryValue Rational))
entryPair = entryPair_ (fmap (EntryValue TFixed) . roundPrecisionCur)

entryPair_
  :: (MonadInsertError m, MonadFinance m)
  => (CurrencyPrec -> v -> v')
  -> TaggedAcnt
  -> TaggedAcnt
  -> CurID
  -> T.Text
  -> v
  -> m (EntrySet AcntID CurrencyPrec TagID Rational v')
entryPair_ f from to curid com val = do
  cp <- lookupCurrency curid
  return $ pair cp from to (f cp val)
  where
    halfEntry :: a -> [t] -> HalfEntrySet a c t v
    halfEntry a ts =
      HalfEntrySet
        { hesPrimary = Entry {eAcnt = a, eValue = (), eComment = com, eTags = ts}
        , hesOther = []
        }
    pair cp (TaggedAcnt fa fts) (TaggedAcnt ta tts) v =
      EntrySet
        { esCurrency = cp
        , esTotalValue = v
        , esFrom = halfEntry fa fts
        , esTo = halfEntry ta tts
        }

sortAllo :: MultiAllocation v -> InsertExcept (DaySpanAllocation v)
sortAllo a@Allocation {alloAmts = as} = do
  bs <- foldSpan [] $ L.sortOn amtWhen as
  return $ a {alloAmts = reverse bs}
  where
    foldSpan acc [] = return acc
    foldSpan acc (x : xs) = do
      let start = amtWhen x
      res <- case xs of
        [] -> resolveDaySpan start
        (y : _) -> resolveDaySpan_ (intStart $ amtWhen y) start
      foldSpan (x {amtWhen = res} : acc) xs

--------------------------------------------------------------------------------
-- Income

-- TODO this will scan the interval allocations fully each time
-- iteration which is a total waste, but the fix requires turning this
-- loop into a fold which I don't feel like doing now :(
insertIncome
  :: (MonadInsertError m, MonadFinance m)
  => CommitRId
  -> T.Text
  -> IntAllocations
  -> Maybe Interval
  -> Income
  -> m [Tx BudgetMeta]
insertIncome
  key
  name
  (intPre, intTax, intPost)
  localInterval
  Income
    { incWhen
    , incCurrency
    , incFrom
    , incPretax
    , incPosttax
    , incTaxes
    , incToBal
    , incGross
    , incPayPeriod
    } =
    combineErrorM
      (combineError incRes nonIncRes (,))
      (combineError precRes dayRes (,))
      $ \_ (precision, days) -> do
        let gross = roundPrecision precision incGross
        concat <$> foldDays (allocate precision gross) start days
    where
      incRes = isIncomeAcnt $ taAcnt incFrom
      nonIncRes =
        mapErrors isNotIncomeAcnt $
          taAcnt incToBal
            : (alloAcnt <$> incPretax)
            ++ (alloAcnt <$> incTaxes)
            ++ (alloAcnt <$> incPosttax)
      precRes = lookupCurrencyPrec incCurrency
      dayRes = askDays incWhen localInterval
      start = fromGregorian' $ pStart incPayPeriod
      pType' = pType incPayPeriod
      meta = BudgetMeta key name
      flatPre = concatMap flattenAllo incPretax
      flatTax = concatMap flattenAllo incTaxes
      flatPost = concatMap flattenAllo incPosttax
      sumAllos = sum . fmap faValue
      -- TODO ensure these are all the "correct" accounts
      allocate precision gross prevDay day = do
        scaler <- liftExcept $ periodScaler pType' prevDay day
        let (preDeductions, pre) =
              allocatePre precision gross $
                flatPre ++ concatMap (selectAllos day) intPre
        let tax =
              allocateTax precision gross preDeductions scaler $
                flatTax ++ concatMap (selectAllos day) intTax
            aftertaxGross = gross - sumAllos (tax ++ pre)
        let post =
              allocatePost precision aftertaxGross $
                flatPost ++ concatMap (selectAllos day) intPost
        let balance = aftertaxGross - sumAllos post
        -- TODO double or rational here?
        primary <-
          entryPair
            incFrom
            incToBal
            incCurrency
            "balance after deductions"
            (fromRational balance)
        allos <- mapErrors (allo2Trans meta day incFrom) (pre ++ tax ++ post)
        let bal =
              Tx
                { txCommit = meta
                , txDate = day
                , txPrimary = primary
                , txOther = []
                , txDescr = "balance after deductions"
                }
        if balance < 0
          then throwError $ InsertException [IncomeError day name balance]
          else return (bal : allos)

periodScaler
  :: PeriodType
  -> Day
  -> Day
  -> InsertExcept PeriodScaler
periodScaler pt prev cur = return scale
  where
    n = fromIntegral $ workingDays wds prev cur
    wds = case pt of
      Hourly HourlyPeriod {hpWorkingDays} -> hpWorkingDays
      Daily ds -> ds
    scale precision x = case pt of
      Hourly HourlyPeriod {hpAnnualHours, hpDailyHours} ->
        fromRational (rnd $ x / fromIntegral hpAnnualHours)
          * fromIntegral hpDailyHours
          * n
      Daily _ -> x * n / 365.25
      where
        rnd = roundPrecision precision

-- ASSUME start < end
workingDays :: [Weekday] -> Day -> Day -> Natural
workingDays wds start end = fromIntegral $ daysFull + daysTail
  where
    interval = diffDays end start
    (nFull, nPart) = divMod interval 7
    daysFull = fromIntegral (length wds') * nFull
    daysTail = fromIntegral $ length $ takeWhile (< nPart) wds'
    startDay = dayOfWeek start
    wds' = L.sort $ (\x -> diff (fromWeekday x) startDay) <$> L.nub wds
    diff a b = fromIntegral $ mod (fromEnum a - fromEnum b) 7

-- ASSUME days is a sorted list
foldDays
  :: MonadInsertError m
  => (Day -> Day -> m a)
  -> Day
  -> [Day]
  -> m [a]
foldDays f start days = case NE.nonEmpty days of
  Nothing -> return []
  Just ds
    | any (start >) ds ->
        throwError $
          InsertException [PeriodError start $ minimum ds]
    | otherwise ->
        combineErrors $
          snd $
            L.mapAccumL (\prevDay day -> (day, f prevDay day)) start days

isIncomeAcnt :: (MonadInsertError m, MonadFinance m) => AcntID -> m ()
isIncomeAcnt = checkAcntType IncomeT

isNotIncomeAcnt :: (MonadInsertError m, MonadFinance m) => AcntID -> m ()
isNotIncomeAcnt = checkAcntTypes (AssetT :| [EquityT, ExpenseT, LiabilityT])

checkAcntType
  :: (MonadInsertError m, MonadFinance m)
  => AcntType
  -> AcntID
  -> m ()
checkAcntType t = checkAcntTypes (t :| [])

checkAcntTypes
  :: (MonadInsertError m, MonadFinance m)
  => NE.NonEmpty AcntType
  -> AcntID
  -> m ()
checkAcntTypes ts i = void $ go =<< lookupAccountType i
  where
    go t
      | t `L.elem` ts = return i
      | otherwise = throwError $ InsertException [AccountError i ts]

flattenAllo :: SingleAllocation v -> [FlatAllocation v]
flattenAllo Allocation {alloAmts, alloCur, alloTo} = fmap go alloAmts
  where
    go Amount {amtValue, amtDesc} =
      FlatAllocation
        { faCur = alloCur
        , faTo = alloTo
        , faValue = amtValue
        , faDesc = amtDesc
        }

-- ASSUME allocations are sorted
selectAllos :: Day -> DaySpanAllocation v -> [FlatAllocation v]
selectAllos day Allocation {alloAmts, alloCur, alloTo} =
  go <$> filter ((`inDaySpan` day) . amtWhen) alloAmts
  where
    go Amount {amtValue, amtDesc} =
      FlatAllocation
        { faCur = alloCur
        , faTo = alloTo
        , faValue = amtValue
        , faDesc = amtDesc
        }

allo2Trans
  :: (MonadInsertError m, MonadFinance m)
  => BudgetMeta
  -> Day
  -> TaggedAcnt
  -> FlatAllocation Rational
  -> m (Tx BudgetMeta)
allo2Trans meta day from FlatAllocation {faValue, faTo, faDesc, faCur} = do
  -- TODO double here?
  p <- entryPair from faTo faCur faDesc (fromRational faValue)
  return
    Tx
      { txCommit = meta
      , txDate = day
      , txPrimary = p
      , txOther = []
      , txDescr = faDesc
      }

allocatePre
  :: Natural
  -> Rational
  -> [FlatAllocation PretaxValue]
  -> (M.Map T.Text Rational, [FlatAllocation Rational])
allocatePre precision gross = L.mapAccumR go M.empty
  where
    go m f@FlatAllocation {faValue} =
      let c = preCategory faValue
          p = preValue faValue
          v =
            if prePercent faValue
              then (roundPrecision 3 p / 100) * gross
              else roundPrecision precision p
       in (mapAdd_ c v m, f {faValue = v})

allocateTax
  :: Natural
  -> Rational
  -> M.Map T.Text Rational
  -> PeriodScaler
  -> [FlatAllocation TaxValue]
  -> [FlatAllocation Rational]
allocateTax precision gross preDeds f = fmap (fmap go)
  where
    go TaxValue {tvCategories, tvMethod} =
      let agi = gross - sum (mapMaybe (`M.lookup` preDeds) tvCategories)
       in case tvMethod of
            TMPercent p ->
              roundPrecision precision $
                fromRational $
                  roundPrecision 3 p / 100 * agi
            TMBracket TaxProgression {tpDeductible, tpBrackets} ->
              let taxDed = roundPrecision precision $ f precision tpDeductible
               in foldBracket f precision (agi - taxDed) tpBrackets

-- | Compute effective tax percentage of a bracket
-- The algorithm can be thought of in three phases:
-- 1. Find the highest tax bracket by looping backward until the AGI is less
--    than the bracket limit
-- 2. Computing the tax in the top bracket by subtracting the AGI from the
--    bracket limit and multiplying by the tax percentage.
-- 3. Adding all lower brackets, which are just the limit of the bracket less
--    the amount of the lower bracket times the percentage.
--
-- In reality, this can all be done with one loop, but it isn't clear these
-- three steps are implemented from this alone.
foldBracket :: PeriodScaler -> Natural -> Rational -> [TaxBracket] -> Rational
foldBracket f precision agi bs = fst $ foldr go (0, agi) $ L.sortOn tbLowerLimit bs
  where
    go TaxBracket {tbLowerLimit, tbPercent} a@(acc, remain) =
      let l = roundPrecision precision $ f precision tbLowerLimit
          p = roundPrecision 3 tbPercent / 100
       in if remain >= l then (acc + p * (remain - l), l) else a

allocatePost
  :: Natural
  -> Rational
  -> [FlatAllocation PosttaxValue]
  -> [FlatAllocation Rational]
allocatePost precision aftertax = fmap (fmap go)
  where
    go PosttaxValue {postValue, postPercent} =
      let v = postValue
       in if postPercent
            then aftertax * roundPrecision 3 v / 100
            else roundPrecision precision v

--------------------------------------------------------------------------------
-- Standalone Transfer

expandTransfers
  :: (MonadInsertError m, MonadSqlQuery m, MonadFinance m)
  => CommitRId
  -> T.Text
  -> Maybe Interval
  -> [BudgetTransfer]
  -> m [Tx BudgetMeta]
expandTransfers key name localInterval ts = do
  txs <-
    fmap (L.sortOn txDate . concat) $
      combineErrors $
        fmap (expandTransfer key name) ts
  case localInterval of
    Nothing -> return txs
    Just i -> do
      bounds <- liftExcept $ resolveDaySpan i
      return $ filter (inDaySpan bounds . txDate) txs

expandTransfer
  :: (MonadInsertError m, MonadSqlQuery m, MonadFinance m)
  => CommitRId
  -> T.Text
  -> BudgetTransfer
  -> m [Tx BudgetMeta]
expandTransfer key name Transfer {transAmounts, transTo, transCurrency, transFrom} = do
  fmap concat $ mapErrors go transAmounts
  where
    go
      Amount
        { amtWhen = pat
        , amtValue = TransferValue {tvVal = v, tvType = t}
        , amtDesc = desc
        } =
        withDates pat $ \day -> do
          let meta = BudgetMeta {bmCommit = key, bmName = name}
          p <- entryPair_ (\_ x -> EntryValue t $ toRational x) transFrom transTo transCurrency desc v
          return
            Tx
              { txCommit = meta
              , txDate = day
              , txPrimary = p
              , txOther = []
              , txDescr = desc
              }

withDates
  :: (MonadSqlQuery m, MonadFinance m, MonadInsertError m)
  => DatePat
  -> (Day -> m a)
  -> m [a]
withDates dp f = do
  bounds <- askDBState kmBudgetInterval
  days <- liftExcept $ expandDatePat bounds dp
  combineErrors $ fmap f days

--------------------------------------------------------------------------------
-- shadow transfers

-- TODO this is going to be O(n*m), which might be a problem?
addShadowTransfers
  :: (MonadInsertError m, MonadFinance m)
  => [ShadowTransfer]
  -> [Tx BudgetMeta]
  -> m [Tx BudgetMeta]
addShadowTransfers ms txs = mapErrors go txs
  where
    go tx = do
      es <- catMaybes <$> mapErrors (fromShadow tx) ms
      return $ tx {txOther = es}

fromShadow
  :: (MonadInsertError m, MonadFinance m)
  => Tx BudgetMeta
  -> ShadowTransfer
  -> m (Maybe (EntrySet AcntID CurrencyPrec TagID Rational (Either Double (EntryValue Rational))))
fromShadow tx ShadowTransfer {stFrom, stTo, stDesc, stRatio, stCurrency, stMatch} = do
  res <- liftExcept $ shadowMatches stMatch tx
  es <- entryPair_ (\_ v -> Left v) stFrom stTo stCurrency stDesc stRatio
  return $ if not res then Nothing else Just es

shadowMatches :: TransferMatcher -> Tx BudgetMeta -> InsertExcept Bool
shadowMatches TransferMatcher {tmFrom, tmTo, tmDate} Tx {txPrimary, txDate} = do
  -- NOTE this will only match against the primary entry set since those
  -- are what are guaranteed to exist from a transfer
  -- valRes <- valMatches tmVal $ esTotalValue $ txPrimary
  return $
    memberMaybe (eAcnt $ hesPrimary $ esFrom txPrimary) tmFrom
      && memberMaybe (eAcnt $ hesPrimary $ esTo txPrimary) tmTo
      && maybe True (`dateMatches` txDate) tmDate
  where
    -- && valRes

    memberMaybe x AcntSet {asList, asInclude} =
      (if asInclude then id else not) $ x `elem` asList

--------------------------------------------------------------------------------
-- random

-- initialCurrency :: TransferCurrency -> CurID
-- initialCurrency (NoX c) = c
-- initialCurrency (X Exchange {xFromCur = c}) = c

alloAcnt :: Allocation w v -> AcntID
alloAcnt = taAcnt . alloTo

data UnbalancedValue = UnbalancedValue
  { cvType :: !TransferType
  , cvValue :: !Rational
  }
  deriving (Show)

-- TODO need to make this into the same ish thing as the Tx/EntrySet structs
-- in the history algorithm, which will entail resolving the budget currency
-- stuff earlier in the chain, and preloading multiple entries into this thing
-- before balancing.
-- type UnbalancedTransfer = FlatTransfer UnbalancedValue

-- ubt2tx :: UnbalancedTransfer -> Tx BudgetMeta
-- ubt2tx
--   FlatTransfer
--     { ftFrom
--     , ftTo
--     , ftValue
--     , ftWhen
--     , ftDesc
--     , ftMeta
--     , ftCur
--     } =
--     Tx
--       { txDescr = ftDesc
--       , txDate = ftWhen
--       , txPrimary = p
--       , txOther = maybeToList os
--       , txCommit = ftMeta
--       }
--     where
--       (p, os) = entries ftCur
--       entries (NoX curid) = (pair curid ftFrom ftTo ftValue, Nothing)
--       entries (X Exchange {xFromCur, xToCur, xAcnt, xRate}) =
--         let middle = TaggedAcnt xAcnt []
--             p1 = pair xFromCur ftFrom middle ftValue
--             p2 = pair xToCur middle ftTo (ftValue * roundPrecision 3 xRate)
--          in (p1, Just p2)
--       pair c (TaggedAcnt fa fts) (TaggedAcnt ta tts) v =
--         EntrySet
--           { esTotalValue = v
--           , esCurrency = c
--           , esFrom =
--               HalfEntrySet
--                 { hesPrimary =
--                     Entry
--                       { eValue = ()
--                       , eComment = ""
--                       , eAcnt = fa
--                       , eTags = fts
--                       }
--                 , hesOther = []
--                 }
--           , esTo =
--               HalfEntrySet
--                 { hesPrimary =
--                     Entry
--                       { eValue = ()
--                       , eComment = ""
--                       , eAcnt = ta
--                       , eTags = tts
--                       }
--                 , hesOther = []
--                 }
--           }

-- type BalancedTransfer = FlatTransfer Rational

-- data FlatTransfer v = FlatTransfer
--   { ftFrom :: !TaggedAcnt
--   , ftTo :: !TaggedAcnt
--   , ftValue :: !v
--   , ftWhen :: !Day
--   , ftDesc :: !T.Text
--   , ftMeta :: !BudgetMeta
--   , ftCur :: !TransferCurrency
--   }
--   deriving (Show)

data BudgetMeta = BudgetMeta
  { bmCommit :: !CommitR
  , bmName :: !T.Text
  }
  deriving (Show)

type IntAllocations =
  ( [DaySpanAllocation PretaxValue]
  , [DaySpanAllocation TaxValue]
  , [DaySpanAllocation PosttaxValue]
  )

type DaySpanAllocation = Allocation DaySpan

type EntryPair = (KeyEntry, KeyEntry)

type PeriodScaler = Natural -> Double -> Double

data FlatAllocation v = FlatAllocation
  { faValue :: !v
  , faDesc :: !T.Text
  , faTo :: !TaggedAcnt
  , faCur :: !CurID
  }
  deriving (Functor, Show)