{-# LANGUAGE BlockArguments #-}
{-# LANGUAGE LambdaCase #-}
{-# LANGUAGE DeriveAnyClass #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE ViewPatterns #-}
{-# LANGUAGE TupleSections #-}
module Elab.Eval where

import Control.Monad.Reader
import Control.Exception

import qualified Data.Map.Strict as Map
import qualified Data.Sequence as Seq
import qualified Data.Set as Set
import qualified Data.Text as T
import Data.Map.Strict (Map)
import Data.Sequence (Seq)
import Data.List (sortOn)
import Data.Traversable
import Data.Set (Set)
import Data.Typeable
import Data.Foldable
import Data.IORef
import Data.Maybe

import Elab.Eval.Formula
import Elab.Monad

import GHC.Stack

import Presyntax.Presyntax (Plicity(..))

import Prettyprinter

import Syntax.Pretty
import Syntax

import System.IO.Unsafe ( unsafePerformIO )

import {-# SOURCE #-} Elab.WiredIn

eval :: HasCallStack => Term -> ElabM Value
eval t = asks (flip eval' t)

-- everywhere force
zonkIO :: Value -> IO Value
zonkIO (VNe hd sp) = do
  sp' <- traverse zonkSp sp
  case hd of
    HMeta (mvCell -> cell) -> do
      solved <- liftIO $ readIORef cell
      case solved of
        Just vl -> zonkIO $ foldl applProj vl sp'
        Nothing -> pure $ VNe hd sp'
    hd -> pure $ VNe hd sp'

zonkIO (GluedVl h sp vl) = GluedVl h <$> traverse zonkSp sp <*> zonkIO vl

zonkIO (VLam p (Closure s k)) = pure $ VLam p (Closure s (zonk . k))
zonkIO (VPi p d (Closure s k)) = VPi p <$> zonkIO d <*> pure (Closure s (zonk . k))
zonkIO (VSigma d (Closure s k)) = VSigma <$> zonkIO d <*> pure (Closure s (zonk . k))
zonkIO (VPair a b) = VPair <$> zonkIO a <*> zonkIO b

zonkIO (VPath line x y) = VPath <$> zonkIO line <*> zonkIO x <*> zonkIO y
zonkIO (VLine line x y f) = VLine <$> zonkIO line <*> zonkIO x <*> zonkIO y <*> zonkIO f

-- Sorts
zonkIO VType = pure VType
zonkIO VTypeω = pure VTypeω

zonkIO VI = pure VI
zonkIO VI0 = pure VI0
zonkIO VI1 = pure VI1

zonkIO (VIAnd x y) = iand <$> zonkIO x <*> zonkIO y
zonkIO (VIOr x y) = ior <$> zonkIO x <*> zonkIO y
zonkIO (VINot x) = inot <$> zonkIO x

zonkIO (VPartial x y)  = VPartial <$> zonkIO x <*> zonkIO y
zonkIO (VPartialP x y) = VPartialP <$> zonkIO x <*> zonkIO y
zonkIO (VSystem fs) = do
  t <- for (Map.toList fs) $ \(a, b) -> (,) <$> zonkIO a <*> zonkIO b
  pure (mkVSystem (Map.fromList t))
zonkIO (VSub a b c) = VSub <$> zonkIO a <*> zonkIO b <*> zonkIO c
zonkIO (VInc a b c) = incS <$> zonkIO a <*> zonkIO b <*> zonkIO c
zonkIO (VComp a b c d) = comp <$> zonkIO a <*> zonkIO b <*> zonkIO c <*> zonkIO d
zonkIO (VHComp a b c d) = hComp <$> zonkIO a <*> zonkIO b <*> zonkIO c <*> zonkIO d

zonkIO (VGlueTy a phi ty e)   = glueType <$> zonkIO a <*> zonkIO phi <*> zonkIO ty <*> zonkIO e
zonkIO (VGlue a phi ty e t x) = glueElem <$> zonkIO a <*> zonkIO phi <*> zonkIO ty <*> zonkIO e <*> zonkIO t <*> zonkIO x
zonkIO (VUnglue a phi ty e x) = unglue   <$> zonkIO a <*> zonkIO phi <*> zonkIO ty <*> zonkIO e <*> zonkIO x
zonkIO (VCase env t x xs) = pure $ VCase env t x xs

zonkIO (VEqStrict a x y) = VEqStrict <$> zonkIO a <*> zonkIO x <*> zonkIO y
zonkIO (VReflStrict a x) = VReflStrict <$> zonkIO a <*> zonkIO x

zonkSp :: Projection -> IO Projection
zonkSp (PApp p x) = PApp p <$> zonkIO x
zonkSp (PIElim l x y i) = PIElim <$> zonkIO l <*> zonkIO x <*> zonkIO y <*> zonkIO i
zonkSp (POuc a phi u) = POuc <$> zonkIO a <*> zonkIO phi <*> zonkIO u
zonkSp (PK a x p pr) = PK <$> zonkIO a <*> zonkIO x <*> zonkIO p <*> zonkIO pr
zonkSp (PJ a x p pr y) = PJ <$> zonkIO a <*> zonkIO x <*> zonkIO p <*> zonkIO pr <*> zonkIO y
zonkSp PProj1 = pure PProj1
zonkSp PProj2 = pure PProj2

zonk :: Value -> Value
zonk = unsafePerformIO . zonkIO

eval' :: HasCallStack => ElabEnv -> Term -> Value
eval' env (Ref x) =
  case Map.lookup x (getEnv env) of
    Just (_, vl) -> vl
    _ -> VNe (HVar x) mempty

eval' env (Con x) =
  case Map.lookup x (getEnv env) of
    Just (ty, _) -> VNe (HCon ty x) mempty
    Nothing -> error $ "constructor " ++ show x ++ " has no type in scope"

eval' env (PCon sys x) =
  case Map.lookup x (getEnv env) of
    Just (ty, _) -> VNe (HPCon (eval' env sys) ty x) mempty
    Nothing -> error $ "constructor " ++ show x ++ " has no type in scope"

eval' _ (Data n x) = VNe (HData n x) mempty
eval' env (App p f x) = vApp p (eval' env f) (eval' env x)

eval' env (Lam p s t) =
  VLam p $ Closure s $ \a ->
    eval' env { getEnv = Map.insert s (error ("type of abs " ++ show (pretty (Lam p s t))), a) (getEnv env) } t

eval' env (Pi p s d t) =
  VPi p (eval' env d) $ Closure s $ \a ->
    eval' env { getEnv = (Map.insert s (error "type of abs", a) (getEnv env))} t

eval' _ (Meta m) = VNe (HMeta m) mempty

eval' env (Sigma s d t) =
  VSigma (eval' env d) $ Closure s $ \a ->
    eval' env { getEnv = Map.insert s (error "type of abs", a) (getEnv env) } t

eval' e (Pair a b) = VPair (eval' e a) (eval' e b)

eval' e (Proj1 a) = vProj1 (eval' e a)
eval' e (Proj2 a) = vProj2 (eval' e a)

eval' _ Type = VType
eval' _ Typeω = VTypeω
eval' _ I = VI
eval' _ I0 = VI0
eval' _ I1 = VI1

eval' e (IAnd x y) = iand (eval' e x) (eval' e y)
eval' e (IOr x y) = ior (eval' e x) (eval' e y)
eval' e (INot x) = inot (eval' e x)

eval' e (PathP l a b) = VPath (eval' e l) (eval' e a) (eval' e b)
eval' e (IElim l x y f i) = ielim (eval' e l) (eval' e x) (eval' e y) (eval' e f) (eval' e i)
eval' e (PathIntro p x y f) = VLine (eval' e p) (eval' e x) (eval' e y) (eval' e f)

eval' e (Partial x y)  = VPartial (eval' e x) (eval' e y)
eval' e (PartialP x y) = VPartialP (eval' e x) (eval' e y)
eval' e (System fs) = mkVSystem (Map.fromList $ map (\(x, y) -> (eval' e x, eval' e y)) $ Map.toList $ fs)

eval' e (Sub a phi u)   = VSub (eval' e a) (eval' e phi) (eval' e u)
eval' e (Inc a phi u)   = incS (eval' e a) (eval' e phi) (eval' e u)
eval' e (Ouc a phi u x) = outS (eval' e a) (eval' e phi) (eval' e u) (eval' e x)

eval' e (Comp a phi u a0) = comp (eval' e a) (eval' e phi) (eval' e u) (eval' e a0)
eval' e (HComp a phi u a0) = hComp (eval' e a) (eval' e phi) (eval' e u) (eval' e a0)

eval' e (GlueTy a phi tys f) = glueType (eval' e a) (eval' e phi) (eval' e tys) (eval' e f)
eval' e (Glue a phi tys eqvs t x) = glueElem (eval' e a) (eval' e phi) (eval' e tys) (eval' e eqvs) (eval' e t) (eval' e x)
eval' e (Unglue a phi tys f x) = unglue (eval' e a) (eval' e phi) (eval' e tys) (eval' e f) (eval' e x)
eval' e (Let ns x) =
  let env' = foldl (\newe (n, ty, x) ->
                      let nft = eval' newe ty
                       in newe { getEnv = Map.insert n (nft, evalFix' newe n nft x) (getEnv newe) })
                    e
                    ns
   in eval' env' x

eval' e (Case range sc xs) = evalCase e (eval' e range @@) (force (eval' e sc)) xs

eval' e (EqS a x y) = VEqStrict (eval' e a) (eval' e x) (eval' e y)
eval' e (Syntax.Refl a x) = VReflStrict (eval' e a) (eval' e x)
eval' e (Syntax.AxK a x p pr eq) = strictK (eval' e a) (eval' e x) (eval' e p) (eval' e pr) (eval' e eq)
eval' e (Syntax.AxJ a x p pr y eq) = strictJ (eval' e a) (eval' e x) (eval' e p) (eval' e pr) (eval' e y) (eval' e eq)

evalCase :: ElabEnv -> (Value -> Value) -> Value -> [(Term, Int, Term)] -> Value
evalCase env rng sc [] = VCase (getEnv env) (fun rng) sc []

evalCase env rng (VSystem fs) cases = VSystem (fmap (flip (evalCase env rng) cases) fs)

evalCase env rng (VHComp a φ u u0) cases =
  comp (fun \i -> rng (v i))
       φ
       (system \i is1 -> α (u @@ i @@ is1))
       (VInc (rng a) φ (α (outS a φ (u @@ VI0) u0)))
  where
    v = Elab.WiredIn.fill (fun (const a)) φ u u0
    α x = evalCase env rng x cases

evalCase env _ sc ((Ref _, _, k):_) = eval' env k @@ sc

evalCase env rng (val@(VNe (HCon _ x) sp)) ((Con x', _, k):xs)
  | x == x'   = foldl applProj (eval' env k) sp
  | otherwise = evalCase env rng val xs

evalCase env rng (val@(VNe (HPCon _ _ x) sp)) ((Con x', _, k):xs)
  | x == x'   = foldl applProj (eval' env k) sp
  | otherwise = evalCase env rng val xs

evalCase env rng sc xs = VCase (getEnv env) (fun rng) sc xs

evalFix' :: HasCallStack => ElabEnv -> Name -> NFType -> Term -> Value
evalFix' env name nft term = fix $ \val -> eval' env{ getEnv = Map.insert name (nft, val) (getEnv env) } term

evalFix :: HasCallStack => Name -> NFType -> Term -> ElabM Value
evalFix name nft term = do
  t <- ask
  pure (evalFix' t name nft term)

data NotEqual = NotEqual Value Value
  deriving (Show, Typeable, Exception)

unify' :: HasCallStack => Value -> Value -> ElabM ()
-- unify' (GluedVl h sp _) (GluedVl h' sp' _)
--   | h == h', length sp == length sp' = traverse_ (uncurry unify'Spine) (Seq.zip sp sp')
unify' topa topb = join $ go <$> forceIO topa <*> forceIO topb where
  go (VNe (HMeta mv) sp) rhs = solveMeta mv sp rhs
  go rhs (VNe (HMeta mv) sp) = solveMeta mv sp rhs

  go (VNe (HPCon _ _ x) sp) (VNe (HPCon _ _ y) sp')
    | x == y = traverse_ (uncurry unify'Spine) (Seq.zip sp sp')
  go (VNe (HPCon s _ _) _) rhs | Just v <- trivialSystem s = go v rhs
  go lhs (VNe (HPCon s _ _) _) | Just v <- trivialSystem s = go lhs v

  go (VCase e _ a b) (VCase e' _ a' b') = do
    env <- ask
    unify' a a'
    let go (_, _, a) (_, _, b) = unify' (eval' env{getEnv=e} a) (eval' env{getEnv=e'} b)
    zipWithM_ go (sortOn (\(x, _, _) -> x) b) (sortOn (\(x, _, _) -> x) b')

  go (VCase e _ _ b) y = do
    env <- ask
    let
      go (_, n, a') = do
        ns <- replicateM n (VVar <$> newName)
        let a = foldl (vApp Ex) (eval' env{getEnv=e} a') ns
        unify' a y
    traverse_ go b

  go (VNe x a) (VNe x' a')
    | x == x', length a == length a' =
      traverse_ (uncurry unify'Spine) (Seq.zip a a')

  go (VLam p (Closure n k)) vl = do
    t <- VVar <$> newName' n
    unify' (k t) (vApp p vl t)

  go vl (VLam p (Closure n k)) = do
    t <- VVar <$> newName' n
    unify' (vApp p vl t) (k t)

  go (VPair a b) vl = unify' a (vProj1 vl) *> unify' b (vProj2 vl)
  go vl (VPair a b) = unify' (vProj1 vl) a *> unify' (vProj2 vl) b

  go (VPi p d (Closure _ k)) (VPi p' d' (Closure _ k')) | p == p' = do
    t <- VVar <$> newName
    unify' d d'
    unify' (k t) (k' t)

  go (VSigma d (Closure _ k)) (VSigma d' (Closure _ k')) = do
    t <- VVar <$> newName
    unify' d d'
    unify' (k t) (k' t)

  go VType VType   = pure ()
  go VTypeω VTypeω = pure ()

  go VI VI   = pure ()

  go (VPath l x y) (VPath l' x' y') = do
    unify' l l'
    unify' x x'
    unify' y y'

  go (VLine l x y p) p' = do
    n <- VVar <$> newName
    unify' (p @@ n) (ielim l x y p' n)

  go p' (VLine l x y p) = do
    n <- VVar <$> newName
    unify' (ielim l x y p' n) (p @@ n)

  go (VPartial phi r) (VPartial phi' r')   = unify' phi phi' *> unify' r r'
  go (VPartialP phi r) (VPartialP phi' r') = unify' phi phi' *> unify' r r'

  go (VSub a phi u) (VSub a' phi' u') = traverse_ (uncurry unify') [(a, a'), (phi, phi'), (u, u')]
  go (VInc a phi u) (VInc a' phi' u') = traverse_ (uncurry unify') [(a, a'), (phi, phi'), (u, u')]

  go (VComp a phi u a0) (VComp a' phi' u' a0') =
    traverse_ (uncurry unify') [(a, a'), (phi, phi'), (u, u'), (a0, a0')]

  go (VGlueTy _ (force -> VI1) u _0) rhs = unify' (u @@ VReflStrict VI VI1) rhs
  go lhs (VGlueTy _ (force -> VI1) u _0) = unify' lhs (u @@ VReflStrict VI VI1)

  go (VGlueTy a phi u a0) (VGlueTy a' phi' u' a0') =
    traverse_ (uncurry unify') [(a, a'), (phi, phi'), (u, u'), (a0, a0')]

  go (VGlue a phi u a0 t x) (VGlue a' phi' u' a0' t' x') =
    traverse_ (uncurry unify') [(a, a'), (phi, phi'), (u, u'), (a0, a0'), (t, t'), (x, x')]

  go (VUnglue a phi u a0 x) (VUnglue a' phi' u' a0' x') =
    traverse_ (uncurry unify') [(a, a'), (phi, phi'), (u, u'), (a0, a0'), (x, x')]

  go (VSystem sys) rhs = goSystem unify' sys rhs
  go rhs (VSystem sys) = goSystem (flip unify') sys rhs

  go (VEqStrict a x y) (VEqStrict a' x' y') = traverse_ (uncurry unify') [(a, a'), (x, x'), (y, y')]
  go (VReflStrict a x) (VReflStrict a' x') = traverse_ (uncurry unify') [(a, a'), (x, x')]
  go _ VReflStrict{} = pure ()
  go VReflStrict{} _ = pure ()

  go x y
    | x == y = pure ()
    | otherwise =
      case (toDnf x, toDnf y) of
        (Just xs, Just ys) -> unify'Formula xs ys
        _ -> fail

  goSystem :: (Value -> Value -> ElabM ()) -> Map.Map Value Value -> Value -> ElabM ()
  goSystem k sys rhs = do
    let rhs_q = quote rhs
    env <- ask
    for_ (Map.toList sys) $ \(f, i) -> do
      let i_q = quote i
      for (truthAssignments f (getEnv env)) $ \e -> do
        k (eval' env{getEnv = e} i_q) (eval' env{getEnv = e} rhs_q)

  fail = throwElab $ NotEqual topa topb

  unify'Formula x y
    | compareDNFs x y = pure ()
    | otherwise = fail

trivialSystem :: Value -> Maybe Value
trivialSystem = go . force where
  go VSystem{} = Nothing
  go x = Just x

unify'Spine :: Projection -> Projection -> ElabM ()
unify'Spine (PApp a v) (PApp a' v')
  | a == a' = unify' v v'

unify'Spine PProj1 PProj1 = pure ()
unify'Spine PProj2 PProj2 = pure ()

unify'Spine (PIElim _ _ _ i) (PIElim _ _ _ j) = unify' i j
unify'Spine (POuc a phi u) (POuc a' phi' u') =
  traverse_ (uncurry unify') [(a, a'), (phi, phi'), (u, u')]

unify'Spine (PK a x p pr) (PK a' x' p' pr') =
  traverse_ (uncurry unify') [(a, a'), (x, x'), (p, p'), (pr, pr')]

unify'Spine (PJ a x p pr y) (PJ a' x' p' pr' y') =
  traverse_ (uncurry unify') [(a, a'), (x, x'), (p, p'), (pr, pr'), (y, y')]

unify'Spine _ _ = throwElab (NotEqual undefined undefined)

unify :: HasCallStack => Value -> Value -> ElabM ()
unify a b = unify' a b `catchElab` \(_ :: SomeException) -> liftIO $ throwIO (NotEqual a b)

isConvertibleTo :: Value -> Value -> ElabM (Term -> Term)
isConvertibleTo a b = isConvertibleTo (force a) (force b) where
  VPi Im d (Closure _v k) `isConvertibleTo` ty = do
    meta <- newMeta d
    cont <- k meta `isConvertibleTo` ty
    pure (\f -> cont (App Im f (quote meta)))
  VType `isConvertibleTo` VTypeω = pure id

  VPi p d (Closure _ k) `isConvertibleTo` VPi p' d' (Closure _ k') | p == p' = do
    wp <- d' `isConvertibleTo` d
    n <- newName
    wp_n <- eval (Lam Ex n (wp (Ref n)))

    wp' <- k (VVar n) `isConvertibleTo` k' (wp_n @@ VVar n)
    pure (\f -> Lam p n (wp' (App p f (wp (Ref n)))))

  VPath a x y `isConvertibleTo` VPi Ex d (Closure _ k') = do
    unify d VI
    nm <- newName
    wp <- isConvertibleTo (a @@ VVar nm) (k' (VVar nm))
    pure (\f -> Lam Ex nm (wp (IElim (quote a) (quote x) (quote y) f (Ref nm))))

  isConvertibleTo a b = do
    unify' a b
    pure id

newMeta :: Value -> ElabM Value
newMeta dom = do
  loc <- liftM2 (,) <$> asks currentFile <*> asks currentSpan
  n <- newName
  c <- liftIO $ newIORef Nothing
  let m = MV (getNameText n) c dom (flatten <$> loc)
      flatten (x, (y, z)) = (x, y, z)

  env <- asks getEnv

  t <- for (Map.toList env) $ \(n, _) -> pure $
    case n of
      Bound{} -> Just (PApp Ex (VVar n))
      _ -> Nothing

  pure (VNe (HMeta m) (Seq.fromList (catMaybes t)))

newName :: MonadIO m => m Name
newName = liftIO $ do
  x <- atomicModifyIORef _nameCounter $ \x -> (x + 1, x + 1)
  pure (Bound (T.pack (show x)) x)

newName' :: Name -> ElabM Name
newName' n = do
  ~(Bound _ x) <- newName
  pure (Bound (getNameText n) x)

_nameCounter :: IORef Int
_nameCounter = unsafePerformIO $ newIORef 0
{-# NOINLINE _nameCounter #-}

solveMeta :: MV -> Seq Projection -> Value -> ElabM ()
solveMeta m@(mvCell -> cell) sp rhs = do
  env <- ask
  names <- tryElab $ checkSpine Set.empty sp
  case names of
    Right names -> do
      checkScope (Set.fromList names) rhs
        `withNote` hsep [prettyTm (quote (VNe (HMeta m) sp)), pretty '≡', prettyTm (quote rhs)]
      let tm = quote rhs
          lam = eval' env $ foldr (Lam Ex) tm names
      liftIO . atomicModifyIORef' cell $ \case
        Just _ -> error "filled cell in solvedMeta"
        Nothing -> (Just lam, ())
    Left (_ :: SpineProjection) -> do
      liftIO . atomicModifyIORef' (unsolvedMetas env) $ \x -> (, ()) $
        case Map.lookup m x of
          Just qs -> Map.insert m ((sp, rhs):qs) x
          Nothing -> Map.insert m [(sp, rhs)] x

checkScope :: Set Name -> Value -> ElabM ()
checkScope scope (VNe h sp) =
  do
    case h of
      HVar v@Bound{} ->
        unless (v `Set.member` scope) . throwElab $
          NotInScope v
      HVar{} -> pure ()
      HCon{} -> pure ()
      HPCon{} -> pure ()
      HMeta{} -> pure ()
      HData{} -> pure ()
    traverse_ checkProj sp
  where
    checkProj (PApp _ t) = checkScope scope t
    checkProj (PIElim l x y i) = traverse_ (checkScope scope) [l, x, y, i]
    checkProj (PK l x y i) = traverse_ (checkScope scope) [l, x, y, i]
    checkProj (PJ l x y i j) = traverse_ (checkScope scope) [l, x, y, i, j]
    checkProj (POuc a phi u) = traverse_ (checkScope scope) [a, phi, u]
    checkProj PProj1 = pure ()
    checkProj PProj2 = pure ()

checkScope scope (GluedVl _ _p vl) = checkScope scope vl

checkScope scope (VLam _ (Closure n k)) =
  checkScope (Set.insert n scope) (k (VVar n))

checkScope scope (VPi _ d (Closure n k)) = do
  checkScope scope d
  checkScope (Set.insert n scope) (k (VVar n))

checkScope scope (VSigma d (Closure n k)) = do
  checkScope scope d
  checkScope (Set.insert n scope) (k (VVar n))

checkScope s (VPair a b) = traverse_ (checkScope s) [a, b]

checkScope _ VType = pure ()
checkScope _ VTypeω = pure ()

checkScope _ VI = pure ()
checkScope _ VI0 = pure ()
checkScope _ VI1 = pure ()

checkScope s (VIAnd x y) = traverse_ (checkScope s) [x, y]
checkScope s (VIOr x y)  = traverse_ (checkScope s) [x, y]
checkScope s (VINot x)   = checkScope s x

checkScope s (VPath line a b) = traverse_ (checkScope s) [line, a, b]
checkScope s (VLine _ _ _ line)   = checkScope s line

checkScope s (VPartial x y)  = traverse_ (checkScope s) [x, y]
checkScope s (VPartialP x y) = traverse_ (checkScope s) [x, y]
checkScope s (VSystem fs) =
  for_ (Map.toList fs) $ \(x, y) -> traverse_ (checkScope s) [x, y]

checkScope s (VSub a b c) = traverse_ (checkScope s) [a, b, c]
checkScope s (VInc a b c) = traverse_ (checkScope s) [a, b, c]
checkScope s (VComp a phi u a0) = traverse_ (checkScope s) [a, phi, u, a0]
checkScope s (VHComp a phi u a0) = traverse_ (checkScope s) [a, phi, u, a0]

checkScope s (VGlueTy a phi ty eq) = traverse_ (checkScope s) [a, phi, ty, eq]
checkScope s (VGlue a phi ty eq inv x) = traverse_ (checkScope s) [a, phi, ty, eq, inv, x]
checkScope s (VUnglue a phi ty eq vl) = traverse_ (checkScope s) [a, phi, ty, eq, vl]

checkScope s (VCase _ _ v _) = checkScope s v

checkScope s (VEqStrict a x y) = traverse_ (checkScope s) [a, x, y]
checkScope s (VReflStrict a x) = traverse_ (checkScope s) [a, x]

checkSpine :: Set Name -> Seq Projection -> ElabM [Name]
checkSpine scope (PApp Ex (VVar n@Bound{}) Seq.:<| xs)
  | n `Set.member` scope = throwElab $ NonLinearSpine n
  | otherwise = (n:) <$> checkSpine scope xs
checkSpine _     (p Seq.:<| _) = throwElab $ SpineProj p
checkSpine _     Seq.Empty = pure []

newtype NonLinearSpine = NonLinearSpine { getDupeName :: Name }
  deriving (Show, Typeable, Exception)

newtype SpineProjection = SpineProj { getSpineProjection :: Projection }
  deriving (Show, Typeable, Exception)

substituteIO :: Map.Map Name Value -> Value -> IO Value
substituteIO sub = substituteIO . force where
  substituteIO (VNe hd sp) = do
    sp' <- traverse (substituteSp sub) sp
    case hd of
      HMeta (mvCell -> cell) -> do
        solved <- liftIO $ readIORef cell
        case solved of
          Just vl -> substituteIO $ foldl applProj vl sp'
          Nothing -> pure $ VNe hd sp'
      HVar v ->
        case Map.lookup v sub of
          Just vl -> substituteIO $ foldl applProj vl sp'
          Nothing -> pure $ VNe hd sp'
      hd -> pure $ VNe hd sp'

  substituteIO (GluedVl h sp vl) = GluedVl h <$> traverse (substituteSp sub) sp <*> substituteIO vl

  substituteIO (VLam p (Closure s k)) = pure $ VLam p (Closure s (substitute (Map.delete s sub) . k))
  substituteIO (VPi p d (Closure s k)) = VPi p <$> substituteIO d <*> pure (Closure s (substitute (Map.delete s sub) . k))
  substituteIO (VSigma d (Closure s k)) = VSigma <$> substituteIO d <*> pure (Closure s (substitute (Map.delete s sub) . k))
  substituteIO (VPair a b) = VPair <$> substituteIO a <*> substituteIO b

  substituteIO (VPath line x y) = VPath <$> substituteIO line <*> substituteIO x <*> substituteIO y
  substituteIO (VLine line x y f) = VLine <$> substituteIO line <*> substituteIO x <*> substituteIO y <*> substituteIO f

  -- Sorts
  substituteIO VType = pure VType
  substituteIO VTypeω = pure VTypeω

  substituteIO VI = pure VI
  substituteIO VI0 = pure VI0
  substituteIO VI1 = pure VI1

  substituteIO (VIAnd x y) = iand <$> substituteIO x <*> substituteIO y
  substituteIO (VIOr x y) = ior <$> substituteIO x <*> substituteIO y
  substituteIO (VINot x) = inot <$> substituteIO x

  substituteIO (VPartial x y)  = VPartial <$> substituteIO x <*> substituteIO y
  substituteIO (VPartialP x y) = VPartialP <$> substituteIO x <*> substituteIO y
  substituteIO (VSystem fs) = do
    t <- for (Map.toList fs) $ \(a, b) -> (,) <$> substituteIO a <*> substituteIO b
    pure (mkVSystem (Map.fromList t))
  substituteIO (VSub a b c) = VSub <$> substituteIO a <*> substituteIO b <*> substituteIO c
  substituteIO (VInc a b c) = incS <$> substituteIO a <*> substituteIO b <*> substituteIO c
  substituteIO (VComp a b c d) = comp <$> substituteIO a <*> substituteIO b <*> substituteIO c <*> substituteIO d
  substituteIO (VHComp a b c d) = hComp <$> substituteIO a <*> substituteIO b <*> substituteIO c <*> substituteIO d

  substituteIO (VGlueTy a phi ty e)   = glueType <$> substituteIO a <*> substituteIO phi <*> substituteIO ty <*> substituteIO e
  substituteIO (VGlue a phi ty e t x) = glueElem <$> substituteIO a <*> substituteIO phi <*> substituteIO ty <*> substituteIO e <*> substituteIO t <*> substituteIO x
  substituteIO (VUnglue a phi ty e x) = unglue   <$> substituteIO a <*> substituteIO phi <*> substituteIO ty <*> substituteIO e <*> substituteIO x
  substituteIO (VCase env t x xs) = VCase env <$> substituteIO t <*> substituteIO x <*> pure xs
  substituteIO (VEqStrict a x y) = VEqStrict <$> substituteIO a <*> substituteIO x <*> substituteIO y
  substituteIO (VReflStrict a x) = VReflStrict <$> substituteIO a <*> substituteIO x

substitute :: Map Name Value -> Value -> Value
substitute sub = unsafePerformIO . substituteIO sub

substituteSp :: Map Name Value -> Projection -> IO Projection
substituteSp sub (PApp p x) = PApp p <$> substituteIO sub x
substituteSp sub (PIElim l x y i) = PIElim <$> substituteIO sub l <*> substituteIO sub x <*> substituteIO sub y <*> substituteIO sub i
substituteSp sub (PK l x y i) = PK <$> substituteIO sub l <*> substituteIO sub x <*> substituteIO sub y <*> substituteIO sub i
substituteSp sub (PJ l x y i j) = PJ <$> substituteIO sub l <*> substituteIO sub x <*> substituteIO sub y <*> substituteIO sub i <*> substituteIO sub j
substituteSp sub (POuc a phi u) = POuc <$> substituteIO sub a <*> substituteIO sub phi <*> substituteIO sub u
substituteSp _ PProj1 = pure PProj1
substituteSp _ PProj2 = pure PProj2

mkVSystem :: Map.Map Value Value -> Value
mkVSystem vals =
  let map' = Map.fromList (Map.toList vals >>= go)
      go (x, y) =
        case (force x, y) of
          (VI0, _) -> []
          (VIOr _ _, VSystem y) -> Map.toList y >>= go
          (a, b) -> [(a, b)]
  in case Map.lookup VI1 map' of
    Just x -> x
    Nothing -> VSystem map'

forceIO :: MonadIO m => Value -> m Value
forceIO mv@(VNe (HMeta (mvCell -> cell)) args) = do
  solved <- liftIO $ readIORef cell
  case solved of
    Just vl -> forceIO (foldl applProj vl args)
    Nothing -> pure mv
forceIO vl@(VSystem fs) =
  case Map.lookup VI1 fs of
    Just x -> forceIO x
    Nothing -> pure vl
forceIO (GluedVl _ _ vl) = forceIO vl
forceIO (VComp line phi u a0) = comp <$> forceIO line <*> forceIO phi <*> pure u <*> pure a0
forceIO (VHComp line phi u a0) = hComp <$> forceIO line <*> forceIO phi <*> pure u <*> pure a0
forceIO (VCase env rng v vs) = do
  env' <- liftIO emptyEnv
  r <- forceIO rng
  evalCase env'{getEnv=env} (r @@) <$> forceIO v <*> pure vs
forceIO x = pure x

force :: Value -> Value
force = unsafePerformIO . forceIO

applProj :: HasCallStack => Value -> Projection -> Value
applProj fun (PApp p arg)     = vApp p fun arg
applProj fun (PIElim l x y i) = ielim l x y fun i
applProj fun (POuc a phi u)   = outS a phi u fun
applProj fun (PK a x p pr)    = strictK a x p pr fun
applProj fun (PJ a x p pr y)  = strictJ a x p pr y fun
applProj fun PProj1           = vProj1 fun
applProj fun PProj2           = vProj2 fun

vApp :: HasCallStack => Plicity -> Value -> Value -> Value
vApp _ (VLam _ k) arg = clCont k arg
vApp p (VNe h sp) arg  = VNe h (sp Seq.:|> PApp p arg)
vApp p (GluedVl h sp vl) arg = GluedVl h (sp Seq.:|> PApp p arg) (vApp p vl arg)
vApp p (VSystem fs) arg = mkVSystem (fmap (flip (vApp p) arg) fs)
vApp p (VCase env rng sc branches) arg =
  VCase env (fun \x -> let VPi _ _ (Closure _ r) = rng @@ x in r arg) sc
    (map (projIntoCase (flip (App p) (quote arg))) branches)
vApp _ x _             = error $ "can't apply " ++ show (prettyTm (quote x))

(@@) :: HasCallStack => Value -> Value -> Value
(@@) = vApp Ex
infixl 9 @@

vProj1 :: HasCallStack => Value -> Value
vProj1 (VPair a _) = a
vProj1 (VNe h sp) = VNe h (sp Seq.:|> PProj1)
vProj1 (GluedVl h sp vl) = GluedVl h (sp Seq.:|> PProj1) (vProj1 vl)
vProj1 (VSystem fs) = VSystem (fmap vProj1 fs)
vProj1 (VInc (VSigma a _) b c) = incS a b (vProj1 c)
vProj1 (VCase env rng sc branches) =
  VCase env rng sc (map (projIntoCase Proj1) branches)
vProj1 x = error $ "can't proj1 " ++ show (prettyTm (quote x))

vProj2 :: HasCallStack => Value -> Value
vProj2 (VPair _ b) = b
vProj2 (VNe h sp)  = VNe h (sp Seq.:|> PProj2)
vProj2 (GluedVl h sp vl) = GluedVl h (sp Seq.:|> PProj2) (vProj2 vl)
vProj2 (VSystem fs) = VSystem (fmap vProj2 fs)
vProj2 (VInc (VSigma _ (Closure _ r)) b c) = incS (r (vProj1 c)) b (vProj2 c)
vProj2 (VCase env rng sc branches) =
  VCase env rng sc (map (projIntoCase Proj2) branches)
vProj2 x = error $ "can't proj2 " ++ show (prettyTm (quote x))