{-|

Budget planning in the Kingdom of Malodivo.

-}
module Malodivo.Budget
  ( billsByMinistry
  , manyToOne
  , manyToMany
  , manyToManyLimited
  , normalizeDown
  ) where

import Control.Arrow ((&&&))

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

{-|
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 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 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.

>>> normalizeDown 10 [10, 20, 30, 40]
[1,2,3,4]

>>> normalizeDown 9 [11]
[9]


If requested maximum sum is larger or equal to the sum of input list, the
list is not changed:

>>> normalizeDown 101 [10, 20, 30, 40]
[10,20,30,40]

__TODO__ It /should/ be generalized to
<https://en.wikipedia.org/wiki/Partition_(number_theory) integer partition>
so that the sums of input and output lists are equal to each other. Currently
it is not guaranteed and the numbers are rounded down if necessary.
-}
normalizeDown ::
     Integer -- ^ Maximum sum of all list items.
  -> [Integer] -- ^ Initial list.
  -> [Integer] -- ^ Normalized list.
normalizeDown maxSum inList
  | inSum <= maxSum = inList
  | otherwise = norm <$> inList
  where
    inSum = sum 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
wants to contribute all its funds.

>>> :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 medium = B.Bill { B.amount = 30, B.name = "A medium bill", B.ministry = M.Science }
>>> let one = B.Bill { B.amount = 1, B.name = "A trivial bill", B.ministry = M.Welfare }


If any district can pay the bill, take funds in proportion.  We use
'HM.lookup', because 'show' of 'HM.HashMap' is not determinate,
and the test can occasionally fail:

>>> let funds = HM.fromList [(D.Palolene, 100), (D.Lakos, 200)]
>>> let contribution = manyToOne funds medium
>>> HM.lookup D.Palolene contribution
Just 10
>>> HM.lookup D.Lakos contribution
Just 20

>>> let funds = HM.fromList [(D.Palolene, 30), (D.Lakos, 30)]
>>> HM.elems $ manyToOne funds medium
[15,15]


It works with a single district:

>>> let funds = HM.fromList [(D.SouthernPalolene, 500)]
>>> HM.elems $ manyToOne funds medium
[30]
>>> HM.elems $ manyToOne funds one
[1]


__TODO__ It /should/ not have rounding issues. In particular,
when the bill's amount is bigger than the number of districts,
/each/ district would contribute some. This problem is known as
<https://en.wikipedia.org/wiki/Partition_(number_theory) integer partition>.

>>> let funds = HM.fromList [(D.Palolene, 10000), (D.Lakos, 1)]
>>> B.amount medium > fromIntegral (HM.size funds)
True
>>> let contribution = manyToOne funds medium
>>> let taken = HM.foldl' (+) 0 contribution


But at the moment we make use of some freedom coming from the fact
that it is possible that a bill will receive less funds than the
Parliament decides, and round down contribution of each district.
/Thus these two tests show 'False', while should show 'True':/

>>> HM.null $ HM.filter (== 0) contribution
False
>>> taken == B.amount medium
False


If all districts together can't pay the bill, take all their money.
Note that due to the principle of proportionality it is impossible
that some districts can pay their shares and others can't:

>>> let low = HM.fromList [(D.Palolene, 10), (D.Lakos, 15)]
>>> manyToOne low medium == low
True


It works without any districts specified:

>>> manyToOne HM.empty medium
fromList []

-}
manyToOne ::
     D.DistrictFunds -- ^ Amounts of available funds per district.
  -> B.Bill -- ^ A bill requiring funding.
  -> D.DistrictFunds -- ^ Contribution of each district.
manyToOne df b = HM.fromList $ zip ds (normalizeDown (B.amount b) fs)
  where
    (ds, fs) = unzip $ HM.toList df

{-|
Districts funding multiple bills. No constraints.

>>> :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 welfareA = B.Bill { B.amount = 100, B.name = "Welfare A", B.ministry = M.Welfare }


It works without any districts specified:

>>> snd . head $  manyToMany HM.empty [scienceA]
fromList []


It works without bills:

>>> manyToMany HM.empty []
[]
>>> let funds = HM.fromList [(D.Palolene, 100), (D.Lakos, 200)]
>>> manyToMany funds []
[]
-}
manyToMany ::
     D.DistrictFunds -- ^ Amounts of available funds per district.
  -> [B.Bill] -- ^ Bills requiring funding.
  -> [(B.Bill, D.DistrictFunds)] -- ^ Contribution of each district to each bill.
manyToMany funds bills = zipWith (\b df -> (b, df)) bills allocated
  where
    fundsTotal = sum $ HM.elems funds
    allocated = manyToOne funds <$> normalizeBills fundsTotal bills

{-|
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 (fst bb) == B.name b)
>>> let amounts budget = map (B.amount . fst . 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.
  -> [(B.Bill, D.DistrictFunds)] -- ^ 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