{-# OPTIONS_GHC -fno-warn-orphans #-}
module Hcompta.Balance where

import           Control.DeepSeq (NFData(..))
import           Control.Exception (assert)
import           Data.Bool
import           Data.Data
import           Data.Either (Either(..))
import           Data.Eq (Eq(..))
import qualified Data.Foldable as Foldable
import           Data.Function (($), (.), const, flip)
import           Data.Map.Strict (Map)
import qualified Data.Map.Strict as Map
import           Data.Monoid (Monoid(..))
import qualified Data.MonoTraversable as MT
import           Data.Ord (Ord(..))
import           Data.Semigroup (Semigroup(..))
import qualified Data.Sequences as Seqs
import qualified Data.Strict.Maybe as Strict
import           Data.TreeMap.Strict (TreeMap(..))
import qualified Data.TreeMap.Strict as TreeMap
import           Data.Tuple (curry, fst, snd)
import           Data.Typeable ()
import           Prelude (seq, undefined)
import           Text.Show (Show(..))

import           Hcompta.Has
import qualified Hcompta.Lib.Foldable as Foldable
import qualified Hcompta.Lib.Strict as Strict
import           Hcompta.Quantity

-- * Type 'Balance_Account'
-- | 'Balance' operations works on this type of 'Account'.
type Balance_Account = TreeMap.Path
instance Get (Balance_Account acct_sect)
             (Balance_Account acct_sect
             ,Balance_Amounts unit qty) where get = fst

-- * Type 'Balance_Amounts'
-- | 'Balance' operations works on this type of 'Amounts'.
type Balance_Amounts = Map
instance Get (Balance_Amounts unit qty)
             (Balance_Account acct_sect
             ,Balance_Amounts unit qty) where get = snd
instance Set (Balance_Amounts unit qty)
             (Balance_Account acct_sect
             ,Balance_Amounts unit qty) where set x (a, _) = (a, x)

-- * Type 'Balance'

-- | 'Balance_Account' and 'BalByUnit' of some @post@s.
--
-- NOTE: to reduce memory consumption
--       when applying 'balance_cons' incrementally,
--       the fields are explicitely stricts.
data Balance acct_sect unit qty
 =   Balance
 {   balByAccount :: !(BalByAccount acct_sect unit qty)
 ,   balByUnit    :: !(BalByUnit    acct_sect unit qty)
 } deriving (Data, Eq, Show, Typeable)
instance -- NFData
 ( NFData acct_sect
 , NFData unit
 , NFData qty
 , Ord acct_sect
 ) => NFData (Balance acct_sect unit qty) where
	rnf (Balance a u) = rnf a `seq` rnf u
instance -- Semigroup
 ( Addable qty
 , Ord     unit
 , Ord     acct_sect
 ) => Semigroup (Balance acct_sect unit qty) where
	(<>) = balance_union
instance -- Monoid
 ( Addable qty
 , Ord     unit
 , Ord     acct_sect
 ) => Monoid (Balance acct_sect unit qty) where
	mempty  = balance_empty
	mappend = (<>)

balance_empty :: Balance acct_sect unit qty
balance_empty =
	Balance
	 { balByAccount = TreeMap.empty
	 , balByUnit    = BalByUnit Map.empty
	 }

-- | Return the first given 'Balance'
--   updated by the second given 'Balance'.
balance_union ::
 (Addable qty, Ord acct_sect, Ord unit)
 => Balance acct_sect unit qty
 -> Balance acct_sect unit qty
 -> Balance acct_sect unit qty
balance_union
 (Balance b0a b0u)
 (Balance b1a b1u) =
	Balance
	 { balByAccount = balByAccount_union b0a b1a
	 , balByUnit    = balByUnit_union    b0u b1u
	 }

-- | Return the given 'Balance'
--   updated by the given @post@.
balance_cons ::
 ( Get (Balance_Account acct_sect) post
 , Get (Balance_Amounts unit qty)  post
 , Addable qty, Ord acct_sect, Ord unit )
 => post
 -> Balance acct_sect unit qty
 -> Balance acct_sect unit qty
balance_cons post bal =
	bal
	 { balByAccount = balByAccount_cons post (balByAccount bal)
	 , balByUnit    = balByUnit_cons    post (balByUnit    bal)
	 }

-- | Return the given 'Balance'
--   updated by the given @post@s.
balance_postings ::
 ( post ~ MT.Element posts
 , MT.MonoFoldable posts
 , Get (Balance_Account acct_sect) post
 , Get (Balance_Amounts unit qty)  post
 , Addable qty, Ord acct_sect, Ord unit )
 => posts
 -> Balance acct_sect unit qty
 -> Balance acct_sect unit qty
balance_postings = flip (MT.ofoldr balance_cons)

-- ** Type 'BalByAccount'
type BalByAccount acct_sect unit qty
 = TreeMap acct_sect (SumByAccount unit qty)

-- | Return the first given 'BalByAccount'
--   updated by the second given 'BalByAccount'.
balByAccount_union ::
 ( Addable qty
 , Ord acct_sect
 , Ord unit )
 => BalByAccount acct_sect unit qty
 -> BalByAccount acct_sect unit qty
 -> BalByAccount acct_sect unit qty
balByAccount_union = TreeMap.union (<>)

-- | Return the given 'BalByAccount'
--   updated by the given @post@.
balByAccount_cons ::
 ( Get (Balance_Account acct_sect) post
 , Get (Balance_Amounts unit qty)  post
 , Ord acct_sect
 , Ord unit
 , Addable qty
 ) => post
 -> BalByAccount acct_sect unit qty
 -> BalByAccount acct_sect unit qty
balByAccount_cons post =
	TreeMap.insert (<>) (get post) (SumByAccount $ get post)

-- *** Type 'SumByAccount'
-- | A sum of @qty@s, concerning a single 'Balance_Account'.
newtype SumByAccount unit qty
 =      SumByAccount (Balance_Amounts unit qty)
 deriving (Data, Eq, NFData, Show, Typeable)
instance -- Semigroup
 (Addable qty, Ord unit) =>
 Semigroup (SumByAccount unit qty) where
	SumByAccount x <> SumByAccount y =
		SumByAccount $ Map.unionWith (flip quantity_add) x y
instance -- Monoid
 (Addable qty, Ord unit) =>
 Monoid (SumByAccount unit qty) where
	mempty  = SumByAccount mempty
	mappend = (<>)

unSumByAccount
 :: SumByAccount unit qty
 -> Map unit qty
unSumByAccount (SumByAccount m) = m

-- ** Type 'BalByUnit'
newtype BalByUnit acct_sect unit qty
 =      BalByUnit (Map unit (SumByUnit (Balance_Account acct_sect) qty))
 deriving (Data, Eq, NFData, Show, Typeable)
instance -- Semigroup
 (Addable qty, Ord acct_sect, Ord unit) =>
 Semigroup (BalByUnit acct_sect unit qty) where
	(<>) = balByUnit_union
instance -- Monoid
 (Addable qty, Ord acct_sect, Ord unit) =>
 Monoid (BalByUnit acct_sect unit qty) where
	mempty  = BalByUnit mempty
	mappend = (<>)

-- | Return the first given 'BalByUnit'
--   updated by the second given 'BalByUnit'.
balByUnit_union
 :: (Addable qty, Ord acct_sect, Ord unit)
 => BalByUnit acct_sect unit qty
 -> BalByUnit acct_sect unit qty
 -> BalByUnit acct_sect unit qty
balByUnit_union
 (BalByUnit a0)
 (BalByUnit a1) =
	BalByUnit $
	Map.unionWith
	 (\new old -> SumByUnit
		 { sumByUnit_quantity = quantity_add
			 (sumByUnit_quantity old)
			 (sumByUnit_quantity new)
		 , sumByUnit_accounts = Map.unionWith
			 (const::()->()->())
			 (sumByUnit_accounts old)
			 (sumByUnit_accounts new)
		 })
	 a0 a1

-- | Return the given 'BalByUnit'
--   updated by the given @post@.
balByUnit_cons ::
 ( Get (Balance_Account acct_sect) post
 , Get (Balance_Amounts unit qty)  post
 , Addable qty
 , Ord acct_sect
 , Ord unit
 ) => post
 -> BalByUnit acct_sect unit qty
 -> BalByUnit acct_sect unit qty
balByUnit_cons post =
	balByUnit_union $
	BalByUnit $
	(`Map.map` get post) $
	 \qty -> SumByUnit
		 { sumByUnit_quantity = qty
		 , sumByUnit_accounts = Map.singleton (get post) ()
		 }

-- | Return the given 'BalByUnit'
--   updated by the given 'BalByAccount'.
balByUnit_of_BalByAccount
 :: (Addable qty, Ord acct_sect, Ord unit)
 => BalByAccount acct_sect unit qty
 -> BalByUnit    acct_sect unit qty
 -> BalByUnit    acct_sect unit qty
balByUnit_of_BalByAccount =
	flip $ TreeMap.foldr_with_Path $ curry balByUnit_cons

instance Get ( Balance_Account acct_sect )
             ( Balance_Account acct_sect
             , SumByAccount unit qty )
 where get = fst
instance Get ( Balance_Amounts unit qty )
             ( Balance_Account acct_sect
             , SumByAccount unit qty )
 where get = unSumByAccount . snd

-- *** Type 'SumByUnit'
-- | A sum of @qty@s with their 'Account's involved,
-- concerning a single @unit@.
data SumByUnit acct qty
 =   SumByUnit
 {   sumByUnit_quantity :: !qty
     -- ^ The sum of @qty@s for a single @unit@.
 ,   sumByUnit_accounts :: !(Map acct ())
     -- ^ The 'Balance_Account's involved to build 'sumByUnit_quantity'.
 } deriving (Data, Eq, Ord, Show, Typeable)
instance -- NFData
 ( NFData acct
 , NFData qty
 ) => NFData (SumByUnit acct qty) where
	rnf (SumByUnit q a) = rnf q `seq` rnf a

-- * Type 'DeviationByUnit'
-- | The 'BalByUnit' whose 'sumByUnit_quantity'
--   is not zero and possible 'Balance_Account' to 'equilibrium'.
newtype DeviationByUnit            acct_sect unit qty
 =      DeviationByUnit (BalByUnit acct_sect unit qty)
 deriving (Data, Eq, NFData, Show, Typeable)

-- | Return the 'balByUnit' of the given 'Balance' with:
--
-- * @unit@s whose 'sumByUnit_quantity' verifying 'quantity_null' are removed,
--
-- * and remaining @unit@s have their 'sumByUnit_accounts'
--   complemented with the 'balByAccount' of the given 'Balance'
--   (i.e. now mapping to the 'Balance_Account's __not__ involved to build the 'SumByUnit').
deviationByUnit
 :: (Zero qty, Addable qty, Ord acct_sect, Ord unit)
 => Balance         acct_sect unit qty
 -> DeviationByUnit acct_sect unit qty
deviationByUnit Balance
 { balByAccount
 , balByUnit=BalByUnit balByUnit
 } =
	let all_accounts = TreeMap.flatten (const ()) balByAccount in
	let max_accounts = Map.size all_accounts in
	DeviationByUnit $
		Map.foldlWithKey
		 (\(BalByUnit m) unit SumByUnit{..} ->
			BalByUnit $
			if quantity_null sumByUnit_quantity
			then m
			else
				case Map.size sumByUnit_accounts of
				 n | n == max_accounts ->
					Map.insert unit SumByUnit
					 { sumByUnit_quantity
					 , sumByUnit_accounts = Map.empty
					 } m
				 _ ->
					let diff = Map.difference all_accounts sumByUnit_accounts in
					Map.insert unit SumByUnit
					 { sumByUnit_quantity
					 , sumByUnit_accounts = diff
					 } m
		 )
		 mempty
		 balByUnit

-- ** Balance equilibrium
-- | Return the 'Balance' (adjusted by inferred @qty@s)
--   of the given @post@s and either:
--
--   * 'Left': the @unit@s which have a non null 'SumByUnit'
--             and for which no equibrating 'Balance_Account' can be inferred.
--   * 'Right': the given @post@s with inferred @qty@s inserted.
equilibrium ::
 ( post ~ MT.Element posts
 , Seqs.IsSequence posts
 , Get (Balance_Account acct_sect) post
 , Has (Balance_Amounts unit qty)  post
 , Zero qty, Addable qty, Negable qty
 , Ord acct, Ord acct_sect, Ord unit
 , Get acct (Balance_Account acct_sect)
 ) => Map acct posts
   -> ( Balance acct_sect unit qty
      , Either [(unit, SumByUnit (Balance_Account acct_sect) qty)]
               (Map acct posts) )
equilibrium posts =
	let bal_initial = MT.ofoldr balance_postings balance_empty posts in
	let DeviationByUnit (BalByUnit dev) = deviationByUnit bal_initial in
	let (bal_adjusted, eithers) =
		Map.foldrWithKey
		 (\unit unit_sum@SumByUnit{..} (bal, lr) ->
			case Map.size sumByUnit_accounts of
			 1 ->
				let acct = fst $ Map.elemAt 0 sumByUnit_accounts in
				let qty  = quantity_neg sumByUnit_quantity in
				let amts = Map.singleton unit qty in
				( balance_cons (acct, SumByAccount amts) bal
				, Right (acct, unit, qty) : lr
				)
			 _ -> (bal, Left [(unit, unit_sum)] : lr))
		 (bal_initial, [])
		 dev in
	let (l, r) = Foldable.accumLeftsAndFoldrRights
		 (\(acct, unit, qty) ->
			Map.insertWith
			 (\_new_ps -> Seqs.fromList . insert_amount (unit, qty) . MT.otoList)
			 (get acct) (assert False undefined))
			 -- NOTE: acct is within bal_initial,
			 -- hence posts already has a mapping for acct.
		 posts eithers in
	case l of
	 [] -> (bal_adjusted, Right r)
	 _  -> (bal_adjusted, Left  l)
	where
		insert_amount
		 :: forall post unit qty.
		 ( Ord unit
		 , Has (Balance_Amounts unit qty) post
		 ) => (unit, qty) -> [post] -> [post]
		insert_amount amt@(unit, qty) l =
			case l of
			 [] -> assert False []
				-- NOTE: the acct being in bal_initial,
				-- hence there was at least one post for this acct.
			 p:ps ->
				let amts :: Balance_Amounts unit qty = get p in
				if Map.notMember unit amts
				then set (Map.insert unit qty amts) p:ps
				else p:insert_amount amt ps

-- | Return 'True' if and only if the given 'DeviationByUnit' maps no @unit@.
is_equilibrium :: DeviationByUnit acct_sect unit qty -> Bool
is_equilibrium (DeviationByUnit (BalByUnit dev)) = Map.null dev

-- | Return 'True' if and only if the given 'DeviationByUnit'
-- maps only to 'SumByUnit's whose 'sumByUnit_accounts'
-- maps exactly one 'Balance_Account'.
is_equilibrium_inferrable :: DeviationByUnit acct_sect unit qty -> Bool
is_equilibrium_inferrable (DeviationByUnit (BalByUnit dev)) =
	Foldable.all ((== 1) . Map.size . sumByUnit_accounts) dev

-- * Type 'ClusiveBalByAccount'

-- | {Ex,In}clusive 'BalByAccount':
-- descending propagation of @qty@s accross 'Account's.
type ClusiveBalByAccount acct_sect unit qty
 = TreeMap acct_sect (ClusiveSumByAccount unit qty)

-- ** Type 'ClusiveSumByAccount'
-- |
-- * 'Strict.exclusive': contains the original 'SumByAccount'.
-- * 'Strict.inclusive': contains 'quantity_add' folded
--   over 'Strict.exclusive' and 'Strict.inclusive'
--   of 'TreeMap.node_descendants'
type ClusiveSumByAccount unit qty
 = Strict.Clusive (SumByAccount unit qty)

-- | Return the given 'BalByAccount' with:
--
-- * all missing 'Account.parent' 'Account's inserted;
-- * and every mapped @qty@ added with any @qty@
--   of the 'Account's for which it is 'Account.parent'.
clusiveBalByAccount
 :: (Addable qty, Ord acct_sect, Ord unit)
 => BalByAccount    acct_sect unit qty
 -> ClusiveBalByAccount acct_sect unit qty
clusiveBalByAccount =
	TreeMap.map_by_depth_first
	 (\descendants value ->
		let exclusive = Strict.fromMaybe mempty value in
		Strict.Clusive
		 { Strict.exclusive
		 , Strict.inclusive =
			Map.foldl'
			 ( flip $ (<>) . Strict.inclusive
			 . Strict.fromMaybe (assert False undefined)
			 . TreeMap.node_value )
			 exclusive $
			TreeMap.nodes descendants
		 })

-- | Return a 'BalByUnit'
--   derived from the given 'ClusiveBalByAccount' balance.
--
--   NOTE: also correct if the 'ClusiveBalByAccount' has been filtered.
balByUnit_of_ClusiveBalByAccount
 :: (Addable qty, Ord acct_sect, Ord unit)
 => ClusiveBalByAccount acct_sect unit qty
 -> BalByUnit acct_sect unit qty
 -> BalByUnit acct_sect unit qty
balByUnit_of_ClusiveBalByAccount =
	go []
	where
	go p (TreeMap nodes) bal =
		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 acct = TreeMap.reverse $ TreeMap.path k p in
				balByUnit_cons (acct, Strict.inclusive a) acc)
		 bal nodes