# LANGUAGE GeneralizedNewtypeDeriving #

-- | Histograms
--
  Indended for qualified import .
--
-- > import Util.Histogram (Histogram)
-- > import qualified Util.Histogram as Histogram
module Util.Histogram (
    -- * Basic definitions
    Bin
  , Count
  , Histogram  -- opaque
    -- * Conversion
  , toMap
  , toList
  , fromMap
  , fromList
    -- * I/O
  , writeFile
  , readFile
    -- * Ranges
  , SplitRanges(..)
  , splitRanges
  , splitGivenRanges
  , range
    -- * Construction
  , BinSize(..)
  , discretize
  , empty
  , insert
  , insertMany
    -- * Combinators
  , max
  , add
  , unionWith
  , filterBins
  , filterCounts
  , nLargestBins
  ) where

import Prelude hiding (empty, max, readFile, toList, writeFile)
import qualified Prelude

import qualified Data.IntMap.Strict as Map
import qualified Data.Set as Set
import qualified Data.String as Base (lines)
import qualified GHC.Show (Show(show))
import qualified System.IO as IO
import Text.Read (read)

import Util.Range (Range(..), Ranges(..), SplitRanges(..))
import qualified Util.Range as Range

{-------------------------------------------------------------------------------
  Basic definitions
-------------------------------------------------------------------------------}

type Bin   = Int
type Count = Int

data Histogram =
    -- | Non-empty histogram (known binsize)
    Histogram !BinSize !(IntMap Count)
  deriving (Show)

toMap :: Histogram -> IntMap Count
toMap (Histogram _ m) = m

toList :: Histogram -> [(Bin, Count)]
toList = toPairs . toMap

-- | Construct a histogram from already-discretized data
fromMap :: BinSize -> IntMap Count -> Histogram
fromMap = Histogram

-- | Convenience wrapper around 'fromMap'. See 'fromMap' for more details.
fromList :: BinSize -> [(Bin, Count)] -> Histogram
fromList bz = fromMap bz . Map.fromList

{-------------------------------------------------------------------------------
  Output
-------------------------------------------------------------------------------}

-- | Write out histogram
--
  Example plot using gnuplot :
--
-- > set grid
  > set xrange [ 5:105 ]   -- for bins 10 , 20 , .. 100
  > set yrange [ 0:10.5 ] -- for counts from 0 .. 10
  > plot ' data.csv ' using 1:2 with boxes
--
-- See 'range' to compute @xrange@ and @yrange@.
writeFile :: FilePath -> Histogram -> IO ()
writeFile fp hist =
    withFile fp WriteMode $ \h ->
      forM_ (toList hist) $ \(step, count) ->
        IO.hPutStrLn h $ show step ++ "\t" ++ show count

-- | Read file previously written by 'writeFile'
--
-- Since we do not store the bin size, this must be passed separately. We
-- cannot verify this value.
readFile :: BinSize -> FilePath -> IO Histogram
readFile binSize fp = parse <$> IO.readFile fp
  where
    parse :: String -> Histogram
    parse = fromList binSize
          . map parseLine
          . Base.lines

    parseLine :: String -> (Bin, Count)
    parseLine str =
        let (bin, '\t':count) = break (== '\t') str
        in (read bin, read count)

{-------------------------------------------------------------------------------
  Construction
-------------------------------------------------------------------------------}

newtype BinSize = BinSize { binSizeToInt :: Int }
  deriving (Eq, Buildable)

instance Show BinSize where
  show = show . binSizeToInt

-- | Construct histogram by counting all the doubles per bin
--
-- See also 'fromMap' if the data is already discretized.
discretize :: BinSize -> [Double] -> Histogram
discretize = flip insertMany . empty

-- | Empty histogram
empty :: BinSize -> Histogram
empty bz = Histogram bz Map.empty

  |  and insert value
insert :: Double -> Histogram -> Histogram
insert d (Histogram (BinSize bz) h) =
    Histogram (BinSize bz) (Map.alter incr bin h)
  where
    bin :: Bin
    bin = floor (d / fromIntegral bz) * bz

    incr :: Maybe Count -> Maybe Count
    incr Nothing  = Just 1
    incr (Just n) = Just (n + 1)

insertMany :: [Double] -> Histogram -> Histogram
insertMany ds h = foldl' (flip insert) h ds

{-------------------------------------------------------------------------------
  Combinators
-------------------------------------------------------------------------------}

-- | Takes maximum across all bins
--
  This only makes sense if the two histograms uses the same bin size .
