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

import           Data.Data
import qualified Data.Foldable
import qualified Data.List
import qualified Data.Map.Strict as Data.Map
import           Data.Map.Strict (Map)
import           Data.Typeable ()
import           Data.Maybe (fromMaybe)
import qualified GHC.Num

import qualified Hcompta.Model as Model ()
import qualified Hcompta.Model.Account as Account
import qualified Hcompta.Lib.TreeMap as Lib.TreeMap
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
import qualified Hcompta.Model.Journal as Journal
import           Hcompta.Model.Journal (Journal)

-- * 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
 = Lib.TreeMap.TreeMap Account.Name Account_Sum
-- | A sum of 'Amount's,
-- concerning a single 'Account'.
type Account_Sum
 = Amount.By_Unit

type By_Unit
 = Map Amount.Unit Unit_Sum
-- | A sum of 'Amount's with their 'Account's involved,
-- concerning a single 'Unit'.
data Unit_Sum
 =   Unit_Sum
 { amount   :: Amount -- ^ The sum of 'Amount's for a single 'Unit'.
 , accounts :: Map Account () -- ^ The 'Account's involved to build 'amount'.
 } deriving (Data, Eq, Read, Show, Typeable)

-- ** Constructors

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

nil_By_Account :: By_Account
nil_By_Account =
	Lib.TreeMap.empty

nil_By_Unit :: By_Unit
nil_By_Unit =
	Data.Map.empty

nil_Sum_by_Account :: Account_Sum
nil_Sum_by_Account =
	Data.Map.empty

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

-- | Return a tuple associating the given 'Unit_Sum' with its 'Unit'.
assoc_unit_sum :: Unit_Sum -> (Unit, Unit_Sum)
assoc_unit_sum s = (Amount.unit $ amount s, s)

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

-- ** Incremental constructors

