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

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

import qualified Hcompta.Lib.Foldable as Lib.Foldable
import qualified Hcompta.Lib.TreeMap as Lib.TreeMap
import           Hcompta.Lib.TreeMap (TreeMap)
import qualified Hcompta.Model.Account as Account
import           Hcompta.Model.Account (Account)

-- * The 'Posting' class

-- | A 'posting' used to produce a 'Balance'
--   must be an instance of this class.
class
 ( Ord (Posting_Unit p)
 , Num (Posting_Amount p)
 ) =>
 Posting p where
	type Posting_Amount p
	type Posting_Unit p
	posting_account :: p -> Account
	posting_amounts :: p -> Map (Posting_Unit p) (Posting_Amount p)
	posting_make    :: Account -> Map (Posting_Unit p) (Posting_Amount p) -> p

-- * The 'Balance' type

-- | Sum by 'Account' and sum by 'unit' of some 'Posting's.
data Balance amount unit
 =   Balance
 { balance_by_account :: Balance_by_Account amount unit
 , balance_by_unit    :: Balance_by_Unit amount unit
 } deriving (Data, Eq, Show, Typeable)

type Balance_by_Account amount unit
 = TreeMap Account.Name
           (Account_Sum amount unit)

-- | A sum of 'amount's,
-- concerning a single 'Account'.
type Account_Sum amount unit
 = Data.Map.Map unit amount

type Balance_by_Unit amount unit
 = Map unit (Unit_Sum amount)

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

-- ** Constructors

balance :: (Num amount, Ord unit) => Balance amount unit
balance =
	Balance
	 { balance_by_account = Lib.TreeMap.empty
	 , balance_by_unit    = Data.Map.empty
	 }

-- | Return the given 'Balance'
--   updated by the given 'Posting'.
posting ::
 ( Posting posting
 , unit   ~ Posting_Unit posting
 , amount ~ Posting_Amount posting
 ) => posting -> Balance amount unit -> Balance amount unit
posting post bal =
	bal
	 { balance_by_account =
		Lib.TreeMap.insert
		 (Data.Map.unionWith (flip (+)))
		 (posting_account post)
		 (posting_amounts post)
		 (balance_by_account bal)
	 , balance_by_unit =
		Data.Map.unionWith
		 (\new old -> Unit_Sum
			 { unit_sum_amount   = (+)
				 (unit_sum_amount old)
				 (unit_sum_amount new)
			 , unit_sum_accounts = Data.Map.unionWith
				 (const::()->()->())
				 (unit_sum_accounts old)
				 (unit_sum_accounts new)
			 })
		 (balance_by_unit bal) $
		Data.Map.map
		 (\amount -> Unit_Sum
			 { unit_sum_amount   = amount
			 , unit_sum_accounts = Data.Map.singleton (posting_account post) ()
			 })
		 (posting_amounts post)
	 }

-- | Return the given 'Balance'
--   updated by the given 'Posting's.
postings ::
 ( Posting posting
 , unit   ~ Posting_Unit   posting
 , amount ~ Posting_Amount posting
 , Foldable foldable )
 => foldable posting
 -> Balance amount unit
 -> Balance amount unit
postings = flip (Data.Foldable.foldr posting)

-- | Return the first given 'Balance'
--   updated by the second given 'Balance'.
union
 :: (Num amount, Ord unit)
 => Balance amount unit
 -> Balance amount unit
 -> Balance amount unit
union b0 b1 =
	b0
	 { balance_by_account =
		Lib.TreeMap.union
		 (Data.Map.unionWith (flip (+)))
		 (balance_by_account b0)
		 (balance_by_account b1)
	 , balance_by_unit =
		Data.Map.unionWith
		 (\new old -> Unit_Sum
			 { unit_sum_amount = (+)
				 (unit_sum_amount old)
				 (unit_sum_amount new)
			 , unit_sum_accounts = Data.Map.unionWith
				 (const::()->()->())
				 (unit_sum_accounts old)
				 (unit_sum_accounts new)
			 })
		 (balance_by_unit b0)
		 (balance_by_unit b1)
	 }

-- * The 'Deviation' type

