module Election(
    Election,

    createElection,
    doCount
    ) where




import qualified System.IO as IO
import qualified Control.Monad as Con
import qualified Control.Monad.Trans.Either as ET
import qualified Control.Monad.IO.Class as MIO
import qualified Data.List as List
import qualified Data.Either.Unwrap as Either
import qualified Counter as Sen
import qualified Candidate as Typ
import qualified CSV as CSV




data Election = Election
    { getEntries     :: [Entry]
    , getCounter     :: Sen.SenateCounter
    , getLogDir      :: FilePath
    , getTotalPapers :: Int
    , getQuota       :: Int
    , getNextLogNum  :: Int
    , getVacancies   :: Int
    , isDone         :: Bool }

data Entry = Entry
    { getID         :: Typ.CandidateID
    , getVoteChange :: Int
    , getTotalVotes :: Int
    , getCritTrace  :: [Trace]
    , getStatus     :: Status
    , getChanged    :: Bool }

data Trace = Trace
    { getCriteria    :: Sen.Criteria
    , getTransferVal :: Float }

data Status = Running | Elected | Eliminated
    deriving (Show, Eq)




createElection :: FilePath -> Sen.SenateCounter -> Int -> IO Election
createElection outDir counter numToElect = do
    entries <- mapM (candToEntry counter) (Sen.getBallot counter)
    let total = Sen.getTotal counter
        quota = 1 + floor ((fromIntegral total) / (fromIntegral (numToElect + 1)))
    return (Election entries counter outDir total quota 1 numToElect False)




candToEntry :: Sen.SenateCounter -> Typ.CandidateID -> IO Entry
candToEntry counter candidate = do
    let criteria = [(1,candidate)]
        trace = Trace criteria 1
    firstPrefs <- Sen.doCount counter criteria
    return (Entry candidate firstPrefs firstPrefs [trace] Running False)




doCount :: Election -> IO ()
doCount e = do
    writeLog e
    --  these following calculations probably aren't the
    --  intended use of Either monads, but the pattern fits
    --  and it's certainly a lot better than a bunch of
    --  if-them-else constructs in haskell
    r <- ET.eitherT return return $
            incrementLog e >>=
            checkNoneLeft >>=
            electSomeone >>=
            excludeSomeone >>=
            checkNoQuota
    Con.when (not (isDone r)) $ doCount r




incrementLog :: Election -> ET.EitherT Election IO Election
incrementLog e =
    ET.right (e { getNextLogNum = 1 + (getNextLogNum e) })




checkNoneLeft :: Election -> ET.EitherT Election IO Election
checkNoneLeft e = do
    let running = filter ((== Running) . getStatus) (getEntries e)
    if (getVacancies e == 0 || length running == 0)
        then ET.left (e { isDone = True })
        else ET.right e




electSomeone :: Election -> ET.EitherT Election IO Election
electSomeone e = do
    let (running, notRunning) = List.partition ((== Running) . getStatus) (getEntries e)
        (reachedQuota, notReached) = List.partition ((>= getQuota e) . getTotalVotes) running
        sortedReached = List.sortBy (\x y -> compare (getTotalVotes y) (getTotalVotes x)) reachedQuota
        (electedEntry, otherEntries) = (head sortedReached, (tail sortedReached) ++ notReached ++ notRunning)

        newTransferValue = (fromIntegral (getTotalVotes electedEntry)) / (fromIntegral (getQuota e))
        revisedElectedEntry = electedEntry
            { getStatus = Elected
            , getChanged = True
            , getTotalVotes = (getQuota e)
            , getVoteChange = (getQuota e) - (getTotalVotes electedEntry) }
    revisedOtherEntries <- MIO.liftIO $ transferVotes (getCounter e) newTransferValue electedEntry otherEntries

    let revisedEntries = revisedElectedEntry:revisedOtherEntries

    if (length reachedQuota > 0)
        then ET.left (e { getEntries = revisedEntries, getVacancies = (getVacancies e) - 1 })
        else ET.right e




excludeSomeone :: Election -> ET.EitherT Election IO Election
excludeSomeone e = do
    let (running, notRunning) = List.partition ((== Running) . getStatus) (getEntries e)
        (reachedQuota, notReached) = List.partition ((>= getQuota e) . getTotalVotes) running
        sortedNotReached = List.sortBy (\x y -> compare (getTotalVotes x) (getTotalVotes y)) notReached
        (excludedEntry, otherEntries) = (head sortedNotReached, (tail sortedNotReached) ++ reachedQuota ++ notRunning)

        revisedExcludedEntry = excludedEntry
            { getStatus = Eliminated
            , getChanged = True
            , getTotalVotes = 0
            , getVoteChange = getTotalVotes excludedEntry }
    revisedOtherEntries <- MIO.liftIO $ transferVotes (getCounter e) 1 excludedEntry otherEntries

    let revisedEntries = revisedExcludedEntry:revisedOtherEntries

    if (length notReached > 0)
        then ET.left (e { getEntries = revisedEntries })
        else ET.right e




transferVotes :: Sen.SenateCounter -> Float -> Entry -> [Entry] -> IO [Entry]
transferVotes counter value from to = do
    let addToCriteria candID crit =
            let maxRank = fst (List.maximumBy (\x y -> compare (fst x) (fst y)) crit)
            in (maxRank + 1, candID):crit
        addToTrace candID trace = trace
            { getCriteria = addToCriteria candID (getCriteria trace)
            , getTransferVal = value * (getTransferVal trace) }

        transferFunc entryFrom entryTo = do
            let newTraces = map (addToTrace (getID entryTo)) (getCritTrace entryFrom)
            voteList <- Con.mapM (Sen.doCount counter) (map getCriteria newTraces)
            let changeList = zipWith (*) (map getTransferVal newTraces) (map fromIntegral voteList)
                totalVoteChange = List.foldl' (+) 0 changeList
            return (entryTo
                { getVoteChange = round totalVoteChange
                , getTotalVotes = (getTotalVotes entryTo) + (round totalVoteChange)
                , getCritTrace = newTraces
                , getChanged = True })

    Con.mapM (transferFunc from) to




checkNoQuota :: Election -> ET.EitherT Election IO Election
checkNoQuota e = do
    let (running, notRunning) = List.partition ((== Running) . getStatus) (getEntries e)
        sortedRunning = List.sortBy (\x y -> compare (getTotalVotes y) (getTotalVotes x)) running
        (topRunning, rest) = List.splitAt (getVacancies e) sortedRunning
        changed = map (\x -> x { getStatus = Elected, getChanged = True }) topRunning
        revisedEntries = map (\x -> x { getVoteChange = 0 }) (changed ++ rest ++ notRunning)
    if (getVacancies e == 1 || getVacancies e == length running)
        then ET.left (e { getEntries = revisedEntries, getVacancies = 0 })
        else ET.right e




writeLog :: Election -> IO ()
writeLog e = do
    let logName = (getLogDir e) ++ "/" ++ (show (getNextLogNum e)) ++ ".csv"
        header = ["Vacancies", "Total Papers", "Quota", "Candidate", "Votes", "Transfer", "Status", "Changed"]
        static = [show (getVacancies e), show (getTotalPapers e), show (getQuota e)]
        dynFunc c = [getID c, show (getTotalVotes c), show (getVoteChange c), show (getStatus c), show (getChanged c)]
        records = map (\x -> static ++ dynFunc x) (getEntries e)
        headerLine = CSV.unParseRecord CSV.defaultSettings header
        recordLines = map (CSV.unParseRecord CSV.defaultSettings) records
        output = unlines (headerLine:recordLines)
    IO.writeFile logName output