{-# LANGUAGE DeriveDataTypeable #-}
{-# LANGUAGE ScopedTypeVariables #-}
module Hcompta.Calc.Balance where

import           Data.Data
import qualified Data.Foldable
import qualified Data.List
import qualified Data.Map
import           Data.Map (Map)
import           Data.Typeable ()
import qualified GHC.Num

import qualified Hcompta.Model as Model
import qualified Hcompta.Model.Account as Account
import           Hcompta.Model.Account (Account)
import qualified Hcompta.Model.Amount as Amount
import           Hcompta.Model.Amount (Amount, Unit)
import qualified Hcompta.Model.Transaction as Transaction
import           Hcompta.Model.Transaction (Transaction, Posting)
import qualified Hcompta.Model.Transaction.Posting as Posting

-- * The 'Balance' type

-- | Sum by 'Account' and sum by 'Unit' of some 'Posting's.
data Balance
 =   Balance
 { by_account :: By_Account
 , by_unit    :: By_Unit
 } deriving (Data, Eq, Read, Show, Typeable)
type By_Account
 = Map Account Sum_by_Account
type By_Unit
 = Map Amount.Unit Sum_by_Unit


-- | A sum by 'Account' of the 'Amount's of some 'Posting's.
type Sum_by_Account
 = Amount.By_Unit

-- | A sum by 'Unit' of the 'Amount's of some 'Posting's,
--   with the 'Account's involved to build that sum.
data Sum_by_Unit
 =   Sum_by_Unit
 { accounts :: Map Account () -- ^ The 'Account's involved to build 'amount'.
 , amount   :: Amount -- ^ The sum of 'Amount's for a same 'Unit'.
 } deriving (Data, Eq, Read, Show, Typeable)

-- ** Constructors

nil :: Balance
nil =
	Balance
	 { by_account = Data.Map.empty
	 , by_unit    = Data.Map.empty
	 }

nil_By_Account :: By_Account
nil_By_Account =
	Data.Map.empty

nil_By_Unit :: By_Unit
nil_By_Unit =
	Data.Map.empty

nil_Sum_by_Account :: Sum_by_Account
nil_Sum_by_Account =
	Data.Map.empty

nil_Sum_by_Unit :: Sum_by_Unit
nil_Sum_by_Unit =
	Sum_by_Unit
	 { accounts = Data.Map.empty
	 , amount = Amount.nil
	 }

-- | Return a tuple associating the given 'Sum_by_Unit' with its 'Unit'.
assoc_by_amount_unit :: Sum_by_Unit -> (Unit, Sum_by_Unit)
assoc_by_amount_unit sum = (Amount.unit $ amount sum, sum)

-- | Return a 'Map' associating the given 'Sum_by_Unit' with their respective 'Unit'.
by_Unit_from_List :: [Sum_by_Unit] -> By_Unit
by_Unit_from_List balances =
	Data.Map.fromListWith
	 (\x y -> Sum_by_Unit
		 { amount=(GHC.Num.+) (amount x) (amount y)
		 , accounts=Data.Map.union (accounts x) (accounts y)
		 }) $
	Data.List.map assoc_by_amount_unit balances

-- ** Incremental constructors

-- | Return the given 'Balance'
--   updated by the given 'Posting'.
posting :: Posting -> Balance -> Balance
posting post balance =
	balance
	 { by_account =
		Data.Map.insertWith
		 (Data.Map.unionWith (GHC.Num.+))
		 (Posting.account post)
		 (Posting.amounts post)
		 (by_account balance)
	 , by_unit =
		Data.Map.unionWith
		 (\x y -> Sum_by_Unit
			 { amount = (GHC.Num.+) (amount x) (amount y)
			 , accounts = Data.Map.union (accounts x) (accounts y)
			 })
		 (by_unit balance) $
		Data.Map.map
		 (\amt -> Sum_by_Unit
			 { amount=amt
			 , accounts=Data.Map.singleton (Posting.account post) ()
			 })
		 (Posting.amounts post)
	 }

-- | Return the given 'Balance'
--   updated by the 'Transaction.postings' of the given 'Transaction'.
transaction :: Transaction -> Balance -> Balance
transaction transaction balance =
	Data.Map.fold (flip (Data.List.foldl (flip posting))) balance
	 (Transaction.postings transaction)

-- | Return the first given 'Balance'
--   updated by the second given 'Balance'.
union :: Balance -> Balance -> Balance
union b0 b1 =
	b0
	 { by_account =
		Data.Map.unionWith
		 (Data.Map.unionWith (GHC.Num.+))
		 (by_account b0)
		 (by_account b1)
	 , by_unit =
		Data.Map.unionWith
		 (\x y -> Sum_by_Unit
			 { amount = (GHC.Num.+) (amount x) (amount y)
			 , accounts = Data.Map.union (accounts x) (accounts y)
			 })
		 (by_unit b0)
		 (by_unit b1)
	 }

-- ** Tests

-- | Return 'True' if and only if the 'Balance'
--   has all its 'by_unit' 'amount's verify 'Amount.is_zero'
--   or exactly one 'Account' of the 'by_unit' 'accounts' is not in 'by_account'.
is_equilibrated :: Balance -> Bool
is_equilibrated balance =
	Data.Foldable.all
	 (\s ->
		(Amount.is_zero $ amount s) ||
		(Data.Map.size (accounts s) ==
		(Data.Map.size (by_account balance) - 1))) -- NOTE: Data.Map.size is O(1)
	 (by_unit balance)

-- * The 'Expanded' type

newtype Expanded
 =      Expanded By_Account
 deriving (Data, Eq, Read, Show, Typeable)

-- | Return the given 'By_Account'
--   with all missing 'Account.ascending' 'Account's inserted,
--   and every mapped Amount.'Amount.By_Unit'
--   added with any 'Account's Amount.'Amount.By_Unit'
--   to which it is 'Account.ascending'.
expand :: By_Account -> Expanded
expand balance =
	-- TODO: because (+) is associative
	--       the complexity could be improved a bit
	--       by only adding to the longest 'Account.ascending'
	--       and reuse this result thereafter,
	--       but coding this requires access
	--       to the hidden constructors of 'Data.Map.Map',
	--       which could be done through TemplateHaskell and lens:
	--       https://stackoverflow.com/questions/24153577/pattern-matching-on-data-map-for-implementing-map
	--
	-- a0' = a0 + a1 + a2 + a3   <-- current calculus
	--     = a0 + a1'            <-- improved calculus todo
	-- a1' =      a1 + a2 + a3
	--     =      a1 + a2'
	-- a2' =           a2 + a3
	-- a3' =                a3
	Expanded $
		Data.Map.foldrWithKey
		 (\account amount ->
			Account.fold (Account.ascending account)
			 (\prefix -> Data.Map.insertWith (+) prefix amount))
		 balance
		 balance