{-# LANGUAGE DeriveDataTypeable #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE NamedFieldPuns #-}
{-# LANGUAGE StandaloneDeriving #-}
{-# 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)

-- * Requirements' interface

-- ** Class 'Amount'
class
 ( Data (Amount_Unit a)
 , Data a
 , Eq a
 , Num a
 , Ord  (Amount_Unit a)
 , Show (Amount_Unit a)
 , Show a
 , Typeable (Amount_Unit a)
 , Typeable a
 ) => Amount a where
	type Amount_Unit a
	amount_sign :: a -> Ordering

-- ** Class 'Posting'

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

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

-- * Type 'Amount_Sum'

-- | 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
	 }

-- * Type 'Balance'

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

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

nil :: Amount amount => Balance amount
nil =
	Balance
	 { balance_by_account = Lib.TreeMap.empty
	 , balance_by_unit    = Data.Map.empty
	 }

-- | Return the given 'Balance'
--   updated by the given 'Posting'.
balance ::
 ( Posting posting
 , balance ~ Balance (Posting_Amount posting) )
 => posting -> balance -> balance
balance 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
 , balance ~ Balance (Posting_Amount posting)
 , Foldable foldable )
 => foldable posting -> balance -> balance
postings = flip (Data.Foldable.foldr balance)

-- | Return the first given 'Balance'
--   updated by the second given 'Balance'.
union :: Amount amount
 => Balance amount -> Balance amount -> Balance amount
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)
	 }

-- * Type 'Deviation'

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

-- | 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
 => Balance amount
 -> Deviation amount
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 )
 => Map Account [posting]
 -> ( Balance (Posting_Amount posting)
    , Either [Unit_Sum (Posting_Amount posting)] (Map Account [posting])
    )
infer_equilibrium posts = do
	let bal_initial = Data.Foldable.foldr postings nil 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
				( balance (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
		 => (Amount_Unit (Posting_Amount 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 :: Amount amount => Deviation amount -> 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 :: Amount amount => Deviation amount -> 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 :: Amount amount => Deviation amount -> Bool
is_equilibrium_non_inferrable (Deviation dev) =
	Data.Foldable.any
	 (\Unit_Sum{unit_sum_accounts} -> Data.Map.size unit_sum_accounts > 1)
	 dev

-- * Type 'Expanded'

-- | Descending propagation of 'Amount's accross 'Account's.
type Expanded amount
 = TreeMap Account.Name (Account_Sum_Expanded amount)
data Amount amount => Account_Sum_Expanded amount
 =   Account_Sum_Expanded
 { exclusive :: Map (Amount_Unit amount) amount
 , inclusive :: Amount_Sum amount (Amount_Unit amount) -- ^ 'amount_sum_add' folded over 'exclusive' and 'inclusive' of 'Lib.TreeMap.node_descendants'
 }
deriving instance Amount amount => Data    (Account_Sum_Expanded amount)
deriving instance Amount amount => Eq      (Account_Sum_Expanded amount)
deriving instance Amount amount => Show    (Account_Sum_Expanded amount)
deriving instance                  Typeable Account_Sum_Expanded

-- | 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
 => Balance_by_Account amount (Amount_Unit amount)
 -> Expanded amount
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
		 })