type Symb = String

infixl 4 :@:
infixr 3 >->

data Expr = Idx Int
          | Ast
          | Box
          | Expr :@: Expr
          | Lmb Decl Expr
          | Pi Decl Expr
    deriving (Read,Show,Eq,Ord)

data Decl = EDecl Symb Expr
    deriving (Read,Show,Ord)

instance Eq Decl where
  EDecl _ t1 == EDecl _ t2 = t1 == t2

type Env = [Decl]

type Rule = (Expr,Expr)

rulesS,rulesF,rulesO,rulesP,rulesFO,rulesPF,rulesPO,rulesPFO :: [Rule]
rulesS   = [(Ast,Ast)]
rulesF   = (Box,Ast) : rulesS
rulesO   = (Box,Box) : rulesS
rulesP   = (Ast,Box) : rulesS
rulesFO  = (Box,Ast) : rulesO
rulesPF  = (Ast,Box) : rulesF
rulesPO  = (Ast,Box) : rulesO
rulesPFO = (Ast,Box) : rulesFO

lE, pE :: Symb -> Expr -> Expr -> Expr
lE v = Lmb . EDecl v
pE v = Pi  . EDecl v

(>->) :: Expr -> Expr -> Expr
a >-> b = pE "_" a (shift 1 b)
----------------------

validEnv :: [Rule] -> Env -> Bool
validEnv rule [] = True
validEnv rule (EDecl v t : ds)
  | Just s <- infer rule ds t >>= isSort = True  -- environment is being checked inside 'infer'
  | otherwise                            = False

infer :: [Rule] -> Env -> Expr -> Maybe Expr
infer rule env expr
  | validEnv rule env = infer' rule env expr
  | otherwise         = Nothing

infer' :: [Rule] -> Env -> Expr -> Maybe Expr
infer' _   _         Ast                   = Just Box
infer' rule env      (Idx n)
  | 0 <= n && n < length env, -- we do not check environment here, because it's being checked only once in 'infer'
    EDecl _ t <- env !! n
                                      = Just (shift (n + 1) t)
infer' rule env      (m :@: n)
  | Just (Pi (EDecl x a) b) <- nf <$> infer' rule env m,
    Just a'                 <- infer' rule env n,
    a' == a || nf a == nf a'
                                      = Just $ subst0 n b
infer' rule env      (Lmb d@(EDecl x a) m)
  | Just s1   <- (infer' rule env a) >>= isSort,
    Just b    <- infer' rule (d : env) m,
    Just s2   <- (infer' rule env (Pi d b)) >>= isSort
                                      = Just $ Pi d b

infer' rule env      (Pi d@(EDecl x a) b)
  | Just s1 <- infer' rule env a,
    isSort s1 /= Nothing,
    Just s2 <- infer' rule (d:env) b,
    isSort s2 /= Nothing,
    (s1, s2) `elem` rule
                                      = Just s2

infer' _ _          _                   = Nothing

infer0 :: [Rule] -> Expr -> Maybe Expr
infer0 rls = infer rls []

isSort :: Expr -> Maybe Expr
isSort Box = Just Box
isSort Ast = Just Ast
isSort _   = Nothing

shift :: Int -> Expr -> Expr
shift val = shiftAbove 0 where
  shiftAbove cut Ast = Ast
  shiftAbove cut Box = Box
  shiftAbove cut (Idx k)
                         | k < cut   = Idx k
                         | otherwise = Idx $ k + val
  shiftAbove cut (u :@: w)           = shiftAbove cut u :@: shiftAbove cut w
  shiftAbove cut (Lmb (EDecl v t) u) = Lmb (EDecl v $ shiftAbove cut t) $ shiftAbove (cut + 1) u
  shiftAbove cut (Pi (EDecl v t) u)  = Pi (EDecl v $ shiftAbove cut t) $ shiftAbove (cut + 1) u

subst :: Int -> Expr -> Expr -> Expr
subst j s (Idx k)
         | k == j    = s
         | otherwise = Idx k
subst j s (Lmb (EDecl x t) u) = Lmb (EDecl x $ subst j s t) $ subst (j + 1) (shift 1 s) u
subst j s (u :@: w) = subst j s u :@: subst j s w
subst j s (Pi (EDecl x t) u) = Pi (EDecl x $ subst j s t) $ subst (j + 1) (shift 1 s) u
subst j s expr = expr

subst0 s u = shift (-1) $ subst 0 (shift 1 s) u

oneStep :: Expr -> Maybe Expr
oneStep Ast             = Nothing
oneStep Box             = Nothing
oneStep (Idx k)         = Nothing

oneStep (Lmb (EDecl x t) u)
  | Just t' <- oneStep t    = Just $ Lmb (EDecl x t') u
  | Just u' <- oneStep u    = Just $ Lmb (EDecl x t ) u'
  | otherwise               = Nothing

oneStep (Lmb _ u :@: w) = Just $ subst0 w u

oneStep (u :@: w)
  | Just u' <- oneStep u    = Just $ u' :@: w
  | Just w' <- oneStep w    = Just $ u  :@: w'
  | otherwise               = Nothing

oneStep (Pi (EDecl x a) m)
  | Just a' <- oneStep a    = Just $ Pi (EDecl x a') m
  | Just m' <- oneStep m    = Just $ Pi (EDecl x a)  m'
  | otherwise               = Nothing

nf :: Expr -> Expr
nf = nfWith oneStep

nfWith :: (Expr -> Maybe Expr) -> Expr -> Expr
nfWith f t | Just r <- f t = nfWith f r
           | otherwise     = t