-- | Return the given 'Balance'
--   updated by the given 'Posting'.
posting :: Posting -> Balance -> Balance
posting post balance =
	balance
	 { by_account =
		Lib.TreeMap.insert
		 (Data.Map.unionWith (GHC.Num.+))
		 (Posting.account post)
		 (Posting.amounts post)
		 (by_account balance)
	 , by_unit =
		Data.Map.unionWith
		 (\x y -> Unit_Sum
			 { amount   = (GHC.Num.+) (amount x) (amount y)
			 , accounts = Data.Map.union (accounts x) (accounts y)
			 })
		 (by_unit balance) $
		Data.Map.map
		 (\amount -> Unit_Sum
			 { amount
			 , 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 tran balance =
	Data.Map.foldr
	 (flip (Data.List.foldl (flip posting)))
	 balance
	 (Transaction.postings tran)

-- | Return the given 'Balance'
--   updated by the 'Transaction.postings'
--   and 'Transaction.virtual_postings'
--   and 'Transaction.balanced_virtual_postings'
--   of the given 'Transaction'.
transaction_with_virtual :: Transaction -> Balance -> Balance
transaction_with_virtual tran balance =
	Data.Map.foldr
	 (flip (Data.List.foldl (flip posting)))
	 balance
	 (Transaction.postings tran)

-- | Return the given 'Balance'
--   updated by the 'Transaction.balanced_virtual_postings' of the given 'Transaction'.
transaction_balanced_virtual :: Transaction -> Balance -> Balance
transaction_balanced_virtual tran balance =
	Data.Map.foldr
	 (flip (Data.List.foldl (flip posting)))
	 balance
	 (Transaction.balanced_virtual_postings tran)

-- | Return the given 'Balance'
--   updated by the 'Journal.transactions'
--   of the given 'Journal',
--   through 'transaction'.
journal :: Journal -> Balance -> Balance
journal jour balance =
	Data.Map.foldl
	 (Data.List.foldl (flip transaction))
	 balance
	 (Journal.transactions jour)

-- | Return the given 'Balance'
--   updated by the 'Journal.transactions'
--   of the given 'Journal',
--   through 'transaction'.
journal_with_virtual :: Journal -> Balance -> Balance
journal_with_virtual jour balance =
	Data.Map.foldl
	 (Data.List.foldl (flip transaction_with_virtual))
	 balance
	 (Journal.transactions jour)

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

-- * The 'Equilibre' type

-- | See 'equilibre'.
newtype Equilibre
 =      Equilibre By_Unit
 deriving (Data, Eq, Read, Show, Typeable)

-- | Return the 'by_unit' of the given 'Balance' with:
--
-- * 'Unit's whose 'Unit_Sum' 'amount' verifying 'Amount.is_zero' removed,
--
-- * and remaining 'Unit's having their 'Unit_Sum' 'accounts'
-- complemented with the 'by_account' of the given 'Balance'
-- (i.e. now mapping to accounts __not__ involved to build the 'Unit_Sum').
equilibre :: Balance -> Equilibre
equilibre balance = do
	let all_accounts = Lib.TreeMap.flatten (const ()) (by_account balance)
	let max_accounts = Data.Map.size all_accounts
	Equilibre $
		Data.Map.foldlWithKey
		 (\m unit Unit_Sum{amount, accounts} ->
			if Amount.is_zero amount
			then m
			else
				case Data.Map.size accounts of
				 n | n == max_accounts ->
					Data.Map.insert unit Unit_Sum{amount, accounts=Data.Map.empty} m
				 _ -> do
					let diff = Data.Map.difference all_accounts accounts
					Data.Map.insert unit Unit_Sum{amount, accounts=diff} m
		 )
		 Data.Map.empty
		 (by_unit balance)

-- ** Tests

-- | Return 'True' if and only if the given 'Equilibre' satisfies:
--
-- * 'is_equilibrated',
-- * or 'is_inferrable'.
is_equilibrable :: Equilibre -> Bool
is_equilibrable e@(Equilibre eq) =
	Data.Map.null eq || is_inferrable e

-- | Return 'True' if and only if the given 'Equilibre' maps no 'Unit'.
is_equilibrated :: Equilibre -> Bool
is_equilibrated (Equilibre eq) = Data.Map.null eq

-- | Return 'True' if and only if the given 'Equilibre'
-- maps only to 'Unit_Sum's whose 'accounts'
-- maps exactly one 'Account'.
is_inferrable :: Equilibre -> Bool
is_inferrable (Equilibre eq) =
	Data.Foldable.all
	 (\Unit_Sum{accounts} -> Data.Map.size accounts == 1)
	 eq

-- | Return 'True' if and only if the given 'Equilibre'
-- maps to at least one 'Unit_Sum's whose 'accounts'
-- maps more than one 'Account'.
is_non_inferrable :: Equilibre -> Bool
is_non_inferrable (Equilibre eq) =
	Data.Foldable.any
	 (\Unit_Sum{accounts} -> Data.Map.size accounts > 1)
	 eq

-- * The 'Expanded' type

-- | See 'expand'.
type Expanded = Lib.TreeMap.TreeMap Account.Name Account_Sum_Expanded
data Account_Sum_Expanded
 =   Account_Sum_Expanded
 { inclusive :: Amount.By_Unit
 , exclusive :: Amount.By_Unit
 }
 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 Amount.'Amount.By_Unit'
-- of the 'Account's’ for which it is 'Account.ascending'.
expand :: By_Account -> Expanded
expand =
	Lib.TreeMap.depth_first_map
	 (\descendants content ->
		let exc = fromMaybe Data.Map.empty content in
		Account_Sum_Expanded
		 { exclusive = exc
		 , inclusive =
			Data.Map.foldr
			 ( Data.Map.unionWith (GHC.Num.+)
			 . ( inclusive
				 . fromMaybe (error "Oops, should not happen")
				 . Lib.TreeMap.node_content) )
			 exc descendants
		 })