module InstantRunoff where
import qualified Data.List as L
import qualified Data.List.NonEmpty as LN
import qualified Data.Map.Strict as M
import qualified Data.Set as S
import qualified Data.Maybe as M
import qualified Data.Foldable as F

solve :: forall a. Ord a => [a] ->  LN.NonEmpty (LN.NonEmpty a) -> [S.Set a]
solve candidates votes = reverse . (consIfAny notVotedForCandidates) . rank $ votes
  where

    consIfAny :: S.Set a -> [S.Set a] -> [S.Set a]
    consIfAny [] xs = xs
    consIfAny x xs = x : xs

    -- candidates not on any submitted ballot
    notVotedForCandidates :: S.Set a
    notVotedForCandidates =
      S.difference (S.fromList candidates)
      -- all candidates on any submited ballot
      . S.unions . LN.map (S.fromList . LN.toList) $ votes


rank :: forall a. Ord a => LN.NonEmpty (LN.NonEmpty a) -> [S.Set a]
rank votes =
  maybe [losers] ((losers :) . rank)
  . filterVotes (`S.notMember` losers) $ votes
  where

    -- losers in terms of the first preference on ballot
    firstChoiceLosers :: S.Set a
    firstChoiceLosers = M.keysSet . M.filter ((==) . L.minimum . M.elems $ voteCounts) $ voteCounts

    -- in the event of multiple first choice losers,
    -- count how many ballots prefer one choice over another
    losers :: S.Set a
    losers = S.fromList
      . maybe [] (smallestSubsetBy preferred)
      . LN.nonEmpty . S.toList
      $ firstChoiceLosers

    preferred :: a -> a -> Ordering
    preferred x0 x1 = compare (numPrefers x0) (numPrefers x1)
      where
        numPrefers x = length . LN.filter (== Just x) $ preferences

        -- which candidate each voter prefers
        preferences :: LN.NonEmpty (Maybe a)
        preferences = LN.map (F.find (`S.member` [x0, x1])) $ votes

    voteCounts :: M.Map a Word
    voteCounts =
        M.unionsWith (+)
      . (M.fromSet (const 0) candidates :)
      . map (`M.singleton` 1)
      . LN.toList .
      LN.map LN.head $ votes

    -- set of all candidates on anybody's ballot
    candidates :: S.Set a
    candidates = S.unions . LN.map (S.fromList . LN.toList) $ votes

filterVotes :: (a -> Bool) -> LN.NonEmpty (LN.NonEmpty a) -> Maybe (LN.NonEmpty (LN.NonEmpty a))
filterVotes f = InstantRunoff.catMaybes . LN.map (LN.nonEmpty . LN.filter f)

catMaybes :: LN.NonEmpty (Maybe a) -> Maybe (LN.NonEmpty a)
catMaybes = LN.nonEmpty . M.catMaybes . LN.toList

-- smallest by the ordering function, not by size
smallestSubsetBy :: forall t. (t -> t -> Ordering) -> LN.NonEmpty t -> [t]
smallestSubsetBy f xs = LN.filter ((== EQ) . f (F.minimumBy f xs)) . LN.sortBy f $ xs