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

import           Control.Exception (assert)
import           Data.Data
import qualified Data.Foldable
import           Data.Foldable (Foldable(..))
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 'Amount_Sum' type

class Num a => Amount a where
	amount_sign :: a -> Ordering

-- | Sum keeping track of negative and positive 'amount's.
data Amount_Sum amount unit
 =   Amount_Sum
 { amount_sum_negative :: Map unit amount
 , amount_sum_positive :: Map unit amount
 , amount_sum_balance  :: Map unit amount
 } deriving (Data, Eq, Show, Typeable)

amount_sum
 :: Amount amount
 => Map unit amount
 -> Amount_Sum amount unit
amount_sum a =
	Amount_Sum
	 { amount_sum_negative = Data.Map.filter ((==) LT . amount_sign) a
	 , amount_sum_positive = Data.Map.filter ((==) GT . amount_sign) a
	 , amount_sum_balance  = a
	 }

amount_sum_add
 :: (Amount amount, Ord unit)
 => Amount_Sum amount unit
 -> Amount_Sum amount unit
 -> Amount_Sum amount unit
amount_sum_add a b =
	let add = Data.Map.unionWith (flip (+)) in
	Amount_Sum
	 { amount_sum_negative = add (amount_sum_negative a) (amount_sum_negative b)
	 , amount_sum_positive = add (amount_sum_positive a) (amount_sum_positive b)
	 , amount_sum_balance  = add (amount_sum_balance  a) (amount_sum_balance  b)
	 }

amount_sum_negate
 :: Amount amount
 => Amount_Sum amount unit
 -> Amount_Sum amount unit
amount_sum_negate a =
	Amount_Sum
	 { amount_sum_negative = amount_sum_positive a
	 , amount_sum_positive = amount_sum_negative a
	 , amount_sum_balance  = Data.Map.map negate $ amount_sum_balance a
	 }

-- * The 'Posting' class

-- | A 'posting' used to produce a 'Balance'
--   must be an instance of this class.
class
 ( Amount (Posting_Amount p)
 , Ord    (Posting_Unit   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_set_amounts :: Map (Posting_Unit p) (Posting_Amount p) -> p -> p

instance (Amount amount, Ord unit)
 => Posting (Account, Map unit amount)
 where
	type Posting_Amount (Account, Map unit amount) = amount
	type Posting_Unit   (Account, Map unit amount) = unit
	posting_account = fst
	posting_amounts = snd
	posting_set_amounts amounts (acct, _) = (acct, amounts)

-- * 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_Sum 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 :: (Amount 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
 , Amount amount
 , 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 = amount_sum_add
				 (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_sum $ Data.Map.singleton () 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
 , Amount amount
 , 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
 :: (Amount 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 = amount_sum_add
				 (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 EQ == amount_sign (amount_sum_balance unit_sum_amount Data.Map.! ())
			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 equilibrium

-- | Return the 'Balance' (adjusted by inferred 'amount's)
--   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 posts = do
	let bal_initial = Data.Foldable.foldr postings balance posts
	let Deviation dev = deviation bal_initial
	let (bal_adjusted, eithers) =
		Data.Map.foldrWithKey
		 (\unit unit_sum@(Unit_Sum{unit_sum_amount, unit_sum_accounts})
		 (bal, lr) ->
			case Data.Map.size unit_sum_accounts of
			 1 ->
				let acct = fst $ Data.Map.elemAt 0 unit_sum_accounts in
				let amt = (amount_sum_balance $ amount_sum_negate unit_sum_amount) Data.Map.! () in
				let amts = Data.Map.singleton unit amt in
				( posting (acct, amts) bal
				, Right (acct, unit, amt) : lr
				)
			 _ -> (bal, Left [unit_sum] : lr))
		 (bal_initial, [])
		 dev
	let (l, r) = Lib.Foldable.accumLeftsAndFoldrRights
		 (\(acct, unit, amt) ->
			Data.Map.insertWith
			 (\_new_ps -> insert_amount (unit, amt))
			 acct (assert False []))
		 posts eithers
	case l of
	 [] -> (bal_adjusted, Right r)
	 _  -> (bal_adjusted, Left  l)
	where
		insert_amount
		 :: Posting posting
		 => (Posting_Unit posting, Posting_Amount posting)
		 -> [posting] -> [posting]
		insert_amount p@(unit, amt) ps =
			case ps of
			 [] -> assert False []
			 (x:xs) | Data.Map.null (posting_amounts x) ->
				posting_set_amounts (Data.Map.singleton unit amt) x:xs
			        | Data.Map.notMember unit (posting_amounts x) ->
				let amts =
					Data.Map.insertWith (assert False undefined)
					 unit amt (posting_amounts x) in
				posting_set_amounts amts x:xs
			 (x:xs) -> x:insert_amount p xs

-- | 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
 { exclusive :: Map unit amount
 , inclusive :: Amount_Sum amount unit -- ^ 'amount_sum_add' folded over 'exclusive' and 'inclusive' of 'Lib.TreeMap.node_descendants'
 }
 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 ::
 ( Amount amount, Ord unit )
 => Balance_by_Account amount unit
 -> Expanded amount unit
expanded =
	let from_value = fromMaybe (assert False undefined) . Lib.TreeMap.node_value in
	Lib.TreeMap.map_by_depth_first
	 (\descendants value ->
		let nodes = Lib.TreeMap.nodes descendants in
		let exclusive = fromMaybe Data.Map.empty value in
		Account_Sum_Expanded
		 { exclusive
		 , inclusive =
			Data.Map.foldr
			 (amount_sum_add . inclusive . from_value)
			 (amount_sum exclusive) nodes
		 })