{-# LANGUAGE NoRebindableSyntax #-}
{-# OPTIONS_GHC -Wno-orphans #-}

module Literate.Accounting.Chart where

import Control.Applicative (Applicative (..))
import Control.DeepSeq (NFData (..))
import Control.Monad (Monad (..))
import Control.Monad.Trans.Writer qualified as MT
import Data.Bool
import Data.Eq (Eq (..))
import Data.Foldable (Foldable (..), all)
import Data.Function (const, flip, ($), (.))
import Data.Functor (Functor (..), (<$>))
import Data.List qualified as List
import Data.List.NonEmpty (NonEmpty (..))
import Data.List.NonEmpty qualified as NonEmpty
import Data.Map.Strict (Map)
import Data.Map.Strict qualified as Map
import Data.Maybe (Maybe (..), isNothing)
import Data.Monoid (Monoid (..))
import Data.Ord (Ord (..))
import Data.Semigroup (Semigroup (..))
import Data.String (String)
import Data.Traversable (Traversable (..))
import Data.Tuple (fst, snd)
import GHC.Exts qualified as GHC
import Text.Show (Show (..))

import Literate.Accounting.Math

-- * Type 'Chart'
newtype Chart k a = Chart {unChart :: Map.Map k (a, Chart k a)}
  deriving newtype (Eq, NFData)
instance (Show k, Show a) => Show (Chart k a) where
  show = List.unlines . drawMap
   where
    drawNode :: (k, (a, Chart k a)) -> [String]
    drawNode (k, (a, ts0)) =
      List.zipWith (<>) (List.lines (show k)) (" " <> show a : List.repeat "")
        <> drawMap ts0
    drawMap = go . Map.toList . unChart
     where
      go [] = []
      go [t] = shift "` " "  " (drawNode t)
      go (t : ts) = shift "+ " "| " (drawNode t) <> go ts
    shift ind0 ind = List.zipWith (<>) (ind0 : List.repeat ind)
instance Functor (Chart k) where
  fmap f = Chart . fmap (\(a, ch) -> (f a, fmap f ch)) . unChart
instance Foldable (Chart k) where
  foldMap f = foldMap (\(a, ch) -> f a <> foldMap f ch) . unChart
instance Traversable (Chart k) where
  traverse f =
    (Chart <$>)
      . traverse (\(a, ch) -> (,) <$> f a <*> traverse f ch)
      . unChart
instance (Semigroup a, Ord k) => Semigroup (Chart k a) where
  x <> y =
    Chart $
      Map.unionWith
        (\new old -> (fst old <> fst new, snd old <> snd new))
        (unChart x)
        (unChart y)
instance (Semigroup a, Ord k) => Monoid (Chart k a) where
  mempty = Chart Map.empty
instance (Ord k, Addable a) => Addable (Chart k a) where
  x + y =
    Chart $
      Map.unionWith
        (\(ym, ya) (xm, xa) -> (xm + ym, xa + ya))
        (unChart x)
        (unChart y)
instance (Ord k, Subable a) => Subable (Chart k a) where
  x - y =
    Chart $
      Map.unionWith
        (\(ym, ya) (xm, xa) -> (xm - ym, xa - ya))
        (unChart x)
        (unChart y)

-- ** Type 'ChartPath'
type ChartPath = NonEmpty.NonEmpty

-- * Type 'Account'
type Account = ChartPath

--newtype Account acct = Account (NonEmpty acct)

{- | @('insert' merge path value chart)@
returns the given @chart@ with a mapping from @path@ to @value@,
using @merge value oldValue@ in case @path@ already map to an @oldValue@.
-}
insert :: Ord k => a -> (a -> a -> a) -> ChartPath k -> a -> Chart k a -> Chart k a
insert init merge p a ch = go ch p
 where
  go (Chart m) = \case
    k :| [] ->
      Chart $
        Map.insertWith
          (\_new (old, c) -> (merge a old, c))
          k
          (a, empty)
          m
    k :| k1 : ks ->
      Chart $
        Map.insertWith
          (\_new (old, c) -> (old, go c (k1 :| ks)))
          k
          (init, go empty (k1 :| ks))
          m

-- | Return the value (if any) associated with the given 'Path'.
lookup :: Ord k => ChartPath k -> Chart k a -> Maybe a
lookup (k :| ks) (Chart m) = do
  (a, ms) <- Map.lookup k m
  case ks of
    [] -> Just a
    (k' : ks') -> lookup (k' :| ks') ms

{- | @('filter' predicate chart)@ returns the given @chart@
with only the paths whose mappings satisfied the given @predicate@.
-}
filter :: Ord k => (a -> Bool) -> Chart k a -> Chart k (Maybe a)
filter f =
  Chart
    . Map.mapMaybe
      ( \(x, m) ->
          let fx = f x
           in let fm = filter f m
               in if not fx && all isNothing fm
                    then Nothing
                    else Just (if fx then Just x else Nothing, fm)
      )
    . unChart

-- | The empty 'Chart'.
empty :: Chart k a
empty = Chart Map.empty

{- | @('singleton' ancestorsValue path leafValue)@
returns a 'Chart' mapping @path@ to @leafValue@
and mapping all its ancestors paths to @ancestorsValue@.
-}
singleton :: Ord k => a -> ChartPath k -> a -> Chart k a
singleton init ks a = insert init const ks a empty

{- | Return a 'Map' associating each 'ChartPath' in the given 'Chart',
with its value mapped by the given function.
-}
flatten :: Ord k => (x -> y) -> Chart k x -> Map (ChartPath k) y
flatten = flattenWithPath . const

flattenWithPath :: Ord k => ([k] -> x -> y) -> Chart k x -> Map (ChartPath k) y
flattenWithPath = go []
 where
  go p f ch =
    Map.unions $
      Map.mapKeysMonotonic
        (NonEmpty.reverse . flip (:|) p)
        ( Map.mapWithKey
            ( \k (a, _children) ->
                f (List.reverse (k : p)) a
            )
            (unChart ch)
        ) :
      Map.foldrWithKey
        (\k (_a, children) -> (go (k : p) f children :))
        []
        (unChart ch)

mapByDepthFirst :: Ord k => (Chart k b -> a -> b) -> Chart k a -> Chart k b
mapByDepthFirst f =
  Chart
    . Map.map
      (\(a, ch) -> let m = mapByDepthFirst f ch in (f m a, m))
    . unChart

foldrPath :: Ord k => (ChartPath k -> a -> acc -> acc) -> ChartPath k -> Chart k a -> acc -> acc
foldrPath f = go [] . NonEmpty.toList
 where
  go _ [] _m acc = acc
  go p (k : ks) (Chart m) acc =
    case Map.lookup k m of
      Just (a, ch) -> f (NonEmpty.reverse (k :| p)) a $ go (k : p) ks ch acc
      Nothing -> acc

foldrWithPath :: Ord k => (ChartPath k -> a -> acc -> acc) -> acc -> Chart k a -> acc
foldrWithPath f = go []
 where
  go p acc =
    Map.foldrWithKey
      ( \k (a, ch) acc' ->
          f
            (NonEmpty.reverse (k :| p))
            a
            (go (k : p) acc' ch)
      )
      acc
      . unChart
-- * Type 'ChartM'

-- | A 'Monad' to construct a 'Chart'.
newtype ChartM k v m a = ChartM
  { unChartM ::
      MT.WriterT (v, Chart k v) m a
  }
  deriving newtype (Functor, Applicative, Monad)

runChartM :: Monad m => ChartM k v m a -> m (Chart k v)
runChartM chM = do
  (_a, (_v, ch)) <- MT.runWriterT (unChartM chM)
  return ch

instance (Ord k, Monoid v, Monad m) => Semigroup (ChartM k v m a) where
  (<>) = (Control.Monad.>>)
instance (Ord k, Monoid v, Monad m, Monoid a) => Monoid (ChartM k v m a) where
  mempty = return mempty
instance
  ( Ord k
  , Monoid v
  , Monad m
  , Monoid a
  ) =>
  GHC.IsList (ChartM k v m a)
  where
  type Item (ChartM k v m a) = ChartM k v m a
  fromList = mconcat
  fromListN _n = mconcat
  toList = return