{-# LANGUAGE TypeFamilies               #-}
{-# LANGUAGE GADTs                      #-}
{-# LANGUAGE RankNTypes                 #-}
{-# LANGUAGE TypeOperators              #-}
{-# LANGUAGE MultiParamTypeClasses      #-}
{-# LANGUAGE FlexibleInstances          #-}
{-# LANGUAGE FlexibleContexts           #-}
{-# LANGUAGE ScopedTypeVariables        #-}
{-# LANGUAGE DataKinds                  #-}
{-# LANGUAGE PolyKinds                  #-}
{-# LANGUAGE ConstraintKinds            #-}

module Generics.MultiRec.Transformations.Main
    ( diff, apply, Ixs, Transformation )
  where

import Generics.MultiRec hiding ( show, foldM )
import Control.Applicative ( (<|>) )
import Control.Monad ( foldM )
import Control.Monad.State hiding ( foldM, mapM )

import Generics.MultiRec.ShallowEq
import Generics.MultiRec.Transformations.Path
import Generics.MultiRec.Transformations.Children
import Generics.MultiRec.Transformations.ZipChildren
import Generics.MultiRec.Transformations.MemoTable

--------------------------------------------------------------------------------
-- Synonyms
--------------------------------------------------------------------------------

-- | A constraint synonym for convenience
type Transform phi = ( Fam phi, HFunctor phi (PF phi), SEq phi (PF phi)
                     , ZipChildren phi (PF phi), MapP phi (PF phi)
                     , Extract (PF phi))

-- | Transformations are just sequences of insertions
type Transformation phi top = [ Insert phi top top ]

--------------------------------------------------------------------------------
-- Applying
--------------------------------------------------------------------------------

-- | Apply the transformation to the given tree
apply :: forall phi ix. (Transform phi)
         => phi ix -> ix -> Transformation phi ix -> Maybe ix
apply p t = foldM apply' t where
  apply' :: ix -> Insert phi ix ix -> Maybe ix
  apply' a (Insert _ loc repl) = mapP p loc (\pt _ -> lookupRefs p t pt repl) a

-- | Look up the references using the original structure
lookupRefs :: forall phi top a. (Transform phi)
              => phi top -> top -> phi a -> HWithRef phi top a -> Maybe a
lookupRefs p r p' (HIn (InR x))   = liftM (to p') $
                                    hmapM (\p'' -> liftM I0 . lookupRefs p r p'') p' x
lookupRefs p r _  (HIn (Ref loc)) = extract p loc r

-- | Extract the subtree at the given path
extract :: forall phi i t. (Transform phi)
        => phi i -> Path phi t i -> i -> Maybe t
extract w1 Empty          x = Just x
extract w1 (Push w2 p ps) x =     fmap unI0 (extract' return w1 p (from w1 x))
                              >>= extract w2 ps

class Extract f where
  extract' :: (r t -> Maybe (r t))
           -> phi ix -> Dir f t ix -> f r ix -> Maybe (r t)

instance (Extract f, Extract g) => Extract (f :+: g) where
  extract' f w (CL p) (L x) = extract' f w p x
  extract' f w (CR p) (R x) = extract' f w p x
  extract' f w _      _     = Nothing

instance (Extract f, Extract g) => Extract (f :*: g) where
  extract' f w (C1 p _) (x :*: _) = extract' f w p x
  extract' f w (C2 _ p) (_ :*: y) = extract' f w p y

instance Extract (I ix) where
  extract' f w CId (I x) = f x

instance Extract U     where extract' _ _ _ _ = Nothing
instance Extract (K a) where extract' _ _ _ _ = Nothing

instance (Extract f) => Extract (f :>: ix) where
  extract' f w (CTag p) (Tag x) = extract' f w p x

instance (Extract f) => Extract (C c f) where
  extract' f w (CC p) (C x) = extract' f w p x

instance (Extract f) => Extract (Maybe :.: f) where
  extract' f w (CCM p) (D x) = x >>= extract' f w p

instance (Extract f) => Extract ([] :.: f) where
  extract' f w (CCL l p _) (D x) = extract' f w p (x !! length l)

--------------------------------------------------------------------------------
-- Diffing
--------------------------------------------------------------------------------
-- | Find a set of insertions to transform the first into the second tree

-- Jeroen says: we could make the code of uses nicer if we define a
-- <|>' :: Monad m => m (Maybe a) -> m (Maybe a) -> m (Maybe a)
-- then it will look similar to the version in the paper
diff :: forall phi top. (Transform phi, Children phi (PF phi) top, EmptyMemo phi top (Ixs phi), ChildrenTable phi top (Ixs phi), GetChildrenTable phi (Ixs phi) top, Eq top)
        => phi top -> top -> top -> Transformation phi top
diff p a b = runMemo (Proxy :: Proxy '(phi,top)) $ build False p a b where
  chTbl :: HList (ChildTable phi top (Ixs phi))
  chTbl = childrenTable (Proxy :: Proxy '(phi, top,Ixs phi)) a
  build :: forall a. (Children phi (PF phi) a, Eq a, GetChildrenTable phi (Ixs phi) a) => Bool -> phi a -> a -> a -> Memo phi top [ Insert phi top a ]
  build False p' a' b' | a' == b' = return []
  build ins   p' a' b' =
    let -- All children of this type
        allChildren :: Children phi (PF phi) a => [(Path phi a top, a)]
        allChildren = getChTable (Proxy :: Proxy '(phi, top, Ixs phi)) chTbl
    in case childLookup p' b' allChildren of
      Just l  -> return [ Insert p' Empty (HIn (Ref l)) ]
      Nothing -> uses >>= maybe insert return where -- Only insert when we cannot reuse
        -- Construct the edits for the children based on a root
        construct :: Bool -> a -> Memo phi top (Maybe [ Insert phi top a ])
        construct ins' c =
          if shallowEq p' (from p' c) (from p' b')
          then do r <- zipChildrenM p'
                       (\p1 l1 c1 c2 -> recMemo build ins' p1 c1 c2
                                        >>= return . map (updatePath l1)
                       ) c b'
                  return $ Just $ concat r
          else return Nothing
        -- Possible edits reusing the existing tree or using a part of
        -- the original tree. The existing tree is only used if we didn't
        -- just insert it, since we want to keep the inserts small
        uses :: Memo phi top (Maybe [ Insert phi top a ])
        uses =  reuses >>= \re -> case re of
          Just r | ins -> return re
          _            -> construct ins a' >>= return . pickBest re
        -- Possible edits that include reusing a part of the original tree
        reuses :: Memo phi top (Maybe [ Insert phi top a ])
        reuses = foldM f Nothing allChildren where
          f :: Maybe [ Insert phi top a ] -> (Path phi a top, a)
               -> Memo phi top (Maybe [ Insert phi top a ])
          f c (l,x) = construct False x >>= return . pickBest c . addRef l
          addRef :: Path phi a top -> Maybe [ Insert phi top a ]
                    -> Maybe [ Insert phi top a ]
          addRef l = liftM ((Insert p' Empty (HIn (Ref l))):)
        -- Best edit including insertion, only chosen if nothing can be reused
        insert :: Memo phi top [ Insert phi top a ]
        insert = do
          Just r <- construct True b'
          let (r',e') = partialApply p' (annotate p' b') r
          return $ (Insert p' Empty r') : e'

-- | Update insert location
updatePath :: Path phi a b -> Insert phi top a -> Insert phi top b
updatePath p (Insert w loc v) = Insert w (p <.> loc) v

-- | Pick the best edit
pickBest :: Maybe [a] -> Maybe [a] -> Maybe [a]
pickBest e1 e2 = case (e1,e2) of
  (Just e1', Just e2') -> Just (pickShortest e1' e2')
  _                    -> e1 <|> e2

-- | Pick the shortest of two lists lazily
pickShortest :: [a] -> [a] -> [a]
pickShortest a b = if f a b then a else b
  where f []     _      = True
        f _      []     = False
        f (_:xs) (_:ys) = f xs ys

-- | Lookup a child with a given type
childLookup :: (Fam phi, Eq t)
             => phi t -> t -> [(Path phi t ix, t)] -> Maybe (Path phi t ix)
childLookup p _ [] = Nothing
childLookup p x ((r,y) : ys) | x == y    = Just r
                             | otherwise = childLookup p x ys

-- | Lift a tree to an edit structure
annotate :: (Fam phi, HFunctor phi (PF phi))
         => phi ix -> ix -> HWithRef phi top ix
annotate p = HIn . InR . hmap (\p' (I0 x) -> annotate p' x) p . from p

-- | Try to apply as much edits to the edit structure as possible
--   to make the final edit smaller
partialApply :: forall phi top a. (Transform phi)
                => phi a
                -> HWithRef phi top a
                -> [Insert phi top a]
                -> (HWithRef phi top a, [Insert phi top a])
partialApply _ a [] = (a, [])
partialApply p a (Insert w l x : xs) = case mapPR p l (\_ _ -> Just x) a of
  Just a' -> partialApply p a' xs
  Nothing -> let (a',xs') = partialApply p a xs in (a', Insert w l x : xs')