{-# 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.DeepSeq (NFData(..))
import           Control.Exception (assert)
import           Data.Bool
import           Data.Data
import           Data.Either (Either(..))
import           Data.Eq (Eq(..))
import           Data.Ord (Ord(..))
import qualified Data.Foldable
import           Data.Foldable (Foldable(..))
import qualified Data.Map.Strict as Data.Map
import           Data.Map.Strict (Map)
import           Data.Monoid (Monoid(..))
import qualified Data.Strict.Maybe as Strict
import           Data.Tuple (fst, snd)
import           Data.Typeable ()
import           Text.Show (Show(..))
import           Prelude (($), (.), const, curry, flip, undefined)

import           Hcompta.Quantity (Zero(..), Addable(..), Negable(..))
import           Hcompta.Account (Account(..), Account_Path)
import qualified Hcompta.Lib.Foldable as Lib.Foldable
import           Hcompta.Lib.TreeMap (TreeMap)
import qualified Hcompta.Lib.TreeMap as TreeMap

-- * Requirements' interface

-- ** Class 'Posting'

-- | A 'posting' used to produce a 'Balance'
--   must be an instance of this class.
class
 ( Account (Posting_Account p)
 ) =>  Posting           p where
	type Posting_Account   p
	type Posting_Quantity  p
	type Posting_Unit      p
	posting_account     :: p -> Posting_Account p
	posting_amounts     :: p -> Map (Posting_Unit p) (Posting_Quantity p)
	posting_set_amounts ::      Map (Posting_Unit p) (Posting_Quantity p) -> p -> p
instance
 ( Account account
 ) => Posting (account, Map unit quantity)
 where
	type Posting_Account  (account, Map unit quantity) = account
	type Posting_Quantity (account, Map unit quantity) = quantity
	type Posting_Unit     (account, Map unit quantity) = unit
	posting_account = fst
	posting_amounts = snd
	posting_set_amounts amounts (acct, _) = (acct, amounts)
instance
 ( Account account
 ) =>  Posting          (account, Account_Sum unit quantity) where
	type Posting_Account  (account, Account_Sum unit quantity) = account
	type Posting_Quantity (account, Account_Sum unit quantity) = quantity
	type Posting_Unit     (account, Account_Sum unit quantity) = unit
	posting_account = fst
	posting_amounts (_, Account_Sum x) = x
	posting_set_amounts amounts (acct, _) = (acct, Account_Sum amounts)

-- * Type 'Balance'

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

instance
 ( Account account
 , Addable quantity
 , Ord     unit
 ) => Monoid (Balance account unit quantity) where
	mempty  = empty
	mappend = union

-- ** Type 'Balance_by_Account'
type Balance_by_Account account_section unit quantity
 = TreeMap account_section
           (Account_Sum unit quantity)

-- *** Type 'Account_Sum'
-- | A sum of 'quantity's, concerning a single 'account'.
newtype Account_Sum unit quantity
 =      Account_Sum (Map unit quantity)
 deriving (Data, Eq, Show, Typeable)
get_Account_Sum
 :: Account_Sum unit quantity
 -> Map unit quantity
get_Account_Sum (Account_Sum m) = m

instance
 ( NFData unit
 , NFData quantity
 ) => NFData (Account_Sum unit quantity) where
	rnf (Account_Sum m) = rnf m
instance
 ( Addable quantity
 , Ord unit
 ) => Monoid (Account_Sum unit quantity) where
	mempty = Account_Sum mempty
	mappend
	 (Account_Sum a0)
	 (Account_Sum a1) =
		Account_Sum $ Data.Map.unionWith quantity_add a0 a1

-- ** Type 'Balance_by_Unit'
newtype
 ( Account account
 ) => Balance_by_Unit account unit quantity
 =    Balance_by_Unit (Map unit (Unit_Sum account quantity))
deriving instance ( Account account
                  , Data    account
                  , Data    unit
                  , Data    quantity
                  , Ord     unit
                  , Typeable unit
                  , Typeable quantity
                  , Data    (Account_Section account)
                  ) => Data (Balance_by_Unit account unit quantity)
deriving instance ( Account account
                  , Eq      account
                  , Eq      unit
                  , Eq      quantity
                  ) => Eq (Balance_by_Unit account unit quantity)
deriving instance ( Account account
                  , Show    account
                  , Show    unit
                  , Show    quantity
                  , Show    (Account_Section account)
                  ) => Show (Balance_by_Unit account unit quantity)
deriving instance Typeable3 Balance_by_Unit
 -- FIXME: use 'Typeable' when dropping GHC-7.6 support

instance
 ( Account account
 , Addable quantity
 , Ord     unit
 ) => Monoid (Balance_by_Unit account unit quantity) where
	mempty = Balance_by_Unit mempty
	mappend = union_by_unit

-- *** Type 'Unit_Sum'

-- | A sum of 'quantity's with their 'Account's involved,
-- concerning a single 'unit'.
data (Account account)
 => Unit_Sum account quantity
 =  Unit_Sum
 { unit_sum_quantity :: !quantity -- ^ The sum of 'quantity's for a single 'unit'.
 , unit_sum_accounts :: !(Map (Account_Path (Account_Section account))
                              ()) -- ^ The 'account's involved to build 'unit_sum_quantity'.
 }
deriving instance ( Account account
                  , Data    account
                  , Data    (Account_Section account)
                  , Data    quantity
                  ) => Data (Unit_Sum account quantity)
deriving instance ( Account account
                  , Eq      account
                  , Eq      quantity
                  ) => Eq (Unit_Sum account quantity)
deriving instance ( Account account
                  , Show    account
                  , Show    (Account_Section account)
                  , Show    quantity
                  ) => Show (Unit_Sum account quantity)
deriving instance Typeable2 Unit_Sum
 -- FIXME: use 'Typeable' when dropping GHC-7.6 support

-- ** Constructors

empty ::
 ( Account account
 , Ord     unit
 , Addable quantity
 ) => Balance account unit quantity
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_Account posting)
                     (Posting_Unit posting)
                     (Posting_Quantity posting)
 , Addable (Posting_Quantity posting)
 , Ord (Posting_Unit 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_Account posting)
                     (Posting_Unit posting)
                     (Posting_Quantity posting)
 , Foldable foldable
 , Addable (Posting_Quantity posting)
 , Ord (Posting_Unit posting)
 ) => foldable posting -> balance -> balance
