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, )


{- $setup
>>> import Control.Applicative ((<$>))
>>> import qualified Test.QuickCheck as QC
>>>
>>> forAllNonNegatives ::
>>>    (Num a, Ord a, Show a, QC.Arbitrary a) => ([a] -> Bool) -> QC.Property
>>> forAllNonNegatives = QC.forAll $
>>>    (map QC.getNonNegative <$> QC.arbitrary) `QC.suchThat` (\xs -> sum xs > 0)
-}

{- |
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 :: [a] -> [Int]
_largestRemainderSort [a]
xs =
   let (Int
d, [(Int, a)]
intFracs) = [a] -> (Int, [(Int, a)])
forall a. RealFrac a => [a] -> Fractions a
fractions [a]
xs
       ([Int]
intUps, [Int]
intDowns) =
          Int -> [Int] -> ([Int], [Int])
forall a. Int -> [a] -> ([a], [a])
splitAt Int
d ([Int] -> ([Int], [Int])) -> [Int] -> ([Int], [Int])
forall a b. (a -> b) -> a -> b
$ ((Int, a) -> Int) -> [(Int, a)] -> [Int]
forall a b. (a -> b) -> [a] -> [b]
map (Int, a) -> Int
forall a b. (a, b) -> a
fst ([(Int, a)] -> [Int]) -> [(Int, a)] -> [Int]
forall a b. (a -> b) -> a -> b
$
          ((Int, a) -> (Int, a) -> Ordering) -> [(Int, a)] -> [(Int, a)]
forall a. (a -> a -> Ordering) -> [a] -> [a]
List.sortBy (((Int, a) -> a) -> (Int, a) -> (Int, a) -> Ordering
forall b a. Ord b => (a -> b) -> a -> a -> Ordering
comparing (a -> a
forall a. Num a => a -> a
negate(a -> a) -> ((Int, a) -> a) -> (Int, a) -> a
forall b c a. (b -> c) -> (a -> b) -> a -> c
.(Int, a) -> a
forall a b. (a, b) -> b
snd)) [(Int, a)]
intFracs
   in  (Int -> Int) -> [Int] -> [Int]
forall a b. (a -> b) -> [a] -> [b]
map (Int
1Int -> Int -> Int
forall a. Num a => a -> a -> a
+) [Int]
intUps [Int] -> [Int] -> [Int]
forall a. [a] -> [a] -> [a]
++ [Int]
intDowns

-- ToDo: generalize to Traversable

-- ToDo: require NonEmpty

-- ToDo: It would be safer to use Integer

{- |
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 [1,2,3::Rational]
[1,2,3]
>>> largestRemainder [1.1,2.2,3.3,4.4::Rational]
[1,2,3,5]

prop> \xs -> xs == largestRemainder (map fromIntegral xs :: [Rational])
prop> forAllNonNegatives $ \xs -> round (sum xs) == sum (largestRemainder (xs :: [Rational]))
-}
largestRemainder :: (RealFrac a) => [a] -> [Int]
largestRemainder :: [a] -> [Int]
largestRemainder = Fractions a -> [Int]
forall a. RealFrac a => Fractions a -> [Int]
largestRemainderCore (Fractions a -> [Int]) -> ([a] -> Fractions a) -> [a] -> [Int]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [a] -> Fractions a
forall a. RealFrac a => [a] -> Fractions a
fractions

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

>>> largestRemainderScaled 100 [1,2,3::Rational]
[17,33,50]

That is, it returns integral percentages almost proportional to 1:2:3.

>>> largestRemainderScaled 100 [1,10,100::Rational]
[1,9,90]

prop> forAllNonNegatives $ \xs -> xs == largestRemainderScaled (sum xs) (map fromIntegral xs :: [Rational])
prop> \(QC.Positive s) -> forAllNonNegatives $ \xs -> s == sum (largestRemainderScaled s (xs :: [Rational]))
-}
largestRemainderScaled :: (RealFrac a) => Int -> [a] -> [Int]
largestRemainderScaled :: Int -> [a] -> [Int]
largestRemainderScaled Int
s = Fractions a -> [Int]
forall a. RealFrac a => Fractions a -> [Int]
largestRemainderCore (Fractions a -> [Int]) -> ([a] -> Fractions a) -> [a] -> [Int]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Int -> [a] -> Fractions a
forall a. RealFrac a => Int -> [a] -> Fractions a
fractionsScaled Int
s


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

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

