Add manyToManyLimited

1 files changed, 71 insertions, 2 deletions

diff --git a/lib/Malodivo/Budget.hs b/lib/Malodivo/Budget.hs
index 6a5e8b1..8d70629 100644
--- a/lib/Malodivo/Budget.hs
+++ b/lib/Malodivo/Budget.hs
@@ -9,8 +9,9 @@ Budget planning in the Kingdom of Malodivo.
 module Malodivo.Budget
   ( BillBudget
   , billsByMinistry
-  , manyToMany
   , manyToOne
+  , manyToMany
+  , manyToManyLimited
   , normalizeDown
   ) where
 
@@ -44,7 +45,7 @@ Group bills by ministry.
 
 >>> 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 scienceC = B.Bill { B.amount = 30, B.name = "Science C", B.ministry = M.Science }
 >>> let scienceBills = [scienceA, scienceB, scienceC]
 
 >>> let welfareA = B.Bill { B.amount = 100, B.name = "Welfare A", B.ministry = M.Welfare }
@@ -97,6 +98,19 @@ normalizeDown maxSum inList
     norm i = maxSum * i `div` inSum
 
 {-|
+Normalize bills proportionally reducing required amount of each one if total
+sum is above the limit.  This a convenient wrapper of 'normalizeDown'.
+-}
+normalizeBills ::
+     Integer -- ^ Maximum amount all the bills can have.
+  -> [B.Bill] -- ^ Initial bills.
+  -> [B.Bill] -- ^ Bills with reduced amounts.
+normalizeBills top bills = zipWith (\b a -> b {B.amount = a}) bills amounts'
+  where
+    amounts = B.amount <$> bills
+    amounts' = normalizeDown top amounts
+
+{-|
 
 Trivial case: many districts, one bill, no constraints (wishes,
 limits).  We assume that, with no explicit wishes, each district
@@ -188,6 +202,61 @@ manyToMany df bills = zipBills bills allocated
     fundsTotal = sum $ HM.elems df
     allocated = flip fundRaising df <$> normalizeDown fundsTotal billAmounts
 
+{-|
+Districts funding multiple bills. But each ministry may have been limited
+in the amount of funds it can get.
+
+>>> :set -XOverloadedStrings
+>>> import Data.List (find)
+>>> import Data.Maybe (fromJust)
+>>> 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 C", 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 funds = HM.fromList [(D.Lakos, 1000)]
+>>> let findBill b = find (\bb -> B.name (bill bb) == B.name b)
+>>> let amounts budget = map (B.amount . bill . fromJust) $ map (\b -> findBill b budget) allBills
+
+
+If all bills of the specific ministry can't get enough funds due to limits imposed on the ministry,
+then each bill gets less funds, but in proportion of the requested amounts:
+
+>>> let limits = HM.fromList [(M.Science, 6)]
+>>> let contribution = manyToManyLimited funds limits allBills
+>>> amounts contribution
+[1,2,3,100,200]
+>>> let limits = HM.fromList [(M.Welfare, 30)]
+>>> let contribution = manyToManyLimited funds limits allBills
+>>> amounts contribution
+[10,20,30,10,20]
+
+-}
+manyToManyLimited ::
+     D.DistrictFunds -- ^ Amounts of available funds per district.
+  -> M.MinistryLimits -- ^ Maximum funds ministries can get.
+  -> [B.Bill] -- ^ Bills requiring funding.
+  -> [BillBudget] -- ^ Contribution of each district to each bill.
+manyToManyLimited df ml bills = manyToMany df . concat $ HM.elems byMinistry'
+  where
+    byMinistry = billsByMinistry bills
+    byMinistry' = HM.mapWithKey normalize byMinistry
+    normalize m bs =
+      case HM.lookup m ml of
+        Nothing -> bs
+        Just l -> normalizeBills l bs
+
 -- | Helper function to maintain DRY and backward compatibility.
 fundRaising ::
      Integer -- ^ Amount to be raised.