{-# LANGUAGE CPP #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE DataKinds #-}
{-# LANGUAGE PolyKinds #-}
{-# LANGUAGE GADTs #-}
{-# LANGUAGE TypeOperators #-}

{-# OPTIONS_GHC -fwarn-incomplete-patterns -fno-warn-orphans #-}

module Main where

-- Dara.Functor.Classes in transformers 0.4.0 are totally different
#if MIN_VERSION_transformers(0,5,0) || !MIN_VERSION_transformers(0,4,0)

import Data.Vector as Vector hiding ((++), map)
import Data.List as List
import Data.Foldable
import Data.Traversable
import Data.Monoid (Monoid(..))
import Control.Monad
import Control.Applicative
import Prelude hiding (foldr)
import Data.Functor.Classes
import Data.Vector.Functor.Classes ()
import Data.Type.Equality
import Bound

infixl 9 :@
infixr 5 :>

data Exp a
  = Var a
  | Exp a :@ Exp a
  | forall (b :: Index). Lam (Pat b Exp a) (Scope (Path b) Exp a)
  | Let (Vector (Scope Int Exp a)) (Scope Int Exp a)

data Index = VarI | WildI | AsI Index | ConI [Index]

data Pat :: Index -> (* -> *) -> * -> * where
  VarP  ::                             Pat 'VarI f a
  WildP ::                             Pat 'WildI f a
  AsP   :: Pat i f a                -> Pat ('AsI i) f a
  ConP  :: String    -> Pats bs f a -> Pat ('ConI bs) f a
  ViewP :: f a       -> Pat b f a   -> Pat b f a -- TODO: allow references to earlier variables

data Pats :: [Index] -> (* -> *) -> * -> * where
  NilP  :: Pats '[] f a
  (:>) :: Pat b f a -> Pats bs f a -> Pats (b ': bs) f a

data Path :: Index -> * where
  V :: Path 'VarI
  L :: Path ('AsI a)
  R :: Path a -> Path ('AsI a)
  C :: MPath as -> Path ('ConI as)

data MPath :: [Index] -> * where
  H :: Path a   -> MPath (a ':as)
  T :: MPath as -> MPath (a ':as)

instance Functor Exp where
  fmap = fmapDefault

instance Foldable Exp where
  foldMap = foldMapDefault

instance Applicative Exp where
  pure = Var
  (<*>) = ap

instance Traversable Exp where
  traverse f (Var a)    = Var <$> f a
  traverse f (x :@ y)   = (:@) <$> traverse f x <*> traverse f y
  traverse f (Lam p e)  = Lam <$> traverse f p <*> traverse f e
  traverse f (Let bs e) = Let <$> traverse (traverse f) bs <*> traverse f e

instance Monad Exp where
  return         = Var
  Var a    >>= f = f a
  (x :@ y) >>= f = (x >>= f) :@ (y >>= f)
  Lam p e  >>= f = Lam (p >>>= f) (e >>>= f)
  Let bs e >>= f = Let (fmap (>>>= f) bs) (e >>>= f)

instance Eq a => Eq (Exp a) where (==) = eq1
instance Eq1 Exp where
  liftEq eq (Var a)    (Var b)     = eq a b
  liftEq eq (a :@ a')  (b :@ b')   = liftEq eq a b && liftEq eq a' b'
  liftEq eq (Lam ps a) (Lam qs b)  =
    case eqPat' eq ps qs of
      Nothing -> False
      Just Refl -> liftEq eq a b

  liftEq eq (Let as a) (Let bs b)  = liftEq (liftEq eq) as bs && liftEq eq a b
  liftEq _  _          _           = False

instance Show a => Show (Exp a) where showsPrec = showsPrec1
instance Show1 Exp where
  liftShowsPrec s _ d (Var a)     = showParen (d > 10) $ showString "Var " . s 11 a
  liftShowsPrec s sl d (a :@ b)   = showParen (d > 9)  $ liftShowsPrec s sl 9 a . showString " :@ " . liftShowsPrec s sl 10 b
  liftShowsPrec s sl d (Lam ps b) = showParen (d > 10) $ showString "Lam " . liftShowsPrec s sl 11 ps . showChar ' ' . liftShowsPrec s sl 11 b
  liftShowsPrec s sl d (Let bs b) = showParen (d > 10) $ showString "Let " . liftShowsPrec (liftShowsPrec s sl) (liftShowList s sl) 11 bs . showChar ' ' . liftShowsPrec s sl 11 b

-- * smart lam

-- ** smart patterns

data P a = forall b. P (Pat b Exp a) [a] (a -> Maybe (Path b))

varp :: Eq a => a -> P a
varp a = P VarP [a] (\v -> if a == v then Just V else Nothing)

wildp :: P a
wildp = P WildP [] (const Nothing)

asp :: Eq a => a -> P a -> P a
asp a (P p as f) = P (AsP p) (a:as) $ \v -> case f v of
  Just b              -> Just (R b)
  Nothing | a == v    -> Just L
          | otherwise -> Nothing

data Ps a = forall bs. Ps (Pats bs Exp a) [a] (a -> Maybe (MPath bs))

conp :: String -> [P a] -> P a
conp g ps = case go ps of
  Ps qs as f -> P (ConP g qs) as (fmap C . f)
  where
    go :: [P a] -> Ps a
    go [] = Ps NilP [] (const Nothing)
    go (P p as f : xs) = case go xs of
      Ps ps' ass g' -> Ps (p :> ps') (as ++ ass) $ \v ->
        T <$> g' v <|> H <$> f v

-- * smart lam
lam :: P a -> Exp a -> Exp a
lam (P p _ f) t = Lam p (abstract f t)

-- * smart let
let_ :: Eq a => [(a, Exp a)] -> Exp a -> Exp a
let_ bs b = Let (Vector.fromList $ map (abstr . snd) bs) (abstr b)
  where vs  = map fst bs
        abstr = abstract (`List.elemIndex` vs)

-- * Pat

-- ** A Kind of Shape

eqPat :: (Eq1 f) => (a -> b -> Bool) -> Pat i f a -> Pat i' f b -> Bool
eqPat _  VarP        VarP        = True
eqPat _  WildP       WildP       = True
eqPat eq (AsP p)     (AsP q)     = eqPat eq p q
eqPat eq (ConP g ps) (ConP h qs) = g == h  && eqPats eq ps qs
eqPat eq (ViewP e p) (ViewP f q) = liftEq eq e f && eqPat eq p q
eqPat _ _ _ = False

-- The same as eqPat, but if the patterns are equal, it returns a
-- proof that their type arguments are the same.
eqPat' :: (Eq1 f) => (a -> a' -> Bool) -> Pat b f a -> Pat b' f a' -> Maybe (b :~: b')
eqPat' _  VarP VarP = Just Refl
eqPat' _  WildP WildP = Just Refl
eqPat' eq (AsP p) (AsP q) = do
  Refl <- eqPat' eq p q
  Just Refl
eqPat' eq (ConP g ps) (ConP h qs) = do
  guard (g == h)
  Refl <- eqPats' eq ps qs
  Just Refl
eqPat' eq (ViewP e p) (ViewP f q) = guard (liftEq eq e f) >> eqPat' eq p q
eqPat' _ _ _ = Nothing

instance Eq1 f   => Eq1 (Pat b f)        where liftEq = eqPat
instance (Eq1 f, Eq a) => Eq (Pat b f a) where (==) = eq1

instance (Show1 f, Show a) => Show (Pat b f a) where showsPrec = showsPrec1

instance Show1 f => Show1 (Pat b f) where
  liftShowsPrec _ _  _ VarP        = showString "VarP"
  liftShowsPrec _ _  _ WildP       = showString "WildP"
  liftShowsPrec s sl d (AsP p)     = showParen (d > 10) $ showString "AsP " . liftShowsPrec s sl 11 p
  liftShowsPrec s sl d (ConP g ps) = showParen (d > 10) $ showString "ConP " . showsPrec 11 g . showChar ' ' . liftShowsPrec s sl 11 ps
  liftShowsPrec s sl d (ViewP e p) = showParen (d > 10) $ showString "ViewP " . liftShowsPrec s sl 11 e . showChar ' ' . liftShowsPrec s sl 11 p

instance Functor f => Functor (Pat b f) where
  fmap _ VarP = VarP
  fmap _ WildP = WildP
  fmap f (AsP p) = AsP (fmap f p)
  fmap f (ConP g ps) = ConP g (fmap f ps)
  fmap f (ViewP e p) = ViewP (fmap f e) (fmap f p)

instance Foldable f => Foldable (Pat b f) where
  foldMap f (AsP p)     = foldMap f p
  foldMap f (ConP _g ps) = foldMap f ps
  foldMap f (ViewP e p) = foldMap f e `mappend` foldMap f p
  foldMap _ _           = mempty

instance Traversable f => Traversable (Pat b f) where
  traverse _ VarP = pure VarP
  traverse _ WildP = pure WildP
  traverse f (AsP p) = AsP <$> traverse f p
  traverse f (ConP g ps) = ConP g <$> traverse f ps
  traverse f (ViewP e p) = ViewP <$> traverse f e <*> traverse f p

instance Bound (Pat b) where
  VarP      >>>= _ = VarP
  WildP     >>>= _ = WildP
  AsP p     >>>= f = AsP (p >>>= f)
  ConP g ps >>>= f = ConP g (ps >>>= f)
  ViewP e p >>>= f = ViewP (e >>= f) (p >>>= f)

-- ** Pats
eqPats :: (Eq1 f) => (a -> b -> Bool) -> Pats bs f a -> Pats bs' f b -> Bool
eqPats _  NilP      NilP      = True
eqPats eq (p :> ps) (q :> qs) = eqPat eq p q && eqPats eq ps qs
eqPats _  _         _         = False

-- Like eqPats, but if the patses are equal, it returns a proof that their
-- type arguments are the same.
eqPats' :: (Eq1 f) => (a -> a' -> Bool) -> Pats bs f a -> Pats bs' f a' -> Maybe (bs :~: bs')
eqPats' _  NilP NilP = Just Refl
eqPats' eq (p :> ps) (q :> qs) = do
  Refl <- eqPat' eq p q
  Refl <- eqPats' eq ps qs
  Just Refl
eqPats' _ _ _ = Nothing

instance Eq1 f         => Eq1 (Pats bs f)   where liftEq = eqPats
instance (Eq1 f, Eq a) => Eq  (Pats bs f a) where (==)  = eq1

instance (Show1 f, Show a) => Show (Pats bs f a) where showsPrec = showsPrec1
instance Show1 f => Show1 (Pats bs f) where
  liftShowsPrec _ _  _ NilP      = showString "NilP"
  liftShowsPrec s sl d (p :> ps) = showParen (d > 5) $
    liftShowsPrec s sl 6 p . showString " :> " . liftShowsPrec s sl 5 ps

instance Functor f => Functor (Pats bs f) where
  fmap _ NilP = NilP
  fmap f (p :> ps) = fmap f p :> fmap f ps

instance Foldable f => Foldable (Pats bs f) where
  foldMap f (p :> ps) = foldMap f p `mappend` foldMap f ps
  foldMap _ _    = mempty

instance Traversable f => Traversable (Pats bs f) where
  traverse _f NilP = pure NilP
  traverse f (p :> ps) = (:>) <$> traverse f p <*> traverse f ps

instance Bound (Pats bs) where
  NilP >>>= _ = NilP
  (p :> ps) >>>= f = (p >>>= f) :> (ps >>>= f)

-- ** Path into Pats
-- Internally, this is only used to implement eqPath, which is only
-- used to implement this.
eqMPath :: MPath is -> MPath js -> Bool
eqMPath (H m) (H n) = eqPath m n
eqMPath (T p) (T q) = eqMPath p q
eqMPath _     _     = False

instance Eq (MPath is) where
    H m == H n = m == n
    T p == T q = p == q
    _   == _   = False

-- Internally, this is only used to define comparePath, which
-- is only used here to define this.
compareMPath :: MPath is -> MPath js -> Ordering
compareMPath (H m) (H n) = comparePath m n
compareMPath (H _) (T _) = LT
compareMPath (T p) (T q) = compareMPath p q
compareMPath (T _) (H _) = GT

instance Ord (MPath is) where
    compare (H m) (H n) = compare m n
    compare (H _) (T _) = LT
    compare (T p) (T q) = compare p q
    compare (T _) (H _) = GT

instance Show (MPath is) where
  showsPrec d (H m) = showParen (d > 10) $ showString "H " . showsPrec 11 m
  showsPrec d (T p) = showParen (d > 10) $ showString "T " . showsPrec 11 p

-- instance Read (MPath is)

-- ** Path into Pat
-- Internally, this is only used to implement eqMPath, which is only used
-- to implement this.
eqPath :: Path i -> Path j -> Bool
eqPath V     V     = True
eqPath L     L     = True
eqPath (R m) (R n) = eqPath m n
eqPath (C p) (C q) = eqMPath p q
eqPath _     _     = False

instance Eq (Path i) where
    p == q = case compare p q of
               EQ -> True
               _ -> False

-- Internally, this is only used to define compareMPath, which
-- is only used here to define this.
comparePath :: Path i -> Path j -> Ordering
comparePath V     V     = EQ
comparePath V     _     = LT
comparePath L     V     = GT
comparePath L     L     = EQ
comparePath L     _     = LT
comparePath (R _) V     = GT
comparePath (R _) L     = GT
comparePath (R m) (R n) = comparePath m n
comparePath (R _) (C _) = LT
comparePath (C p) (C q) = compareMPath p q
comparePath (C _) _     = GT

instance Ord (Path i) where
    compare V y = case (y :: Path 'VarI) of V -> EQ
    compare L y = cpL y
        where
          cpL :: Path ('AsI a) -> Ordering
          cpL L = EQ
          cpL (R _) = LT
    compare (R r) y = cpR r y
        where
          cpR :: Path a -> Path ('AsI a) -> Ordering
          cpR _ L = GT
          cpR m (R n) = compare m n
    compare (C c) y = cpC c y
        where
          cpC :: MPath as -> Path ('ConI as) -> Ordering
          cpC p (C q) = compare p q

instance Show (Path i) where
  showsPrec _ V     = showString "V"
  showsPrec _ L     = showString "L"
  showsPrec d (R m) = showParen (d > 10) $ showString "R " . showsPrec 11 m
  showsPrec d (C p) = showParen (d > 10) $ showString "C " . showsPrec 11 p

-- |
-- >>> let_ [("x",Var "y"),("y",Var "x" :@ Var "y")] $ lam (varp "z") (Var "z" :@ Var "y")
-- Let (fromList [Scope (Var (B 1)),Scope (Var (B 0) :@ Var (B 1))]) (Scope (Lam VarP (Scope (Var (B V) :@ Var (F (Var (B 1)))))))
--
-- >>> lam (varp "x") (Var "x")
-- Lam VarP (Scope (Var (B V)))
--
-- >>> lam (conp "Hello" [varp "x", wildp]) (Var "y")
-- Lam (ConP "Hello" (VarP :> WildP :> NilP)) (Scope (Var (F (Var "y"))))
#endif
main :: IO ()
main = return ()