{-# LANGUAGE DeriveAnyClass #-}
{-# LANGUAGE DeriveGeneric #-}
{-# LANGUAGE DerivingStrategies #-}
{-# LANGUAGE GeneralizedNewtypeDeriving #-}
{-# LANGUAGE TupleSections #-}
{-# OPTIONS_GHC -Wno-unrecognised-pragmas #-}
module Main (main) where
import Control.Monad
import qualified Data.ByteString as BS
import Data.Char (isSpace)
import Data.Foldable (toList)
import Data.List (dropWhileEnd, sortOn)
import Data.Map (Map)
import qualified Data.Map as Map
import Data.Maybe (mapMaybe)
import Data.Serialize (Serialize, decode, encode)
import Data.Word (Word64)
import GHC.Generics (Generic)
import Options.Applicative hiding (str)
import System.AtomicWrite.Writer.ByteString (atomicWriteFile)
import System.Directory (createDirectoryIfMissing, doesFileExist)
import System.Exit (die)
import System.FilePath (takeDirectory)
import System.Process (readProcess)
import Text.Printf (printf)
main :: IO ()
main = do
now <- getUTC
cmd <- customExecParser (prefs $ showHelpOnEmpty <> showHelpOnError) (info (pCommand <**> helper) mempty)
case cmd of
Bump str thresh fa -> withFrecencies fa $ expire thresh . bump str . decay now
Delete str fa -> withFrecencies fa $ delete str
View augmentArgs weights fa -> do
frecs <- loadFrecencies fa
fAugment <- augment augmentArgs
putStr . unlines . view weights . fAugment . decay now $ frecs
Scores augmentArgs weights fa -> do
frecs <- loadFrecencies fa
fAugment <- augment augmentArgs
printScores weights . fAugment . decay now $ frecs
Touch expireArgs fileArgs -> do
withFrecencies fileArgs $ expire expireArgs . decay now
data Command
= Bump NEString ExpireArgs FileArgs
| View AugmentArgs Weights FileArgs
| Delete NEString FileArgs
| Scores AugmentArgs Weights FileArgs
| Touch ExpireArgs FileArgs
pCommand :: Parser Command
pCommand =
subparser $
command
"bump"
( info
(withFile $ Bump <$> pStringArg (help "The entry to bump") <*> pExpireArgs)
( progDesc "Bump a single entry and update the database"
<> footer "Bumping adds 1 to the entry's hourly/weekly/monthly energy, updates all entries' energy, and removes entries whose monthly energy has dropped below the threshold."
)
)
<> command
"view"
( info
(withFile $ View <$> pAugmentArgs <*> pWeights)
( progDesc "View the history"
<> footer "When used with --augment and/or --restrict, this blocks on stdin."
)
)
<> command
"delete"
( info
(withFile $ Delete <$> pStringArg (help "The entry to delete"))
(progDesc "Delete an entry from the history")
)
<> command
"scores"
( info
(withFile $ Scores <$> pAugmentArgs <*> pWeights)
( progDesc "View score table"
<> footer "The hourly/daily/weekly energies are presented unweighted. Supports the same stdin functionality as the view command."
)
)
<> command
"touch"
( info
(withFile $ Touch <$> pExpireArgs)
(progDesc "Create and/or update a history file")
)
withFile :: Parser (FileArgs -> a) -> Parser a
withFile inner =
(\fp a err -> a (FileArgs fp err))
<$> strArgument (help "History file to use" <> metavar "FILE")
<*> inner
<*> flag False True (long "missing-file-error" <> short 'e' <> help "Throw an error if the file is missing, instead of treating it as an empty history")
data FileArgs = FileArgs
{ faPath :: FilePath,
_faErrorIfMissing :: Bool
-- TODO different output file
}
newtype Weights = Weights Energy
pStringArg :: Mod ArgumentFields NEString -> Parser NEString
pStringArg extraInfo = argument (maybeReader $ \str -> guard ('\n' `notElem` str) >> stripWhitespace str) (metavar "KEY" <> extraInfo)
pWeights :: Parser Weights
pWeights =
fmap Weights $
Energy
<$> option auto (short 'h' <> long "hourly" <> metavar "FLOAT" <> help "Hourly energy weight" <> showDefault <> value 720)
<*> option auto (short 'd' <> long "daily" <> metavar "FLOAT" <> help "Daily energy weight" <> showDefault <> value 30)
<*> option auto (short 'm' <> long "monthly" <> metavar "FLOAT" <> help "Monthly energy weight" <> showDefault <> value 1)
-- TODO no-decay
newtype ExpireArgs = ExpireArgs {_uaThreshold :: Double}
pExpireArgs :: Parser ExpireArgs
pExpireArgs =
ExpireArgs
<$> option
auto
( long "threshold"
<> short 't'
<> help "Expiration threshold. Entries with a monthly energy below this will be removed."
<> metavar "FLOAT"
<> value 0.1
<> showDefault
)
readInput :: IO [NEString]
readInput = mapMaybe stripWhitespace . lines <$> getContents
augment :: AugmentArgs -> IO (Frecencies -> Frecencies)
augment (AugmentArgs False False) = pure id
augment (AugmentArgs aug res) = do
strs <- readInput
pure $ \(Frecencies t fs) -> Frecencies t $
case (aug, res) of
(False, True) -> Map.fromList [(str, nrg) | str <- strs, nrg <- toList (Map.lookup str fs)]
(True, False) -> foldr (\str -> Map.insertWith (<>) str (Energy 0 0 0)) fs strs
(True, True) -> Map.fromList $ (\key -> maybe (key, Energy 0 0 0) (key,) (Map.lookup key fs)) <$> strs
_ -> fs
data AugmentArgs = AugmentArgs
{ _aaAugment :: Bool,
_aaRestrict :: Bool
}
pAugmentArgs :: Parser AugmentArgs
pAugmentArgs =
AugmentArgs
<$> flag False True (long "augment" <> short 'a' <> help "Augment the output with entries read from stdin, treating them as if they had a score of 0 if they are not present in the history")
<*> flag False True (long "restrict" <> short 'r' <> help "Only output entries present in the keys read from stdin")
type Time = Word64
newtype NEString = NEString {unNEString :: String}
deriving newtype (Eq, Ord, Show, Serialize)
getUTC :: IO Time
getUTC = read <$> readProcess "date" ["+%s"] ""
data Energy = Energy
{ hourly :: Double,
daily :: Double,
monthly :: Double
}
deriving stock (Generic)
deriving anyclass (Serialize)
instance Semigroup Energy where Energy h d m <> Energy h' d' m' = Energy (h + h') (d + d') (m + m')
data Frecencies = Frecencies
{ lastUpdate :: Time,
energies :: Map NEString Energy
}
deriving stock (Generic)
deriving anyclass (Serialize)
decay :: Time -> Frecencies -> Frecencies
decay tNow (Frecencies tUpdate frecs) = Frecencies tNow (fmap f frecs)
where
f (Energy h d m) = Energy (alphaHour * h) (alphaDay * d) (alphaMonth * m)
deltaSec = fromIntegral $ tNow - tUpdate
alphaHour = 0.5 ** (deltaSec / 3600)
alphaDay = 0.5 ** (deltaSec / 86400)
alphaMonth = 0.5 ** (deltaSec / 2592000)
score :: Weights -> Energy -> Double
score (Weights (Energy wh wd wm)) (Energy h d m) = wh * h + wd * d + wm * m
delete :: NEString -> Frecencies -> Frecencies
delete str (Frecencies t fs) = Frecencies t (Map.delete str fs)
expire :: ExpireArgs -> Frecencies -> Frecencies
expire (ExpireArgs threshold) (Frecencies t fs) = Frecencies t (Map.filter ((> threshold) . monthly) fs)
-- TODO bump by 1 1 1, multiply scores before presenting
bump :: NEString -> Frecencies -> Frecencies
bump str (Frecencies t fs) = Frecencies t (Map.insertWith (<>) str (Energy 1 1 1) fs)
view :: Weights -> Frecencies -> [String]
view weights = fmap (unNEString . fst) . sortOn (negate . score weights . snd) . Map.toList . energies
stripWhitespace :: String -> Maybe NEString
stripWhitespace str = if null str' then Nothing else Just (NEString str)
where
str' = (dropWhileEnd isSpace . dropWhile isSpace) str
loadFrecencies :: FileArgs -> IO Frecencies
loadFrecencies (FileArgs fp errorOnMissing) = do
exists <- doesFileExist fp
if exists
then BS.readFile fp >>= either die pure . decode
else
if errorOnMissing
then die $ "Error: missing history file " <> fp
else pure $ Frecencies 0 mempty
withFrecencies :: FileArgs -> (Frecencies -> Frecencies) -> IO ()
withFrecencies fa f = loadFrecencies fa >>= writeFrecencies (faPath fa) . f
writeFrecencies :: FilePath -> Frecencies -> IO ()
writeFrecencies path fs = do
createDirectoryIfMissing True (takeDirectory path)
atomicWriteFile path (encode fs)
printScores :: Weights -> Frecencies -> IO ()
printScores weights (Frecencies _ fs) = do
printf "weighted score\thourly\t\tdaily\t\tmonthly\n"
forM_ (sortOn (negate . score weights . snd) $ Map.toList fs) $ \(str, nrg@(Energy h d m)) ->
printf "%12.6f\t%10.6f\t%10.6f\t%10.6f\t%s\n" (score weights nrg) h d m (unNEString str)