{-
ToDo:
only read log file if present
translation from country code to country name
   including warning for unknown countries
-}
{-# LANGUAGE GeneralizedNewtypeDeriving #-}
module Main where

import qualified Control.Monad.Trans.State as SM
import qualified Control.Monad.Trans.Reader as RM
import Control.Monad.IO.Class (MonadIO, liftIO)
import Control.Monad.Trans.Class (lift)
import Control.Applicative (Applicative)

import Data.Tuple.HT (swap, mapSnd, )
import Data.List.HT (mapAdjacent, )
import Data.Maybe.HT (toMaybe, )
import Data.Maybe (mapMaybe, )
import Data.Char (isAlpha, toUpper, )
import Data.List (sort, mapAccumL, )
import qualified Data.List.Key as Key

import Data.Time.LocalTime (timeToTimeOfDay, timeOfDayToTime, )
import qualified Data.Time.Format as TimeFmt
import qualified Data.Time.Clock as Clock

import qualified Data.Map as Map
import Data.Map (Map)

import qualified System.IO.Strict as SIO
import qualified System.IO as IO


type Country = String
data Contestant =
   Contestant {
      contestantCountry :: Country,
      contestantNumber :: Int
   }
   deriving (Eq, Show)
type Time = Clock.DiffTime

type Entries = Map Country (Int, Time)


newtype M a = M (SM.StateT Entries (RM.ReaderT IO.Handle IO) a)
   deriving (Functor, Applicative, Monad, MonadIO)

gets :: (Entries -> a) -> M a
gets f =
   M (SM.gets f)

modify :: (Entries -> Entries) -> M ()
modify f =
   M (SM.modify f)


logFileName :: FilePath
logFileName = "toilet.log"


parseContestant :: String -> Either Country Contestant
parseContestant str =
   let (country, number) = span isAlpha str
   in  case reads number of
         [(n, "")] -> Right $ Contestant country n
         _ -> Left country

output :: String -> M ()
output str =
   liftIO $ putStrLn $ "       " ++ str

data Direction = Enter | Leave
   deriving (Eq, Show, Enum)

writeLog :: Contestant -> Time -> Direction -> M ()
writeLog cnt time dir =
   M (lift RM.ask >>= \h ->
      liftIO $
        (IO.hPutStrLn h $
         formatContestant cnt ++ " " ++
         formatDirection dir ++ " " ++
         formatTime time) >>
        IO.hFlush h)

logToEntries :: String -> Entries
logToEntries =
   foldl (\ent (Contestant country number, dirStr, time) ->
      case dirStr of
         Enter -> Map.insert country (number, time) ent
         Leave -> Map.delete country ent)
      Map.empty .
   parseLog

parseLog :: String -> [(Contestant, Direction, Time)]
parseLog =
   map (\[cntStr,dirStr,timeStr] ->
      (either
         (error $ "invalid contestant identifier " ++ cntStr)
         id $ parseContestant cntStr,
       case dirStr of
          "enters" -> Enter
          "leaves" -> Leave
          _ -> error $ "invalid direction " ++ dirStr,
       parseTime timeStr)) .
   map words .
   lines

timeZoneOffset :: Time
timeZoneOffset = 2*60*60

formatDirection :: Direction -> String
formatDirection dir =
   case dir of
      Enter -> "enters"
      Leave -> "leaves"

formatContestant :: Contestant -> String
formatContestant (Contestant country number) =
   country ++ show number

parseTime :: String -> Time
parseTime =
   subtract timeZoneOffset .
   timeOfDayToTime .
   TimeFmt.parseTimeOrError True TimeFmt.defaultTimeLocale "%R"

formatTime :: Time -> String
formatTime =
   TimeFmt.formatTime TimeFmt.defaultTimeLocale "%R" .
   timeToTimeOfDay .
   (timeZoneOffset+)

blockedMsg :: Contestant -> Time -> M ()
blockedMsg cnt time =
   output $
      "toilet blocked by " ++ formatContestant cnt ++
      " since " ++ formatTime time

caseToilet ::
   Country ->
   (Int -> Time -> M a) ->
   (M a) ->
   M a
caseToilet country ifBlocked ifFree =
   maybe
      ifFree
      (\(number, time) -> ifBlocked number time)
     =<< gets (Map.lookup country)

listToilet :: M ()
listToilet =
   mapM_ (\(country, (number, time)) ->
      output $ country ++ show number ++ " since " ++ formatTime time) =<<
   gets Map.toAscList

countryFrequencies :: String -> [(Int, String)]
countryFrequencies =
   reverse . sort .
   map swap .
   Map.toList .
   Map.fromListWith (+) .
   map (\(cnt,_,_) -> (contestantCountry cnt, 1)) .
   filter (\(_,dir,_) -> dir == Enter) .
   parseLog


type Minute = Int

countryTotalTimes :: String -> [(Minute, String)]
countryTotalTimes =
   reverse . sort .
   map swap .
   Map.toList .
   Map.fromListWith (+) .
   map (\(cnt,dir,time) ->
      (contestantCountry cnt,
       minutesFromTime $
       case dir of
          Enter -> -time
          Leave -> time)) .
   parseLog


countryMaximumTimes :: String -> [(Minute, String)]
countryMaximumTimes =
   reverse . sort .
   map swap .
   Map.toList .
   accumulateDurations max .
   map (\(cnt,dir,time) ->
      (contestantCountry cnt,
       dir,
       minutesFromTime time)) .
   parseLog

accumulateDurations ::
   (Num time, Ord key) =>
   (time -> time -> time) ->
   [(key, Direction, time)] -> Map key time
accumulateDurations acc =
   snd .
   foldl
      (\(startTimes, maxTimes) (cnt,dir,time) ->
          case dir of
             Enter ->
                (Map.insertWith (error "entered twice") cnt time startTimes,
                 maxTimes)
             Leave ->
                (Map.delete cnt startTimes,
                 let duration =
                        time -
                        Map.findWithDefault (error "never entered") cnt startTimes
                 in  Map.insertWith acc cnt duration maxTimes))
      (Map.empty, Map.empty)

minutesFromTime :: Time -> Int
minutesFromTime =
   round . (/ (60::Double)) . realToFrac

{- |
Contestants that left toilet immediately
are probably actually corrections by the registrars.
-}
contestantImmediateLeave :: String -> [(Contestant, Time)]
contestantImmediateLeave =
   mapMaybe
      (\((cnt0,dir0,time0), (cnt1,dir1,time1)) ->
         toMaybe
            (cnt0==cnt1 && dir0==Enter && dir1==Leave && time1-time0 <= 60)
            (cnt0, time0)) .
   mapAdjacent (,) .
   parseLog


loadPerTime :: String -> [(Time, Int)]
loadPerTime =
   snd .
   mapAccumL
      (\oldLoad (time,loadDiff) ->
         let newLoad = oldLoad + loadDiff
         in  (newLoad, (time, newLoad))) 0 .
   map (mapSnd $
      sum . map (\dir -> case dir of Enter -> 1; Leave -> -1)) .
   buckets (take (4*60+30) $ iterate (60+) (parseTime "09:00")) .
   map (\(_cnt,dir,time) -> (time, dir)) .
   parseLog


buckets :: (Eq a) =>
   [a] -> [(a,b)] -> [(a,[b])]
buckets as bs =
   snd $
   mapAccumL
      (\bs0 a ->
         let (bucket, bs1) = span ((a==) . fst) bs0
         in  (bs1, (a, map snd bucket)))
      bs as


generateStatistics :: IO ()
generateStatistics =
   do str <- readFile "toilet-correct.log"
      writeFile "country-frequency.csv" . unlines .
         map (\(n,country) -> show n ++ " " ++ country) . countryFrequencies $ str
      writeFile "country-totaltimes.csv" . unlines .
         map (\(n,country) -> show n ++ " " ++ country) . countryTotalTimes $ str
      writeFile "country-maxtimes.csv" . unlines .
         map (\(n,country) -> show n ++ " " ++ country) . countryMaximumTimes $ str
      writeFile "time-load.csv" . unlines .
         map (\(time,load) -> formatTime time ++ " " ++ show load) . loadPerTime $ str


{-
writeFile "toilet-by-country.log" . sortForCountry =<< readFile "toilet-correct.log"
-}
sortForCountry :: String -> String
sortForCountry =
   unlines .
   map unlines .
   Key.group (take 3) .
   Key.sort (take 3) .
   lines


{-
writeFile "toilet-by-contestant.log" . sortForContestant =<< readFile "toilet-correct.log"
-}
sortForContestant :: String -> String
sortForContestant =
   unlines .
   map unlines .
   Key.group (take 4) .
   Key.sort (take 4) .
   lines


loop :: M ()
loop =
   do ln <- fmap (map toUpper) $ liftIO getLine
      if ln=="LIST"
        then listToilet
        else
          case parseContestant ln of
             Left country ->
                caseToilet country
                   (blockedMsg . Contestant country)
                   (output "toilet free")
             Right cnt@(Contestant country number) ->
                fmap Clock.utctDayTime (liftIO Clock.getCurrentTime) >>= \time ->
                caseToilet country
                   (\oldNumber oldTime ->
                       let oldCnt = Contestant country oldNumber
                       in  if number==oldNumber
                             then
                                output (ln ++ " leaves toilet at " ++ formatTime time) >>
                                writeLog oldCnt time Leave >>
                                modify (Map.delete country)
                             else blockedMsg oldCnt oldTime)
                   (do writeLog cnt time Enter
                       modify (Map.insert country (number, time))
                       output (ln ++ " enters toilet at " ++ formatTime time))
      loop

main :: IO ()
main =
   let M lp = loop
   in  SIO.readFile logFileName >>= \logContent ->
       IO.withFile logFileName IO.AppendMode $ \h ->
       flip RM.runReaderT h $ flip SM.evalStateT
          (logToEntries logContent) lp