{-# 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.Foldable
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.Foldable as Lib.Foldable
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_Account_Sum :: Account_Sum
nil_Account_Sum =
	Data.Map.empty

nil_Unit_Sum :: Unit_Sum
nil_Unit_Sum =
	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 given 'Posting's.
postings :: (Foldable to, Foldable ti) => to (ti Posting) -> Balance -> Balance
postings = flip $ Data.Foldable.foldr (flip (Data.Foldable.foldr posting))

-- | Return the given 'Balance'
--   updated by the 'Transaction.postings'
--   of the given 'Transaction'.
transaction :: Transaction -> Balance -> Balance
transaction = postings . Transaction.postings

-- | 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 tr =
	postings (Transaction.balanced_virtual_postings tr) .
	postings (Transaction.virtual_postings tr) .
	postings (Transaction.postings tr)

-- | Return the given 'Balance'
--   updated by the 'Transaction.balanced_virtual_postings' of the given 'Transaction'.
transaction_balanced_virtual :: Transaction -> Balance -> Balance
transaction_balanced_virtual =
	postings . Transaction.balanced_virtual_postings

-- | 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 'Deviation' type

-- | The 'By_Unit' whose 'Unit_Sum's’ 'amount'
--   is not zero and possible 'Account' to 'infer_equilibrium'.
newtype Deviation
 =      Deviation 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').
deviation :: Balance -> Deviation
deviation balance = do
	let all_accounts = Lib.TreeMap.flatten (const ()) (by_account balance)
	let max_accounts = Data.Map.size all_accounts
	Deviation $
		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)

-- | Return the 'Balance' of the given 'Posting's and either:
--
--   * 'Left': the 'Posting's that cannot be inferred.
--   * 'Right': the given 'Posting's with inferred 'Amount's inserted.
infer_equilibrium
 :: Posting.By_Account
 -> (Balance, Either [Unit_Sum] Posting.By_Account)
infer_equilibrium ps = do
	let bal = postings ps nil
	let Deviation dev = deviation bal
	(\(l, r) -> (bal, case l of { [] -> Right r; _ -> Left l })) $ do
	Lib.Foldable.accumLeftsAndFoldrRights
	 (\p -> Data.Map.insertWith
		 (\pn -> (++) pn . Data.List.filter (not . Data.Map.null . Posting.amounts))
		 (Posting.account p) [p])
	 ps $ do
	Data.Foldable.foldr
	 (\unit_sum@(Unit_Sum{ amount=amt, accounts }) acc ->
		case Data.Map.size accounts of
		 1 -> (Right $ (Posting.nil $ fst $ Data.Map.elemAt 0 accounts)
			 { Posting.amounts = Amount.from_List [negate amt] }):acc
		 _ -> Left [unit_sum]:acc)
	 []
	 dev

-- ** Tests

-- | Return 'True' if and only if the given 'Deviation' maps no 'Unit'.
is_at_equilibrium :: Deviation -> Bool
is_at_equilibrium (Deviation dev) = Data.Map.null dev

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

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

-- * The 'Expanded' type

-- | Descending propagation of 'Amount's accross 'Account's.
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'.
expanded :: By_Account -> Expanded
expanded =
	Lib.TreeMap.map_by_depth_first
	 (\descendants value ->
		let exc = fromMaybe Data.Map.empty value in
		Account_Sum_Expanded
		 { exclusive = exc
		 , inclusive =
			Data.Map.foldr
			 ( Data.Map.unionWith (GHC.Num.+)
			 . ( inclusive
				 . fromMaybe (error "Oops, should not happen in: Hcompta.Calc.Balance.expanded")
				 . Lib.TreeMap.node_value) )
			 exc $ Lib.TreeMap.nodes $ descendants
		 })