{-# LANGUAGE FlexibleContexts      #-}
{-# LANGUAGE StandaloneDeriving    #-}
{-# LANGUAGE TypeFamilies          #-}
{-# LANGUAGE PatternSignatures     #-}
{-# LANGUAGE ScopedTypeVariables   #-}
{-# LANGUAGE GADTs                 #-}
{-# LANGUAGE TypeSynonymInstances  #-}
{-# LANGUAGE EmptyDataDecls        #-}
{-# LANGUAGE TypeOperators         #-}
{-# LANGUAGE Rank2Types            #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE FlexibleInstances     #-}
{-# LANGUAGE KindSignatures        #-}
{-# LANGUAGE EmptyDataDecls        #-}

module Zipper where

import Control.Monad
import Data.Maybe

import Base

infixr 5 :.

data Zipper :: (* -> *) -> * -> * where
  Zipper :: (Ix l ixh) => ixh -> CList l ix ixh -> Zipper l ix

data CList :: (* -> *) -> * -> * -> * where
  CNil  :: CList l ix ix
  (:.)  :: (Ix l ix) => D (PF l) l ix ixh -> CList l ix' ix -> CList l ix' ixh 


data family D f :: (* -> *) -> * -> * -> *
data instance D (K a)     l ix ixh
data instance D (Id x)    l ix ixh = DId                  (x :=: ixh)
data instance D (f ::: x) l ix ixh = DTag   (D f l x ixh) (x :=: ix)
data instance D (f :*: g) l ix ixh = DProdL (D f l ix ixh) (g l ix)
                                   | DProdR (f l ix) (D g l ix ixh)
data instance D (f :+: g) l ix ixh = DSumL  (D f l ix ixh)
                                   | DSumR  (D g l ix ixh)

data (:=:) :: * -> * -> * where
  Refl :: a :=: a

data DPair :: ((* -> *) -> * -> *) -> (* -> *) -> * -> * where
  DPair :: (Ix l ixh) => D f l ix ixh -> ixh -> DPair f l ix

mapDFst :: (forall ixh. Ix l ixh => D f l ix ixh -> D f' l ix' ixh) ->
           DPair f l ix -> DPair f' l ix'
mapDFst f (DPair ctx x) = DPair (f ctx) x

class Zipping_ f where
  firstf :: f l ix -> Maybe (DPair f l ix)
  upf    :: forall l ix ixh . (Ix l ixh) => ixh -> D f l ix ixh -> f l ix

class (Ix l ix, Zipping_ (PF l)) => Zipping l ix
  
instance Zipping_ (K a) where
  firstf (K _) = Nothing
  upf ixh zero = error "impossible"

-- The type annotations below are necessary for the GADT pattern match to succeed.
instance Zipping_ (Id x) where
  firstf (Id x)                              = Just (DPair (DId Refl) x)
  upf (ixh :: ixh) (DId (Refl :: x :=: ixh)) = Id ixh

instance Zipping_ f => Zipping_ (f ::: x) where
  firstf (Tag x) = liftM (mapDFst (\ ctx -> DTag ctx Refl)) (firstf x)
  upf ixh (DTag ctx prf) = upf' prf ixh ctx
    where
      upf' :: Ix l ixh => (x :=: ix) -> ixh -> D f l x ixh -> (f ::: x) l ix
      upf' Refl ixh ctx = Tag (upf ixh ctx)

instance (Zipping_ f, Zipping_ g) => Zipping_ (f :+: g) where
  firstf (L x) = liftM (mapDFst DSumL) (firstf x)
  firstf (R x) = liftM (mapDFst DSumR) (firstf x)
  upf h (DSumL ctx) = L (upf h ctx)
  upf h (DSumR ctx) = R (upf h ctx)

instance (Zipping_ f, Zipping_ g) => Zipping_ (f :*: g) where
  firstf (x :*: y) = liftM (mapDFst (\ ctx -> DProdL ctx y)) (firstf x)
             `mplus` liftM (mapDFst (\ cty -> DProdR x cty)) (firstf y)
  upf h (DProdL ctx y) = upf h ctx :*: y
  upf h (DProdR x cty) = x         :*: upf h cty

down :: Zipping l ix => Zipper l ix -> Maybe (Zipper l ix)
down (Zipper x ctxs) = 
  do
    DPair ctx x' <- firstf (from x)
    return (Zipper x' (ctx :. ctxs))

up :: Zipping l ix => Zipper l ix -> Maybe (Zipper l ix)
up (Zipper _ CNil)          = Nothing
up (Zipper x (ctx :. ctxs)) = Just (Zipper (to (upf x ctx)) ctxs)

enter :: Zipping l ix => l ix -> ix -> Maybe (Zipper l ix)
enter _ x = return (Zipper x CNil)

leave :: Zipping l ix => Zipper l ix -> Maybe ix
leave (Zipper x CNil) = return x
leave x               = leave <=< up $ x

update :: (forall ix. Ix l ix => ix -> ix) -> Zipper l ix -> Maybe (Zipper l ix)
update f (Zipper x ctx) = return (Zipper (f x) ctx)

applyZipper :: (forall ix . l ix -> ix -> a) -> Zipper l ix -> a
applyZipper f (Zipper x _) = f ix x