```{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE BangPatterns #-}
-- |
-- Module     : Data.Histogram.Bin
-- Maintainer : Alexey Khudyakov <alexey.skladnoy@gmail.com>
-- Stability  : experimental
--
-- Binning algorithms. This is mapping from set of interest to integer
-- indices and approximate reverse.

module Data.Histogram.Bin ( -- * Type classes
Bin(..)
, Bin1D(..)
, Indexable(..)
, Indexable2D(..)
-- * Bin types
-- ** Integer bins
, BinI(..)
, binI0
-- ** Integer bins with non-1 size
, BinInt
, binInt
-- ** Indexed bins
, BinIx(BinIx,unBinIx)
, binIx
-- ** Floating point bins
, BinF
, binF
, binFn
, binI2binF
, scaleBinF
-- *** Specialized for Double
, BinD
, binD
, binDn
, binI2binD
, scaleBinD
-- ** Log scale point
, LogBinD
, logBinD
-- ** 2D bins
, Bin2D(..)
, (><)
, nBins2D
, toIndex2D
, binX
, binY
, fmapBinX
, fmapBinY
-- ** 2D indexed bins
, BinIx2D (unBinIx2D)
, binIx2D
) where

import Data.Histogram.Parse

import GHC.Float (double2Int)
----------------------------------------------------------------
-- | Abstract binning algorithm. It provides way to map some values
-- onto continous range of integer values starting from zero.
--
-- Following invariant is expected to hold:
--
-- > toIndex . fromIndex == id
--
-- Reverse is not nessearily true.
class Bin b where
-- | Type of value to bin
type BinValue b
-- | Convert from value to index. No bound checking performed
toIndex :: b -> BinValue b -> Int
-- | Convert from index to value.
fromIndex :: b -> Int -> BinValue b
-- | Check whether value in range.
inRange :: b -> BinValue b -> Bool
-- | Total number of bins
nBins :: b -> Int

----------------------------------------------------------------
-- | One dimensional binning algorithm. It means that bin values have
-- some inherent ordering. For example all binning algorithms for real
-- numbers could be members or this type class whereas binning
-- algorithms for R^2 could not.
class Bin b => Bin1D b where
-- | List of center of bins in ascending order.
binsList :: b -> [BinValue b]
-- | List of bins in ascending order.
binsListRange :: b -> [(BinValue b, BinValue b)]

----------------------------------------------------------------
-- | Indexable is value which could be converted to and from Int
-- without information loss.
--
-- Always true
--
-- > deindex . index = id
--
-- Only if Int is in range
--
-- > index . deindex = id
class Indexable a where
-- | Convert value to index
index :: a -> Int
-- | Convert index to value
deindex :: Int -> a

instance Indexable Int where
index   = id
deindex = id

----------------------------------------------------------------
-- | This type class is same as Indexable but for 2D values.
class Indexable2D a where
-- | Convert value to index
index2D :: a -> (Int,Int)
-- | Convert index to value
deindex2D :: (Int,Int) -> a

instance (Indexable a, Indexable b) => Indexable2D (a,b) where
index2D   (x,y) = (index x,   index y)
deindex2D (i,j) = (deindex i, deindex j)

----------------------------------------------------------------
-- Integer bin
----------------------------------------------------------------
-- | Simple binning algorithm which map continous range of bins onto
-- indices. Each number correcsponds to different bin
data BinI = BinI {-# UNPACK #-} !Int {-# UNPACK #-} !Int
deriving Eq

-- | Construct BinI with n bins. Indexing starts from 0
binI0 :: Int -> BinI
binI0 n = BinI 0 (n-1)

instance Bin BinI where
type BinValue BinI = Int
toIndex   !(BinI base _) !x = x - base
{-# INLINE toIndex #-}
fromIndex !(BinI base _) !x = x + base
inRange   !(BinI x y) i     = i>=x && i<=y
{-# INLINE inRange #-}
nBins     !(BinI x y) = y - x + 1

instance Bin1D BinI where
binsList (BinI lo hi) = [lo .. hi]
binsListRange b = zip (binsList b) (binsList b)

instance Show BinI where
show (BinI lo hi) = unlines [ "# BinI"
, "# Low  = " ++ show lo
, "# High = " ++ show hi
]
readPrec = keyword "BinI" >> liftM2 BinI (value "Low") (value "High")

----------------------------------------------------------------
-- Another form of Integer bin
----------------------------------------------------------------

-- | Integer bins with size which differ from 1.
data BinInt = BinInt
{-# UNPACK #-} !Int -- Low bound
{-# UNPACK #-} !Int -- Bin size
{-# UNPACK #-} !Int -- Number of bins
deriving Eq

-- | Construct BinInt.
binInt :: Int                   -- ^ Lower bound
-> Int                   -- ^ Bin size
-> Int                   -- ^ Upper bound
-> BinInt
binInt lo n hi = BinInt lo n nb
where
nb = (hi-lo) `div` n

instance Bin BinInt where
type BinValue BinInt = Int
toIndex   !(BinInt base sz _) !x = (x - base) `div` sz
{-# INLINE toIndex #-}
fromIndex !(BinInt base sz _) !x = x * sz + base
inRange   !(BinInt base sz n) i  = i>=base && i<(base+n*sz)
{-# INLINE inRange #-}
nBins     !(BinInt _ _ n) = n

instance Show BinInt where
show (BinInt base sz n) =
unlines [ "# BinInt"
, "# Base = " ++ show base
, "# Step = " ++ show sz
, "# Bins = " ++ show n
]

readPrec = keyword "BinInt" >> liftM3 BinInt (value "Base") (value "Step") (value "Bins")

----------------------------------------------------------------
-- Bins for indexables
----------------------------------------------------------------

-- | Binning for indexable values
newtype BinIx i = BinIx { unBinIx :: BinI }
deriving Eq

-- | Construct indexed bin
binIx :: Indexable i => i -> i -> BinIx i
binIx lo hi = BinIx \$ BinI (index lo) (index hi)

instance Indexable i => Bin (BinIx i) where
type BinValue (BinIx i) = i
toIndex   (BinIx b) x = toIndex b (index x)
fromIndex (BinIx b) i = deindex (fromIndex b i)
inRange   (BinIx b) x = inRange b (index x)
nBins (BinIx b) = nBins b

instance Indexable i => Bin1D (BinIx i) where
binsList (BinIx b) = map deindex (binsList b)
binsListRange b    = let bins = binsList b in zip bins bins

instance (Show i, Indexable i) => Show (BinIx i) where
show (BinIx (BinI lo hi)) = unlines [ "# BinIx"
, "# Low  = " ++ show (deindex lo :: i)
, "# High = " ++ show (deindex hi :: i)
]
readPrec = keyword "BinIx" >> liftM2 binIx (value "Low") (value "High")

----------------------------------------------------------------
-- Floating point bin
----------------------------------------------------------------
-- | Floaintg point bins with equal sizes.
data BinF f where
BinF :: RealFrac f => !f -> !f -> !Int -> BinF f

instance Eq f => Eq (BinF f) where
(BinF lo hi n) == (BinF lo' hi' n') = lo == lo'  && hi == hi' && n == n'

-- | Create bins.
binF :: RealFrac f =>
f   -- ^ Lower bound of range
-> Int -- ^ Number of bins
-> f   -- ^ Upper bound of range
-> BinF f
binF from n to = BinF from ((to - from) / fromIntegral n) n

-- | Create bins. Note that actual upper bound can differ from specified.
binFn :: RealFrac f =>
f -- ^ Begin of range
-> f -- ^ Size of step
-> f -- ^ Approximation of end of range
-> BinF f
binFn from step to = BinF from step (round \$ (to - from) / step)

-- | Convert BinI to BinF
binI2binF :: RealFrac f => BinI -> BinF f
binI2binF b@(BinI i _) = BinF (fromIntegral i) 1 (nBins b)

-- | 'scaleBinF a b' scales BinF using linear transform 'a+b*x'
scaleBinF :: RealFrac f => f -> f -> BinF f -> BinF f
scaleBinF a b (BinF base step n)
| b > 0     = BinF (a + b*base) (b*step) n
| otherwise = error \$ "scaleBinF: b must be positive (b = "++show b++")"

instance Bin (BinF f) where
type BinValue (BinF f) = f
toIndex   !(BinF from step _) !x = floor \$ (x-from) / step
{-# INLINE toIndex #-}
fromIndex !(BinF from step _) !i = (step/2) + (fromIntegral i * step) + from
inRange   !(BinF from step n) x  = x > from && x < from + step*fromIntegral n
{-# INLINE inRange #-}
nBins     !(BinF _ _ n) = n

instance Bin1D (BinF f) where
binsList b@(BinF _ _ n) = map (fromIndex b) [0..n-1]
binsListRange b@(BinF _ step _) = map toPair (binsList b)
where
toPair x = (x - step/2, x + step/2)

instance Show f => Show (BinF f) where
show (BinF base step n) = unlines [ "# BinF"
, "# Base = " ++ show base
, "# Step = " ++ show step
, "# N    = " ++ show n
]
keyword "BinF"
base <- value "Base"
step <- value "Step"
n    <- value "N"
return \$ BinF base step n

----------------------------------------------------------------
-- Floating point bin /Specialized for Double
----------------------------------------------------------------
-- | Floaintg point bins with equal sizes. If you work with Doubles
-- this data type should be used instead of BinF.
data BinD = BinD {-# UNPACK #-} !Double {-# UNPACK #-} !Double {-# UNPACK #-} !Int

instance Eq BinD where
(BinD lo hi n) == (BinD lo' hi' n') = lo == lo'  && hi == hi' && n == n'

-- | Create bins.
binD :: Double -- ^ Lower bound of range
-> Int    -- ^ Number of bins
-> Double -- ^ Upper bound of range
-> BinD
binD from n to = BinD from ((to - from) / fromIntegral n) n

-- | Create bins. Note that actual upper bound can differ from specified.
binDn :: Double -- ^ Begin of range
-> Double -- ^ Size of step
-> Double -- ^ Approximation of end of range
-> BinD
binDn from step to = BinD from step (round \$ (to - from) / step)

-- | Convert BinI to BinF
binI2binD :: BinI -> BinD
binI2binD b@(BinI i _) = BinD (fromIntegral i) 1 (nBins b)

-- | 'scaleBinF a b' scales BinF using linear transform 'a+b*x'
scaleBinD :: Double -> Double -> BinD -> BinD
scaleBinD a b (BinD base step n)
| b > 0     = BinD (a + b*base) (b*step) n
| otherwise = error \$ "scaleBinF: b must be positive (b = "++show b++")"

-- Fast variant of flooor
floorD :: Double -> Int
floorD x | x < 0     = double2Int x - 1
| otherwise = double2Int x
{-# INLINE floorD #-}

instance Bin BinD where
type BinValue BinD = Double
toIndex   !(BinD from step _) !x = floorD \$ (x-from) / step
{-# INLINE toIndex #-}
fromIndex !(BinD from step _) !i = (step/2) + (fromIntegral i * step) + from
inRange   !(BinD from step n) x  = x > from && x < from + step*fromIntegral n
{-# INLINE inRange #-}
nBins     !(BinD _ _ n) = n

instance Bin1D BinD where
binsList b@(BinD _ _ n) = map (fromIndex b) [0..n-1]
binsListRange b@(BinD _ step _) = map toPair (binsList b)
where
toPair x = (x - step/2, x + step/2)

instance Show BinD where
show (BinD base step n) = unlines [ "# BinD"
, "# Base = " ++ show base
, "# Step = " ++ show step
, "# N    = " ++ show n
]
keyword "BinD"
base <- value "Base"
step <- value "Step"
n    <- value "N"
return \$ BinD base step n

----------------------------------------------------------------
-- Log-scale bin
----------------------------------------------------------------
-- | Logarithmic scale bins.
data LogBinD = LogBinD
Double -- Low border
Double -- Hi border
Double -- Increment ratio
Int    -- Number of bins
deriving Eq

-- | Create log-scale bins.
logBinD :: Double -> Int -> Double -> LogBinD
logBinD lo n hi = LogBinD lo hi ((hi/lo) ** (1 / fromIntegral n)) n

instance Bin LogBinD where
type BinValue LogBinD = Double
toIndex   !(LogBinD base _ step _) !x = floorD \$ logBase step (x / base)
{-# INLINE toIndex #-}
fromIndex !(LogBinD base _ step _) !i = base * step ^ i
inRange   !(LogBinD lo hi _ _) x  = x >= lo && x < hi
{-# INLINE inRange #-}
nBins     !(LogBinD _ _ _ n) = n

instance Show LogBinD where
show (LogBinD lo hi step n) =
unlines [ "# LogBinD"
, "# Lo   = " ++ show lo
, "# Hi   = " ++ show hi
, "# Step = " ++ show step
, "# N    = " ++ show n
]

----------------------------------------------------------------
-- 2D bin
----------------------------------------------------------------

-- | 2D bins. binX is binning along X axis and binY is one along Y axis.
data Bin2D binX binY = Bin2D !binX !binY
deriving Eq

-- | Alias for 'Bin2D'.
(><) :: binX -> binY -> Bin2D binX binY
(><) = Bin2D

-- | Get binning algorithm along X axis
binX :: Bin2D bx by -> bx
binX !(Bin2D bx _) = bx

-- | Get binning algorithm along Y axis
binY :: Bin2D bx by -> by
binY !(Bin2D _ by) = by

instance (Bin binX, Bin binY) => Bin (Bin2D binX binY) where
type BinValue (Bin2D binX binY) = (BinValue binX, BinValue binY)
toIndex b@(Bin2D bx by) (x,y)
| inRange b (x,y) = toIndex bx x + (toIndex by y)*(fromIntegral \$ nBins bx)
| otherwise       = maxBound
{-# INLINE toIndex #-}
fromIndex b@(Bin2D bx by) i = let (ix,iy) = toIndex2D b i
in  (fromIndex bx ix, fromIndex by iy)
inRange (Bin2D bx by) (x,y) = inRange bx x && inRange by y
{-# INLINE inRange #-}
nBins (Bin2D bx by) = (nBins bx) * (nBins by)

toIndex2D :: (Bin binX, Bin binY) => Bin2D binX binY -> Int -> (Int,Int)
toIndex2D b i = let (iy,ix) = divMod i (nBins \$ binX b) in (ix,iy)

-- | 2-dimensional size of binning algorithm
nBins2D :: (Bin bx, Bin by) => Bin2D bx by -> (Int,Int)
nBins2D (Bin2D bx by) = (nBins bx, nBins by)

-- | Apply function to X binning algorithm. If new binning algorithm
--   have different number of bins will fail.
fmapBinX :: (Bin bx, Bin bx') => (bx -> bx') -> Bin2D bx by -> Bin2D bx' by
fmapBinX f (Bin2D bx by)
| nBins bx' /= nBins bx = error "fmapBinX: new binnig algorithm has different number of bins"
| otherwise             = Bin2D bx' by
where
bx' = f bx

-- | Apply function to Y binning algorithm. If new binning algorithm
--   have different number of bins will fail.
fmapBinY ::(Bin by, Bin by') => (by -> by') -> Bin2D bx by -> Bin2D bx by'
fmapBinY f (Bin2D bx by)
| nBins by' /= nBins by = error "fmapBinY: new binnig algorithm has different number of bins"
| otherwise             = Bin2D bx by'
where
by' = f by

instance (Show b1, Show b2) => Show (Bin2D b1 b2) where
show (Bin2D b1 b2) = concat [ "# Bin2D\n"
, "# X\n"
, show b1
, "# Y\n"
, show b2
]
keyword "Bin2D"
keyword "X"
keyword "Y"
return \$ Bin2D b1 b2

----------------------------------------------------------------
-- Indexed 2D bins
----------------------------------------------------------------
-- | Binning for 2D indexable value
newtype BinIx2D i = BinIx2D {unBinIx2D :: (Bin2D BinI BinI) }

-- | Construct indexed bin
binIx2D :: Indexable2D i => i -> i -> BinIx2D i
binIx2D lo hi = let (ix,iy) = index2D lo
(jx,jy) = index2D hi
in BinIx2D \$ BinI ix jx >< BinI iy jy

instance Indexable2D i => Bin (BinIx2D i) where
type BinValue (BinIx2D i) = i
toIndex   (BinIx2D b) x = toIndex b (index2D x)
fromIndex (BinIx2D b) i = deindex2D \$ fromIndex b i
inRange   (BinIx2D b) x = inRange b (index2D x)
nBins     (BinIx2D b)   = nBins b

instance (Show i, Indexable2D i) => Show (BinIx2D i) where
show (BinIx2D b) = unlines [ "# BinIx2D"
, "# Low  = " ++ show (deindex2D (fromIndex b 0            ) :: i)
, "# High = " ++ show (deindex2D (fromIndex b (nBins b - 1)) :: i)
]