module Election(
    Election,

    createElection,
    doCount
    ) where




--  This source is licensed under Creative Commons CC0 v1.0.

--  To read the full text, see license.txt in the main directory of this repository
--  or go to https://creativecommons.org/publicdomain/zero/1.0/legalcode.txt

--  For a human readable summary, go to https://creativecommons.org/publicdomain/zero/1.0/




import qualified System.IO as IO
import qualified System.Exit as Ex
import qualified Control.Monad as Con
import qualified Control.Monad.Trans.Either as ET
import qualified Control.Monad.IO.Class as MIO
import Data.List ( (\\) )
import qualified Data.List as List
import qualified Data.Maybe as Maybe
import qualified Data.Either.Unwrap as Either
import Data.Ratio ( (%) )
import qualified Counter as Sen
import qualified Candidate as Typ
import qualified CSV as CSV
import Miscellaneous ( (?) )
import qualified Miscellaneous as Misc




data Election = Election
    { getEntries       :: [Entry]
    , getCounter       :: Sen.SenateCounter
    , getLogDir        :: FilePath
    , getTotalPapers   :: Int
    , getQuota         :: Int
    , getMainLog       :: String
    , getNextLogNum    :: Int
    , getSeats         :: Int
    , getVacancies     :: Int
    , getTransferQueue :: [Transfer]
    , getNextToElect   :: Int
    , isDone           :: Bool
    , isVerbose        :: Bool }

data Entry = Entry
    { getID           :: Typ.CandidateID
    , getVoteChange   :: Int
    , getTotalVotes   :: Int
    , getCritTrace    :: [Trace]
    , getStatus       :: Status
    , getChanged      :: Bool
    , getOrderElected :: Maybe Int }
    deriving (Eq)

data Trace = Trace
    { getCriteria    :: Sen.Criteria
    , getTransferVal :: Rational }
    deriving (Eq)

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

data Transfer = Transfer
    { getWhoFrom    :: Typ.CandidateID
    , getVoteAmount :: Int
    , getNewValue   :: Rational
    , getWhatToDist :: [Trace] }
    deriving (Eq)




createElection :: FilePath -> FilePath -> Sen.SenateCounter -> Int -> Bool -> IO Election
createElection outDir mainLog counter numToElect verbosity = do
    entries <- mapM (candToEntry counter) (Sen.getBallot counter)
    return (Election
        { getEntries       = entries
        , getCounter       = counter
        , getLogDir        = outDir
        , getTotalPapers   = Sen.getTotal counter
        , getQuota         = droopQuota (Sen.getTotal counter) numToElect
        , getMainLog       = mainLog
        , getNextLogNum    = 1
        , getSeats         = numToElect
        , getVacancies     = numToElect
        , getTransferQueue = []
        , getNextToElect   = 1
        , isDone           = False
        , isVerbose        = verbosity })




droopQuota :: Int -> Int -> Int
droopQuota votes seats =
    1 + floor ((fromIntegral votes) / (fromIntegral (seats + 1)))




compareVotes :: Entry -> Entry -> Ordering
compareVotes x y = compare (getTotalVotes x) (getTotalVotes y)




candToEntry :: Sen.SenateCounter -> Typ.CandidateID -> IO Entry
candToEntry counter candidate = do
    let criteria = [Sen.matchID candidate]
        trace = Trace criteria 1
    firstPrefs <- Sen.doCount counter criteria
    return (Entry
        { getID           = candidate
        , getVoteChange   = firstPrefs
        , getTotalVotes   = firstPrefs
        , getCritTrace    = [trace]
        , getStatus       = Running
        , getChanged      = False
        , getOrderElected = Nothing })




doCount :: Election -> IO ()
doCount e =
    Con.when (not (isDone e)) $ do
        writeLog e
        let e1 = e { getNextLogNum = 1 + getNextLogNum e }
        let e2 = e1 { getEntries = map clearChange (getEntries e1) }

        --  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-then-else constructs in haskell
        r <- ET.eitherT return return $
            electCandidates e2 >>=
            checkIfDone >>=
            transferVotes >>=
            checkNoQuota >>=
            excludeCandidates

        --  this should never happen unless there's a bug somewhere
        Con.when (getEntries e2 == getEntries r && getTransferQueue e2 == getTransferQueue r && not (isDone r)) $
            Ex.die "Infinite loop detected in election counting"

        doCount r




