{-# LANGUAGE CPP #-}
{-|


Data values for zero or more report periods, and for the pre-report period.
Report periods are assumed to be contiguous, and represented only by start dates
(as keys of an IntMap).

-}
module Hledger.Data.PeriodData
( periodDataFromList

, lookupPeriodData
, insertPeriodData
, opPeriodData
, mergePeriodData
, padPeriodData

, tests_PeriodData
) where

#if MIN_VERSION_base(4,18,0)
import Data.Foldable1 (Foldable1(..))
#else
import Control.Applicative (liftA2)
#endif
import qualified Data.IntMap.Strict as IM
import qualified Data.IntSet as IS
#if !MIN_VERSION_base(4,20,0)
import Data.List (foldl')
#endif
import Data.Time (Day(..), fromGregorian)

import Test.Tasty (testGroup)
import Test.Tasty.HUnit ((@?=), testCase)

import Hledger.Data.Amount
import Hledger.Data.Types


instance Show a => Show (PeriodData a) where
  showsPrec :: Int -> PeriodData a -> ShowS
showsPrec Int
d (PeriodData a
h IntMap a
ds) =
    Bool -> ShowS -> ShowS
showParen (Int
d Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Int
10) (ShowS -> ShowS) -> ShowS -> ShowS
forall a b. (a -> b) -> a -> b
$
        String -> ShowS
showString String
"PeriodData"
      ShowS -> ShowS -> ShowS
forall b c a. (b -> c) -> (a -> b) -> a -> c
. String -> ShowS
showString String
"{ pdpre = " ShowS -> ShowS -> ShowS
forall b c a. (b -> c) -> (a -> b) -> a -> c
. a -> ShowS
forall a. Show a => a -> ShowS
shows a
h
      ShowS -> ShowS -> ShowS
forall b c a. (b -> c) -> (a -> b) -> a -> c
. String -> ShowS
showString String
", pdperiods = "
      ShowS -> ShowS -> ShowS
forall b c a. (b -> c) -> (a -> b) -> a -> c
. String -> ShowS
showString String
"fromList " ShowS -> ShowS -> ShowS
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [(Day, a)] -> ShowS
forall a. Show a => a -> ShowS
shows (((Int, a) -> (Day, a)) -> [(Int, a)] -> [(Day, a)]
forall a b. (a -> b) -> [a] -> [b]
map (\(Int
day, a
x) -> (Integer -> Day
ModifiedJulianDay (Integer -> Day) -> Integer -> Day
forall a b. (a -> b) -> a -> b
$ Int -> Integer
forall a. Integral a => a -> Integer
toInteger Int
day, a
x)) ([(Int, a)] -> [(Day, a)]) -> [(Int, a)] -> [(Day, a)]
forall a b. (a -> b) -> a -> b
$ IntMap a -> [(Int, a)]
forall a. IntMap a -> [(Int, a)]
IM.toList IntMap a
ds)
      ShowS -> ShowS -> ShowS
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Char -> ShowS
showChar Char
'}'

instance Foldable PeriodData where
  foldr :: forall a b. (a -> b -> b) -> b -> PeriodData a -> b
foldr a -> b -> b
f b
z (PeriodData a
h IntMap a
as) = (a -> b -> b) -> b -> IntMap a -> b
forall a b. (a -> b -> b) -> b -> IntMap a -> b
forall (t :: * -> *) a b.
Foldable t =>
(a -> b -> b) -> b -> t a -> b
foldr a -> b -> b
f (a -> b -> b
f a
h b
z) IntMap a
as
  foldl :: forall b a. (b -> a -> b) -> b -> PeriodData a -> b
foldl b -> a -> b
f b
z (PeriodData a
h IntMap a
as) = (b -> a -> b) -> b -> IntMap a -> b
forall b a. (b -> a -> b) -> b -> IntMap a -> b
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl b -> a -> b
f (b -> a -> b
f b
z a
h) IntMap a
as
  foldl' :: forall b a. (b -> a -> b) -> b -> PeriodData a -> b
foldl' b -> a -> b
f b
z (PeriodData a
h IntMap a
as) = let fzh :: b
fzh = b -> a -> b
f b
z a
h in b
fzh b -> b -> b
forall a b. a -> b -> b
`seq` (b -> a -> b) -> b -> IntMap a -> b
forall b a. (b -> a -> b) -> b -> IntMap a -> b
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl' b -> a -> b
f b
fzh IntMap a
as

#if MIN_VERSION_base(4,18,0)
instance Foldable1 PeriodData where
  foldrMap1 :: forall a b. (a -> b) -> (a -> b -> b) -> PeriodData a -> b
foldrMap1 a -> b
f a -> b -> b
g (PeriodData a
h IntMap a
as) = (a -> b -> b) -> b -> IntMap a -> b
forall a b. (a -> b -> b) -> b -> IntMap a -> b
forall (t :: * -> *) a b.
Foldable t =>
(a -> b -> b) -> b -> t a -> b
foldr a -> b -> b
g (a -> b
f a
h) IntMap a
as
  foldlMap1 :: forall a b. (a -> b) -> (b -> a -> b) -> PeriodData a -> b
foldlMap1 a -> b
f b -> a -> b
g (PeriodData a
h IntMap a
as) = (b -> a -> b) -> b -> IntMap a -> b
forall b a. (b -> a -> b) -> b -> IntMap a -> b
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl b -> a -> b
g (a -> b
f a
h) IntMap a
as
  foldlMap1' :: forall a b. (a -> b) -> (b -> a -> b) -> PeriodData a -> b
foldlMap1' a -> b
f b -> a -> b
g (PeriodData a
h IntMap a
as) = let fh :: b
fh = a -> b
f a
h in b
fh b -> b -> b
forall a b. a -> b -> b
`seq` (b -> a -> b) -> b -> IntMap a -> b
forall b a. (b -> a -> b) -> b -> IntMap a -> b
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl' b -> a -> b
g b
fh IntMap a
as
#endif

instance Traversable PeriodData where
  traverse :: forall (f :: * -> *) a b.
Applicative f =>
(a -> f b) -> PeriodData a -> f (PeriodData b)
traverse a -> f b
f (PeriodData a
h IntMap a
as) = (b -> IntMap b -> PeriodData b)
-> f b -> f (IntMap b) -> f (PeriodData b)
forall a b c. (a -> b -> c) -> f a -> f b -> f c
forall (f :: * -> *) a b c.
Applicative f =>
(a -> b -> c) -> f a -> f b -> f c
liftA2 b -> IntMap b -> PeriodData b
forall a. a -> IntMap a -> PeriodData a
PeriodData (a -> f b
f a
h) (f (IntMap b) -> f (PeriodData b))
-> f (IntMap b) -> f (PeriodData b)
forall a b. (a -> b) -> a -> b
$ (a -> f b) -> IntMap a -> f (IntMap b)
forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
(a -> f b) -> t a -> f (t b)
forall (f :: * -> *) a b.
Applicative f =>
(a -> f b) -> IntMap a -> f (IntMap b)
traverse a -> f b
f IntMap a
as

-- | The Semigroup instance for 'AccountBalance' will simply take the union of
-- keys in the date map section. This may not be the result you want if the
-- keys are not identical.
instance Semigroup a => Semigroup (PeriodData a) where
  PeriodData a
h1 IntMap a
as1 <> :: PeriodData a -> PeriodData a -> PeriodData a
<> PeriodData a
h2 IntMap a
as2 = a -> IntMap a -> PeriodData a
forall a. a -> IntMap a -> PeriodData a
PeriodData (a
h1 a -> a -> a
forall a. Semigroup a => a -> a -> a
<> a
h2) (IntMap a -> PeriodData a) -> IntMap a -> PeriodData a
forall a b. (a -> b) -> a -> b
$ (a -> a -> a) -> IntMap a -> IntMap a -> IntMap a
forall a. (a -> a -> a) -> IntMap a -> IntMap a -> IntMap a
IM.unionWith a -> a -> a
forall a. Semigroup a => a -> a -> a
(<>) IntMap a
as1 IntMap a
as2

instance Monoid a => Monoid (PeriodData a) where
  mempty :: PeriodData a
mempty = a -> IntMap a -> PeriodData a
forall a. a -> IntMap a -> PeriodData a
PeriodData a
forall a. Monoid a => a
mempty IntMap a
forall a. Monoid a => a
mempty

-- | Construct an 'PeriodData' from a list.
periodDataFromList :: a -> [(Day, a)] -> PeriodData a
periodDataFromList :: forall a. a -> [(Day, a)] -> PeriodData a
periodDataFromList a
h = a -> IntMap a -> PeriodData a
forall a. a -> IntMap a -> PeriodData a
PeriodData a
h (IntMap a -> PeriodData a)
-> ([(Day, a)] -> IntMap a) -> [(Day, a)] -> PeriodData a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [(Int, a)] -> IntMap a
forall a. [(Int, a)] -> IntMap a
IM.fromList ([(Int, a)] -> IntMap a)
-> ([(Day, a)] -> [(Int, a)]) -> [(Day, a)] -> IntMap a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ((Day, a) -> (Int, a)) -> [(Day, a)] -> [(Int, a)]
forall a b. (a -> b) -> [a] -> [b]
map (\(Day
d, a
a) -> (Integer -> Int
forall a. Num a => Integer -> a
fromInteger (Integer -> Int) -> Integer -> Int
forall a b. (a -> b) -> a -> b
$ Day -> Integer
toModifiedJulianDay Day
d, a
a))

-- | Get account balance information to the period containing a given 'Day'.
lookupPeriodData :: Day -> PeriodData a -> a
lookupPeriodData :: forall a. Day -> PeriodData a -> a
lookupPeriodData Day
d (PeriodData a
h IntMap a
as) =
    a -> ((Int, a) -> a) -> Maybe (Int, a) -> a
forall b a. b -> (a -> b) -> Maybe a -> b
maybe a
h (Int, a) -> a
forall a b. (a, b) -> b
snd (Maybe (Int, a) -> a) -> Maybe (Int, a) -> a
forall a b. (a -> b) -> a -> b
$ Int -> IntMap a -> Maybe (Int, a)
forall a. Int -> IntMap a -> Maybe (Int, a)
IM.lookupLE (Integer -> Int
forall a. Num a => Integer -> a
fromInteger (Integer -> Int) -> Integer -> Int
forall a b. (a -> b) -> a -> b
$ Day -> Integer
toModifiedJulianDay Day
d) IntMap a
as

-- | Add account balance information to the appropriate location in 'PeriodData'.
insertPeriodData :: Semigroup a => Maybe Day -> a -> PeriodData a -> PeriodData a
insertPeriodData :: forall a.
Semigroup a =>
Maybe Day -> a -> PeriodData a -> PeriodData a
insertPeriodData Maybe Day
mday a
b PeriodData a
balances = case Maybe Day
mday of
    Maybe Day
Nothing  -> PeriodData a
balances{pdpre = pdpre balances <> b}
    Just Day
day -> PeriodData a
balances{pdperiods = IM.insertWith (<>) (fromInteger $ toModifiedJulianDay day) b $ pdperiods balances}

-- | Merges two 'PeriodData', using the given operation to combine their balance information.
--
-- This will drop keys if they are not present in both 'PeriodData'.
opPeriodData :: (a -> b -> c) -> PeriodData a -> PeriodData b -> PeriodData c
opPeriodData :: forall a b c.
(a -> b -> c) -> PeriodData a -> PeriodData b -> PeriodData c
opPeriodData a -> b -> c
f (PeriodData a
h1 IntMap a
as1) (PeriodData b
h2 IntMap b
as2) =
    c -> IntMap c -> PeriodData c
forall a. a -> IntMap a -> PeriodData a
PeriodData (a -> b -> c
f a
h1 b
h2) (IntMap c -> PeriodData c) -> IntMap c -> PeriodData c
forall a b. (a -> b) -> a -> b
$ (a -> b -> c) -> IntMap a -> IntMap b -> IntMap c
forall a b c. (a -> b -> c) -> IntMap a -> IntMap b -> IntMap c
IM.intersectionWith a -> b -> c
f IntMap a
as1 IntMap b
as2

-- | Merges two 'PeriodData', using the given operations for balance
-- information only in the first, only in the second, or in both
-- 'PeriodData', respectively.
mergePeriodData :: (a -> c) -> (b -> c) -> (a -> b -> c)
                -> PeriodData a -> PeriodData b -> PeriodData c
mergePeriodData :: forall a c b.
(a -> c)
-> (b -> c)
-> (a -> b -> c)
-> PeriodData a
-> PeriodData b
-> PeriodData c
mergePeriodData a -> c
only1 b -> c
only2 a -> b -> c
f = \(PeriodData a
h1 IntMap a
as1) (PeriodData b
h2 IntMap b
as2) ->
    c -> IntMap c -> PeriodData c
forall a. a -> IntMap a -> PeriodData a
PeriodData (a -> b -> c
f a
h1 b
h2) (IntMap c -> PeriodData c) -> IntMap c -> PeriodData c
forall a b. (a -> b) -> a -> b
$ IntMap a -> IntMap b -> IntMap c
merge IntMap a
as1 IntMap b
as2
  where
    merge :: IntMap a -> IntMap b -> IntMap c
merge = (Int -> a -> b -> Maybe c)
-> (IntMap a -> IntMap c)
-> (IntMap b -> IntMap c)
-> IntMap a
-> IntMap b
-> IntMap c
forall a b c.
(Int -> a -> b -> Maybe c)
-> (IntMap a -> IntMap c)
-> (IntMap b -> IntMap c)
-> IntMap a
-> IntMap b
-> IntMap c
IM.mergeWithKey (\Int
_ a
x b
y -> c -> Maybe c
forall a. a -> Maybe a
Just (c -> Maybe c) -> c -> Maybe c
forall a b. (a -> b) -> a -> b
$ a -> b -> c
f a
x b
y) ((a -> c) -> IntMap a -> IntMap c
forall a b. (a -> b) -> IntMap a -> IntMap b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap a -> c
only1) ((b -> c) -> IntMap b -> IntMap c
forall a b. (a -> b) -> IntMap a -> IntMap b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap b -> c
only2)

-- | Pad out the datemap of an 'PeriodData' so that every key from a set is present.
padPeriodData :: Monoid a => IS.IntSet -> PeriodData a -> PeriodData a
padPeriodData :: forall a. Monoid a => IntSet -> PeriodData a -> PeriodData a
padPeriodData IntSet
keys PeriodData a
bal = PeriodData a
bal{pdperiods = pdperiods bal <> IM.fromSet (const mempty) keys}


-- tests

tests_PeriodData :: TestTree
tests_PeriodData =
  let
    dayMap :: PeriodData MixedAmount
dayMap  = MixedAmount -> [(Day, MixedAmount)] -> PeriodData MixedAmount
forall a. a -> [(Day, a)] -> PeriodData a
periodDataFromList ([Amount] -> MixedAmount
forall (t :: * -> *). Foldable t => t Amount -> MixedAmount
mixed [DecimalRaw Integer -> Amount
usd DecimalRaw Integer
1]) [(Integer -> Int -> Int -> Day
fromGregorian Integer
2000 Int
01 Int
01, [Amount] -> MixedAmount
forall (t :: * -> *). Foldable t => t Amount -> MixedAmount
mixed [DecimalRaw Integer -> Amount
usd DecimalRaw Integer
2]), (Integer -> Int -> Int -> Day
fromGregorian Integer
2004 Int
02 Int
28, [Amount] -> MixedAmount
forall (t :: * -> *). Foldable t => t Amount -> MixedAmount
mixed [DecimalRaw Integer -> Amount
usd DecimalRaw Integer
3])]
    dayMap2 :: PeriodData MixedAmount
dayMap2 = MixedAmount -> [(Day, MixedAmount)] -> PeriodData MixedAmount
forall a. a -> [(Day, a)] -> PeriodData a
periodDataFromList ([Amount] -> MixedAmount
forall (t :: * -> *). Foldable t => t Amount -> MixedAmount
mixed [DecimalRaw Integer -> Amount
usd DecimalRaw Integer
2]) [(Integer -> Int -> Int -> Day
fromGregorian Integer
2000 Int
01 Int
01, [Amount] -> MixedAmount
forall (t :: * -> *). Foldable t => t Amount -> MixedAmount
mixed [DecimalRaw Integer -> Amount
usd DecimalRaw Integer
4]), (Integer -> Int -> Int -> Day
fromGregorian Integer
2004 Int
02 Int
28, [Amount] -> MixedAmount
forall (t :: * -> *). Foldable t => t Amount -> MixedAmount
mixed [DecimalRaw Integer -> Amount
usd DecimalRaw Integer
6])]
  in String -> [TestTree] -> TestTree
testGroup String
"PeriodData" [

  String -> Assertion -> TestTree
testCase String
"periodDataFromList" (Assertion -> TestTree) -> Assertion -> TestTree
forall a b. (a -> b) -> a -> b
$ do
    PeriodData MixedAmount -> Int
forall a. PeriodData a -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length PeriodData MixedAmount
dayMap Int -> Int -> Assertion
forall a. (Eq a, Show a, HasCallStack) => a -> a -> Assertion
@?= Int
3,

  String -> Assertion -> TestTree
testCase String
"Semigroup instance" (Assertion -> TestTree) -> Assertion -> TestTree
forall a b. (a -> b) -> a -> b
$ do
    PeriodData MixedAmount
dayMap PeriodData MixedAmount
-> PeriodData MixedAmount -> PeriodData MixedAmount
forall a. Semigroup a => a -> a -> a
<> PeriodData MixedAmount
dayMap PeriodData MixedAmount -> PeriodData MixedAmount -> Assertion
forall a. (Eq a, Show a, HasCallStack) => a -> a -> Assertion
@?= PeriodData MixedAmount
dayMap2,

  String -> Assertion -> TestTree
testCase String
"Monoid instance" (Assertion -> TestTree) -> Assertion -> TestTree
forall a b. (a -> b) -> a -> b
$ do
    PeriodData MixedAmount
dayMap PeriodData MixedAmount
-> PeriodData MixedAmount -> PeriodData MixedAmount
forall a. Semigroup a => a -> a -> a
<> PeriodData MixedAmount
forall a. Monoid a => a
mempty PeriodData MixedAmount -> PeriodData MixedAmount -> Assertion
forall a. (Eq a, Show a, HasCallStack) => a -> a -> Assertion
@?= PeriodData MixedAmount
dayMap
  ]