max :: Histogram -> Histogram -> Histogram
max = unionWith Prelude.max

-- | Sum across all bins
add :: Histogram -> Histogram -> Histogram
add = unionWith (+)

  | Like ' Map.unionWith '
unionWith :: (Count -> Count -> Count) -> Histogram -> Histogram -> Histogram
unionWith f (Histogram bz a) (Histogram bz' b) =
    if bz /= bz'
      then error "Cannot union two histograms with different bin sizes"
      else Histogram bz (Map.unionWith f a b)

-- | Filter bins
filterBins :: (Bin -> Bool) -> Histogram -> Histogram
filterBins p (Histogram bz m) = Histogram bz (aux m)
  where
    aux :: IntMap Count -> IntMap Count
    aux = Map.filterWithKey $ \b _c -> p b

-- | Filter counts
filterCounts :: (Count -> Bool) -> Histogram -> Histogram
filterCounts p (Histogram bz m) = Histogram bz (aux m)
  where
    aux :: IntMap Count -> IntMap Count
    aux = Map.filter p

-- | Keep a percentage of the bins, prioritizing bins with more elements
nLargestBins :: Double -> Histogram -> Histogram
nLargestBins p (Histogram bz m) = Histogram bz (aux m)
  where
    allBins  = sortBy (flip (comparing snd)) $ toPairs m
    keep     = ceiling $ p * fromIntegral (length allBins)
    nLargest = Set.fromList $ map fst $ take keep allBins
    aux      = Map.filterWithKey $ \b _c -> b `Set.member` nLargest

{-------------------------------------------------------------------------------
  Split ranges
-------------------------------------------------------------------------------}

-- | Compute the range in bins and counts
--
-- @splitRanges n m@ splits the histogram into different ranges at " markers ",
-- where a marker is a consecutive series of bins of exactly length @m@ where
  each of the  has a value at or less than the threshold @n@.
--
  We make sure that the first range starts at bin 0 .
splitRanges :: Int -> Int -> Histogram -> SplitRanges Bin Count
splitRanges n m (Histogram _ h) =
    case toPairs h of
      []           -> error "splitRanges: empty histogram"
      (b,c) : bins -> go ([(Range 0 b, 1)], Range.singleton c) 0 False bins
  where
    go :: ([(Range Bin, Int)], Range Count)  -- Accumulator
       -> Int  -- Number of consecutive bins below the threshold
       -> Bool -- Have we ever exceeded the threshold?
       -> [(Bin, Count)]                     -- To do
       -> SplitRanges Bin Count
    go (xRanges, yRange) _ _ [] = SplitRanges {
          _splitXRanges = reverse xRanges
        , _splitYRange  = yRange
        }
    go (xRanges, yRange) numBelowThreshold haveExceeded ((b, c) : bins) =
        go (xRanges'', yRange') numBelowThreshold' haveExceeded' bins
      where
        numBelowThreshold' = if c <= n
                               then numBelowThreshold + 1
                               else 0
        haveExceeded'      = haveExceeded || c > n

        (xRange, numBins) : xRanges' = xRanges -- cannot be empty
        xRanges'' = if haveExceeded && numBelowThreshold' == m
                     then (Range.singleton b, 1)             : xRanges
                     else (Range.with b xRange, numBins + 1) : xRanges'
        yRange'   = Range.with c yRange

-- | Like 'splitRanges', but with explicitly specified ranges
splitGivenRanges :: [Bin] -> Histogram -> SplitRanges Bin Count
splitGivenRanges = \splits (Histogram _ h) ->
    case toPairs h of
      []            -> error "splitGivenRanges: empty histogram"
      (b, c) : bins -> go ([(Range 0 b, 1)], Range.singleton c) splits bins
  where
    go :: ([(Range Bin, Int)], Range Count)  -- Accumulator
       -> [Bin]                              -- Unhandled splits
       -> [(Bin, Count)]                     -- To do
       -> SplitRanges Bin Count
    go (xRanges, yRange) _ [] = SplitRanges {
          _splitXRanges = reverse xRanges
        , _splitYRange  = yRange
        }
    go (xRanges, yRange) splits ((b, c) : bins) =
        case splits of
          s:splits' | b >= s ->
            let xRanges'' = (Range.singleton b, 1) : xRanges in
            go (xRanges'', yRange') splits' bins
          _otherwise ->
            let xRanges'' = (Range.with b xRange, numBins + 1) : xRanges' in
            go (xRanges'', yRange') splits bins
      where
        (xRange, numBins) : xRanges' = xRanges -- cannot be empty
        yRange' = Range.with c yRange

-- | Specialization of 'splitRanges' with no maximum gap.
range :: Histogram -> Ranges Bin Count
range h =
    let SplitRanges [(xRange, _numBins)] yRange = splitRanges 0 maxBound h
    in Ranges xRange yRange

 | Histograms


 > import Util.Histogram (Histogram)
 > import qualified Util.Histogram as Histogram
 * Basic definitions
 opaque
 * Conversion
 * I/O
 * Ranges
 * Construction
 * Combinators
------------------------------------------------------------------------------
  Basic definitions
------------------------------------------------------------------------------
 | Non-empty histogram (known binsize)
 | Construct a histogram from already-discretized data
 | Convenience wrapper around 'fromMap'. See 'fromMap' for more details.
------------------------------------------------------------------------------
  Output
------------------------------------------------------------------------------
 | Write out histogram


 > set grid
 for bins 10 , 20 , .. 100
 for counts from 0 .. 10

 See 'range' to compute @xrange@ and @yrange@.
 | Read file previously written by 'writeFile'

 Since we do not store the bin size, this must be passed separately. We
 cannot verify this value.
------------------------------------------------------------------------------
  Construction
------------------------------------------------------------------------------
 | Construct histogram by counting all the doubles per bin

 See also 'fromMap' if the data is already discretized.
 | Empty histogram
------------------------------------------------------------------------------
  Combinators
------------------------------------------------------------------------------
 | Takes maximum across all bins

 | Sum across all bins
 | Filter bins
 | Filter counts
 | Keep a percentage of the bins, prioritizing bins with more elements
------------------------------------------------------------------------------
  Split ranges
------------------------------------------------------------------------------
 | Compute the range in bins and counts

 @splitRanges n m@ splits the histogram into different ranges at " markers ",
 where a marker is a consecutive series of bins of exactly length @m@ where

 Accumulator
 Number of consecutive bins below the threshold
 Have we ever exceeded the threshold?
 To do
 cannot be empty
 | Like 'splitRanges', but with explicitly specified ranges
 Accumulator
 Unhandled splits
 To do
 cannot be empty
 | Specialization of 'splitRanges' with no maximum gap.
# LANGUAGE GeneralizedNewtypeDeriving #

  Indended for qualified import .
module Util.Histogram (
    Bin
  , Count
  , toMap
  , toList
  , fromMap
  , fromList
  , writeFile
  , readFile
  , SplitRanges(..)
  , splitRanges
  , splitGivenRanges
  , range
  , BinSize(..)
  , discretize
  , empty
  , insert
  , insertMany
  , max
  , add
  , unionWith
  , filterBins
  , filterCounts
  , nLargestBins
  ) where

import Prelude hiding (empty, max, readFile, toList, writeFile)
import qualified Prelude

import qualified Data.IntMap.Strict as Map
import qualified Data.Set as Set
import qualified Data.String as Base (lines)
import qualified GHC.Show (Show(show))
import qualified System.IO as IO
import Text.Read (read)

import Util.Range (Range(..), Ranges(..), SplitRanges(..))
import qualified Util.Range as Range


type Bin   = Int
type Count = Int

data Histogram =
    Histogram !BinSize !(IntMap Count)
  deriving (Show)

toMap :: Histogram -> IntMap Count
toMap (Histogram _ m) = m

toList :: Histogram -> [(Bin, Count)]
toList = toPairs . toMap

fromMap :: BinSize -> IntMap Count -> Histogram
fromMap = Histogram

fromList :: BinSize -> [(Bin, Count)] -> Histogram
fromList bz = fromMap bz . Map.fromList


  Example plot using gnuplot :
  > plot ' data.csv ' using 1:2 with boxes
writeFile :: FilePath -> Histogram -> IO ()
writeFile fp hist =
    withFile fp WriteMode $ \h ->
      forM_ (toList hist) $ \(step, count) ->
        IO.hPutStrLn h $ show step ++ "\t" ++ show count

readFile :: BinSize -> FilePath -> IO Histogram
readFile binSize fp = parse <$> IO.readFile fp
  where
    parse :: String -> Histogram
    parse = fromList binSize
          . map parseLine
          . Base.lines

    parseLine :: String -> (Bin, Count)
    parseLine str =
        let (bin, '\t':count) = break (== '\t') str
        in (read bin, read count)


newtype BinSize = BinSize { binSizeToInt :: Int }
  deriving (Eq, Buildable)

instance Show BinSize where
  show = show . binSizeToInt

discretize :: BinSize -> [Double] -> Histogram
discretize = flip insertMany . empty

empty :: BinSize -> Histogram
empty bz = Histogram bz Map.empty

  |  and insert value
insert :: Double -> Histogram -> Histogram
insert d (Histogram (BinSize bz) h) =
    Histogram (BinSize bz) (Map.alter incr bin h)
  where
    bin :: Bin
    bin = floor (d / fromIntegral bz) * bz

    incr :: Maybe Count -> Maybe Count
    incr Nothing  = Just 1
    incr (Just n) = Just (n + 1)

insertMany :: [Double] -> Histogram -> Histogram
insertMany ds h = foldl' (flip insert) h ds


  This only makes sense if the two histograms uses the same bin size .
max :: Histogram -> Histogram -> Histogram
max = unionWith Prelude.max

add :: Histogram -> Histogram -> Histogram
add = unionWith (+)

  | Like ' Map.unionWith '
unionWith :: (Count -> Count -> Count) -> Histogram -> Histogram -> Histogram
unionWith f (Histogram bz a) (Histogram bz' b) =
    if bz /= bz'
      then error "Cannot union two histograms with different bin sizes"
      else Histogram bz (Map.unionWith f a b)

filterBins :: (Bin -> Bool) -> Histogram -> Histogram
filterBins p (Histogram bz m) = Histogram bz (aux m)
  where
    aux :: IntMap Count -> IntMap Count
    aux = Map.filterWithKey $ \b _c -> p b

filterCounts :: (Count -> Bool) -> Histogram -> Histogram
filterCounts p (Histogram bz m) = Histogram bz (aux m)
  where
    aux :: IntMap Count -> IntMap Count
    aux = Map.filter p

nLargestBins :: Double -> Histogram -> Histogram
nLargestBins p (Histogram bz m) = Histogram bz (aux m)
  where
    allBins  = sortBy (flip (comparing snd)) $ toPairs m
    keep     = ceiling $ p * fromIntegral (length allBins)
    nLargest = Set.fromList $ map fst $ take keep allBins
    aux      = Map.filterWithKey $ \b _c -> b `Set.member` nLargest


  each of the  has a value at or less than the threshold @n@.
  We make sure that the first range starts at bin 0 .
splitRanges :: Int -> Int -> Histogram -> SplitRanges Bin Count
splitRanges n m (Histogram _ h) =
    case toPairs h of
      []           -> error "splitRanges: empty histogram"
      (b,c) : bins -> go ([(Range 0 b, 1)], Range.singleton c) 0 False bins
  where
       -> SplitRanges Bin Count
    go (xRanges, yRange) _ _ [] = SplitRanges {
          _splitXRanges = reverse xRanges
        , _splitYRange  = yRange
        }
    go (xRanges, yRange) numBelowThreshold haveExceeded ((b, c) : bins) =
        go (xRanges'', yRange') numBelowThreshold' haveExceeded' bins
      where
        numBelowThreshold' = if c <= n
                               then numBelowThreshold + 1
                               else 0
        haveExceeded'      = haveExceeded || c > n

        xRanges'' = if haveExceeded && numBelowThreshold' == m
                     then (Range.singleton b, 1)             : xRanges
                     else (Range.with b xRange, numBins + 1) : xRanges'
        yRange'   = Range.with c yRange

splitGivenRanges :: [Bin] -> Histogram -> SplitRanges Bin Count
splitGivenRanges = \splits (Histogram _ h) ->
    case toPairs h of
      []            -> error "splitGivenRanges: empty histogram"
      (b, c) : bins -> go ([(Range 0 b, 1)], Range.singleton c) splits bins
  where
       -> SplitRanges Bin Count
    go (xRanges, yRange) _ [] = SplitRanges {
          _splitXRanges = reverse xRanges
        , _splitYRange  = yRange
        }
    go (xRanges, yRange) splits ((b, c) : bins) =
        case splits of
          s:splits' | b >= s ->
            let xRanges'' = (Range.singleton b, 1) : xRanges in
            go (xRanges'', yRange') splits' bins
          _otherwise ->
            let xRanges'' = (Range.with b xRange, numBins + 1) : xRanges' in
            go (xRanges'', yRange') splits bins
      where
        yRange' = Range.with c yRange

range :: Histogram -> Ranges Bin Count
range h =
    let SplitRanges [(xRange, _numBins)] yRange = splitRanges 0 maxBound h
    in Ranges xRange yRange