writeLog :: Election -> IO ()
writeLog e = do
    let logName = (getLogDir e) ++ "/" ++ (show (getNextLogNum e)) ++ ".csv"

        header =
            [ "Seats"
            , "Vacancies"
            , "Total Papers"
            , "Quota"
            , "Candidate"
            , "Votes"
            , "Transfer"
            , "Status"
            , "Changed"
            , "Order Elected" ]
        static =
            [ show (getSeats e)
            , show (getVacancies e)
            , show (getTotalPapers e)
            , show (getQuota e)]
        dynFunc c =
            [ getID c
            , show (getTotalVotes c)
            , show (getVoteChange c)
            , show (getStatus c)
            , if (getChanged c) then show (getChanged c) else ""
            , if (Maybe.isJust (getOrderElected c)) then show (Maybe.fromJust (getOrderElected c)) else "" ]

        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




clearChange :: Entry -> Entry
clearChange entry = entry
    { getChanged = False
    , getVoteChange = 0 }




electCandidates :: Election -> ET.EitherT Election IO Election
electCandidates e = do
    let oldToElectNum = getNextToElect e
        electLoop x = ET.eitherT return electLoop $ doElectCandidate x
    r <- MIO.liftIO $ electLoop e
    if (getNextToElect r > oldToElectNum)
        then ET.left r
        else ET.right r




--  needs to be modified to take into account ties
--  may be prudent to put this just in the IO monad instead of EitherT
doElectCandidate :: Election -> ET.EitherT Election IO Election
doElectCandidate e = do
    let running = filter ((== Running) . getStatus) (getEntries e)
        electedEntry = List.maximumBy compareVotes running
        (beforeEntries, afterEntries) = Misc.partBeforeAfter electedEntry (getEntries e)

        newTransfer = Transfer
            { getWhoFrom = getID electedEntry
            , getVoteAmount = getTotalVotes electedEntry - getQuota e
            , getNewValue = (fromIntegral (getTotalVotes electedEntry - getQuota e)) %
                                (fromIntegral (getTotalVotes electedEntry))
            , getWhatToDist = getCritTrace electedEntry }

        revisedElectedEntry = electedEntry
            { getStatus = Elected
            , getChanged = True
            , getOrderElected = Just (getNextToElect e) }
        allRevised = beforeEntries ++ [revisedElectedEntry] ++ afterEntries

    if (getTotalVotes electedEntry >= getQuota e)
        then do
            let logmsg = show (getID electedEntry) ++ " elected at logfile #" ++ show (getNextLogNum e)
            MIO.liftIO $ IO.appendFile (getMainLog e) (logmsg ++ "\n")
            MIO.liftIO $ Con.when (isVerbose e) (IO.hPutStrLn IO.stderr logmsg)
            ET.right (e
                { getEntries = allRevised
                , getTransferQueue = (getTransferQueue e) ++ [newTransfer]
                , getNextToElect = getNextToElect e + 1
                , getVacancies = getVacancies e - 1 })
        else ET.left e




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




--  redistributing votes in STV is annoying as hell
transferVotes :: Election -> ET.EitherT Election IO Election
transferVotes e =
    if (length (getTransferQueue e) > 0)
        then (MIO.liftIO $ doVoteTransfer e) >>= ET.left
        else ET.right e




doVoteTransfer :: Election -> IO Election
doVoteTransfer e = do
    let (currentTransfer:remainingTransfers) = getTransferQueue e
        fromEntry = Maybe.fromJust (List.find ((== getWhoFrom currentTransfer) . getID) (getEntries e))
        (beforeEntries, afterEntries) = Misc.partBeforeAfter fromEntry (getEntries e)
        notRunningIDs = map getID (filter ((/= Running) . getStatus) (getEntries e))

        reviseTrace candID trace = trace
            { getCriteria = getCriteria trace ++ [Sen.matchList notRunningIDs, Sen.matchID candID]
            , getTransferVal = getTransferVal trace * getNewValue currentTransfer }

        reviseFunc entry = do
            let newTraces = map (reviseTrace (getID entry)) (getWhatToDist currentTransfer)
            rawVoteChanges <- mapM (Sen.doCount (getCounter e)) (map getCriteria newTraces)
            let tracesAndChanges = zip (map fromIntegral rawVoteChanges) newTraces
                adjustedChanges = map (\(r,t) -> (floor (r * getTransferVal t), t)) tracesAndChanges
                filteredChanges = filter ((/= 0) . fst) adjustedChanges
                totalVoteChange = sum (map fst filteredChanges)
                addedTraces = map snd filteredChanges
            return (entry
                { getVoteChange = totalVoteChange
                , getTotalVotes = getTotalVotes entry + totalVoteChange
                , getCritTrace = getCritTrace entry ++ addedTraces })

        revisedFromEntry = fromEntry
            { getVoteChange = -(getVoteAmount currentTransfer)
            , getTotalVotes = getTotalVotes fromEntry - getVoteAmount currentTransfer
            , getCritTrace = getCritTrace fromEntry \\ getWhatToDist currentTransfer }
    revisedBeforeEntries <- mapM reviseFunc beforeEntries
    revisedAfterEntries <- mapM reviseFunc afterEntries
    let allRevised = revisedBeforeEntries ++ [revisedFromEntry] ++ revisedAfterEntries

    return (e
        { getEntries = allRevised
        , getTransferQueue = remainingTransfers })




