{-# OPTIONS_GHC -fglasgow-exts #-}

module Base where

import Control.Monad

import Indices

class Zip f where
  zipM :: Monad m => (forall ix. Ix l ix => l ix -> r ix -> s ix -> m (t ix))
                  -> f l r ix -> f l s ix -> m (f l t ix)
  zipW :: Monad m => (forall ix. Ix l ix => l ix -> r ix -> s ix -> t ix)
                  -> f l r ix -> f l s ix -> m (f l t ix)
  zipW f = zipM (\w x y -> return (f w x y))

  zip  :: (forall ix. r ix -> s ix -> t ix)
       -> f l r ix -> f l s ix -> f l t ix
  zip f a b = case zipW (\ w x y -> f x y) a b of
                Nothing -> error "generic zip failed"
                Just res -> res

newtype K' a ix = K' { unK' :: a }

from'' :: forall l ix pfl . (Ix l ix,pfl ~ PF l) => ix -> pfl l I ix
from'' = from :: ix -> pfl l I ix

geq :: (Zip (PF l), Ix l ix) => l ix -> ix -> ix -> Bool
geq ix x y = maybe False (const True) (geq' ix (I x) (I y))

geq' :: (Monad m,Zip (PF l), Ix l ix) => l ix -> I ix -> I ix -> m ()
geq' ix (I x) (I y) = combine ix geq' (from'' x) (from'' y)

combine :: (Monad m,Zip f) => l ix -> (forall ix. Ix l ix => l ix -> r ix -> s ix -> m ())
        -> f l r ix -> f l s ix -> m ()
combine ix f x y = zipM wrapf x y >> return ()
  where
    wrapf ix x y = f ix x y >>= return . K'

instance Zip (Id ix') where
  zipM f (Id x) (Id y) = liftM Id (f ix x y)

instance (Zip a,Zip b) => Zip (Prod a b) where
  zipM f (Prod x1 y1) (Prod x2 y2) = liftM2 Prod (zipM f x1 x2) (zipM f y1 y2)

instance Zip Unit where
  zipM f Unit Unit = return Unit

instance (Zip a,Zip b) => Zip (Sum a b) where
  zipM f (Inl x) (Inl y) = liftM Inl (zipM f x y)
  zipM f (Inr x) (Inr y) = liftM Inr (zipM f x y)
  zipM f _ _ = fail "zip failed"

instance Zip f => Zip (Tag ix' f) where
  zipM f (Tag x) (Tag y) = liftM Tag (zipM f x y)

instance Eq a => Zip (K a) where
  zipM f (K x) (K y) | x == y = return (K x)
                     | otherwise = fail "zip failed in K"

instance Monad I where
  return = I
  (I x) >>= f = f x

class GMap f where
  gmapM :: Monad m => (forall ix. Ix l ix => l ix -> r ix -> m (s ix))
                   -> f l r ix -> m (f l s ix)
  gmapW  :: (forall ix. Ix l ix => l ix -> r ix -> s ix) -> f l r ix -> f l s ix
  gmapW f = unI . gmapM (\w x -> I (f w x))
  gmap  :: (forall ix. Ix l ix => r ix -> s ix) -> f l r ix -> f l s ix
  gmap f = gmapW (const f)

instance GMap (Id ix') where
  gmapM f (Id x) = liftM Id (f ix x)

instance GMap (K a) where
  gmapM f (K x) = return (K x)

instance (GMap a,GMap b) => GMap (Prod a b) where
  gmapM f (Prod x y) = liftM2 Prod (gmapM f x) (gmapM f y)

instance GMap Unit where
  gmapM f Unit = return Unit

instance (GMap a,GMap b) => GMap (Sum a b) where
  gmapM f (Inl x) = liftM Inl (gmapM f x)
  gmapM f (Inr x) = liftM Inr (gmapM f x)

instance GMap f => GMap (Tag ix' f) where
  gmapM f (Tag x) = liftM Tag (gmapM f x)

from' :: (Ix l ix,PF l ~ pfl,GMap pfl) => ix -> HFix (pfl l) ix
from' = HIn . gmapW (\ix -> from' . unI) . from

fold :: (Ix l ix, GMap (PF l))
     => (forall ix . Ix l ix => l ix -> PF l l r ix -> r ix) -> ix -> r ix
fold f = f ix . gmapW (\ix -> fold f . unI) . from


--data Box l r = forall ix. Ix l ix => Box (l ix) (r ix)
--
--flatten :: Crush f => f l a ix -> [Box l a]
--flatten = crush (\ ix -> (:) . Box ix) []

class Crush f where
  crush :: (forall ix. Ix l ix => l ix -> a ix -> b -> b) -> b -> f l a ix -> b

instance Crush (Id ix') where
  crush f e (Id x) = f ix x e

instance (Crush a,Crush b) => Crush (Prod a b) where
  crush f e (Prod x y) = crush f (crush f e y) x

instance (Crush a,Crush b) => Crush (Sum a b) where
  crush f e (Inl x) = crush f e x
  crush f e (Inr x) = crush f e x

instance (Crush f) => Crush (Tag ix' f) where
  crush f e (Tag x) = crush f e x

instance Crush (K a) where
  crush f e _ = e

class LR (f :: (* -> *) -> (* -> *) -> * -> *) l where
  leftf  :: l ix -> (forall ix . Ix l ix => l ix -> r ix) -> [f l r ix]
  rightf :: l ix -> (forall ix . Ix l ix => l ix -> r ix) -> [f l r ix]

instance (LR a l,LR b l) => LR (Sum a b) l where
  leftf  w f = map Inl (leftf  w f) ++ map Inr (leftf  w f)
  rightf w f = map Inr (rightf w f) ++ map Inl (rightf w f)

instance (LR a l, LR b l) => LR (Prod a b) l where
  leftf  w f = zipWith Prod (leftf w f) (leftf w f)
  rightf w f = zipWith Prod (rightf w f) (rightf w f)

instance Ix l ixr => LR (Id ixr) l where
  leftf   _ f = [Id (f ix)]
  rightf  _ f = [Id (f ix)]

instance LRA a => LR (K a) l where
  leftf  w f = [K lefta]
  rightf w f = [K righta]

instance LR Unit l where
  leftf  w f = [Unit]
  rightf w f = [Unit]

class LRA a where
  lefta  :: a
  righta :: a


{-
 -
 - This should be more efficient (compile time checking of Tag ls)
instance (Zero (f r l ix)) => LR (Tag ix f r l ix) where
  left  = [Tag zero]
  right = [Tag zero]

instance LR (Tag ix' f r l ix) where
  left  = []
  right = []
-}

-- The instance below can be used when no undecidable instances are desired.
-- In that case type comparisons are carried out at runtime
--

newtype WWrap f (l :: * -> *) ix = WWrap { wwrap :: [f l ix] }

instance (Repr l,Ix l ixt,LR f l) => LR (Tag ixt f) l where
  leftf w f
    = case coerce w' w (WWrap (map Tag (leftf w' f))) of
        Nothing -> []
        Just (WWrap xs) -> xs
    where
      w' :: l ixt
      w' = ix
  rightf w f
    = case coerce w' w (WWrap (map Tag (rightf w' f))) of
        Nothing -> []
        Just (WWrap xs) -> xs
    where
      w' :: l ixt
      w' = ix

-- Ad-hoc left/right
instance LRA Char where
  lefta  = 'L'
  righta = 'R'

instance LRA a => LRA [a] where
  lefta  = []
  righta = [righta]

instance LRA Int where
  lefta  = 0
  righta = 1

safeHead [] = error "Internal error, left or right returned []"
safeHead (x:xs) = x


left :: (Ix l ix, LR (PF l) l) => l ix -> ix
left w = to $ safeHead $ leftf w (I . left)

right :: (Ix l ix, LR (PF l) l) => l ix -> ix
right w = to $ safeHead $ rightf w (I . right)