{-# 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 {-# SOURCE #-} Elab.Eval.Formula
import Elab.Monad

import GHC.Stack

import Presyntax.Presyntax (Plicity(..))

import Syntax.Pretty
import Syntax

import System.IO.Unsafe ( unsafePerformIO )

import {-# SOURCE #-} Elab.WiredIn
import Debug (traceM, traceDocM)
import Prettyprinter (pretty, (<+>))

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

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

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) = pure $ VComp a b c d
zonkIO (VHComp a b c d) = pure $ VHComp a b c 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 (idkT, 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 (idkT, 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 (idkT, 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)

idkT :: NFType
idkT = VVar (Defined (T.pack "dunno") (negate 1))

isIdkT :: NFType -> Bool
isIdkT (VVar (Defined (T.unpack -> "dunno") (negate -> 1))) = True
isIdkT _ = False

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 _ _ (VVar ((== trueCaseSentinel) -> True)) _ = VI1

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

-- This is a great big HACK; When we see a system [ case x of ... => p
-- ], we somehow need to make the 'case x of ...' become VI1. The way we
-- do this is by substituting x/trueCaseSentinel in truthAssignments,
-- and then making case trueCaseSentinel of ... => VI1 always.
trueCaseSentinel :: Name
trueCaseSentinel = Bound (T.pack "sentinel for true cases") (-1000)

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

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

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

unify' :: HasCallStack => Bool -> Value -> Value -> ElabM ()

unify' cs topa@(GluedVl h sp a) topb@(GluedVl h' sp' b)
  | h == h', length sp == length sp' =
    traverse_ (uncurry (unify'Spine cs topa topb)) (Seq.zip sp sp')
      `catchElab` \(_ :: SomeException) -> unify' cs a b

unify' canSwitch 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 topa@(VNe (HPCon _ _ x) sp) topb@(VNe (HPCon _ _ y) sp')
    | x == y = traverse_ (uncurry (unify'Spine canSwitch topa topb)) (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 _ _ b) (VCase e' _ _ b') = do
    env <- ask
    let
      go (_, _, a) (_, _, b)
        | a == b = pure ()
        | otherwise = unify' canSwitch (eval' env{getEnv=moreDefinedFrom e e' <$> e} a) (eval' env{getEnv=moreDefinedFrom e e' <$> 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' canSwitch a y
    traverse_ go b

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

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

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

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

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

  go (VSigma d (Closure n k)) (VSigma d' (Closure _ k')) = do
    t <- VVar <$> newName' n
    unify' canSwitch d d'
    unify' canSwitch (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' canSwitch l l'
    unify' canSwitch x x'
    unify' canSwitch y y'

  go (VLine l x y p) p' = do
    n <- VVar <$> newName' (Bound (T.singleton 'i') (- 1))
    unify' canSwitch (p @@ n) (ielim l x y p' n)

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

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

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

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

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

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

  go (VGlueTy a phi u a0) (VGlueTy a' phi' u' a0') =
    traverse_ (uncurry (unify' canSwitch)) [(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' canSwitch)) [(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' canSwitch)) [(a, a'), (phi, phi'), (u, u'), (a0, a0'), (x, x')]

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

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

  go (VINot x) (VINot y) = unify' canSwitch x y

  go x y =
    case (toDnf x, toDnf y) of
      (Just xs, Just ys) -> unify'Formula xs ys
      _ ->
        if canSwitch
          then goDumb x y
          else fail

  goDumb (VIOr a b) (VIOr a' b') = unify' canSwitch a a' *> goDumb b b'
  goDumb (VIAnd a b) (VIAnd a' b') = unify' canSwitch a a' *> goDumb b b'
  goDumb x y = switch $ unify' False x y

  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

moreDefinedFrom :: Map Name (NFType, Value) -> Map Name (NFType, Value) -> (NFType, Value) -> (NFType, Value)
moreDefinedFrom map1 map2 ours@(_, VNe (HVar name) _) =
  case Map.lookup name map1 of
    Just (_, VNe HVar{} _) -> map2's
    Just (ty, x) -> (ty, x)
    Nothing -> map2's
  where
    map2's = case Map.lookup name map2 of
      Just (_, VNe HVar{} _) -> ours
      Just (ty, x) -> (ty, x)
      Nothing -> ours
moreDefinedFrom _ _ ours = ours

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

unify'Spine :: Bool -> Value -> Value -> Projection -> Projection -> ElabM ()
unify'Spine cs _ _ (PApp a v) (PApp a' v')
  | a == a' = unify' cs v v'

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

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

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

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

unify'Spine _ x y _ _ = throwElab (NotEqual x y)

unify :: HasCallStack => Value -> Value -> ElabM ()
unify x y = shallowly $ go x y where
  go topa@(GluedVl h sp a) topb@(GluedVl h' sp' b)
    | h == h', length sp == length sp' =
      traverse_ (uncurry (unify'Spine True topa topb)) (Seq.zip sp sp')
        `catchElab` \(_ :: SomeException) -> unify' True a b
  go a b = unify' True 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' True a b
    pure id

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

  env <- asks getEnv

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

  let
    ts = Map.fromList $ fmap (\(_, n, (t, _)) -> (n, t)) t
    t' = fmap (\(x, _, _) -> x) t

  um <- asks unsolvedMetas
  liftIO . atomicModifyIORef um $ \um -> (Map.insert (m ts) [] um, ())

  pure (VNe (HMeta (m ts)) (Seq.fromList t'))

newMeta :: Value -> ElabM Value
newMeta = newMeta' False

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 Seq.Empty (VNe (HMeta m') Seq.Empty) | m == m' = pure ()
solveMeta m@(mvCell -> cell) sp rhs = do
  env <- ask
  names <- tryElab $ checkSpine Set.empty sp
  case names of
    Right names -> do
      scope <- tryElab $ checkScope m (Set.fromList names) rhs
      case scope of
        Right () -> do
          let tm = quote rhs
              lam = eval' env $ foldr (Lam Ex) tm names
          liftIO . atomicModifyIORef (unsolvedMetas env) $ \mp -> (Map.delete m mp, ())
          liftIO . atomicModifyIORef' cell $ \case
            Just _ -> error "filled cell in solvedMeta"
            Nothing -> (Just lam, ())
        Left (_ :: MetaException) -> abort env
    Left (_ :: MetaException) -> abort env
  where
    abort env =
      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 :: MV -> Set Name -> Value -> ElabM ()
checkScope mv scope (VNe h sp) =
  do
    case h of
      HVar v@Bound{} ->
        unless (v `Set.member` scope) . throwElab $
          ScopeCheckingFail v
      HVar{} -> pure ()
      HCon{} -> pure ()
      HPCon{} -> pure ()
      HMeta m' -> when (mv == m') $ throwElab $ CircularSolution mv
      HData{} -> pure ()
    traverse_ checkProj sp
  where
    checkProj (PApp _ t) = checkScope mv scope t
    checkProj (PIElim l x y i) = traverse_ (checkScope mv scope) [l, x, y, i]
    checkProj (PK l x y i) = traverse_ (checkScope mv scope) [l, x, y, i]
    checkProj (PJ l x y i j) = traverse_ (checkScope mv scope) [l, x, y, i, j]
    checkProj (POuc a phi u) = traverse_ (checkScope mv scope) [a, phi, u]
    checkProj PProj1 = pure ()
    checkProj PProj2 = pure ()

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

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

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

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

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

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

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

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

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

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

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

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

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

checkScope mv s (VEqStrict a x y) = traverse_ (checkScope mv s) [a, x, y]
checkScope mv s (VReflStrict a x) = traverse_ (checkScope mv 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 []

data MetaException = NonLinearSpine { getDupeName :: Name }
                   | SpineProj { getSpineProjection :: Projection } 
                   | CircularSolution { getCycle :: MV }
                   | ScopeCheckingFail { outOfScope :: Name }
  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 $ foldl applProj (VNe hd mempty) 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

  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 _ (VLine _ _ _ (VLam _ k)) arg = clCont k arg
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))