content: {-# LANGUAGE DeriveFunctor, FlexibleInstances #-}
  
  module Ask where
  
  import Control.Applicative
  import Control.Monad
  import Control.Monad.Fail
  
  data Ki
    = Ty (Bwd Ki)
    | Si
    deriving (Show, Eq)
  
  newtype Sc x = Sc x deriving (Show, Eq)
  
  data Ty
    = D String [Ty]
    | Ty :> Ty
    | Tup [Ty]
    | All Ki (Sc Ty)
    | Abs (Sc Ty)
    | Int :$ Bwd Ty
    | Gdd Bool Int
    | Any
    deriving (Show, Eq)
  
  data DataType
    = DPar (Ordering, Ki) (Sc DataType)
    | DCons [(String, [Ty])]
    deriving Show
  
  type DataTypes = [(String, DataType)]
  
  data Bwd x = B0 | Bwd x :< x deriving (Show, Eq, Functor)
  
  (<><) :: Bwd x -> [x] -> Bwd x
  xz <>< [] = xz
  xz <>< (x : xs) = (xz :< x) <>< xs
  
  (<>>) :: Bwd x -> [x] -> [x]
  B0 <>> xs = xs
  (xz :< x) <>> xs = xz <>> (x : xs)
  
  prj :: Bwd x -> Int -> Either String x
  prj B0 _ = Left "missing projection"
  prj (_ :< x)  0 = return x
  prj (xz :< _) i = prj xz (i - 1)
  
  class Ren t where
    ren :: (Int -> Int) -> t -> t
  
  instance Ren t => Ren (Sc t) where
    ren f (Sc t) = Sc (ren wf t) where
      wf 0 = 0
      wf i = f (i - 1) + 1
  
  instance Ren t => Ren [t] where
    ren = fmap . ren
  
  instance Ren t => Ren (Bwd t) where
    ren = fmap . ren
  
  instance Ren Ty where
    ren f (D d ts)  = D d (ren f ts)
    ren f (s :> t)  = ren f s :> ren f t
    ren f (Tup ts)  = Tup (ren f ts)
    ren f (All k t) = All k (ren f t)
    ren f (Abs t)   = Abs (ren f t)
    ren f (i :$ ts) = f i :$ ren f ts
    ren f (Gdd b i) = Gdd b (f i)
    ren f Any       = Any
  
  isTy :: DataTypes -> Bwd Ki -> Ty -> Either String ()
  isTy ds ga (D d ts) = do
    da <- case lookup d ds of
      Just da -> return da
      _ -> Left $ "unknown " ++ d
    dPars ds ga da ts
  isTy ds ga (s :> t) = do
    isTy ds ga s
    isTy ds ga t
  isTy ds ga (Tup ts) = do
    traverse (isTy ds ga) ts
    return ()
  isTy ds ga (All k (Sc t)) = do
    isTy ds (ga :< k) t
  isTy ds ga (i :$ tz) = do
    Ty kz <- prj ga i
    myzw (kiTy ds ga) kz tz
    return ()
  isTy _ _ _ = Left "not a type"
  
  myzw :: Alternative m => (a -> b -> m c) -> Bwd a -> Bwd b -> m (Bwd c)
  myzw f B0 B0 = pure B0
  myzw f (az :< a) (bz :< b) = (:<) <$> myzw f az bz <*> f a b
  myzw _ _ _ = empty
  
  dPars :: DataTypes -> Bwd Ki -> DataType -> [Ty] -> Either String ()
  dPars ds ga (DCons _) [] = return ()
  dPars ds ga (DPar (_, k) (Sc da)) (t : ts) = do
    kiTy ds ga k t
    dPars ds ga da ts
  dPars _ _ _ _ = Left "invalid data type"
  
  kiTy :: DataTypes -> Bwd Ki -> Ki -> Ty -> Either String ()
  kiTy ds ga (Ty kz) t = go ga (kz <>> []) t where
    go ga [] t = isTy ds ga t
    go ga (k : ks) (Abs (Sc t)) = go (ga :< k) ks t
    go _ _ _ = Left "abstraction arity mismatch"
  kiTy ds ga Si t = isSi ga t
  
  isSi :: Bwd Ki -> Ty -> Either String ()
  isSi _ Any = return ()
  isSi ga (Gdd b i) = do
    Si <- prj ga i
    return ()
  isSi _ _ = Left "not a size"
  
  subTy :: DataTypes -> Bwd Ki -> (Ty, Ty) -> Either String ()
  subTy ds ga (s0 :> t0, s1 :> t1) = do
    subTy ds ga (s1, s0)
    subTy ds ga (t0, t1)
  subTy ds ga (D d0 ss, D d1 ts) = do
    guard (d0 == d1)
    da <- case lookup d0 ds of
      Just da -> return da
      _ -> Left $ "unknown " ++ d0
    subPars ds ga da ss ts
  subTy ds ga (s, t)
    | s == t = return ()
    | otherwise = Left $ "subTy " ++ show (s, t)
  
  subPars :: DataTypes -> Bwd Ki -> DataType -> [Ty] -> [Ty] -> Either String ()
  subPars ds ga (DCons _) [] [] = return ()
  subPars ds ga (DPar (o, k) (Sc da)) (s : s0) (t : t0) = do
    case o of
      EQ -> guard (s == t)
      LT -> subKiTy ds ga k s t
      GT -> subKiTy ds ga k t s
    subPars ds ga da s0 t0
  subPars _ _ _ _ _ = Left "arity mismatch"
  
  subKiTy :: DataTypes -> Bwd Ki -> Ki -> Ty -> Ty -> Either String ()
  subKiTy ds ga Si _ Any = return ()
  subKiTy ds ga Si (Gdd a i) (Gdd b j) = do
    guard (i == j)
    guard (a || not b)
  subKiTy ds ga Si Any (Gdd _ _) = Left "unguarded"
  subKiTy ds ga (Ty kz) s t = go ga (kz <>> []) s t where
    go ga [] s t = subTy ds ga (s, t)
    go ga (k : ks) (Abs (Sc s)) (Abs (Sc t)) = go (ga :< k) ks s t
    go _ _ _ _ = Left "arity mismatch"
  
  class Sbs t where
    sbs, sbw :: (Bwd Ty, Int) -> t -> t
    sbs (B0, _) = id
    sbs sg = sbw sg
  
  instance Sbs t => Sbs (Sc t) where
    sbw (tz, n) (Sc t) = Sc (sbw (tz, n + 1) t)
  
  instance Sbs t => Sbs [t] where
    sbw = fmap . sbw
  
  instance Sbs t => Sbs (Bwd t) where
    sbw = fmap . sbw
  
  instance Sbs Ty where
    sbw sg (D d ts) = D d (sbw sg ts)
    sbw sg (s :> t)  = sbw sg s :> sbw sg t
    sbw sg (Tup ts)  = Tup (sbw sg ts)
    sbw sg (All k t) = All k (sbw sg t)
    sbw sg (Abs t)   = Abs (sbw sg t)
    sbw sg@(_, n) (i :$ sz) | i < n = i :$ sbw sg sz
    sbw sg@(tz, n) (i :$ sz) = go tz (i - n) (sbw sg sz) where
      go B0 i uz = (i + n) :$ uz
      go (_ :< t) 0 uz = app (if n == 0 then t else ren (n +) t) uz
      go (tz :< _) i uz = go tz (i - 1) uz
      app (Abs (Sc t)) uz = app t uz
      app t uz = sbs (uz, 0) t
    sbw (_, n) (Gdd b i) | i < n = Gdd b i
    sbw (tz, n) (Gdd b i) = go tz (i - n) where
      go B0 i = Gdd b (i + n)
      go (_ :< t) 0 = case t of
        Any      -> Any
        Gdd c j  -> Gdd (b || c) j
      go (tz :< _) i = go tz (i - 1)
    sbw sg Any       = Any
  
  data Tm
    = E El
    | A (Sc Tm)
    | L Tm
    | Int :! [Tm]
    deriving Show
  
  data El
    = X Int
    | Tm ::: Ty
    | El :/ Ty
    | El :- Tm
    | El :. Int
    | El :? (Ty, [Tm])
    | El :* (Ty, Tm)
    deriving Show
  
  
  daCons :: DataType -> [Ty] -> Either String [(String, [Ty])]
  daCons da ts = go B0 da ts where
    go tz (DCons css) [] = do
      return [(c, sbs (tz, 0) s) | (c, s) <- css]
    go tz (DPar _ (Sc da)) (t : ts) = go (tz :< t) da ts
    go _ _ _ = Left "bad data type"
  
  chk :: DataTypes -> Bwd Ki -> Bwd Ty -> Ty -> Tm -> Either String ()
  chk ds ga de tT (E e) = do
    sS <- syn ds ga de e
    subTy ds ga (sS, tT)
  chk ds ga de (All k (Sc tT)) (A (Sc t)) = do
    chk ds (ga :< k) de tT t
  chk ds ga de (sS :> tT) (L t) = do
    chk ds ga (de :< sS) tT t
  chk ds ga de (Tup tTs) (0 :! ts) = chks ds ga de tTs ts
  chk ds ga de (D d (Any : sSs)) (i :! ts) = do
    da <- case lookup d ds of
      Just da -> return da
      _ -> Left $ "unknown " ++ d
    cTs <- daCons da (Any : sSs)
    guard (i < length cTs)
    chks ds ga de (snd (cTs !! i)) ts
  chk _ ga de tT t = Left $ "chk" ++ show (ga, de, tT, t)
  
  chks :: DataTypes -> Bwd Ki -> Bwd Ty -> [Ty] -> [Tm] -> Either String ()
  chks ds ga de [] [] = return ()
  chks ds ga de (tT : tTs) (t : ts) = do
      chk  ds ga de tT t
      chks ds ga de tTs ts
  chks _ _ _ _ _  = Left "length mismatch"
  
  cases :: DataTypes -> Bwd Ki -> Bwd Ty ->
    [(String, [Ty])] -> Ty -> [Tm] -> Either String ()
  cases _ _ _ [] _ [] = return ()
  cases ds ga de ((_, sSs) : cTs) tT (t : ts) = do
    chk ds ga (de <>< sSs) tT t
    cases ds ga de cTs tT ts
  cases _ _ _ _ _ _ = Left "wrong number of cases"
  
  syn :: DataTypes -> Bwd Ki -> Bwd Ty -> El -> Either String Ty
  syn ds ga de (X i) = prj de i
  syn ds ga de (t ::: tT) = do
    isTy ds ga tT
    chk ds ga de tT t
    return tT
  syn ds ga de (f :/ sS) = do
    All k (Sc tT) <- syn ds ga de f
    kiTy ds ga k sS
    return (sbs (B0 :< sS, 0) tT)
  syn ds ga de (f :- s) = do
    sS :> tT <- syn ds ga de f
    chk ds ga de sS s
    return tT
  syn ds ga de (e :. i) = do
    Tup tTs <- syn ds ga de e
    guard (i < length tTs)
    return (tTs !! i)
  syn ds ga de (e :? (tT, ts)) = do
    D d sSs <- syn ds ga de e
    da <- case lookup d ds of
      Just da -> return da
      _ -> Left $ "unknown " ++ d
    cTs <- daCons da sSs
    cases ds ga de cTs tT ts
    return tT
  syn ds ga de (e :* (tT, t)) = do
    D d (Any : sSs) <- syn ds ga de e
    let sSs' = ren (1 +) sSs
    let gD b = D d (Gdd b 0 : sSs')
    let tT' = ren (1 +) tT
    chk ds ga de (All Si (Sc ((gD True :> tT') :> (gD False :> tT')))) t
    return tT
  
  data Va
    = Bwd Va :\: Tm
    | Int :!: [Va]
    | Mu Va
    deriving Show
  
  evTm :: Bwd Va -> Tm -> Va
  evTm rh (E e) = evEl rh e
  evTm rh (A (Sc t)) = evTm rh t
  evTm rh (L t) = rh :\: t
  evTm rh (i :! ts) = i :!: fmap (evTm rh) ts
  
  evEl :: Bwd Va -> El -> Va
  evEl rh (X i) = case prj rh i of
    Right v -> v
  evEl rh (t ::: _) = evTm rh t
  evEl rh (e :/ _) = evEl rh e
  evEl rh (f :- s) = apply (evEl rh f) (evTm rh s)
  evEl rh (e :. i) = case evEl rh e of
    _ :!: vs -> vs !! i
  evEl rh (e :? (_, ts)) = case evEl rh e of
    i :!: vs -> evTm (rh <>< vs) (ts !! i)
  evEl rh (e :* (tT, t)) = apply (Mu (evTm rh t)) (evEl rh e)
  
  apply :: Va -> Va -> Va
  apply (rh :\: t) s = evTm (rh :< s) t
  apply r@(Mu f) v@(i :!: ts) = apply (apply f r) v
  
  myNat :: (String, DataType)
  myNat = (,) "Nat" $
    DPar (LT, Si) . Sc $
    DCons
      [ ("Z", [])
      , ("S", [D "Nat" [Gdd True 0]])
      ]
  
  myPlus :: Tm
  myPlus = L (L (E (X 1 :* (D "Nat" [Any],
    A (Sc (L (L (E (X 0 :? (D "Nat" [Any],
      [ E (X 2)
      , 1 :! [E (X 2 :- E (X 0))]
      ]))))
    ))))))
  
  tyPlus :: Ty
  tyPlus = D "Nat" [Any] :> (D "Nat" [Any] :> D "Nat" [Any])
  
  myTwo :: Tm
  myTwo = 1 :! [1 :! [0 :! []]]
  
  myTest :: El
  myTest = ((myPlus ::: tyPlus) :- myTwo) :- myTwo
  
  myBadPlus :: Tm
  myBadPlus = L (L (E (X 1 :* (D "Nat" [Any],
    A (Sc (L (L (E (X 3 :? (D "Nat" [Any],
      [ E (X 2)
      , 1 :! [E (X 2 :- E (X 0))]
      ]))))
    ))))))
  
  instance MonadFail (Either String) where
    fail = Left
  
  instance Alternative (Either String) where
    empty = Left "it went wrong"
    Left e <|> r = r
    r <|> _ = r