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 Data.Function (($), (.), const, flip)
import Data.Map.Strict (Map)
import Data.Ord (Ord(..))
import Data.TreeMap.Strict (TreeMap(..))
import Data.Tuple (fst)
import Data.Typeable ()
import Prelude (seq, undefined)
import Text.Show (Show(..))
import qualified Data.Foldable as Fold
import qualified Data.Map.Strict as Map
import qualified Data.MonoTraversable as MT
import qualified Data.Sequences as Seqs
import qualified Data.Strict.Maybe as Strict
import qualified Data.TreeMap.Strict as TM

import Hcompta.Data
import Hcompta.Quantity
import qualified Hcompta.Lib.Foldable as Fold
import qualified Hcompta.Lib.Strict as Strict

-- * Type 'Balance'

-- | 'BalByAccount' and 'BalByUnit' of some @post@s.
--
-- NOTE: to reduce memory consumption
--       when applying ('+=') incrementally,
--       the fields are explicitely stricts.
data Balance name unit qty
 =   Balance
 {   balByAccount :: !(BalByAccount name unit qty)
 ,   balByUnit    :: !(BalByUnit    name unit qty)
 } deriving (Data, Eq, Show, Typeable)

instance (NFData name, NFData unit, NFData qty, Ord name) => NFData (Balance name unit qty) where
	rnf (Balance a u) = rnf a `seq` rnf u
instance Zeroable (Balance name unit qty) where
	zero = Balance TM.empty Map.empty
instance Nullable qty => Nullable (Balance name unit qty) where
	null (Balance a u) = TM.null a && null u
instance (Ord name, Ord unit, Addable qty) =>
         Addable (Balance name unit qty) where
	Balance xa xu + Balance ya yu = Balance (xa + ya) (xu + yu)
instance (Ord name, Ord unit, Addable qty) =>
         Sumable (Balance name unit qty)
                 (TM.Path name, SumByAccount unit qty) where
	Balance a u += x = Balance (a += x) (u += x)

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

-- ** Type 'SumByAccount'
-- | A sum of quantities, concerning a single account.
type SumByAccount = Map

-- * Type 'BalByUnit'
type BalByUnit name unit qty
 =   Map unit (SumByUnit (TM.Path name) qty)

-- ** Type 'SumByUnit'
-- | A sum of quantities with their accounts involved,
-- concerning a single @unit@.
data SumByUnit acct qty
 =   SumByUnit
 {   sumByUnit_quantity :: !qty
     -- ^ The sum of quantities for a single @unit@.
 ,   sumByUnit_accounts :: !(Map acct ())
     -- ^ The accounts involved to build 'sumByUnit_quantity'.
 } deriving (Data, Eq, Ord, Show, Typeable)
instance (NFData acct, NFData qty) => NFData (SumByUnit acct qty) where
	rnf (SumByUnit q a) = rnf q `seq` rnf a
instance Zeroable qty => Zeroable (SumByUnit acct qty) where
	zero = SumByUnit zero zero
instance Nullable qty => Nullable (SumByUnit acct qty) where
	null (SumByUnit q a) = null q && Map.null a
instance (Ord acct, Addable qty) => Addable (SumByUnit acct qty) where
	x + y = SumByUnit
	 (sumByUnit_quantity x + sumByUnit_quantity y)
	 (sumByUnit_accounts x + sumByUnit_accounts y)
instance (Ord name, Ord unit, Addable qty) =>
         Sumable (BalByUnit name unit qty)
                 (TM.Path name, Map unit qty) where
	bal += (acct, amts) =
		(+) ((`Map.map` amts) $ \qty ->
		SumByUnit
		 { sumByUnit_quantity = qty
		 , sumByUnit_accounts = Map.singleton acct ()
		 }) bal

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

-- | Return the 'balByUnit' of the given 'Balance' with:
--
-- * @unit@s whose 'sumByUnit_quantity' verifying 'zero' are removed,
--
-- * and remaining @unit@s have their 'sumByUnit_accounts'
--   complemented with the 'balByAccount' of the given 'Balance'
--   (i.e. now mapping to the accounts __not__ involved to build the 'SumByUnit').
deviationByUnit ::
 (Ord name, Ord unit, Addable qty, Nullable qty) =>
 Balance         name unit qty ->
 DeviationByUnit name unit qty