-- | The 'Balance_by_Unit' whose 'unit_sum_amount'
--   is not zero and possible 'Account' to 'infer_equilibrium'.
newtype Deviation amount unit
 =      Deviation (Balance_by_Unit amount unit)
 deriving (Data, Eq, Show, Typeable)

-- | Return the 'balance_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 'balance_by_account' of the given 'Balance'
-- (i.e. now mapping to 'Account's __not__ involved to build the 'Unit_Sum').
deviation
 :: (Amount amount, Ord unit)
 => Balance amount unit
 -> Deviation amount unit
deviation bal = do
	let all_accounts = Lib.TreeMap.flatten (const ()) (balance_by_account bal)
	let max_accounts = Data.Map.size all_accounts
	Deviation $
		Data.Map.foldlWithKey
		 (\m unit Unit_Sum{unit_sum_amount, unit_sum_accounts} ->
			if amount_is_zero unit_sum_amount
			then m
			else
				case Data.Map.size unit_sum_accounts of
				 n | n == max_accounts ->
					Data.Map.insert unit Unit_Sum
					 { unit_sum_amount
					 , unit_sum_accounts = Data.Map.empty
					 } m
				 _ -> do
					let diff = Data.Map.difference all_accounts unit_sum_accounts
					Data.Map.insert unit Unit_Sum
					 { unit_sum_amount
					 , unit_sum_accounts = diff
					 } m
		 )
		 Data.Map.empty
		 (balance_by_unit bal)

-- ** The 'Amount' class

-- | An 'amount' used to produce a 'Deviation'
--   must be an instance of this class.
class
 Num a =>
 Amount a where
	amount_is_zero :: a -> Bool

-- ** The equilibrium

-- | 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 posting
 , Amount amount
 , Ord unit
 , amount ~ Posting_Amount posting
 , unit   ~ Posting_Unit   posting )
 => Map Account [posting]
 -> ( Balance amount unit
    , Either [Unit_Sum amount] (Map Account [posting])
    )
infer_equilibrium ps = do
	let bal = flip (Data.Foldable.foldr postings) ps balance
	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.Map.foldrWithKey
	 (\unit unit_sum@(Unit_Sum{ unit_sum_amount=amt, unit_sum_accounts }) acc ->
		case Data.Map.size unit_sum_accounts of
		 1 -> (Right $ (posting_make $ fst $ Data.Map.elemAt 0 unit_sum_accounts) (Data.Map.singleton unit (negate amt))):acc
		 _ -> Left [unit_sum]:acc)
	 []
	 dev

-- | Return 'True' if and only if the given 'Deviation' maps no 'unit'.
is_at_equilibrium :: Deviation amount unit -> 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 'unit_sum_accounts'
-- maps exactly one 'Account'.
is_equilibrium_inferrable :: Deviation amount unit -> Bool
is_equilibrium_inferrable (Deviation dev) =
	Data.Foldable.all
	 (\Unit_Sum{unit_sum_accounts} -> Data.Map.size unit_sum_accounts == 1)
	 dev

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

-- * The 'Expanded' type

-- | Descending propagation of 'amount's accross 'Account's.
type Expanded amount unit
 = TreeMap Account.Name
           (Account_Sum_Expanded amount unit)
data Account_Sum_Expanded amount unit
 =   Account_Sum_Expanded
 { inclusive :: Map unit amount
 , exclusive :: Map unit amount
 }
 deriving (Data, Eq, Show, Typeable)

-- | Return the given 'Balance_by_Account' with:
--
-- * all missing 'Account.ascending' 'Account's inserted,
--
-- * and every mapped 'amount'
-- added with any 'amount'
-- of the 'Account's’ for which it is 'Account.ascending'.
expanded ::
 ( Num amount, Ord unit )
 => Balance_by_Account amount unit
 -> Expanded amount unit
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 (flip (+))
			 . ( inclusive
				 . fromMaybe (error "Oops, should not happen in: Hcompta.Calc.Balance.expanded")
				 . Lib.TreeMap.node_value) )
			 exc $ Lib.TreeMap.nodes $ descendants
		 })