module Prometheus.Metric.Summary (
    Summary
,   Quantile
,   summary
,   defaultQuantiles
,   observe
,   observeDuration
,   getSummary

,   dumpEstimator
,   emptyEstimator
,   Estimator (..)
,   Item (..)
,   insert
,   compress
,   query
,   invariant
) where

import Prometheus.Info
import Prometheus.Metric
import Prometheus.Metric.Observer
import Prometheus.MonadMonitor

import Data.Int (Int64)
import Data.Foldable (foldr')
import qualified Control.Concurrent.STM as STM
import qualified Data.ByteString.UTF8 as BS


newtype Summary = MkSummary (STM.TVar Estimator)

-- | Creates a new summary metric with a given name, help string, and a list of
-- quantiles. A reasonable set set of quantiles is provided by
-- 'defaultQuantiles'.
summary :: Info -> [Quantile] -> IO (Metric Summary)
summary info quantiles = do
    valueTVar <- STM.newTVarIO (emptyEstimator quantiles)
    return Metric {
            handle = MkSummary valueTVar
        ,   collect = collectSummary info valueTVar
        }

withSummary :: MonadMonitor m
            => Metric Summary -> (Estimator -> Estimator) -> m ()
withSummary (Metric {handle = MkSummary valueTVar}) f =
    doIO $ STM.atomically $ do
        STM.modifyTVar' valueTVar compress
        STM.modifyTVar' valueTVar f

instance Observer Summary where
    -- | Adds a new observation to a summary metric.
    observe v s = withSummary s (insert v)

-- | Retrieves a list of tuples containing a quantile and its associated value.
getSummary :: Metric Summary -> IO [(Rational, Double)]
getSummary (Metric {handle = MkSummary valueTVar}) = do
    estimator <- STM.atomically $ do
        STM.modifyTVar' valueTVar compress
        STM.readTVar valueTVar
    let quantiles = map fst $ estQuantiles estimator
    let values = map (query estimator) quantiles
    return $ zip quantiles values

collectSummary :: Info -> STM.TVar Estimator -> IO [SampleGroup]
collectSummary info valueTVar = STM.atomically $ do
    STM.modifyTVar' valueTVar compress
    estimator@(Estimator count itemSum _ _) <- STM.readTVar valueTVar
    let quantiles = map fst $ estQuantiles estimator
    let samples =  map (toSample estimator) quantiles
    let sumSample = Sample (metricName info ++ "_sum") [] (bsShow itemSum)
    let countSample = Sample (metricName info ++ "_count") [] (bsShow count)
    return [SampleGroup info SummaryType $ samples ++ [sumSample, countSample]]
    where
        bsShow :: Show s => s -> BS.ByteString
        bsShow = BS.fromString . show

        toSample estimator q =
            Sample (metricName info) [("quantile", show $ toDouble q)] $
                bsShow $ query estimator q

        toDouble :: Rational -> Double
        toDouble = fromRational

dumpEstimator :: Metric Summary -> IO Estimator
dumpEstimator (Metric {handle = MkSummary valueTVar}) =
    STM.atomically $ STM.readTVar valueTVar

-- | A quantile is a pair of a quantile value and an associated acceptable error
-- value.
type Quantile = (Rational, Rational)

data Item = Item {
    itemValue :: Double
,   itemG     :: !Int64
,   itemD     :: !Int64
} deriving (Eq, Show)

instance Ord Item where
    compare a b = itemValue a `compare` itemValue b

data Estimator = Estimator {
    estCount      :: !Int64
,   estSum        :: !Double
,   estQuantiles  :: [Quantile]
,   estItems      :: [Item]
} deriving (Show)

defaultQuantiles :: [Quantile]
defaultQuantiles = [(0.5, 0.05), (0.9, 0.01), (0.99, 0.001)]

emptyEstimator :: [Quantile] -> Estimator
emptyEstimator quantiles = Estimator 0 0 quantiles []

insert :: Double -> Estimator -> Estimator
insert value estimator@(Estimator oldCount oldSum quantiles items) =
        newEstimator $ insertItem 0 items
    where
        newEstimator = Estimator (oldCount + 1) (oldSum + value) quantiles

        insertItem _ [] = [Item value 1 0]
        insertItem r [x]
            -- The first two cases cover the scenario where the initial size of
            -- the list is one.
            | r == 0 && value < itemValue x = Item value 1 0 : [x]
            | r == 0                        = x : [Item value 1 0]
            -- The last case covers the scenario where the we have walked off
            -- the end of a list with more than 1 element in the final case of
            -- insertItem in which case we already know that x < value.
            | otherwise                     = x : [Item value 1 0]
        insertItem r (x:y:xs)
            -- This first case only covers the scenario where value is less than
            -- the first item in a multi-item list. For subsequent steps of
            -- a multi valued list, this case cannot happen as it would have
            -- fallen through to the case below in the previous step.
            | value <= itemValue x = Item value 1 0 : x : y : xs
            | value <= itemValue y = x : Item value 1 (calcD $ r + itemG x)
                                       : y : xs
            | otherwise            = x : insertItem (itemG x + r) (y : xs)

        calcD r = max 0
                $ floor (invariant estimator (fromIntegral r)) - 1


compress :: Estimator -> Estimator
compress est@(Estimator _ _ _ [])    = est
compress est@(Estimator _ _ _ items) = est {
        estItems = (minItem :)
                 $ foldr' compressPair []
                 $ drop 1  -- The exact minimum item must be kept exactly.
                 $ zip items
                 $ scanl (+) 0 (map itemG items)
    }
    where
        minItem = head items
        compressPair (a, _) [] = [a]
        compressPair (a@(Item _ aG _), r) (b@(Item bVal bG bD):bs)
            | bD == 0             = a : b : bs
            | aG + bG + bD <= inv = Item bVal (aG + bG) bD : bs
            | otherwise           = a : b : bs
            where
                inv = floor $ invariant est (fromIntegral r)

query :: Estimator -> Rational -> Double
query est@(Estimator count _ _ items) q = findQuantile allRs items
    where
        allRs = scanl (+) 0 $ map itemG items

        n = fromIntegral count
        f = invariant est

        rank  = q * n
        bound = rank + (f rank / 2)

        findQuantile _        []   = 0 / 0  -- NaN
        findQuantile _        [a]  = itemValue a
        findQuantile (_:bR:rs) (a@(Item{}):b@(Item _ bG bD):xs)
            | fromIntegral (bR + bG + bD) > bound = itemValue a
            | otherwise            = findQuantile (bR:rs) (b:xs)
        findQuantile _        _    = error "Query impossibility"

invariant :: Estimator -> Rational -> Rational
invariant (Estimator count _ quantiles _) r = max 1
                                            $ minimum $ map fj quantiles
    where
        n = fromIntegral count
        fj (q, e) | q * n <= r = 2 * e * r / q
                  | otherwise  = 2 * e * (n - r) / (1 - q)