{-# LANGUAGE DataKinds #-}
{-# LANGUAGE PartialTypeSignatures #-}
{-# LANGUAGE RankNTypes #-}
{-# LANGUAGE StandaloneDeriving #-}
{-# LANGUAGE UndecidableInstances #-}
{-# LANGUAGE NoMonomorphismRestriction #-}
{-# LANGUAGE NoOverloadedLists #-}

module Literate.Accounting.Balance where

import Control.DeepSeq (NFData)
import Data.Either (Either (..), lefts, rights)
import Data.Eq (Eq (..))
import Data.Function (id, ($), (.))
import Data.Functor (Functor (..), (<$), (<$>))
import Data.List qualified as List
import Data.Map.Strict (Map)
import Data.Map.Strict qualified as Map
import Data.Maybe (Maybe (..), catMaybes, fromMaybe, isJust)
import Data.Monoid (Monoid (..))
import Data.Ord (Ord (..))
import Data.Semigroup (Semigroup (..))
import Data.String (String)
import Data.Typeable ()
import GHC.Generics (Generic)
import Lens.Micro
import Literate.Accounting.Chart as Chart
import Literate.Accounting.Flow
import Literate.Accounting.Math
import Literate.Accounting.Move
import Literate.Accounting.Unit
import Text.Show (Show (..))

--type Lens' a b = forall f . Functor f => (b -> f b) -> (a -> f a)

-- * Type 'Balance'
data Balance acct unit qty sum = Balance
  { balanceByAccount :: Chart acct (sum (Amounts unit qty))
  , balanceByUnit :: Amounts unit (SumForUnit (Account acct) qty)
  }
  deriving (Generic)
deriving instance (Eq acct, Eq unit, Eq qty, Eq (sum (Amounts unit qty))) => Eq (Balance acct unit qty sum)
deriving instance (Show acct, Show unit, Show qty, Show (sum (Amounts unit qty))) => Show (Balance acct unit qty sum)
deriving instance (NFData acct, NFData unit, NFData qty, NFData (sum (Amounts unit qty))) => NFData (Balance acct unit qty sum)
instance Zeroable (Balance acct unit qty sum) where
  zero = Balance (Chart Map.empty) zero
instance (Ord acct, Ord unit, Addable qty, Addable (sum (Amounts unit qty))) => Addable (Balance acct unit qty sum) where
  Balance xa xu + Balance ya yu = Balance (xa + ya) (xu + yu)

tableBalance ::
  Ord acct =>
  Show acct =>
  Show unit =>
  Show qty =>
  Addable qty =>
  Balance acct unit (Flow qty) Trickle ->
  [[String]]
tableBalance Balance{..} =
  Chart.foldrWithPath
    ( \acct Trickle{inclusive = Amounts 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 'Account's involved,
 concerning a single @unit@.
-}
data SumForUnit acct qty = SumForUnit
  { -- | The sum of quantities for a single @unit@.
    sumForUnitQuantity :: qty
  , -- | The accounts either involved to build 'sumForUnitQuantity',
    -- or *not* involved when inside a 'DeviationByUnit'.
    sumForUnitAccounts :: Map acct ()
  }
  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
      }

-- * Class 'Balanceable'
class Balanceable a acct unit qty f where
  balance :: a -> Balance acct unit qty f
instance
  ( Balanceable a acct unit qty f
  , Addable qty
  , Addable (f (Amounts unit qty))
  , Ord acct
  , Ord unit
  ) =>
  Balanceable [a] acct unit qty f
  where
  balance = List.foldr (\a acc -> acc + balance a) zero
instance Ord acct => Balanceable (Account acct, Amounts unit qty) acct unit qty Maybe where
  balance (acct, amt) =
    Balance
      { balanceByAccount = Chart.singleton Nothing acct (Just amt)
      , balanceByUnit =
          Amounts $
            Map.map
              ( \qty ->
                  SumForUnit
                    { sumForUnitQuantity = qty
                    , sumForUnitAccounts = Map.singleton acct ()
                    }
              )
              $ unAmounts amt
      }
instance Ord acct => Postable acct unit qty (Balance acct unit qty Maybe) where
  acct += amt =
    Balance
      { balanceByAccount = Chart.singleton Nothing acct (Just amt)
      , balanceByUnit =
          Amounts $
            Map.map
              ( \qty ->
                  SumForUnit
                    { sumForUnitQuantity = qty
                    , sumForUnitAccounts = Map.singleton acct ()
                    }
              )
              $ unAmounts amt
      }

data Post acct amt = Post
  { _postAccount :: acct
  , postAmounts :: amt
  }

class Accountable a section where
  account :: Lens' a (ChartPath section)
class Amountable a unit qty where
  amounts :: Lens' a (Amounts unit qty)

--accountOf :: Functor f => (acct -> f acct) -> (Post acct amt) -> f (Post acct amt)
-- accountOf k post = fmap (\newAccount -> post { _postAccount = newAccount }) (k (_postAccount post))

-- postAccount :: Lens' (Post acct amt) acct
-- postAccount :: Functor f => (acct -> f acct) -> (Post acct amt) -> f (Post acct amt)
-- postAccount k post = fmap (\newAccount -> post { _postAccount = newAccount }) (k (_postAccount post))

-- type instance AccountOf (Post acct amt) = acct
type instance AccountSectionOf (Post acct amt) = AccountSectionOf acct
type instance AmountOf (Post acct amt) = amt
type instance UnitOf (Post acct amt) = UnitOf amt
type instance QuantityOf (Post acct amt) = QuantityOf amt
type instance QuantityOf (Maybe a) = QuantityOf a
type instance UnitOf (Maybe a) = UnitOf a

{-
newtype Amount unit qty = Amount (unit, qty)
type instance UnitOf (Amount unit qty) = unit
type instance QuantityOf (Amount unit qty) = qty
-}

-- type instance AccountSectionOf (Account accountSection) = accountSection

-- ** 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 post unit qty acct.
  Balanceable post acct unit qty Maybe =>
  Nullable qty =>
  Addable qty =>
  Negable qty =>
  Ord unit =>
  Ord acct =>
  Accountable post acct =>
  Amountable post unit qty =>
  [post] ->
  Either
    [(unit, SumForUnit (Account acct) qty)]
    [post]
equilibrium posts =
  let Balance{..} :: Balance acct unit qty Maybe = balance posts
   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 $
                                  Amounts (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)
                  )
                  []
                  (unAmounts balanceByUnit)
           in case lefts eithers of
                [] ->
                  Right $
                    let inferredPosts = Map.fromListWith (+) $ rights eithers
                     in ( \post ->
                            case Map.lookup (post ^. account) inferredPosts of
                              Nothing -> post
                              Just inferredAmounts -> post & amounts %~ (+ inferredAmounts)
                        )
                          <$> posts
                ls -> Left ls

{-
-- | 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)
type instance QuantityOf (Trickle amt) = QuantityOf amt
type instance UnitOf (Trickle amt) = UnitOf amt
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 'Balance' with:

 * all missing parent accounts inserted;
 * and every mapped 'Amounts' added with any 'Amounts'
   of the account for which it is a parent.
-}
trickleBalance ::
  Ord acct =>
  Ord unit =>
  Addable qty =>
  Balance acct unit qty Maybe ->
  Balance acct unit qty Trickle
trickleBalance bal =
  bal
    { balanceByAccount =
        Chart.mapByDepthFirst
          ( \(Chart ch) a ->
              let exclusive = fromMaybe zero a
               in Trickle
                    { exclusive
                    , inclusive = Map.foldr (\(sba, _ch) -> (+ inclusive sba)) exclusive ch
                    }
          )
          (balanceByAccount bal)
    }