{-# LANGUAGE DeriveDataTypeable #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE NamedFieldPuns #-}
{-# LANGUAGE StandaloneDeriving #-}
{-# LANGUAGE TypeFamilies #-}
{-# OPTIONS_GHC -fno-warn-deprecations #-} -- FIXME: to be removed when dropping GHC-7.6 support
module Hcompta.Balance where

-- import           Control.Applicative (Const(..))
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 qualified Data.Strict.Maybe as Strict
import           Data.Typeable ()

-- import           Hcompta.Lib.Consable (Consable(..))
import qualified Hcompta.Lib.Foldable as Lib.Foldable
import qualified Hcompta.Lib.TreeMap as Lib.TreeMap
import           Hcompta.Lib.TreeMap (TreeMap)
import qualified Hcompta.Account as Account
import           Hcompta.Account (Account)

-- * Requirements' interface

-- ** Class 'Amount'
class
 ( Data     (Amount_Unit a)
 , Ord      (Amount_Unit a)
 , Show     (Amount_Unit a)
 , Typeable (Amount_Unit a)
 ) => Amount a where
	type Amount_Unit a
	amount_null   :: a -> Bool
	amount_add    :: a -> a -> a
	amount_negate :: a -> a

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

{- NOTE: not needed so far.
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)
-}

instance (Amount amount)
 => Posting (Account, Account_Sum amount)
 where
	type Posting_Amount (Account, Account_Sum amount) = amount
	posting_account = fst
	posting_amounts (_, Account_Sum x) = x
	posting_set_amounts amounts (acct, _) = (acct, Account_Sum amounts)

-- * Type 'Balance'

-- | Sum by 'Account' and sum by 'unit' of some 'Posting's.
--
-- NOTE: to reduce memory consumption when 'cons'ing iteratively,
--       the fields are explicitely stricts.
data Amount  amount
 =>  Balance amount
 =   Balance
 { balance_by_account :: !(Balance_by_Account amount)
 , balance_by_unit    :: !(Balance_by_Unit    amount)
 }
deriving instance ( Amount amount
                  , Data   amount
                  ) => Data (Balance amount)
deriving instance ( Amount amount
                  , Eq     amount
                  ) => Eq (Balance amount)
deriving instance ( Amount amount
                  , Show   amount
                  ) => Show (Balance amount)
deriving instance Typeable1 Balance
 -- FIXME: use 'Typeable' when dropping GHC-7.6 support

instance Amount amount => Monoid (Balance amount) where
	mempty  = empty
	mappend = union

-- ** Type 'Balance_by_Account'
type Balance_by_Account amount
 = TreeMap Account.Name
           (Account_Sum amount)

-- *** Type 'Account_Sum'
-- | A sum of 'amount's,
-- concerning a single 'Account'.
newtype Amount amount
 => Account_Sum amount
 =  Account_Sum (Map (Amount_Unit amount) amount)
get_Account_Sum :: Amount amount => Account_Sum amount -> Map (Amount_Unit amount) amount
get_Account_Sum (Account_Sum m) = m
deriving instance ( Amount amount
                  , Data   amount
                  ) => Data (Account_Sum amount)
deriving instance ( Amount amount
                  , Eq     amount
                  ) => Eq (Account_Sum amount)
deriving instance ( Amount amount
                  , Show   amount
                  ) => Show (Account_Sum amount)
deriving instance Typeable1 Account_Sum
 -- FIXME: use 'Typeable' when dropping GHC-7.6 support

instance Amount amount
 => Monoid (Account_Sum amount) where
	mempty = Account_Sum mempty
	mappend
	 (Account_Sum a0)
	 (Account_Sum a1) =
		Account_Sum $ Data.Map.unionWith amount_add a0 a1

-- ** Type 'Balance_by_Unit'
newtype Amount      amount
 => Balance_by_Unit amount
 =  Balance_by_Unit (Map (Amount_Unit amount) (Unit_Sum amount))
deriving instance ( Amount amount
                  , Data   amount
                  ) => Data (Balance_by_Unit amount)
deriving instance ( Amount amount
                  , Eq     amount
                  ) => Eq (Balance_by_Unit amount)
deriving instance ( Amount amount
                  , Show   amount
                  ) => Show (Balance_by_Unit amount)
deriving instance Typeable1 Balance_by_Unit
 -- FIXME: use 'Typeable' when dropping GHC-7.6 support

instance Amount amount
 => Monoid (Balance_by_Unit amount) where
	mempty = Balance_by_Unit mempty
	mappend = union_by_unit

-- *** Type 'Unit_Sum'

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

empty :: Amount amount => Balance amount
empty =
	Balance
	 { balance_by_account = mempty
	 , balance_by_unit    = mempty
	 }

-- | Return the given 'Balance'
--   updated by the given 'Posting'.
cons ::
 ( Posting posting
 , balance ~ Balance (Posting_Amount posting) )
 => posting -> balance -> balance
cons post bal =
	bal
	 { balance_by_account = cons_by_account post (balance_by_account bal)
	 , balance_by_unit    = cons_by_unit    post (balance_by_unit    bal)
	 }

-- | 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 cons)

-- | Return the first given 'Balance'
--   updated by the second given 'Balance'.
union :: Amount amount
 => Balance amount -> Balance amount -> Balance amount
union
 (Balance b0a b0u)
 (Balance b1a b1u) =
	Balance
	 { balance_by_account = union_by_account b0a b1a
	 , balance_by_unit    = union_by_unit    b0u b1u
	 }

-- | Return the given 'Balance_by_Account'
--   updated by the given 'Posting'.
cons_by_account ::
 ( Posting posting
 , amount ~ Posting_Amount posting
 , unit   ~ Amount_Unit amount
 )
 => posting
 -> Balance_by_Account amount
 -> Balance_by_Account amount
cons_by_account post =
	Lib.TreeMap.insert mappend
	 (posting_account post)
	 (Account_Sum $ posting_amounts post)

-- | Return the given 'Balance_by_Unit'
--   updated by the given 'Posting'.
cons_by_unit ::
 ( Posting posting
 , amount ~ Posting_Amount posting
 , unit   ~ Amount_Unit amount )
 => posting
 -> Balance_by_Unit amount
 -> Balance_by_Unit amount
cons_by_unit post =
	union_by_unit $
	Balance_by_Unit $
	Data.Map.map
	 (\amount -> Unit_Sum
		 { unit_sum_amount   = amount
		 , unit_sum_accounts = Data.Map.singleton (posting_account post) ()
		 })
	 (posting_amounts post)

-- | Return a 'Balance_by_Unit'
--   derived from the given 'Balance_by_Account'.
by_unit_of_by_account ::
 ( Amount amount
 , unit ~ Amount_Unit amount
 )
 => Balance_by_Account amount
 -> Balance_by_Unit    amount
 -> Balance_by_Unit    amount
by_unit_of_by_account =
	flip $ Lib.TreeMap.foldr_with_Path $ curry cons_by_unit

-- | Return the first given 'Balance_by_Account'
--   updated by the second given 'Balance_by_Account'.
union_by_account :: Amount amount
 => Balance_by_Account amount
 -> Balance_by_Account amount
 -> Balance_by_Account amount
union_by_account = Lib.TreeMap.union mappend

-- | Return the first given 'Balance_by_Unit'
--   updated by the second given 'Balance_by_Unit'.
union_by_unit :: (Amount amount, unit ~ Amount_Unit amount)
 => Balance_by_Unit amount
 -> Balance_by_Unit amount
 -> Balance_by_Unit amount
union_by_unit
 (Balance_by_Unit a0)
 (Balance_by_Unit a1) =
	Balance_by_Unit $
	Data.Map.unionWith
	 (\new old -> Unit_Sum
		 { unit_sum_amount = amount_add
			 (unit_sum_amount old)
			 (unit_sum_amount new)
		 , unit_sum_accounts = Data.Map.unionWith
			 (const::()->()->())
			 (unit_sum_accounts old)
			 (unit_sum_accounts new)
		 })
	 a0 a1

-- * 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)
deriving instance ( Amount amount
                  , Data   amount
                  ) => Data (Deviation amount)
deriving instance ( Amount amount
                  , Eq     amount
                  ) => Eq (Deviation amount)
deriving instance ( Amount amount
                  , Show   amount
                  ) => Show (Deviation amount)
deriving instance Typeable1 Deviation
 -- FIXME: use 'Typeable' when dropping GHC-7.6 support

-- | 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 Balance
 { balance_by_account=ba
 , balance_by_unit=Balance_by_Unit bu
 } = do
	let all_accounts = Lib.TreeMap.flatten (const ()) ba
	let max_accounts = Data.Map.size all_accounts
	Deviation $
		Data.Map.foldlWithKey
		 (\(Balance_by_Unit m) unit Unit_Sum{unit_sum_amount, unit_sum_accounts} ->
			Balance_by_Unit $
			if amount_null 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
		 )
		 mempty
		 bu

-- ** 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 empty posts
	let Deviation (Balance_by_Unit 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_negate unit_sum_amount in
				let amts = Data.Map.singleton unit amt in
				( cons (acct, Account_Sum 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 (Balance_by_Unit 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 (Balance_by_Unit 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 (Balance_by_Unit 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 :: !(Account_Sum amount)
 , inclusive :: !(Account_Sum amount) -- ^ 'amount_add' folded over 'exclusive' and 'inclusive' of 'Lib.TreeMap.node_descendants'
 }
deriving instance ( Amount amount
                  , Data   amount
                  ) => Data (Account_Sum_Expanded amount)
deriving instance ( Amount amount
                  , Eq     amount
                  ) => Eq (Account_Sum_Expanded amount)
deriving instance ( Amount amount
                  , Show   amount
                  ) => Show (Account_Sum_Expanded amount)
deriving instance Typeable1 Account_Sum_Expanded
 -- FIXME: use 'Typeable' when dropping GHC-7.6 support

instance Amount amount => Monoid (Account_Sum_Expanded amount) where
	mempty = Account_Sum_Expanded mempty mempty
	mappend
	 (Account_Sum_Expanded e0 i0)
	 (Account_Sum_Expanded e1 i1) =
		Account_Sum_Expanded
		 (mappend e0 e1)
		 (mappend i0 i1)

-- | 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
 -> Expanded amount
expanded =
	Lib.TreeMap.map_by_depth_first
	 (\descendants value ->
		let exclusive = Strict.fromMaybe mempty value in
		Account_Sum_Expanded
		 { exclusive
		 , inclusive =
			Data.Map.foldl'
			 ( flip $ mappend . inclusive
			 . Strict.fromMaybe (assert False undefined)
			 . Lib.TreeMap.node_value)
			 exclusive $
			Lib.TreeMap.nodes descendants
		 })

-- | Return a 'Balance_by_Unit'
--   derived from the given 'Expanded' balance.
--
--   NOTE: also correct if the 'Expanded' has been filtered.
by_unit_of_expanded ::
 ( Amount amount
 , unit ~ Amount_Unit amount
 )
 => Expanded amount
 -> Balance_by_Unit amount
 -> Balance_by_Unit amount
by_unit_of_expanded =
	go []
	where
		go p (Lib.TreeMap.TreeMap m) bal =
			Data.Map.foldrWithKey
			 (\k Lib.TreeMap.Node{Lib.TreeMap.node_value, Lib.TreeMap.node_descendants} acc ->
				case node_value of
				 Strict.Nothing -> go (k:p) node_descendants acc
				 Strict.Just a  ->
					let account = Lib.TreeMap.reverse $ Lib.TreeMap.path k p in
					cons_by_unit (account, inclusive a) acc)
			 bal m