import Data.Maybe (maybe)

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]

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)
----------------------

isConst e = e == Ast || e == Box

validEnv :: Env -> Bool
validEnv [] = True
validEnv (EDecl _ expr:ctx) = maybe False isConst (infer ctx expr) && validEnv ctx

shift :: Int -> Expr -> Expr
shift = shift' 0

shift' k v Ast = Ast
shift' k v Box = Box
shift' k v (Idx i) = Idx $ if i < k then i else i + v
shift' k v (t1 :@: t2) = shift' k v t1 :@: shift' k v t2
shift' k v (Lmb (EDecl s e) t) = Lmb (EDecl s $ shift' k v e) $ shift' (k + 1) v t
shift' k v (Pi (EDecl n t1) t2) = Pi (EDecl n $ shift' k v t1) (shift' (k + 1) v t2)

subst :: Int -> Expr -> Expr -> Expr
subst _ _ Ast = Ast
subst _ _ Box = Box
subst i e idx@(Idx j) = if i == j then e else idx
subst i e (t1 :@: t2) = subst i e t1 :@: subst i e t2
subst i e (Lmb (EDecl n ex) t) = Lmb (EDecl n $ subst i e ex) $ subst (i + 1) (shift 1 e) t
subst j e (Pi (EDecl n ex) t) = Pi (EDecl n $ subst j e ex) $ subst (j + 1) (shift 1 e) t

infer :: Env -> Expr -> Maybe Expr
infer en ex = if validEnv en then infer' en ex else Nothing

infer' _ Ast = Just Box
infer' en (Idx i)
  | 0 <= i && i < length en,
    (EDecl _ ex) <- en !! i
  = Just $ shift (i + 1) ex
infer' en (Pi d@(EDecl _ ex) b)
  | Just Ast <- infer' en ex,
    Just s   <- infer' (d : en) b,
    isConst s
  = Just s
infer' en (t1 :@: t2)
  | Just (Pi (EDecl _ ex) b) <- nf <$> infer' en t1,
    Just ex' <- infer' en t2,
    nf ex == nf ex'
  = Just $ shift (-1) $ subst 0 (shift 1 t2) b
infer' en (Lmb d@(EDecl _ e) m)
  | Just Ast <- infer' en e,
    Just b <- infer (d : en) m
  = Just $ Pi d b
infer' _ _ = Nothing


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

isNF (Lmb (EDecl _ ex) t) = isNF ex && isNF t
isNF x = isNANF x

isNANF (Idx _) = True
isNANF (t1 :@: t2) = isNANF t1 && isNF t2
isNANF (Pi (EDecl _ ex) b) = isNF ex && isNF b
isNANF e = isConst e

isNA (Idx _)   = True
isNA (_ :@: _) = True
isNA (Pi _ _) = True
isNA _ = False

oneStep :: Expr -> Maybe Expr
oneStep (t1 :@: t2)
  | isNA t1,
    Just t1' <- oneStep t1
  = Just $ t1' :@: t2
  | isNANF t1,
    Just t2' <- oneStep t2
  = Just $ t1 :@: t2'
oneStep (Lmb (EDecl name a) m)
  | Just a' <- oneStep a = Just $ Lmb (EDecl name a') m
  | isNF a, Just m' <- oneStep m = Just $ Lmb (EDecl name a)  m'
oneStep (Lmb (EDecl _ _) t1 :@: t2) = Just $ shift (-1) $ subst 0 (shift 1 t2) t1
oneStep (Pi (EDecl s ex) b)
  | Just ex' <- oneStep ex
  = Just $ Pi (EDecl s ex') b
  | isNF ex,
    Just b' <- oneStep b
  = Just $ Pi (EDecl s ex) b'


oneStep _ = Nothing

nf :: Expr -> Expr
nf t = case oneStep t of
  Just t1 -> nf t1
  Nothing -> t