module Dodge.DoubleTree where

import Dodge.Data.DoubleTree
import Control.Lens
import qualified Data.IntMap.Strict as IM
import Data.Monoid

singleDT :: a -> DTree a
singleDT x = DT x [] []

dtStartPropagate :: (a -> c) -> (c -> a -> c) -> DTree a -> DTree c
dtStartPropagate g f (DT x l r) = DT z 
    (fmap (dtPropagate' f z) l)
    (fmap (dtPropagate' f z) r)
  where
    z = g x

dtPropagate' :: (c -> a -> c) -> c -> DTree a -> DTree c
dtPropagate' f x (DT y l r) = DT z 
    (fmap (dtPropagate' f z) l)
    (fmap (dtPropagate' f z) r)
  where
    z = f x y

-- conceptually, in a tree growing from left to right,
-- bottom -> top is equated with left -> right.
-- this does not match with thinking of a list as top -> bottom, so take care
doubleTreeToIndentList :: DTree a -> [(a, Int, DoubleTreeNodeType)]
doubleTreeToIndentList = dtIL DTRootNode

dtIL :: DoubleTreeNodeType -> DTree a -> [(a, Int, DoubleTreeNodeType)]
dtIL nt (DT x l r) =
    map doindent (concat (headMap (dtIL DTBottomNode) (dtIL DTMidBelowNode) l))
        ++ [(x, 0, nt)]
        ++ map doindent (concat (lastMap (dtIL DTTopNode) (dtIL DTMidAboveNode) r))
  where
    doindent (a, b, c) = (a, b + 1, c)

dtToAdjacency :: (a -> Int) -> DTree a -> IM.IntMap [Int]
dtToAdjacency f (DT x l r) =
    IM.insert (f x) (map g l <> map g r)
        . IM.unions
        $ map (dtToAdjacency f) $ l <> r
  where
    g = f . _dtValue

dtToIntMapWithRoot :: (a -> Int) -> DTree a -> IM.IntMap (Maybe Int, DTree a)
dtToIntMapWithRoot f t@(DT x l r) =
    IM.insert (f x) (Nothing, t) $
        foldMap (dtToRootIntMap' (f x) f) $ l <> r

dtToRootIntMap' :: Int -> (a -> Int) -> DTree a -> IM.IntMap (Maybe Int, DTree a)
dtToRootIntMap' root f t@(DT x l r) =
    IM.insert (f x) (Just root, t) $
        foldMap (dtToRootIntMap' root f) $ l <> r

dtToUpDownAdj :: (a -> Int) -> DTree a -> IM.IntMap ([Int], [Int])
dtToUpDownAdj f (DT x l r) =
    IM.insert (f x) (map g l, map g r)
        . IM.unions
        $ map (dtToUpDownAdj f) $ l <> r
  where
    g = f . _dtValue

-- returns an adjacency map with oldest ancestor and direct parent if they exist
-- and any left and right children
-- this should be all involving invids
dtToLRAdj :: (a -> Int) -> DTree a -> IM.IntMap (Maybe (Int, Int), [Int], [Int])
dtToLRAdj f (DT x l r) =
    IM.insert i (Nothing, map g l, map g r)
        . IM.unions
        $ map (dtToAdjRootParent i i f) $ l <> r
  where
    i = f x
    g = f . _dtValue

-- returns an adjacency map with oldest ancestor and direct parent if they exist
-- and any left and right children
-- allows to propagate failure in the index discovery
dtToLRAdjEither ::
    (a -> Either String Int) ->
    DTree a ->
    Either String (IM.IntMap (Maybe (Int, Int), [Int], [Int]))
dtToLRAdjEither f (DT x l r) = do
    i <- f x
    l' <- mapM g l
    r' <- mapM g r
    childrenasnodes <- mapM (dtToAdjRootParentEither i i f) $ l <> r
    return $
        IM.insert i (Nothing, l', r') $
            IM.unions childrenasnodes
  where
    g = f . _dtValue

dtToAdjRootParent :: Int -> Int -> (a -> Int) -> DTree a -> IM.IntMap (Maybe (Int, Int), [Int], [Int])
dtToAdjRootParent root par f (DT x l r) =
    IM.insert (f x) (Just (root, par), map g l, map g r)
        . IM.unions
        $ map (dtToAdjRootParent root (f x) f) $ l <> r
  where
    g = f . _dtValue

dtToAdjRootParentEither ::
    Int ->
    Int ->
    (a -> Either String Int) ->
    DTree a ->
    Either String (IM.IntMap (Maybe (Int, Int), [Int], [Int]))
dtToAdjRootParentEither root par f (DT x l r) = do
    i <- f x
    l' <- mapM g l
    r' <- mapM g r
    childrenasnodes <- mapM (dtToAdjRootParentEither root i f) $ l <> r
    return $ IM.insert i (Just (root, par), l', r') $ IM.unions childrenasnodes
  where
    g = f . _dtValue

headMap :: (a -> b) -> (a -> b) -> [a] -> [b]
headMap f g (x : xs) = f x : map g xs
headMap _ _ [] = []

lastMap :: (a -> b) -> (a -> b) -> [a] -> [b]
lastMap _ _ [] = []
lastMap f _ [x] = [f x]
lastMap f g (x : xs) = g x : lastMap f g xs

prettyDT :: (a -> String) -> DTree a -> [String]
prettyDT f (DT x l r) =
    concatMap (map ('/' :) . prettyDT f) r
        ++ (f x : concatMap (map ('\\' :) . prettyDT f) l)

locUp' :: LocationDT a -> Maybe (LocationDT a)
locUp' (LocDT TopDT _) = Nothing
locUp' (LocDT c@LeftwardDT{} t) =
    Just $
        LocDT
            (_cdtUp c)
            (DT (_cdtParent c) (_cdtCloseLeft c ++ ( t : _cdtCloseRight c)) (_cdtFarRight c))
locUp' (LocDT c@RightwardDT{} t) =
    Just $
        LocDT
            (_cdtUp c)
            (DT (_cdtParent c) (_cdtFarLeft c) (_cdtCloseLeft c ++ ( t : _cdtCloseRight c)))

locToTop :: LocationDT a -> LocationDT a
locToTop loc = maybe loc locToTop $ locUp' loc

locDTLeftmost :: LocationDT a -> LocationDT a
locDTLeftmost loc = maybe loc locDTLeftmost $ alaf First foldMap Just $ locDTGoLeft loc

locDTRightmost :: LocationDT a -> LocationDT a
locDTRightmost loc = maybe loc locDTRightmost $ alaf Last foldMap Just $ locDTGoRight loc

-- should probably do tests for these
locDTGoLeft :: LocationDT a -> [LocationDT a]
locDTGoLeft (LocDT c (DT v l r)) =
    [LocDT (LeftwardDT c closel v closer r) t | (closel,  t, closer) <- locDTGoHelp id l]

-- should probably do tests for these
locDTGoRight :: LocationDT a -> [LocationDT a]
locDTGoRight (LocDT c (DT v l r)) =
    [LocDT (RightwardDT c l closel v closer) t | (closel,  t, closer) <- locDTGoHelp id r]

-- this seems like it might be very inefficient for large lists
-- difference lists?
locGoHelp :: (a -> b) -> [a] -> [([a], b, [a])]
locGoHelp f = go []
  where
    go cleft (y : ys) = (cleft, f y, ys) : go (cleft <> [y]) ys
    go _ [] = []

-- this seems like it might be very inefficient for large lists
-- difference lists?
locDTGoHelp :: (a -> b) -> [a] -> [([a], b, [a])]
locDTGoHelp f = go []
  where
    go cleft (y : ys) = (cleft, f y, ys) : go (cleft <> [y]) ys
    go _ [] = []

reduceLocDT :: Monoid m => (LocationDT a -> m) -> LocationDT a -> m
reduceLocDT f x =
    foldMap (reduceLocDT f) (locDTGoLeft x) <> f x
        <> foldMap (reduceLocDT f) (locDTGoRight x)

-- Propgates a value (of type c) down the branches of the ContextDT.
-- The value is updated according a "left" or "right" function (lf or rf),
-- that acts on the parent value and the child value.
-- For each context node, the updated value is used to update a final value (of type d).
cdtPropagateFold ::
    (c -> a -> a -> c) ->
    (c -> a -> a -> c) ->
    (c -> LocationDT a -> d -> d) ->
    c ->
    LocationDT a ->
    d ->
    d
cdtPropagateFold lf rf up x loc =
    alaf
        Endo
        foldMap
        ( \(LocDT con' t') ->
            cdtPropagateFold
                lf
                rf
                up
                (lf x (_cdtParent con') (_dtValue t'))
                (LocDT con' t')
        )
        (locDTGoLeft loc)
        . alaf
            Endo
            foldMap
            ( \(LocDT con' t') ->
                cdtPropagateFold
                    lf
                    rf
                    up
                    (rf x (_cdtParent con') (_dtValue t'))
                    (LocDT con' t')
            )
            (locDTGoRight loc)
        . up x loc