module Zips where

import Data.Maybe

-- | Zipper for lists
data ListZip a = LZip { backward :: [a], forward :: [a] }
               deriving ( Eq, Show, Read )

listToZip :: [a] -> ListZip a
listToZip = LZip []

zipToList :: ListZip a -> [a]
zipToList (LZip [] xs) = xs
zipToList (LZip (b:bs) xs) = zipToList $ LZip bs (b:xs)

goForward :: ListZip a -> Maybe (ListZip a)
goForward (LZip bs []) = Nothing
goForward (LZip bs (x:xs)) = Just $ LZip (x:bs) xs

goBackward :: ListZip a -> Maybe (ListZip a)
goBackward (LZip [] xs) = Nothing
goBackward (LZip (b:bs) xs) = Just $ LZip bs (b:xs)

modifyLZip :: (a -> a) -> ListZip a -> ListZip a
modifyLZip f z@(LZip bs []) = z
modifyLZip f (LZip bs (x:xs)) = LZip bs (f x : xs)

data BinTree a = BNode (BinTree a) a (BinTree a)
               | BEmpty
               deriving ( Eq, Show, Read )

data TreeDir a = TLeft a (BinTree a)
               | TRight (BinTree a) a
               deriving ( Eq, Show, Read )

-- | Zipper for binary trees
data TreeZip a = TZip [TreeDir a] (BinTree a)
               deriving ( Eq, Show, Read )

treeToZip :: BinTree a -> TreeZip a
treeToZip tree = TZip [] tree

goLeft :: TreeZip a -> Maybe (TreeZip a)
goLeft (TZip ds BEmpty) = Nothing
goLeft (TZip ds (BNode l v r)) = Just $ TZip (TLeft v r : ds) l

goRight :: TreeZip a -> Maybe (TreeZip a)
goRight (TZip ds BEmpty) = Nothing
goRight (TZip ds (BNode l v r)) = Just $ TZip (TRight l v : ds) r

goUp :: TreeZip a -> Maybe (TreeZip a)
goUp (TZip [] _) = Nothing
goUp (TZip (TLeft v r : ds) tree) = Just $ TZip ds (BNode tree v r)
goUp (TZip (TRight l v : ds) tree) = Just $ TZip ds (BNode l v tree)

zipToTree :: TreeZip a -> BinTree a
zipToTree (TZip [] tree) = tree
zipToTree tree = zipToTree . fromJust $ goUp tree

modifyTZip :: (a -> a) -> TreeZip a -> TreeZip a
modifyTZip f z@(TZip ds BEmpty) = z
modifyTZip f (TZip ds (BNode l v r)) = TZip ds (BNode l (f v) r)

testTree :: BinTree Int
testTree = BNode (BNode (BNode BEmpty 4 BEmpty)
                        2
                        (BNode BEmpty 5 BEmpty))
                 1
                 (BNode (BNode BEmpty 6 BEmpty)
                        3
                        (BNode BEmpty 7 BEmpty))