{-# LANGUAGE NoRebindableSyntax #-}
{-# OPTIONS_GHC -Wno-orphans #-}
module Symantic.Compta.Calc.Chart where

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

--import Symantic.Compta.Lang.Rebindable
import Symantic.Compta.Lang.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

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 :: 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

empty :: Chart k a
empty = Chart Map.empty

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