fractions :: (RealFrac a) => [a] -> Fractions a
fractions :: [a] -> Fractions a
fractions [a]
xs =
   let xsum :: Int
xsum = a -> Int
forall a b. (RealFrac a, Integral b) => a -> b
round (a -> Int) -> a -> Int
forall a b. (a -> b) -> a -> b
$ [a] -> a
forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
sum [a]
xs
       intFracs :: [(Int, a)]
intFracs = (a -> (Int, a)) -> [a] -> [(Int, a)]
forall a b. (a -> b) -> [a] -> [b]
map a -> (Int, a)
forall a b. (RealFrac a, Integral b) => a -> (b, a)
properFraction [a]
xs
       isum :: Int
isum = [Int] -> Int
forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
sum ([Int] -> Int) -> [Int] -> Int
forall a b. (a -> b) -> a -> b
$ ((Int, a) -> Int) -> [(Int, a)] -> [Int]
forall a b. (a -> b) -> [a] -> [b]
map (Int, a) -> Int
forall a b. (a, b) -> a
fst [(Int, a)]
intFracs
   in  (Int
xsumInt -> Int -> Int
forall a. Num a => a -> a -> a
-Int
isum, [(Int, a)]
intFracs)

fractionsScaled :: (RealFrac a) => Int -> [a] -> Fractions a
fractionsScaled :: Int -> [a] -> Fractions a
fractionsScaled Int
xsum [a]
xs =
   let c :: a
c = Int -> a
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
xsum a -> a -> a
forall a. Fractional a => a -> a -> a
/ [a] -> a
forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
sum [a]
xs
       intFracs :: [(Int, a)]
intFracs = (a -> (Int, a)) -> [a] -> [(Int, a)]
forall a b. (a -> b) -> [a] -> [b]
map (a -> (Int, a)
forall a b. (RealFrac a, Integral b) => a -> (b, a)
properFraction (a -> (Int, a)) -> (a -> a) -> a -> (Int, a)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (a -> a -> a
forall a. Num a => a -> a -> a
* a
c)) [a]
xs
       isum :: Int
isum = [Int] -> Int
forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
sum ([Int] -> Int) -> [Int] -> Int
forall a b. (a -> b) -> a -> b
$ ((Int, a) -> Int) -> [(Int, a)] -> [Int]
forall a b. (a -> b) -> [a] -> [b]
map (Int, a) -> Int
forall a b. (a, b) -> a
fst [(Int, a)]
intFracs
   in  (Int
xsumInt -> Int -> Int
forall a. Num a => a -> a -> a
-Int
isum, [(Int, a)]
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 dHondtDivisors 100 [1,2,3::Rational]
[17,33,50]
>>> highestAveragesScaled dHondtDivisors 100 [1,10,100::Rational]
[0,9,91]

>>> highestAveragesScaled sainteLagueDivisors 100 [1,2,3::Rational]
[17,33,50]
>>> highestAveragesScaled sainteLagueDivisors 100 [1,10,100::Rational]
[1,9,90]

prop> forAllNonNegatives $ \xs -> xs == highestAveragesScaled dHondtDivisors (sum xs) (map fromIntegral xs :: [Rational])
prop> forAllNonNegatives $ \xs -> xs == highestAveragesScaled sainteLagueDivisors (sum xs) (map fromIntegral xs :: [Rational])

prop> \(QC.Positive s) -> forAllNonNegatives $ \xs -> s == sum (highestAveragesScaled dHondtDivisors s (xs :: [Rational]))
prop> \(QC.Positive s) -> forAllNonNegatives $ \xs -> s == sum (highestAveragesScaled sainteLagueDivisors s (xs :: [Rational]))
-}
highestAveragesScaled :: (RealFrac a) => [a] -> Int -> [a] -> [Int]
highestAveragesScaled :: [a] -> Int -> [a] -> [Int]
highestAveragesScaled [a]
divs Int
s [a]
xs =
   let m :: Map Int a
m = [(Int, a)] -> Map Int a
forall k a. Ord k => [(k, a)] -> Map k a
Map.fromList ([(Int, a)] -> Map Int a) -> [(Int, a)] -> Map Int a
forall a b. (a -> b) -> a -> b
$ [Int] -> [a] -> [(Int, a)]
forall a b. [a] -> [b] -> [(a, b)]
zip [(Int
0::Int) ..] [a]
xs
   in  Map Int Int -> [Int]