--  needs to be modified to take into account ties
checkNoQuota :: Election -> ET.EitherT Election IO Election
checkNoQuota e = do
    let running = filter ((== Running) . getStatus) (getEntries e)
        sorted = reverse (List.sortBy compareVotes running)

    if (length running <= getVacancies e + 1)
        then do
            let makeElect entry n = do
                    let logmsg = show (getID entry) ++ " elected at logfile #" ++ show (getNextLogNum e)
                    IO.appendFile (getMainLog e) (logmsg ++ "\n")
                    Con.when (isVerbose e) (IO.hPutStrLn IO.stderr logmsg)
                    return (entry
                        { getStatus = Elected
                        , getChanged = True
                        , getOrderElected = Just n })

                reviseFunc input output toChange n =
                    if (length input == 0)
                        then return (reverse output)
                        else if ((head input) `elem` toChange)
                                then do
                                    r <- makeElect (head input) n
                                    reviseFunc (tail input) (r:output) toChange (n + 1)
                                else reviseFunc (tail input) ((head input):output) toChange n

                toChange = if (length running <= getVacancies e) then sorted else init sorted

            allRevised <- MIO.liftIO $ reviseFunc (getEntries e) [] toChange (getNextToElect e)

            ET.left (e
                { getEntries = allRevised
                , getVacancies = 0
                , getNextToElect = getNextToElect e + length toChange })
        else ET.right e




excludeCandidates :: Election -> ET.EitherT Election IO Election
excludeCandidates e = do
    let running = filter ((== Running) . getStatus) (getEntries e)
        sorted = reverse (List.sortBy compareVotes running)
        appliedBreakpoint = getQuota e - getTotalVotes (head sorted)

        excludeLoop n v e = do
            (i,r) <- MIO.liftIO $ excludeSomeone e
            let v1 = v + i
                n1 = n + 1
            if (v1 > appliedBreakpoint)
                then if (n > 0)
                        then do
                            MIO.liftIO $ Con.when (n > 1) $ do
                                let logmsg = "Bulk exclusion of " ++ show n ++
                                        " candidates at logfile #" ++ show (getNextLogNum e)
                                IO.appendFile (getMainLog e) (logmsg ++ "\n")
                                Con.when (isVerbose e) (IO.hPutStrLn IO.stderr logmsg)
                            ET.left e
                        else ET.left r
                else excludeLoop n1 v1 r

    if (length running > 0 && all (< getQuota e) (map getTotalVotes running))
        then excludeLoop 0 0 e
        else ET.right e




--  needs to be modified to take into account ties
--  this function is still in the IO monad in case I want to log something in verbose mode later
excludeSomeone :: Election -> IO (Int, Election)
excludeSomeone e = do
    let running = filter ((== Running) . getStatus) (getEntries e)
        excludedEntry = List.minimumBy compareVotes running
        (beforeEntries, afterEntries) = Misc.partBeforeAfter excludedEntry (getEntries e)

        newTransfer = Transfer
            { getWhoFrom = getID excludedEntry
            , getVoteAmount = getTotalVotes excludedEntry
            , getNewValue = 1
            , getWhatToDist = getCritTrace excludedEntry }

        revisedExcludedEntry = excludedEntry
            { getStatus = Eliminated
            , getChanged = True }
        allRevised = beforeEntries ++ [revisedExcludedEntry] ++ afterEntries

    return (getTotalVotes excludedEntry, e
        { getEntries = allRevised
        , getTransferQueue = (getTransferQueue e) ++ [newTransfer] })