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 Data.List ( (\\) )
import qualified Data.List as List
import qualified Data.Maybe as Maybe
import qualified Data.Either.Unwrap as Either
import qualified Counter as Sen
import qualified Candidate as Typ
import qualified Criteria as Crit
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    :: Crit.Criteria
    , getTransferVal :: Float }

--  can't check for equality on functions, unfortunately,
--  so this'll have to do
instance Eq Trace where
    a == b =
        (length (getCriteria a) == length (getCriteria b))
        && (getTransferVal a == getTransferVal b)

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

data Transfer = Transfer
    { getWhoFrom    :: Typ.CandidateID
    , getVoteAmount :: Int
    , getNewValue   :: Float
    , getWhatToDist :: [Trace] }




createElection :: FilePath -> Sen.SenateCounter -> Int -> Bool -> IO Election
createElection outDir 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       = outDir ++ "/" ++ "log.txt"
        , 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 = [Crit.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 }

        --  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 $
            clearChanges e1 >>=
            electCandidates >>=
            checkIfDone >>=
            transferVotes >>=
            checkNoQuota >>=
            excludeSomeone

        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




clearChanges :: Election -> ET.EitherT Election IO Election
clearChanges e = do
    let clear entry = entry
            { getChanged = False
            , getVoteChange = 0 }
    ET.right (e { getEntries = map clear (getEntries e) })




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
doElectCandidate :: Election -> ET.EitherT Election IO Election
doElectCandidate e = do
    let (running, notRunning) = List.partition ((== 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 doVoteTransfer e
        else ET.right e




doVoteTransfer :: Election -> ET.EitherT 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 ++
                    [Crit.many (Crit.matchFromList notRunningIDs), Crit.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 totalVoteChange = sum (zipWith (floor .: (*)) (map fromIntegral rawVoteChanges) (map getTransferVal newTraces))
            return (entry
                { getVoteChange = totalVoteChange
                , getTotalVotes = getTotalVotes entry + totalVoteChange
                , getCritTrace = getCritTrace entry ++ newTraces })

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

    ET.left (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, notRunning) = List.partition ((== Running) . getStatus) (getEntries e)
        minimumEntry = List.minimumBy compareVotes running
        (beforeEntries, afterEntries) = Misc.partBeforeAfter minimumEntry (getEntries e)

        makeElect x = x
            { getStatus = Elected
            , getChanged = True }
        reviseFunc entry =
            if ((getStatus entry == Running) && (entry /= minimumEntry))
                then makeElect entry
                else entry
        revisedMinEntry =
            if (length running <= getVacancies e)
                then makeElect minimumEntry
                else minimumEntry
        allRevised = (map reviseFunc beforeEntries) ++ [revisedMinEntry] ++ (map reviseFunc afterEntries)

    if (length running <= getVacancies e + 1)
        then ET.left (e
            { getEntries = allRevised
            , getVacancies = 0 })
        else ET.right e




--  needs to be modified to take into account ties
excludeSomeone :: Election -> ET.EitherT Election IO Election
excludeSomeone e = do
    let (running, notRunning) = List.partition ((== 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

    if (getTotalVotes excludedEntry < getQuota e)
        then ET.left (e
            { getEntries = allRevised
            , getTransferQueue = (getTransferQueue e) ++ [newTransfer] })
        else ET.right e