forall k a. Map k a -> [a]
Map.elems (Map Int Int -> [Int]) -> Map Int Int -> [Int]
forall a b. (a -> b) -> a -> b
$ (Map Int Int -> Map Int Int -> Map Int Int)
-> Map Int Int -> Map Int Int -> Map Int Int
forall a b c. (a -> b -> c) -> b -> a -> c
flip Map Int Int -> Map Int Int -> Map Int Int
forall k a. Ord k => Map k a -> Map k a -> Map k a
Map.union ((a -> Int) -> Map Int a -> Map Int Int
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (Int -> a -> Int
forall a b. a -> b -> a
const Int
0) Map Int a
m) (Map Int Int -> Map Int Int) -> Map Int Int -> Map Int Int
forall a b. (a -> b) -> a -> b
$
       (Int -> Int -> Int) -> [(Int, Int)] -> Map Int Int
forall k a. Ord k => (a -> a -> a) -> [(k, a)] -> Map k a
Map.fromListWith Int -> Int -> Int
forall a. Num a => a -> a -> a
(+) ([(Int, Int)] -> Map Int Int) -> [(Int, Int)] -> Map Int Int
forall a b. (a -> b) -> a -> b
$ ((Int, a) -> (Int, Int)) -> [(Int, a)] -> [(Int, Int)]
forall a b. (a -> b) -> [a] -> [b]
map ((a -> Int) -> (Int, a) -> (Int, Int)
forall b c a. (b -> c) -> (a, b) -> (a, c)
mapSnd (Int -> a -> Int
forall a b. a -> b -> a
const Int
1)) ([(Int, a)] -> [(Int, Int)]) -> [(Int, a)] -> [(Int, Int)]
forall a b. (a -> b) -> a -> b
$
       Int -> [(Int, a)] -> [(Int, a)]
forall a. Int -> [a] -> [a]
take Int
s ([(Int, a)] -> [(Int, a)]) -> [(Int, a)] -> [(Int, a)]
forall a b. (a -> b) -> a -> b
$ ([(Int, a)] -> [(Int, a)] -> [(Int, a)])
-> [(Int, a)] -> Map Int [(Int, a)] -> [(Int, a)]
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
Fold.foldl (((Int, a) -> (Int, a) -> Bool)
-> [(Int, a)] -> [(Int, a)] -> [(Int, a)]
forall a. (a -> a -> Bool) -> [a] -> [a] -> [a]
ListHT.mergeBy ((a -> a -> Bool) -> ((Int, a) -> a) -> (Int, a) -> (Int, a) -> Bool
forall b c a. (b -> b -> c) -> (a -> b) -> a -> a -> c
compose2 a -> a -> Bool
forall a. Ord a => a -> a -> Bool
(>=) (Int, a) -> a
forall a b. (a, b) -> b
snd)) [] (Map Int [(Int, a)] -> [(Int, a)])
-> Map Int [(Int, a)] -> [(Int, a)]
forall a b. (a -> b) -> a -> b
$
       (Int -> [a] -> [(Int, a)]) -> Map Int [a] -> Map Int [(Int, a)]
forall k a b. (k -> a -> b) -> Map k a -> Map k b
Map.mapWithKey ((a -> (Int, a)) -> [a] -> [(Int, a)]
forall a b. (a -> b) -> [a] -> [b]
map ((a -> (Int, a)) -> [a] -> [(Int, a)])
-> (Int -> a -> (Int, a)) -> Int -> [a] -> [(Int, a)]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (,)) (Map Int [a] -> Map Int [(Int, a)])
-> Map Int [a] -> Map Int [(Int, a)]
forall a b. (a -> b) -> a -> b
$ (a -> a -> a) -> Map Int a -> [a] -> Map Int [a]
forall (f :: * -> *) (g :: * -> *) a b c.
(Functor f, Functor g) =>
(a -> b -> c) -> f a -> g b -> f (g c)
outerProduct a -> a -> a
forall a. Fractional a => a -> a -> a
(/) Map Int a
m [a]
divs

-- ToDo: use Stream
dHondtDivisors :: Num a => [a]
dHondtDivisors :: [a]
dHondtDivisors = (a -> a) -> a -> [a]
forall a. (a -> a) -> a -> [a]
iterate (a
1a -> a -> a
forall a. Num a => a -> a -> a
+) a
1

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