postings = flip (Data.Foldable.foldr cons)

-- | Return the first given 'Balance'
--   updated by the second given 'Balance'.
union ::
 ( Account account
 , Addable quantity
 , Ord unit
 , balance ~ Balance account unit quantity
 ) => balance -> balance -> balance
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
 , account  ~ Posting_Account  posting
 , quantity ~ Posting_Quantity posting
 , unit     ~ Posting_Unit     posting
 , Addable (Posting_Quantity posting)
 , Ord unit
 )
 => posting
 -> Balance_by_Account (Account_Section account) unit quantity
 -> Balance_by_Account (Account_Section account) unit quantity
cons_by_account post =
	TreeMap.insert mappend
	 (account_path $ posting_account post)
	 (Account_Sum $ posting_amounts post)

-- | Return the given 'Balance_by_Unit'
--   updated by the given 'Posting'.
cons_by_unit ::
 ( Posting posting
 , account  ~ Posting_Account   posting
 , quantity ~ Posting_Quantity  posting
 , unit     ~ Posting_Unit      posting
 , Addable quantity
 , Ord unit
 ) => posting
 -> Balance_by_Unit account unit quantity
 -> Balance_by_Unit account unit quantity
cons_by_unit post =
	union_by_unit $
	Balance_by_Unit $
	Data.Map.map
	 (\quantity -> Unit_Sum
		 { unit_sum_quantity = quantity
		 , unit_sum_accounts = Data.Map.singleton (account_path $ posting_account post) ()
		 })
	 (posting_amounts post)

-- | Return a 'Balance_by_Unit'
--   derived from the given 'Balance_by_Account'.
by_unit_of_by_account ::
 ( Account account
 , account ~ Account_Path (Account_Section account)
 , Addable quantity
 , Ord unit
 )
 => Balance_by_Account (Account_Section account) unit quantity
 -> Balance_by_Unit    account unit quantity
 -> Balance_by_Unit    account unit quantity
by_unit_of_by_account =
	flip $ 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 ::
 ( Addable quantity
 , Ord account_section
 , Ord unit
 ) => Balance_by_Account account_section unit quantity
 -> Balance_by_Account account_section unit quantity
 -> Balance_by_Account account_section unit quantity
union_by_account = TreeMap.union mappend

-- | Return the first given 'Balance_by_Unit'
--   updated by the second given 'Balance_by_Unit'.
union_by_unit ::
 ( Account account
 , Addable quantity
 , Ord unit
 ) => Balance_by_Unit account unit quantity
 -> Balance_by_Unit account unit quantity
 -> Balance_by_Unit account unit quantity
