{-# LANGUAGE BangPatterns #-}
{-# 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

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

instance
 ( Posting posting
 , amount ~ Posting_Amount posting
 )
 => Consable (Const (Balance_by_Account amount)) posting where
	mcons p (Const !bal) = Const $ cons_by_account p bal
instance
 ( Foldable foldable
 , Posting posting
 , amount ~ Posting_Amount posting
 )
 => Consable (Const (Balance_by_Account amount))
 (foldable posting) where
	mcons ps (Const !bal) =
		Const $ Data.Foldable.foldr cons_by_account bal ps

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