1 files changed, 33 insertions, 0 deletions
diff --git a/lib/Malodivo/Budget.hs b/lib/Malodivo/Budget.hs
index 657cf14..6a5e8b1 100644
--- a/lib/Malodivo/Budget.hs
+++ b/lib/Malodivo/Budget.hs
@@ -8,11 +8,13 @@ Budget planning in the Kingdom of Malodivo.
 
 module Malodivo.Budget
   ( BillBudget
+  , billsByMinistry
   , manyToMany
   , manyToOne
   , normalizeDown
   ) where
 
+import Control.Arrow ((&&&))
 import GHC.Generics (Generic)
 
 import Data.Aeson (ToJSON)
@@ -20,6 +22,7 @@ import qualified Data.HashMap.Strict as HM
 
 import qualified Malodivo.Types.Bill as B
 import qualified Malodivo.Types.District as D
+import qualified Malodivo.Types.Ministry as M
 
 -- | This is output type. It describes contribution of each district into a bill.
 data BillBudget = BillBudget
@@ -31,6 +34,36 @@ zipBills :: [B.Bill] -> [D.DistrictFunds] -> [BillBudget]
 zipBills = zipWith $ \b df -> BillBudget {bill = b, districts = D.df2di df}
 
 {-|
+Group bills by ministry.
+
+>>> :set -XOverloadedStrings
+>>> import qualified Data.HashMap.Strict as HM
+>>> import qualified Malodivo.Types.Bill as B
+>>> import qualified Malodivo.Types.District as D
+>>> import qualified Malodivo.Types.Ministry as M
+
+>>> let scienceA = B.Bill { B.amount = 10, B.name = "Science A", B.ministry = M.Science }
+>>> let scienceB = B.Bill { B.amount = 20, B.name = "Science B", B.ministry = M.Science }
+>>> let scienceC = B.Bill { B.amount = 30, B.name = "Science B", B.ministry = M.Science }
+>>> let scienceBills = [scienceA, scienceB, scienceC]
+
+>>> let welfareA = B.Bill { B.amount = 100, B.name = "Welfare A", B.ministry = M.Welfare }
+>>> let welfareB = B.Bill { B.amount = 200, B.name = "Welfare B", B.ministry = M.Welfare }
+>>> let welfareBills = [welfareA, welfareB]
+
+>>> let allBills = scienceBills ++ welfareBills
+>>> let groups = billsByMinistry allBills
+>>> HM.size groups
+2
+>>> length <$> HM.lookup M.Science groups
+Just 3
+>>> length <$> HM.lookup M.Welfare groups
+Just 2
+-}
+billsByMinistry :: [B.Bill] -> HM.HashMap M.Ministry [B.Bill]
+billsByMinistry = HM.fromListWith (++) . map (B.ministry &&& return)
+
+{-|
 Normalize list of integers, i. e. proportionally decrease each list element
 so that the sum of all elements does not exceed given limit.