{-# LANGUAGE StandaloneDeriving #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE UndecidableInstances #-}
{-# LANGUAGE DataKinds #-}
{-# LANGUAGE ConstraintKinds #-}
{-# LANGUAGE PolyKinds #-}
{-# LANGUAGE KindSignatures #-}
{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE GADTs #-}
{-# LANGUAGE ExistentialQuantification #-}
{-# LANGUAGE TypeOperators #-}
module Exp where

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.Monad.Trans.Class
import Control.Applicative
import Prelude hiding (foldr)
import Prelude.Extras
import GHC.Prim (Constraint(..))
import Unsafe.Coerce
import Bound

infixl 9 :@
infixr 5 :>

-- little orphan instances
instance Show1 Vector where showsPrec1 = showsPrec
instance Eq1 Vector where (==#) = (==)

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 (==) = (==#)
instance Eq1 Exp where
  Var a     ==# Var b     = a == b
  (a :@ b)  ==# (c :@ d)  = a ==# c && b ==# d
  Lam ps a  ==# Lam qs b  = eqPat ps qs && a ==# unsafeCoerce b -- eqPat proves equal shape
  Let as a  ==# Let bs b  = as == bs && a ==# b
  _         ==# _         = False

instance Show a => Show (Exp a) where showsPrec = showsPrec1
instance Show1 Exp where
  showsPrec1 d (Var a)    = showParen (d > 10) $ showString "Var " . showsPrec 11 a
  showsPrec1 d (a :@ b)   = showParen (d > 9) $ showsPrec1 9 a . showString " :@ " . showsPrec1 10 b
  showsPrec1 d (Lam ps b) = showParen (d > 10) $ showString "Lam " . showsPrec1 11 ps . showChar ' ' . showsPrec1 11 b
  showsPrec1 d (Let bs b) = showParen (d > 10) $ showString "Let " . showsPrec1 11 bs . showChar ' ' . showsPrec1 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, Eq a) => Pat b f a -> Pat b' f a -> Bool
eqPat VarP        VarP        = True
eqPat WildP       WildP       = True
eqPat (AsP p)     (AsP q)     = eqPat p q
eqPat (ConP g ps) (ConP h qs) = g == h  && eqPats ps qs
eqPat (ViewP e p) (ViewP f q) = e ==# f && eqPat p q

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

instance Show1 f => Show1 (Pat b f) where showsPrec1 = showsPrec
instance (Show1 f, Show a) => Show (Pat b f a) where
  showsPrec _ VarP        = showString "VarP"
  showsPrec _ WildP       = showString "WildP"
  showsPrec d (AsP p)     = showParen (d > 10) $ showString "AsP " . showsPrec 11 p
  showsPrec d (ConP g ps) = showParen (d > 10) $ showString "ConP " . showsPrec 11 g . showChar ' ' . showsPrec 11 ps
  showsPrec d (ViewP e p) = showParen (d > 10) $ showString "ViewP " . showsPrec1 11 e . showChar ' ' . showsPrec 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, Eq a) => Pats bs f a -> Pats bs' f a -> Bool
eqPats NilP      NilP      = True
eqPats (p :> ps) (q :> qs) = eqPat p q && eqPats ps qs
eqPats _         _         = False

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

instance (Show1 f, Show a) => Show (Pats bs f a) where showsPrec = showsPrec1
instance Show1 f => Show1 (Pats bs f) where
  showsPrec1 _ NilP      = showString "NilP"
  showsPrec1 d (p :> ps) = showParen (d > 5) $
    showsPrec1 6 p . showString " :> " . showsPrec1 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

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 (==) = eqMPath

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 = compareMPath

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


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 (==) = eqPath

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     V     = EQ
  compare L     L     = EQ
  compare L     _     = LT
  compare (R _) L     = GT
  compare (R m) (R n) = compare m n
  compare (C 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

-- ghci> let_ [("x",Var "y"),("y",Var "x" :@ Var "y")] $ lam (varp "z") (Var "z" :@ Var "y")
-- ghci> lam (varp "x") (Var "x")
-- ghci> lam (conp "Hello" [varp "x", wildp])) (Var "y")