{-# LANGUAGE TypeFamilies               #-}
{-# LANGUAGE GADTs                      #-}
{-# LANGUAGE RankNTypes                 #-}
{-# LANGUAGE TypeOperators              #-}
{-# LANGUAGE MultiParamTypeClasses      #-}
{-# LANGUAGE FlexibleInstances          #-}
{-# LANGUAGE FlexibleContexts           #-}
{-# LANGUAGE UndecidableInstances       #-}
{-# LANGUAGE ScopedTypeVariables        #-}
{-# LANGUAGE GeneralizedNewtypeDeriving #-}

module Generics.MultiRec.Transformations.Path (
    module Generics.MultiRec.Transformations.Path
  , Ctxs(..), Ctx(..)
  ) where

import Prelude as P hiding ( mapM, sequence )
import Generics.MultiRec hiding ( show, foldM )
import Generics.MultiRec.Zipper ( Ctxs(..), Ctx(..) )
import Control.Monad.State hiding ( foldM, mapM, sequence )
import Data.Traversable ( Traversable, mapM, sequence )

import Generics.MultiRec.CountIs

--------------------------------------------------------------------------------
-- Paths, annotations and edits
--------------------------------------------------------------------------------
-- | A path is a list of connecting directions on a datatype. This is equivalent
-- to a zipper context where the recursive positions are ignored (set to a
-- constant type).
type Path  phi t i = Ctxs phi t (K0 ()) i

-- | A direction points to a single recursive position in a datatype.
type Dir f t i = Ctx f t (K0 ()) i

-- | The type of pattern functors extended with references
data WithRef phi top r a = InR { unInR :: PF phi r a }
                         | Ref { unRef :: Path phi a top }

-- | Closed functors extended with references
type HWithRef phi top t = HFix (WithRef phi top) t

-- | Insertions contain a path of where to insert, and what to insert
data Insert phi top ix where
  Insert :: phi t -> Path phi t ix -> HWithRef phi top t -> Insert phi top ix

--------------------------------------------------------------------------------
-- Path helpers
--------------------------------------------------------------------------------

-- | Concatenate two paths
(<.>) :: forall phi a b c. Path phi b a -> Path phi c b -> Path phi c a
Empty         <.> p2 = p2
(Push p x xs) <.> p2 = Push p x (xs <.> p2)
--------------------------------------------------------------------------------
-- Show instances
--------------------------------------------------------------------------------

newtype ConIndex = CI Int deriving (Eq, Num)

instance Show ConIndex where
  show (CI (-1)) = ""
  show (CI n   ) = "_" ++ show n ++ " "

class ShowPath f where
  showsPrecPath :: ShowS -> ConIndex -> Int -> Dir f i t -> ShowS

instance (ShowPath f, ShowPath g) => ShowPath (f :+: g) where
  showsPrecPath r d n (CL p) = showsPrecPath r d n p
  showsPrecPath r d n (CR p) = showsPrecPath r d n p

instance (ShowPath f, ShowPath g, CountIs f) => ShowPath (f :*: g) where
  -- Going left on a product is unproblematic
  showsPrecPath r d n (C1 p _) = showsPrecPath r d n p
  -- Going right, however, we have to increase |d| by the number of children to
  -- our left, unless |d == -1|, which means we are under a composition, and
  -- shouldn't print any indices anymore.
  showsPrecPath r d n (C2 _ p) =
    let newd = if d == -1 then -1 else d + CI (countIs (undefined :: f r ix))
    in showsPrecPath r newd n p

instance (ShowPath f) => ShowPath (f :>: ix) where
  showsPrecPath r d n (CTag p) = showsPrecPath r d n p

instance (ShowPath f, Constructor c) => ShowPath (C c f) where
  showsPrecPath r d n (CC p) = let name = conName (undefined :: C c f r ix)
                               in showParen (n > 10) $ showString name
                                                     . showsPrecPath r 0 11 p

instance ShowPath (K a) where showsPrecPath _ _ _ _ = id
instance ShowPath U     where showsPrecPath _ _ _ _ = id

instance ShowPath (I ix) where
  showsPrecPath r d n CId = shows d . r

instance (ShowPath f) => ShowPath (Maybe :.: f) where
  showsPrecPath r d n (CCM p) = shows d
                              . showParen (n > 10)
                                  ( showString "Maybe_0 "
                                  . showsPrecPath r (-1) 11 p)

instance (ShowPath f) => ShowPath ([] :.: f) where
  showsPrecPath r d n (CCL l1 p l2) = shows d
                                    . showParen (n > 10)
                                        ( showString "List_"
                                        . P.showsPrec 11 (length l1)
                                        . showChar ' '
                                        . showsPrecPath r (-1) 11 p)

showsPrecPathC :: (ShowPath (PF phi))
               => ConIndex -> Int -> Path phi t i -> ShowS
showsPrecPathC d n Empty         = showString "End"
showsPrecPathC d n (Push w p ps) = showsPrecPath (showsPrecPathC d n ps) d n p

instance (ShowPath (PF phi)) => Show (Path phi t i) where
  showsPrec = showsPrecPathC 0

instance (HFunctor phi (PF phi), HShow phi (PF phi), El phi ix, ShowPath (PF phi))
         => Show (HWithRef phi top ix) where
  showsPrec = showWR proof

instance (HFunctor phi (PF phi), HShow phi (PF phi), El phi ix, ShowPath (PF phi))
         => Show (Insert phi top ix) where
  showsPrec n (Insert w p r) = showParen (n > 10) $
    showString "Insert " . spaces [P.showsPrec 11 p, showWR w 11 r]

-- showsPrec for WithRef working with type index
showWR :: forall phi top ix.
          (HFunctor phi (PF phi), HShow phi (PF phi), ShowPath (PF phi))
       => phi ix -> Int -> HWithRef phi top ix -> ShowS
showWR w n (HIn (InR p)) = showParen (n > 10) $ spaces (("InR"++) : map ($ 11) x) where
    f :: forall ix. phi ix -> HWithRef phi top ix -> K0 [Int -> ShowS] ix
    f w wr = K0 [\n -> showWR w n wr]
    r :: PF phi (K0 [Int -> ShowS]) ix
    r = hmap f w p
    x :: [Int -> ShowS]
    x = hShowsPrecAlg w r
showWR w n (HIn (Ref p)) = showParen (n > 10) $ showString "Ref " . P.showsPrec 11 p

{-
showsPrecPath :: forall phi f i t. ConIndex -> Int -> Dir phi f i t -> ShowS
showsPrecPath d _ End        = showString "End"
showsPrecPath d n (PlusL p)  = showsPrecPath d     n p
showsPrecPath d n (PlusR p)  = showsPrecPath d     n p

-- Going left on a product is unproblematic
showsPrecPath d n (ProdL p)  = showsPrecPath d     n p

-- Going right, however, we have to increase |d| by the number of children to
-- our left, unless |d == -1|, which means we are under a composition, and
-- shouldn't print any indices anymore.
showsPrecPath d n (ProdR (p :: Dir phi g i t))  =
  let newd = if d == -1 then -1 else d + CI (countIs (undefined :: g r ix))
  in showsPrecPath newd n p

showsPrecPath d n (TagP  p)  = showsPrecPath d     n p
showsPrecPath d n (ConsP p)  = let name = conName (undefined :: f r i)
                               in showParen (n > 10) $ showString name
                                                     . showsPrecPath 0 11 p
showsPrecPath d n (RecP _ p) = shows d . showsPrecPath d n p
showsPrecPath d n (TrvI i p) = shows d
                             . showParen (n > 10) (spaces
                                 [ showString "TrvI"
                                 , P.showsPrec 11 i
                                 , showsPrecPath (-1) 11 p])

instance Show (Dir phi f i t) where
  showsPrec = showsPrecPath 0

instance (HFunctor phi (PF phi), HShow phi (PF phi), El phi ix)
         => Show (WithRef phi top ix) where
  showsPrec = showWR proof

instance (HFunctor phi (PF phi), HShow phi (PF phi), El phi ix)
         => Show (Insert phi top ix) where
  showsPrec n (Insert w p r) = showParen (n > 10) $
    showString "Insert " . spaces [P.showsPrec 11 p, showWR w 11 r]

-- showsPrec for WithRef working with type index
showWR :: forall phi top ix. (HFunctor phi (PF phi), HShow phi (PF phi))
       => phi ix -> Int -> WithRef phi top ix -> ShowS
showWR w n (InR p) = showParen (n > 10) $ spaces (("InR"++) : map ($ 11) x) where
    f :: forall ix. phi ix -> WithRef phi top ix -> K0 [Int -> ShowS] ix
    f w wr = K0 [\n -> showWR w n wr]
    r :: PF phi (K0 [Int -> ShowS]) ix
    r = hmap f w p
    x :: [Int -> ShowS]
    x = hShowsPrecAlg w r
showWR w n (Ref p) = showParen (n > 10) $ showString "Ref " . P.showsPrec 11 p
-}

--------------------------------------------------------------------------------
-- MapP
--------------------------------------------------------------------------------
mapP :: forall m phi i t. (Monad m, Fam phi, MapP phi (PF phi))
     => phi i -> Path phi t i -> (phi t -> t -> m t) -> i -> m i
mapP w1 Empty         f = f w1
mapP w1 (Push w2 y p) f =
  liftM (to w1) . mapP' (\w -> liftM I0 . mapP w p f . unI0) w1 y . from w1

class MapP phi f where
  mapP' :: Monad m
        => (phi t -> r t -> m (r t))
        -> phi ix -> Dir f t ix -> f r ix -> m (f r ix)

instance MapP phi U     where mapP' f phi p = return
instance MapP phi (K a) where mapP' f phi p = return

instance (El phi ix) => MapP phi (I ix) where
  mapP' f phi CId (I x) = liftM I (f proof x)

instance (MapP phi f, MapP phi g) => MapP phi (f :+: g) where
  mapP' f phi (CL p) (L x) = liftM L (mapP' f phi p x)
  mapP' f phi (CR p) (R x) = liftM R (mapP' f phi p x)
  mapP' _ _   _      _     = fail "mapP': inconsistent sum"

instance (MapP phi f, MapP phi g) => MapP phi (f :*: g) where
  mapP' f phi (C1 p _) (x :*: y) = liftM2 (:*:) (mapP' f phi p x) (return y)
  mapP' f phi (C2 _ p) (x :*: y) = liftM2 (:*:) (return x) (mapP' f phi p y)

instance (MapP phi f) => MapP phi (C c f) where
  mapP' f phi (CC p) (C x) = liftM C (mapP' f phi p x)

instance (MapP phi f) => MapP phi (f :>: ix) where
  mapP' f phi (CTag p) (Tag x) = liftM Tag (mapP' f phi p x)

instance (MapP phi f) => MapP phi (Maybe :.: f) where
  mapP' f phi (CCM p) = liftM D . sequence . liftM (mapP' f phi p) . unD

instance (MapP phi f) => MapP phi ([] :.: f) where
  mapP' f phi (CCL x p _) = liftM D . mapMwithI (\i ->
                                                  if i == length x
                                                  then mapP' f phi p
                                                  else return) . unD


mapPR :: forall phi top t a. (Fam phi, MapP phi (PF phi))
         => phi a -> Path phi t a
         -> (phi t -> HWithRef phi top t -> Maybe (HWithRef phi top t))
         -> HWithRef phi top a -> Maybe (HWithRef phi top a)
mapPR _  _             _  (HIn (Ref _)) = Nothing
mapPR w1 Empty         f       a        = f w1 a
mapPR w1 (Push w2 y p) f  (HIn (InR a)) =
  liftM (HIn . InR) . mapP' (\w -> mapPR w p f) w1 y $ a

mapMwithI :: (Monad m, Traversable t) => (Int -> a -> m b) -> t a -> m (t b)
mapMwithI f ta = evalStateT (mapM g ta) 0 where
  g a = do i <- get
           put $ i+1
           lift $ f i a