union_by_unit
 (Balance_by_Unit a0)
 (Balance_by_Unit a1) =
	Balance_by_Unit $
	Data.Map.unionWith
	 (\new old -> Unit_Sum
		 { unit_sum_quantity = quantity_add
			 (unit_sum_quantity old)
			 (unit_sum_quantity 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_quantity'
--   is not zero and possible 'account' to 'infer_equilibrium'.
newtype
 ( Account account
 ) => Deviation account unit quantity
 =  Deviation (Balance_by_Unit account unit quantity)
deriving instance ( Account account
                  , Data    account
                  , Data    unit
                  , Data    quantity
                  , Ord     unit
                  , Typeable unit
                  , Typeable quantity
                  , Data    (Account_Section account)
                  ) => Data (Deviation account unit quantity)
deriving instance ( Account account
                  , Eq      account
                  , Eq      unit
                  , Eq      quantity
                  ) => Eq (Deviation account unit quantity)
deriving instance ( Account account
                  , Show    account
                  , Show    (Account_Section account)
                  , Show    unit
                  , Show    quantity
                  ) => Show (Deviation account unit quantity)
deriving instance Typeable3 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_quantity' is verifying 'quantity_null' 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 ::
 ( Account account
 , Zero    quantity
 , Addable quantity
 , Ord     unit
 ) => Balance account unit quantity
 -> Deviation account unit quantity
deviation Balance
 { balance_by_account=ba
 , balance_by_unit=Balance_by_Unit bu
 } = do
	let all_accounts = 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_quantity, unit_sum_accounts} ->
			Balance_by_Unit $
			if quantity_null unit_sum_quantity
			then m
			else
				case Data.Map.size unit_sum_accounts of
				 n | n == max_accounts ->
					Data.Map.insert unit Unit_Sum
					 { unit_sum_quantity
					 , 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_quantity
					 , unit_sum_accounts = diff
					 } m
		 )
		 mempty
		 bu

-- ** The equilibrium

-- | Return the 'Balance' (adjusted by inferred 'quantity's)
--   of the given 'Posting's and either:
--
--   * 'Left': the 'Posting's that cannot be inferred.
--   * 'Right': the given 'Posting's with inferred 'quantity's inserted.
infer_equilibrium ::
 ( Posting posting
 , account  ~ Posting_Account  posting
 , unit     ~ Posting_Unit     posting
 , quantity ~ Posting_Quantity posting
 , account ~ Account_Path (Account_Section account)
 , Addable quantity
 , Negable quantity
 , Zero    quantity
 , Ord     unit
 ) => Map account [posting]
 -> ( Balance account unit quantity
    , Either [(unit, Unit_Sum account quantity)]
             (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_quantity, 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 qty  = quantity_neg unit_sum_quantity in
				let amts = Data.Map.singleton unit qty in
				( cons (acct, Account_Sum amts) bal
				, Right (acct, unit, qty) : lr
				)
			 _ -> (bal, Left [(unit, unit_sum)] : lr))
		 (bal_initial, [])
		 dev
	let (l, r) = Lib.Foldable.accumLeftsAndFoldrRights
		 (\(acct, unit, qty) ->
			Data.Map.insertWith
			 (\_new_ps -> insert_amount (unit, qty))
			 acct (assert False []))
		 posts eithers
	case l of
	 [] -> (bal_adjusted, Right r)
	 _  -> (bal_adjusted, Left  l)
	where
		insert_amount ::
		 ( Posting posting
		 , Ord     (Posting_Unit posting)
		 ) => ( Posting_Unit     posting
		      , Posting_Quantity posting )
		 -> [posting] -> [posting]
		insert_amount p@(unit, qty) ps =
			case ps of
			 [] -> assert False []
			 (x:xs) | Data.Map.null (posting_amounts x) ->
				posting_set_amounts (Data.Map.singleton unit qty) x:xs
			        | Data.Map.notMember unit (posting_amounts x) ->
				let amts = Data.Map.insertWith
					 (assert False undefined)
					 unit qty (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 :: (Account account) => Deviation account unit quantity -> 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 :: (Account account) => Deviation account unit quantity -> 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 :: (Account account) => Deviation account unit quantity -> 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 'quantity's accross 'Account's.
type Expanded account_section unit quantity
 = TreeMap account_section (Account_Sum_Expanded unit quantity)
data Account_Sum_Expanded unit quantity
 =   Account_Sum_Expanded
 { exclusive :: !(Account_Sum unit quantity)
 , inclusive :: !(Account_Sum unit quantity) -- ^ 'quantity_add' folded over 'exclusive' and 'inclusive' of 'TreeMap.node_descendants'
 } deriving (Data, Eq, Show, Typeable)

instance
 ( Addable quantity
 , Ord unit
 ) => Monoid (Account_Sum_Expanded unit quantity) 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 'quantity' added with any 'quantity'
--   of the 'Account's for which it is 'Account.ascending'.
expanded ::
 ( Addable quantity
 , Ord account_section
 , Ord unit
 ) => Balance_by_Account account_section unit quantity
 -> Expanded account_section unit quantity
expanded =
	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)
			 . TreeMap.node_value )
			 exclusive $
			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 ::
 ( Account account
 , account ~ Account_Path (Account_Section account)
 , Addable quantity
 , Ord unit
 ) => Expanded (Account_Section account) unit quantity
 -> Balance_by_Unit account unit quantity
 -> Balance_by_Unit account unit quantity
by_unit_of_expanded =
	go []
	where
		go p (TreeMap.TreeMap nodes) bal =
			Data.Map.foldrWithKey
			 (\k TreeMap.Node{TreeMap.node_value, TreeMap.node_descendants} acc ->
				case node_value of
				 Strict.Nothing -> go (k:p) node_descendants acc
				 Strict.Just a  ->
					let account = TreeMap.reverse $ TreeMap.path k p in
					cons_by_unit (account, inclusive a) acc)
			 bal nodes