deviationByUnit Balance{balByAccount, balByUnit} =
	let all_accounts = TM.flatten (const ()) balByAccount in
	let max_accounts = Map.size all_accounts in
	DeviationByUnit $
		Map.foldlWithKey
		 (\m unit SumByUnit{..} ->
			if 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 = all_accounts `Map.difference` sumByUnit_accounts in
					Map.insert unit SumByUnit
					 { sumByUnit_quantity
					 , sumByUnit_accounts = diff
					 } m
		 )
		 zero
		 balByUnit

-- ** Balance 'equilibrium'
-- | Return the 'Balance' (adjusted by inferred quantities)
--   of the given @post@s and either:
--
--   * 'Left': the @unit@s which have a non null 'SumByUnit'
--             and for which no equibrating account can be inferred.
--   * 'Right': the given @post@s with inferred quantities inserted.
equilibrium ::
 ( post ~ MT.Element posts
 
 , Seqs.IsSequence posts
 , Sumable (Balance name unit qty) post
 
 , Get acct post
 , Has (Map unit qty) post
 , To acct (TM.Path name)
 
 , Ord acct
 , Ord name, Ord unit, Addable qty
 , Nullable qty, Negable qty
 ) =>
 Map acct posts ->
 ( Balance name unit qty
 , Either [(unit, SumByUnit (TM.Path name) qty)]
          (Map acct posts)
 )
equilibrium postsByAcct =
	let bal_initial = Fold.foldr (flip $ MT.ofoldr (flip (+=))) zero postsByAcct in
	let DeviationByUnit dev = deviationByUnit bal_initial in
	let (bal_adjusted, eithers) =
		Map.foldrWithKey
		 (\unit sbu@SumByUnit{..} (bal, lr) ->
			case Map.size $ sumByUnit_accounts of
			 1 ->
				let acct = fst $ Map.elemAt 0 sumByUnit_accounts in
				let qty  = neg sumByUnit_quantity in
				let amts = Map.singleton unit qty in
				( bal += (acct, amts)
				, Right (acct, unit, qty) : lr
				)
			 _ -> (bal, Left [(unit, sbu)] : lr))
		 (bal_initial, [])
		 dev in
	let (l, r) = Fold.accumLeftsAndFoldrRights
		 (\(acct, unit, qty) ->
			Map.adjust (Seqs.fromList . insAmount (unit, qty) . MT.otoList) (to acct))
			 -- NOTE: acct is within bal_initial,
			 -- hence postsByAcct already has a mapping for acct.
		 postsByAcct eithers in
	case l of
	 [] -> (bal_adjusted, Right r)
	 _  -> (bal_adjusted, Left  l)
	where
	insAmount ::
	 forall post unit qty.
	 (Ord unit, Has (Map unit qty) post) =>
	 (unit, qty) -> [post] -> [post]
	insAmount 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 :: Map unit qty = get p in
			if unit `Map.notMember` amts
			then set (Map.insert unit qty amts) p:ps
			else p:insAmount amt ps

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

-- | Return 'True' if and only if the given 'DeviationByUnit'
-- maps only to 'SumByUnit's whose 'sumByUnit_accounts'
-- maps exactly one account.
isEquilibriumInferrable :: DeviationByUnit name unit qty -> Bool
isEquilibriumInferrable (DeviationByUnit dev) =
	Fold.all ((== 1) . Map.size . sumByUnit_accounts) dev

-- * Type 'ClusiveBalByAccount'

-- | {Ex,In}clusive 'BalByAccount':
-- descending propagation of quantities accross accounts.
type ClusiveBalByAccount name unit qty
 = TreeMap name (ClusiveSumByAccount unit qty)

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

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

-- | NOTE: also correct if the 'ClusiveBalByAccount' has been filtered.
instance (Ord name, Ord unit, Addable qty) =>
         Sumable (BalByUnit name unit qty)
                 (ClusiveBalByAccount name unit qty) where
	(+=) = go []
		where
		go p bal (TreeMap nodes) =
			Map.foldrWithKey
			 (\k TM.Node{TM.node_value, TM.node_descendants} acc ->
				case node_value of
				 Strict.Nothing -> go (k:p) acc node_descendants
				 Strict.Just Strict.Clusive{Strict.inclusive=amts}  ->
					let acct = Seqs.reverse $ TM.path k p in
					acc += (acct, amts))
			 bal nodes