module Math.Apportionment (
   largestRemainder,
   largestRemainderScaled,

   highestAveragesScaled,
   dHondtDivisors,
   sainteLagueDivisors,
   ) where

import Control.Functor.HT (outerProduct, )

import qualified Data.Foldable as Fold
import qualified Data.List.HT as ListHT
import qualified Data.List as List
import qualified Data.Map as Map
import Data.Function.HT (compose2, )
import Data.Tuple.HT (mapSnd, )
import Data.Ord.HT (comparing, )


{- |
Like 'largestRemainder' but result values
are sorted with respect to descending fractional parts of the inputs.
This is an artifact of the used algorithm.
The result still depends on the input order,
especially on the order of numbers with equal fractional part.
-}
_largestRemainderSort :: (RealFrac a) => [a] -> [Int]
_largestRemainderSort xs =
   let (d, intFracs) = fractions xs
       (intUps, intDowns) =
          splitAt d $ map fst $
          List.sortBy (comparing (negate.snd)) intFracs
   in  map (1+) intUps ++ intDowns

{- |
This function rounds values
such that the sum of the rounded values
matches the rounded sum of the original values.

Also known as Hare-Niemeyer method.
<https://en.wikipedia.org/wiki/Largest_remainder_method>

Input values must be non-negative, otherwise 'properFraction' bites us.
-}
largestRemainder :: (RealFrac a) => [a] -> [Int]
largestRemainder = largestRemainderCore . fractions

{- |
@largestRemainderScaled s xs@
scales and rounds the values in @xs@ such that their sum becomes @s@.

E.g. @largestRemainderScaled 100 [1,2,3]@
returns integral percentages proportional to 1:2:3.
-}
largestRemainderScaled :: (RealFrac a) => Int -> [a] -> [Int]
largestRemainderScaled s = largestRemainderCore . fractionsScaled s


type Fractions a = (Int, [(Int, a)])

largestRemainderCore :: (RealFrac a) => Fractions a -> [Int]
largestRemainderCore (d, intFracs) =
   let (intUps, intDowns) =
          splitAt d $ map (mapSnd fst) $
          List.sortBy (comparing (negate . snd . snd)) $
          zip [(0::Int) .. ] intFracs
   in  map snd $ List.sortBy (comparing fst) $
       map (mapSnd (1+)) intUps ++ intDowns

fractions :: (RealFrac a) => [a] -> Fractions a
fractions xs =
   let xsum = round $ sum xs
       intFracs = map properFraction xs
       isum = sum $ map fst intFracs
   in  (xsum-isum, intFracs)

fractionsScaled :: (RealFrac a) => Int -> [a] -> Fractions a
fractionsScaled xsum xs =
   let c = fromIntegral xsum / sum xs
       intFracs = map (properFraction . (* c)) xs
       isum = sum $ map fst intFracs
   in  (xsum-isum, intFracs)


{- |
<https://en.wikipedia.org/wiki/Highest_averages_method>

In @highestAveragesScaled divs s xs@,
@divs@ must be an infinite list of strictly increasing positive numbers.
E.g. @highestAveragesScaled dHondtDivisors s xs@ runs the d'Hondt method.
-}
highestAveragesScaled :: (RealFrac a) => [a] -> Int -> [a] -> [Int]
highestAveragesScaled divs s xs =
   let m = Map.fromList $ zip [(0::Int) ..] xs
   in  Map.elems $ flip Map.union (fmap (const 0) m) $
       Map.fromListWith (+) $ map (mapSnd (const 1)) $
       take s $ Fold.foldl (ListHT.mergeBy (compose2 (>=) snd)) [] $
       Map.mapWithKey (map . (,)) $ outerProduct (/) m divs

dHondtDivisors :: Num a => [a]
dHondtDivisors = iterate (1+) 1

sainteLagueDivisors :: Num a => [a]
sainteLagueDivisors = iterate (2+) 1