{-# 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 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
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
 = Map Account 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 = 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 :: 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 =
		Data.Map.insertWith
		 (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 'transactions'.
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 'transactions'.
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 =
		Data.Map.unionWith
		 (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 max_accounts = Data.Map.size $ by_account balance
	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 d = Data.Map.map (const ()) $
					Data.Map.difference (by_account balance) accounts
				Data.Map.insert unit Unit_Sum{amount, accounts=d} m
	 )
	 Data.Map.empty
	 (by_unit balance)

-- ** Tests

-- | 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' 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 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'.
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 Amount.'Amount.By_Unit'
-- of the 'Account's’ for 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 amt ->
			Account.fold (Account.ascending account)
			 (\prefix -> Data.Map.insertWith (+) prefix amt))
		 balance
		 balance