{-# LANGUAGE PartialTypeSignatures #-}
{-# LANGUAGE NoOverloadedLists #-}
{-# LANGUAGE DataKinds #-}
--{-# LANGUAGE ConstraintKinds #-}
{-# LANGUAGE UndecidableInstances #-}
module Symantic.Compta.Calc.Balance where

import Control.DeepSeq (NFData)
import Data.Kind (Type)
import Data.Bool
import Data.Coerce (coerce)
import Data.Either (Either(..), rights, lefts)
import Data.Eq (Eq(..))
import Data.Function (($), (.), id, const)
import Data.Functor ((<$), (<$>))
import Data.Map.Strict (Map)
import Data.Maybe (Maybe(..), maybe, isJust, fromMaybe, catMaybes)
import Data.Monoid (Monoid(..))
import Data.Ord (Ord(..))
import Data.Semigroup (Semigroup(..))
import Data.String (String)
import Data.Typeable ()
import GHC.Generics (Generic)
import Text.Show (Show(..))
import qualified Data.List as List
import qualified Data.List.NonEmpty as NonEmpty
import qualified Data.Map.Strict as Map

import Symantic.Compta.Lang
import Symantic.Compta.Calc.Flow
import qualified Symantic.Compta.Calc.Chart as Chart

-- * Type 'BalanceRepr'
data Balance section unit qty amt = Balance
  { balanceByAccount :: !(Chart.Chart section amt)
  , balanceByUnit    :: !(Map unit (SumForUnit (Chart.ChartPath section) qty))
  } deriving (Eq, Show, Generic, NFData)

-- | 'BalanceReprByAccount' and 'BalanceReprByUnit' of some 'TyPost'.
data BalanceRepr f (repr::Type->Type) a where
  BalanceReprAny ::
    Balanceable a ~ 'False =>
    Ty (BalanceRepr f repr) a -> BalanceRepr f repr a
  BalanceRepr ::
    Balanceable a ~ 'True =>
    { unBal :: Balance (Ty repr TyAccountSection)
                       (Ty repr TyUnit)
                       (Ty repr TyQuantity)
                       (f (Ty (BalanceRepr f repr) TyAmount))
    } -> BalanceRepr f repr a
--type instance Tr (BalanceRepr f repr) TyAmount  = Map (Tr repr TyUnit) (Tr repr TyQuantity)
--type instance Tr (BalanceRepr f repr) TyAccount = Chart.ChartPath (Tr repr TyAccountSection)
instance
  ( unit ~ Ty repr TyUnit
  , qty ~ Ty repr TyQuantity
  ) => Inject (Map unit qty) (BalanceRepr f repr) TyAmount where
  inject = BalanceReprAny
instance
  ( section ~ Ty repr TyAccountSection
  ) => Inject (Chart.ChartPath section) (BalanceRepr f repr) TyAccount where
  inject = BalanceReprAny

runBalanceRepr :: forall f repr a.
  Balanceable a ~ 'True =>
  BalanceRepr f repr a ->
  Balance (Ty repr TyAccountSection)
          (Ty repr TyUnit)
          (Ty repr TyQuantity)
          (f (Ty (BalanceRepr f repr) TyAmount))
runBalanceRepr = unBal

type instance Ty (BalanceRepr f repr) TyAmount = Map (Ty repr TyUnit) (Ty repr TyQuantity)
type instance Ty (BalanceRepr f repr) TyAccount = Chart.ChartPath (Ty repr TyAccountSection)
type instance Ty (BalanceRepr f repr) TyAccountSection = Ty repr TyAccountSection

instance Amountable (BalanceRepr f repr) where
  amount = BalanceReprAny
instance Accountable (BalanceRepr f repr) where
  account = BalanceReprAny

type family Balanceable a :: Bool
type instance Balanceable TyQuantity = 'False
type instance Balanceable TyAccount  = 'False
type instance Balanceable TyAmount   = 'False
type instance Balanceable (Map k a)  = 'False -- Balanceable a
type instance Balanceable TyMove     = 'True
type instance Balanceable TyPost     = 'True
type instance Balanceable [a]        = 'True -- Balanceable a

{-
deriving instance
  ( Eq (Ty repr TyUnit)
  , Eq (Ty repr TyQuantity)
  , Eq (Ty repr TyAccountSection)
  , Eq (f (AmountOf (BalanceRepr f repr)))
  ) => Eq (BalanceRepr f repr a)
deriving instance
  ( Show (Ty repr TyUnit)
  , Show (Ty repr TyQuantity)
  , Show (Ty repr TyAccountSection)
  , Show (f (AmountOf (BalanceRepr f repr)))
  ) => Show (BalanceRepr f repr a)
--deriving instance (NFData (Ty repr TyUnit), NFData (Ty repr TyQuantity)) => NFData (BalanceRepr section repr a)
-}

instance Zeroable (Balance acct unit aty amt) where
  zero = Balance (Chart.Chart Map.empty) Map.empty
instance (Addable amt, Addable qty, Ord acct, Ord unit) => Addable (Balance acct unit qty amt) where
  Balance xa xu + Balance ya yu = Balance (xa + ya) (xu + yu)
instance
  ( Addable (Ty repr TyQuantity)
  , Addable (f (Map (Ty repr TyUnit) (Ty repr TyQuantity)))
  , Ord (Ty repr TyAccountSection)
  , Ord (Ty repr TyUnit)
  ) => Listable (BalanceRepr f repr) where
  nil = BalanceRepr zero
  cons (BalanceRepr x) (BalanceRepr y) = BalanceRepr (x + y)
  concat (BalanceRepr x) (BalanceRepr y) = BalanceRepr (x + y)
instance
  ( Ord (Ty repr TyUnit)
  , Ord (Ty repr TyAccountSection)
  , Addable (Ty repr TyQuantity)
  , Addable (f (Ty (BalanceRepr f repr) TyAmount))
  ) => Addable (BalanceRepr f repr TyPost) where
  BalanceRepr (Balance xa xu) + BalanceRepr (Balance ya yu) =
    BalanceRepr (Balance (xa + ya) (xu + yu))
instance
  ( Ord (Ty repr TyUnit)
  , Addable (Ty repr TyQuantity)
  ) => Addable (BalanceRepr f repr [a]) where
  x + y = coerce x + y
instance Balanceable a ~ 'True => Zeroable (BalanceRepr f repr a) where
  zero = BalanceRepr zero
instance
  ( Ord (Ty repr TyAccountSection)
  ) => Postable (BalanceRepr Maybe repr) where
  post (BalanceReprAny acct) (BalanceReprAny amt) = BalanceRepr Balance
    { balanceByAccount = Chart.singleton Nothing acct (Just amt)
    , balanceByUnit = Map.map (\qty -> SumForUnit
      { sumForUnitQuantity = qty
      , sumForUnitAccounts = Map.singleton acct ()
      }) amt
    }
instance Moveable (BalanceRepr f repr) where
  move (BalanceRepr bal) = BalanceRepr bal
  --move = coerce

{-
instance Nullable qty => Nullable (BalanceRepr section unit qty) where
  null (BalanceRepr a u) = TM.null a && null u
instance (Ord section, Ord unit, Addable qty) =>
         Sumable (BalanceRepr section unit qty)
                 (TM.Path section, SumByAccount unit qty) where
  BalanceRepr a u += x = BalanceRepr (a += x) (u += x)
-}

tableBalanceRepr ::
  Ord section =>
  Show section =>
  Show unit =>
  Show qty =>
  Addable qty =>
  Balance section unit qty (Trickle (Map unit (Flow qty))) ->
  [[String]]
tableBalanceRepr Balance{..} =
  Chart.foldrWithPath
    (\acct Trickle{inclusive=amt} -> (
      [ show acct
      , List.intercalate "+" $ catMaybes $ (\(u,q) -> show . (u,) <$> flowIn  q) <$> Map.toList amt
      , List.intercalate "+" $ catMaybes $ (\(u,q) -> show . (u,) <$> flowOut q) <$> Map.toList amt
      --, show (unFlow <$> amt)
      ] :))
    []
    balanceByAccount

{-
-- * Type 'BalanceReprByAccount'
type BalanceReprByAccount f section unit qty =
  Chart.Chart section (f (SumByAccount unit qty))

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

-- * Type 'BalanceReprByUnit'
type BalanceReprByUnit section unit qty =
  Map unit (SumForUnit (Chart.ChartPath section) qty)
-}

-- ** Type 'SumForUnit'
-- | A sum of quantities with their accounts involved,
-- concerning a single @unit@.
data SumForUnit acct qty
 =   SumForUnit
 {   sumForUnitQuantity :: !qty
     -- ^ The sum of quantities for a single @unit@.
 ,   sumForUnitAccounts :: !(Map acct ())
     -- ^ The accounts either involved to build 'sumForUnitQuantity',
     -- or *not* involved when inside a 'DeviationByUnit'.
 } deriving (Eq, Ord, Show, Generic, NFData)
instance Zeroable qty => Zeroable (SumForUnit acct qty) where
  zero = SumForUnit zero Map.empty
instance (Ord acct, Addable qty) => Addable (SumForUnit acct qty) where
  x + y = SumForUnit
    (sumForUnitQuantity x + sumForUnitQuantity y)
    (sumForUnitAccounts x + sumForUnitAccounts y)
instance (Ord acct, Negable qty) => Negable (SumForUnit acct qty) where
  negate x = SumForUnit
    { sumForUnitQuantity = negate (sumForUnitQuantity x)
    , sumForUnitAccounts = negate (sumForUnitAccounts x)
    }
instance (Ord acct, Subable qty) => Subable (SumForUnit acct qty) where
  x - y = SumForUnit
    { sumForUnitQuantity = sumForUnitQuantity x - sumForUnitQuantity y
    , sumForUnitAccounts = sumForUnitAccounts x - sumForUnitAccounts y
    }

-- ** BalanceRepr 'equilibrium'
-- | Return the 'BalanceRepr' (adjusted by inferred quantities)
--   of the given @post@s and either:
--
--   * 'Left': the @unit@s which have a non null 'SumForUnit'
--             and for which no equibrating account can be inferred.
--   * 'Right': the given @post@s with inferred quantities inserted.
equilibrium ::
  forall repr.
  Ord (Ty repr TyUnit) =>
  Ord (Ty repr TyAccountSection) =>
  Nullable (Ty repr TyQuantity) =>
  Addable (Ty repr TyQuantity) =>
  Negable (Ty repr TyQuantity) =>
  Show (Ty repr TyUnit) =>
  Show (Ty repr TyQuantity) =>
  Trans repr (BalanceRepr Maybe repr) =>
  Trans repr (InferPost repr) =>
  Postable repr =>
  repr [TyPost] ->
  Either
    [(Ty repr TyUnit, SumForUnit (Chart.ChartPath (Ty repr TyAccountSection))
                                 (Ty repr TyQuantity))]
    (repr [TyPost])
equilibrium posts =
  let BalanceRepr Balance{..} = move (trans posts) :: BalanceRepr Maybe repr TyMove in
  let balanceAccounts = Map.filter isJust $ Chart.flatten id balanceByAccount in
  let eithers = Map.foldrWithKey
        (\unt sfu@SumForUnit{..} ->
          let unusedAccounts = balanceAccounts `Map.difference` sumForUnitAccounts in
          case Map.size unusedAccounts of
            0 | null sumForUnitQuantity -> id
            1 ->
              -- The quantity can be inferred since having an equilibrated balance
              -- means it must be the opposite of the quantity for that unit on other accounts.
              (:) $ Right $ Map.singleton unt (negate sumForUnitQuantity)
                <$ Map.elemAt 0 unusedAccounts
            _ ->
              -- There is more than one account not specifying a quantity for that unit
              -- hence those cannot be inferred.
              (:) $ Left (unt, sfu)
        ) [] balanceByUnit in
  case lefts eithers of
    [] -> Right $ (`unInferPost` Map.fromListWith (+) (rights eithers)) $ trans posts
    ls  -> Left ls

-- *** Type 'InferPost'
data InferPost repr a where
  InferPostAccount :: Ty (BalanceRepr Maybe repr) TyAccount -> InferPost repr TyAccount
  InferPostAmount  :: Ty (BalanceRepr Maybe repr) TyAmount -> InferPost repr TyAmount
  InferPost ::
    Balanceable a ~ 'True =>
    { unInferPost :: Map (Ty (BalanceRepr Maybe repr) TyAccount)
                               (Ty (BalanceRepr Maybe repr) TyAmount) -> repr a } ->
    InferPost repr a
instance
  ( Postable repr
  , Accountable repr
  , Amountable repr
  , Addable (Ty repr TyAmount)
  , Ord (Ty repr TyAccountSection)
  , Ty repr TyAccount ~ Chart.ChartPath (Ty repr TyAccountSection)
  , Ty repr TyAmount ~ Map (Ty repr TyUnit) (Ty repr TyQuantity)
  ) => Postable (InferPost repr) where
  post (InferPostAccount acct) (InferPostAmount amt) = InferPost $ \env ->
    post (account acct) $ amount $ maybe amt (amt +) $ Map.lookup acct env
instance Listable repr => Listable (InferPost repr) where
  nil = InferPost (const nil)
  cons x xs = InferPost $ \env ->
    cons
      (unInferPost x env)
      (unInferPost xs env)
  concat xs ys = InferPost $ \env ->
    concat
      (unInferPost xs env)
      (unInferPost ys env)
instance
  ( Moveable repr
  ) => Moveable (InferPost repr) where
  move ps = InferPost $ \env ->
    move (unInferPost ps env)
instance
  ( unit ~ Ty repr TyUnit
  , qty ~ Ty repr TyQuantity
  ) => Inject (Map unit qty) (InferPost repr) TyAmount where
  inject = InferPostAmount
instance
  ( section ~ Ty repr TyAccountSection
  ) => Inject (Chart.ChartPath section) (InferPost repr) TyAccount where
  inject = InferPostAccount

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

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

-- | {Ex,In}clusive 'BalanceReprByAccount':
-- descending propagation of quantities accross accounts.
-}

-- * Type 'Trickle'
-- A data type to calculate an 'inclusive' value
-- (through some propagation mecanism,
-- eg. incorporating the values of the children of a tree node),
-- while keeping the original 'exclusive' value
-- (eg. the original value of a tree node).
--
-- * 'exclusive': contains the original 'SumByAccount'.
-- * 'inclusive': contains ('+') folded
--   over 'exclusive' and 'inclusive' of children.
data Trickle amt = Trickle
 { exclusive :: !amt
 , inclusive :: !amt
 } deriving (Eq, Show, Generic, NFData)
instance Semigroup a => Semigroup (Trickle a) where
  Trickle e0 i0 <> Trickle e1 i1 =
    Trickle (e0<>e1) (i0<>i1)
instance Monoid a => Monoid (Trickle a) where
  mempty = Trickle mempty mempty
  mappend = (<>)
--type instance Ty (Trickle amt) TyUnit = Ty amt TyUnit
--type instance Ty (Trickle amt) TyQuantity = Ty amt TyQuantity
instance Addable amt => Addable (Trickle amt) where
  x + y = Trickle
    { exclusive = exclusive x + exclusive y
    , inclusive = inclusive x + inclusive y
    }

-- | Return the given 'BalanceReprByAccount' with:
--
-- * all missing parent accounts inserted;
-- * and every mapped @qty@ added with any @qty@
--   of the account for which it is a parent.
trickleBalanceRepr ::
  forall repr a.
  Balanceable a ~ 'True =>
  Addable (Ty repr TyQuantity) =>
  Ord (Ty repr TyAccountSection) =>
  Ord (Ty repr TyUnit) =>
  BalanceRepr Maybe repr a ->
  BalanceRepr Trickle repr a
trickleBalanceRepr (BalanceRepr bal) = BalanceRepr Balance
  { balanceByAccount = balByAccount
  , balanceByUnit = balByUnit [] (balanceByUnit bal) balByAccount
  } where
  balByUnit ks ini =
    Map.foldrWithKey
      (\k (amt, ch) acc ->
        let acct = NonEmpty.reverse (k NonEmpty.:| ks) in
        acc + balanceByUnit (unBal (post (account acct) (amount (inclusive amt)) ::
          BalanceRepr Maybe repr TyPost))
      ) ini . Chart.unChart
  balByAccount =
    Chart.mapByDepthFirst (\ch a ->
      let exclusive = fromMaybe Map.empty a in
      Trickle
        { exclusive
        , inclusive =
          Map.foldl'
            (\acc (sba, _ch) -> acc + inclusive sba)
            exclusive
            (Chart.unChart ch)
        }
      ) (balanceByAccount bal)