module IRV where

import Control.Monad (guard)
import Data.List (nub, sortBy)
import Data.Maybe (catMaybes, fromMaybe)
import Data.Ord (Down(..), comparing)

winner :: Eq a => [[a]] -> Maybe a
winner votes = do
  let freqs = frequencies votes
  candidate <- majorityCandidate freqs
  if hasAbsoluteMajority candidate freqs
     then return candidate
     else do
       last <- lastCandidate freqs
       winner (removeFirstPref last votes)

frequencies :: Eq a => [[a]] -> [(a, Int)]
frequencies votes = group firstPreferences
  where firstPreferences = catMaybes (map maybeHead votes)

group :: Eq a => [a] -> [(a, Int)]
group xs = map count (nub xs)
  where count x = (x, length (filter (x ==) xs))

maybeHead :: [a] -> Maybe a
maybeHead [] = Nothing
maybeHead xs = Just (head xs)

majorityCandidate :: Eq a => [(a, Int)] -> Maybe a
majorityCandidate freqs = fst <$> maybeHead (sortBy (comparing (Down . snd)) freqs)

hasAbsoluteMajority :: Eq a => a -> [(a, Int)] -> Bool
hasAbsoluteMajority candidate freqs = candidateVotes > totalVotes `div` 2
  where candidateVotes = fromMaybe 0 (lookup candidate freqs)
        totalVotes = sum (map snd freqs)

lastCandidate :: Eq a => [(a, Int)] -> Maybe a
lastCandidate []  = Nothing
lastCandidate [x] = Nothing
lastCandidate [x, y]
  | snd x == snd y = Nothing -- It's a tie!
lastCandidate freqs = fst <$> maybeHead (sortBy (comparing snd) freqs)

removeFirstPref :: Eq a => a -> [[a]] -> [[a]]
removeFirstPref candidate = map removeFirstPref'
  where removeFirstPref' [] = []
        removeFirstPref' prefs@(x:xs)
          | x == candidate = xs
          | otherwise = prefs