{-# 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 Prelude (error)

import GHC.Int (Int)
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 = showChartAsTree

showChartAsTree :: Show k => Show a => Chart k a -> String
showChartAsTree = List.unlines . drawMap
 where
  -- drawNode :: (k, (a, Chart k a)) -> [String]
  drawNode (k, (a, ts0)) =
    List.zipWith
      (<>)
      (List.lines (showsPrec 11 k ""))
      (List.lines (" " <> showsPrec 11 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)

table :: [[String]] -> String
table cells = List.unlines ((\row -> List.foldr (\cell acc -> "|" <> cell <> acc) "|" row) <$> rows)
 where
  maxCols :: Int
  maxWidths :: [Int] -- for each row
  rows :: [[String]]
  (maxCols, maxWidths, rows) = List.foldr go (0, List.repeat 0, []) cells

  go :: [String] -> (Int, [Int], [[String]]) -> (Int, [Int], [[String]])
  go row (accMaxCols, accMaxWidths, accRows) =
    ( max accMaxCols (List.length row)
    , List.zipWith max accMaxWidths (List.map List.length row <> List.repeat 0)
    , List.take maxCols (List.zipWith alignLeft row maxWidths) : accRows
    )
  alignRight cell maxWidth = List.replicate (maxWidth - List.length cell) ' ' <> cell
  alignLeft cell maxWidth = List.take maxWidth $ cell <> List.repeat ' '

{- >>> error $ table [["toto", "titi"], ["123", "4", "567890"], ["", "", "0"]]
|toto|titi|
|123 |4   |567890|
|    |    |0     |
-}

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