{-# LANGUAGE LambdaCase #-}
module Dodge.DoubleTree where

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

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

singleLDT :: a -> LDTree b a
singleLDT x = LDT x [] []

ldtToDT :: LDTree b a -> DTree a
ldtToDT (LDT x l r) = DT x (map (ldtToDT . snd) l) (map (ldtToDT . snd) r)

locLDTToLocDT :: LocationLDT b a -> LocationDT a
locLDTToLocDT (LocLDT con t) = LocDT (conLDTToConDT con) (ldtToDT t)

conLDTToConDT :: ContextLDT b a -> ContextDT a
conLDTToConDT = \case
    TopLDT -> TopDT
    LeftwardLDT u cl p _ cr fr -> LeftwardDT 
        (conLDTToConDT u)
        (map (ldtToDT . snd) cl)
        p
        (map (ldtToDT . snd) cr)
        (map (ldtToDT . snd) fr)
    RightwardLDT u fl cl p _ cr -> RightwardDT 
        (conLDTToConDT u)
        (map (ldtToDT . snd) fl)
        (map (ldtToDT . snd) cl)
        p
        (map (ldtToDT . snd) cr)
        
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

ldtStartPropagate :: (a -> c) -> (c -> b -> a -> c) -> LDTree b a -> LDTree b c
ldtStartPropagate g f (LDT x l r) = LDT z 
    (fmap (\(j,t) -> (j,ldtPropagate' z j f t)) l)
    (fmap (\(j,t) -> (j,ldtPropagate' z j f t)) r)
  where
    z = g x
 

ldtPropagate' :: c -> b -> (c -> b -> a -> c) -> LDTree b a -> LDTree b c
ldtPropagate' x i f (LDT y l r) = LDT z 
    (fmap (\(j,t) -> (j,ldtPropagate' z j f t)) l)
    (fmap (\(j,t) -> (j,ldtPropagate' z j f t)) r)
  where
    z = f x i y

-- propagate two functions down the links of an LDT tree
-- which function is chosen depends on whether it is a left or right branch
ldtPropagate ::
    (c -> b -> c) ->
    (c -> b -> c) ->
    c ->
    LDTree b a ->
    LDTree c a
ldtPropagate lf rf = ildtPropagate (const lf) (const rf)

-- Propgates a value (of type c) down the branches of the LDT.
-- The value is updated according a "left" or "right" function (lf or rf),
-- that acts on the parent value, the link, and the child value.
-- For each node, the updated value is used to update a final value (of type d).
ldtPropagateFold ::
    (c -> a -> b -> a -> c) ->
    (c -> a -> b -> a -> c) ->
    (c -> a -> d -> d) ->
    c ->
    LDTree b a ->
    d ->
    d
ldtPropagateFold lf rf up x (LDT v l r) =
    alaf Endo foldMap (\(s, y) -> ldtPropagateFold lf rf up (rf x v s (_ldtValue y)) y) r
        . alaf Endo foldMap (\(s, y) -> ldtPropagateFold lf rf up (lf x v s (_ldtValue y)) y) l
        . up x v

-- Propgates a value (of type c) down the branches of the LDT.
-- The value is updated according a "left" or "right" function (lf or rf),
-- that acts on the parent value, the link, and the child value.
-- For each node-tree, the updated value is used to update a final value (of type d).
ldtPropagateFoldTree ::
    (c -> a -> b -> a -> c) ->
    (c -> a -> b -> a -> c) ->
    (c -> LDTree b a -> d -> d) ->
    c ->
    LDTree b a ->
    d ->
    d
ldtPropagateFoldTree lf rf up x t@(LDT v l r) =
    alaf Endo foldMap (\(s, y) -> ldtPropagateFoldTree lf rf up (rf x v s (_ldtValue y)) y) r
        . alaf Endo foldMap (\(s, y) -> ldtPropagateFoldTree lf rf up (lf x v s (_ldtValue y)) y) l
        . up x t

ildtPropagate ::
    (Int -> c -> b -> c) ->
    (Int -> c -> b -> c) ->
    c ->
    LDTree b a ->
    LDTree c a
ildtPropagate lf rf x (LDT v l r) = LDT v (imap (go lf x) l) (imap (go rf x) r)
  where
    go f y i (z, t) = (f i y z, ildtPropagate lf rf (f i y z) t)

ldtPropagateIndices :: LDTree b a -> LDTree b (a, [Either Int Int])
ldtPropagateIndices (LDT x l r) = LDT (x, []) (imap (f Left) l) (imap (f Right) r)
  where
    f e i (y, t) = (y, second (e i :) <$> ldtPropagateIndices t)

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

ldtToIM :: (a -> Int) -> LDTree b a -> IM.IntMap (LDTree b a)
ldtToIM f t@(LDT x l r) = IM.insert (f x) t $ IM.unions $ map (ldtToIM f . snd) $ l <> r

ldtToIndentList :: LDTree b a -> [(a, Int, LabelDoubleTreeNodeType b)]
ldtToIndentList = ldtIL LDTRootNode

ldtIL :: LabelDoubleTreeNodeType b -> LDTree b a -> [(a, Int, LabelDoubleTreeNodeType b)]
ldtIL nt (LDT x l r) =
    map
        doindent
        ( concat
            ( headMap
                (\(lab, c) -> ldtIL (LDTBottomNode lab) c)
                (\(lab, c) -> ldtIL (LDTMidBelowNode lab) c)
                l
            )
        )
        ++ [(x, 0, nt)]
        ++ map
            doindent
            ( concat
                ( lastMap
                    (\(lab, c) -> ldtIL (LDTTopNode lab) c)
                    (\(lab, c) -> ldtIL (LDTMidAboveNode lab) c)
                    r
                )
            )
  where
    doindent (a, b, c) = (a, b + 1, c)

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)

prettyLDT :: (a -> String) -> LDTree b a -> [String]
prettyLDT f (LDT x l r) =
    concatMap (map ('/' :) . prettyLDT f . snd) r
        ++ (f x : concatMap (map ('\\' :) . prettyLDT f . snd) l)

ldtToLoc :: LDTree b a -> LocationLDT b a
ldtToLoc = LocLDT TopLDT

-- should probably do tests for these
locUp :: LocationLDT b a -> Maybe (LocationLDT b a)
locUp (LocLDT TopLDT _) = Nothing
locUp (LocLDT c@LeftwardLDT{} t) =
    Just $
        LocLDT
            (_cldtUp c)
            (LDT (_cldtParent c) (_cldtCloseLeft c ++ ((_cldtLink c, t) : _cldtCloseRight c)) (_cldtFarRight c))
locUp (LocLDT c@RightwardLDT{} t) =
    Just $
        LocLDT
            (_cldtUp c)
            (LDT (_cldtParent c) (_cldtFarLeft c) (_cldtCloseLeft c ++ ((_cldtLink c, t) : _cldtCloseRight c)))

--locToTop :: LocationLDT b a -> LocationLDT b a
--locToTop loc = maybe loc locToTop $ locUp loc

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

--locToTop = fix $ \x -> fromMaybe x $ locUp x

locLeftmost :: LocationLDT b a -> LocationLDT b a
locLeftmost loc = maybe loc locLeftmost $ alaf First foldMap Just $ locGoLeft loc

locRightmost :: LocationLDT b a -> LocationLDT b a
locRightmost loc = maybe loc locRightmost $ alaf Last foldMap Just $ locGoRight 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
locGoLeft :: LocationLDT b a -> [LocationLDT b a]
locGoLeft (LocLDT c (LDT v l r)) =
    [LocLDT (LeftwardLDT c closel v link closer r) t | (closel, (link, t), closer) <- locGoHelp id l]

-- should probably do tests for these
locGoRight :: LocationLDT b a -> [LocationLDT b a]
locGoRight (LocLDT c (LDT v l r)) =
    [LocLDT (RightwardLDT c l closel v link closer) t | (closel, (link, t), closer) <- locGoHelp id r]

-- 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 _ [] = []

-- Propgates a value (of type c) down the branches of the ContextLDT.
-- The value is updated according a "left" or "right" function (lf or rf),
-- that acts on the parent value, the link, and the child value.
-- For each context node, the updated value is used to update a final value (of type d).
cldtPropagateFold ::
    (c -> a -> b -> a -> c) ->
    (c -> a -> b -> a -> c) ->
    (c -> LocationLDT b a -> d -> d) ->
    c ->
    LocationLDT b a ->
    d ->
    d
cldtPropagateFold lf rf up x loc =
    alaf
        Endo
        foldMap
        ( \(LocLDT con' t') ->
            cldtPropagateFold
                lf
                rf
                up
                (lf x (_cldtParent con') (_cldtLink con') (_ldtValue t'))
                (LocLDT con' t')
        )
        (locGoLeft loc)
        . alaf
            Endo
            foldMap
            ( \(LocLDT con' t') ->
                cldtPropagateFold
                    lf
                    rf
                    up
                    (rf x (_cldtParent con') (_cldtLink con') (_ldtValue t'))
                    (LocLDT con' t')
            )
            (locGoRight loc)
        . up x loc

reduceLocLDT :: Monoid m => (LocationLDT b a -> m) -> LocationLDT b a -> m
reduceLocLDT f x =
    foldMap (reduceLocLDT f) (locGoLeft x) <> f x
        <> foldMap (reduceLocLDT f) (locGoRight x)

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 ContextLDT.
-- 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