{-# LANGUAGE DeriveDataTypeable #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE FlexibleInstances #-}
{-# 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.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.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

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

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

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 = by_account post (balance_by_account bal)
	 , balance_by_unit    = 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 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 = union_by_account
		 (balance_by_account b0)
		 (balance_by_account b1)
	 , balance_by_unit = union_by_unit
		 (balance_by_unit b0)
		 (balance_by_unit b1)
	 }

-- | Return the given 'Balance_by_Account'
--   updated by the given 'Posting'.
by_account ::
 ( Posting posting
 , amount ~ Posting_Amount posting
 , unit   ~ Amount_Unit amount )
 => posting
 -> Balance_by_Account amount unit
 -> Balance_by_Account amount unit
by_account post =
	Lib.TreeMap.insert
	 (Data.Map.unionWith (flip amount_add))
	 (posting_account post)
	 (posting_amounts post)

-- | Return the given 'Balance_by_Unit'
--   updated by the given 'Posting'.
by_unit ::
 ( Posting posting
 , amount ~ Posting_Amount posting
 , unit   ~ Amount_Unit amount )
 => posting
 -> Balance_by_Unit amount unit
 -> Balance_by_Unit amount unit
by_unit post bal =
	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)
		 })
	 bal $
	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 unit
 -> Balance_by_Unit    amount unit
 -> Balance_by_Unit    amount unit
by_unit_of_by_account =
	flip $ Lib.TreeMap.foldr_with_Path $ curry by_unit

-- | Return the first given 'Balance_by_Account'
--   updated by the second given 'Balance_by_Account'.
union_by_account :: (Amount amount, unit ~ Amount_Unit amount)
 => Balance_by_Account amount unit
 -> Balance_by_Account amount unit
 -> Balance_by_Account amount unit
union_by_account =
	Lib.TreeMap.union
	 (Data.Map.unionWith (flip amount_add))

-- | 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 unit
 -> Balance_by_Unit amount unit
 -> Balance_by_Unit amount unit
union_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)
		 })

-- * 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 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 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_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
		 )
		 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_negate unit_sum_amount 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 :: Map (Amount_Unit amount) 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

-- | 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
			 (Data.Map.unionWith amount_add . inclusive . from_value)
			 exclusive nodes
		 })

-- | 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 unit
 -> Balance_by_Unit amount unit
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
				 Nothing -> go (k:p) node_descendants acc
				 Just a  ->
					let account = Lib.TreeMap.reverse $ Lib.TreeMap.path k p in
					by_unit (account, inclusive a) acc)
			 bal m