module Obs.NormalForm.Normalise

import Data.Vect

import Obs.Logging
import Obs.NormalForm
import Obs.Sort
import Obs.Substitution

-- Utilities -------------------------------------------------------------------

mergeName : String -> String -> String
mergeName "" s' = s'
mergeName "_" s' = s'
mergeName s s' = s

wkn : Vect k String -> (Fin n -> Logging ann (NormalForm m)) -> Fin (k + n) -> Logging ann (NormalForm (k + m))
wkn []            f = f
wkn (var :: vars) f =
  add
    (Logging ann . NormalForm)
    (pure $ Ntrl $ Var var FZ)
    (\i => pure $ rename !(wkn vars f i) FS)

export
expandContainerTy : Sort -> NormalForm n -> Sort -> Sort -> Constructor n
expandContainerTy s a s' s'' =
  let tagTy : Constructor n
      tagTy = Pi s (suc s') "i" a (cast s')
  in
  let posTy : Constructor (S n)
      posTy =
        Pi s (s' ~> suc s'') "i"
          (weaken 1 a)
          (Cnstr $ Pi s' (suc s'') "t"
            (Ntrl $ App (Var "tag" 1) (Ntrl $ Var "i" 0))
            (cast s''))
  in
  let nextTy : Constructor (2 + n)
      nextTy =
        Pi s (s' ~> suc s'' ~> s) "i"
          (weaken 2 a)
          (Cnstr $ Pi s' (suc s'' ~> s) "t"
            (Ntrl $ App (Var "tag" 2) (Ntrl $ Var "i" 0))
            (Cnstr $ Pi (suc s'') s "p"
              (Ntrl $ App (App (Var "pos" 2) (Ntrl $ Var "i" 1)) (Ntrl $ Var "p" 0))
              (weaken 5 a)))
  in
  Sigma (s ~> suc s') (lub (s ~> s' ~> suc s'') (s ~> s' ~> s'' ~> s)) "tag"
    (Cnstr tagTy) $
  Cnstr $ Sigma (s ~> s' ~> suc s'') (s ~> s' ~> s'' ~> s) "pos"
    (Cnstr posTy)
    (Cnstr nextTy)

-- Substitution ----------------------------------------------------------------

partial
substCnstr : Constructor n -> (Fin n -> Logging ann (NormalForm m)) -> Logging ann (Constructor m)
partial
substNtrl  : Neutral n -> (Fin n -> Logging ann (NormalForm m)) -> Logging ann (NormalForm m)
partial export
subst      : NormalForm n -> (Fin n -> Logging ann (NormalForm m)) -> Logging ann (NormalForm m)
partial export
subst1     : NormalForm n -> NormalForm (S n) -> Logging ann (NormalForm n)
partial export
doApp      : NormalForm n -> NormalForm n -> Logging ann (NormalForm n)
partial export
doFst      : NormalForm n -> Logging ann (NormalForm n)
partial export
doSnd      : NormalForm n -> Logging ann (NormalForm n)
partial export
doCase     : NormalForm n -> NormalForm n -> NormalForm n -> Logging ann (NormalForm n)
partial export
doTag      : NormalForm n -> Logging ann (NormalForm n)
partial export
doPosition : NormalForm n -> Logging ann (NormalForm n)
partial export
doNext     : NormalForm n -> Logging ann (NormalForm n)
partial export
doEqual    : NormalForm n -> NormalForm n -> NormalForm n -> Logging ann (NormalForm n)
partial
doEqualL   : Nat -> NormalForm n -> NormalForm n -> Logging ann (NormalForm n)
partial
doEqualR   : Nat -> Constructor n -> NormalForm n -> Logging ann (NormalForm n)
partial
doEqualU   : Nat -> Constructor n -> Constructor n -> Logging ann (NormalForm n)
partial export
doCastL    : NormalForm n -> NormalForm n -> NormalForm n -> Logging ann (NormalForm n)
partial
doCastR    : Constructor n -> NormalForm n -> NormalForm n -> Logging ann (NormalForm n)
partial
doCastU    : Constructor n -> Constructor n -> NormalForm n -> Logging ann (NormalForm n)
partial export
expandContainer : NormalForm n -> Logging ann (Constructor n)
partial export
contractContainer : NormalForm n -> Logging ann (Constructor n)

substCnstr (Sort s) f = pure $ Sort s
substCnstr (Pi s s' var a b) f = do
  a <- subst a f
  b <- subst b (wkn [var] f)
  pure (Pi s s' var a b)
substCnstr (Lambda var t) f = do
  t <- subst t (wkn [var] f)
  pure (Lambda var t)
substCnstr (Sigma s s' var a b) f = do
  a <- subst a f
  b <- subst b (wkn [var] f)
  pure (Sigma s s' var a b)
substCnstr (Pair t u) f = do
  t <- subst t f
  u <- subst u f
  pure (Pair t u)
substCnstr (Sum s s' a b) f = do
  a <- subst a f
  b <- subst b f
  pure (Sum s s' a b)
substCnstr (Left t) f = do
  t <- subst t f
  pure (Left t)
substCnstr (Right t) f = do
  t <- subst t f
  pure (Right t)
substCnstr (Container s a s' s'') f = do
  a <- subst a f
  pure (Container s a s' s'')
substCnstr (MkContainer t p u) f = do
  t <- subst t f
  p <- subst p f
  u <- subst u f
  pure (MkContainer t p u)
substCnstr Top f = pure $ Top
substCnstr Bottom  f = pure $ Bottom

substNtrl (Var var i) f = do
  Ntrl (Var var' j) <- f i
    | t => pure t
  pure (Ntrl (Var (mergeName var' var) j))
substNtrl (App t u) f = do
  t <- substNtrl t f
  u <- subst u f
  doApp t u
substNtrl (Fst t) f = do
  t <- substNtrl t f
  doFst t
substNtrl (Snd t) f = do
  t <- substNtrl t f
  doSnd t
substNtrl (Case t u t') f = do
  t <- substNtrl t f
  u <- subst u f
  t' <- subst t' f
  doCase t u t'
substNtrl (Tag t) f = do
  t <- substNtrl t f
  doTag t
substNtrl (Position t) f = do
  t <- substNtrl t f
  doPosition t
substNtrl (Next t) f = do
  t <- substNtrl t f
  doNext t
substNtrl Absurd f = pure $ Ntrl Absurd
substNtrl (Equal a t u) f = do
  a <- substNtrl a f
  t <- subst t f
  u <- subst u f
  doEqual a t u
substNtrl (EqualL i t u) f = do
  t <- substNtrl t f
  u <- subst u f
  doEqualL i t u
substNtrl (EqualR i t u) f = do
  t <- substCnstr t f
  u <- substNtrl u f
  doEqualR i t u
substNtrl (EqualU i t u) f = do
  t <- substCnstr t f
  u <- substCnstr u f
  doEqualU i t u
substNtrl (CastL a b t) f = do
  a <- substNtrl a f
  b <- subst b f
  t <- subst t f
  doCastL a b t
substNtrl (CastR a b t) f = do
  a <- substCnstr a f
  b <- substNtrl b f
  t <- subst t f
  doCastR a b t
substNtrl (CastU a b t) f = do
  a <- substCnstr a f
  b <- substCnstr b f
  t <- subst t f
  doCastU a b t

subst (Ntrl t)  f = substNtrl t f
subst (Cnstr t) f = map Cnstr $ substCnstr t f
subst Irrel     f = pure Irrel

subst1 t u = subst u (add (Logging ann . NormalForm) (pure t) (pure . point))

doApp (Ntrl t) u = pure $ Ntrl (App t u)
doApp Irrel u = pure $ Irrel
doApp (Cnstr (Lambda var t)) u = subst t (add (Logging ann . NormalForm) (pure u) (pure . point))
doApp (Cnstr t) u = inScope "bug" $ inScope "wrong constructor in apply" $ fatal t

doFst (Ntrl t) = pure $ Ntrl (Fst t)
doFst Irrel = pure $ Irrel
doFst (Cnstr (Pair t u)) = pure $ t
doFst (Cnstr t) = inScope "bug" $ inScope "wrong constructor in fst" $ fatal t

doSnd (Ntrl t) = pure $ Ntrl (Snd t)
doSnd Irrel = pure $ Irrel
doSnd (Cnstr (Pair t u)) = pure $ u
doSnd (Cnstr t) = inScope "bug" $ inScope "wrong constructor in snd" $ fatal t

doCase (Ntrl t) f g = pure $ Ntrl (Case t f g)
doCase Irrel f g = inScope "bug" $ inScope "wrong constructor in case" $ fatal "Irrel"
doCase (Cnstr (Left t)) f g = doApp f t
doCase (Cnstr (Right t)) f g = doApp g t
doCase (Cnstr t) f g = inScope "bug" $ inScope "wrong constructor in case" $ fatal t

doTag (Ntrl t) = pure $ Ntrl (Tag t)
doTag Irrel = pure $ Irrel
doTag (Cnstr (MkContainer t p f)) = pure t
doTag (Cnstr t) = inScope "bug" $ inScope "wrong constructor in tag" $ fatal t

doPosition (Ntrl t) = pure $ Ntrl (Position t)
doPosition Irrel = pure $ Irrel
doPosition (Cnstr (MkContainer t p f)) = pure p
doPosition (Cnstr t) = inScope "bug" $ inScope "wrong constructor in position" $ fatal t

doNext (Ntrl t) = pure $ Ntrl (Next t)
doNext Irrel = pure $ Irrel
doNext (Cnstr (MkContainer t p f)) = pure f
doNext (Cnstr t) = inScope "bug" $ inScope "wrong constructor in next" $ fatal t

doEqual (Ntrl a) t u = pure $ Ntrl (Equal a t u)
doEqual Irrel t u = inScope "bug" $ inScope "wrong type over equal" $ fatal "Irrel"
doEqual (Cnstr (Sort Prop)) t u =
  pure $
  Cnstr (Sigma Prop Prop "_"
    (Cnstr $ Pi Prop Prop "_" t (weaken 1 u))
    (weaken 1 $ Cnstr $ Pi Prop Prop "_" u (weaken 1 t)))
doEqual (Cnstr (Sort (Set k))) t u = doEqualL k t u
doEqual (Cnstr (Pi s s' var a b)) t u = do
  eqLhs <- doApp (weaken 1 t) (Ntrl $ Var var FZ)
  eqRhs <- doApp (weaken 1 u) (Ntrl $ Var var FZ)
  eq <- doEqual b eqLhs eqRhs -- b in Set because Pi in Set.
  pure $ Cnstr (Pi s Prop var a eq)
doEqual (Cnstr (Sigma s s' var a b)) t u = do
  t1 <- doFst t
  u1 <- doFst u
  eq1 <- case s of
    Prop => pure $ Cnstr Top
    (Set i) => doEqual a t1 u1
  eq2 <- case s' of
    Prop => pure $ Cnstr Top
    (Set i) => do
      bt <- subst1 t1 b
      bu <- subst1 u1 b
      t2 <- doSnd t
      t2 <- doCastL bt bu t2
      u2 <- doSnd u
      doEqual bu t2 u2
  pure $ Cnstr (Sigma Prop Prop "_" eq1 (weaken 1 eq2))
doEqual (Cnstr (Sum s s' a b)) t u = do
  let x = Ntrl $ Var "x" 1
  let y = Ntrl $ Var "y" 0
  ll <- doEqual (weaken 2 a) x y
  lr <- pure $ Cnstr Bottom
  rl <- pure $ Cnstr Bottom
  rr <- doEqual (weaken 2 b) x y
  fBody <- doCase (weaken 1 u) (Cnstr (Lambda "y" ll)) (Cnstr (Lambda "y" lr))
  gBody <- doCase (weaken 1 u) (Cnstr (Lambda "y" rl)) (Cnstr (Lambda "y" rr))
  doCase t (Cnstr (Lambda "x" fBody)) (Cnstr (Lambda "x" gBody))
doEqual (Cnstr (Container s a s' s'')) t u = do
  let containerTy = expandContainerTy s a s' s''
  t <- expandContainer t
  u <- expandContainer u
  doEqual (Cnstr containerTy) (Cnstr t) (Cnstr u)
doEqual (Cnstr Top) t u = pure $ Cnstr Top
doEqual (Cnstr Bottom) t u = pure $ Cnstr Top
doEqual (Cnstr a) t u = inScope "bug" $ inScope "wrong type under equal" $ fatal a

doEqualL i (Ntrl t) u = pure $ Ntrl (EqualL i t u)
doEqualL i Irrel u = inScope "bug" $ inScope "wrong type under equalL" $ fatal "Irrel"
doEqualL i (Cnstr t) u = doEqualR i t u

doEqualR i t (Ntrl u) = pure $ Ntrl (EqualR i t u)
doEqualR i t Irrel = inScope "bug" $ inScope "wrong type under equalR" $ fatal "Irrel"
doEqualR i t (Cnstr u) = doEqualU i t u

doEqualU i (Sort s) (Sort s') = pure $ Cnstr Top -- have suc s = i = suc s', and suc injective
doEqualU i (Pi s s' _ a b) (Pi l l' var a' b') = case (s == s' && l == l') of
  False => pure $ Cnstr Bottom
  True => do
    eqLhs <- doEqual (cast s) a' a
    let x = Ntrl $ Var var FZ
    b <- subst b (add (Logging ann . NormalForm) (doCastL (weaken 1 a') (weaken 1 a) x) (pure . Ntrl . Var "" . FS))
    eqRhs <- doEqual (cast s') b b'
    pure $ Cnstr (Sigma Prop Prop "" eqLhs $ weaken 1 $ Cnstr (Pi s Prop var a' eqRhs))
doEqualU i (Sigma s s' _ a b) (Sigma l l' var a' b') = case (s == s' && l == l') of
  False => pure $ Cnstr Bottom
  True => do
    eqLhs <- doEqual (cast s) a' a
    let x = Ntrl $ Var var FZ
    b <- subst b (add (Logging ann . NormalForm) (doCastL (weaken 1 a') (weaken 1 a) x) (pure . Ntrl . Var "" . FS))
    eqRhs <- doEqual (cast s') b b'
    pure $ Cnstr (Sigma Prop Prop "" eqLhs $ weaken 1 $ Cnstr (Pi s Prop var a' eqRhs))
doEqualU i (Sum s s' a b) (Sum l l' a' b') = case (s == l && s' == l') of
  False => pure $ Cnstr Bottom
  True => do
    eqLhs <- doEqual (cast s) a a'
    eqRhs <- doEqual (cast s) b b'
    pure $ Cnstr (Sigma Prop Prop "" eqLhs (weaken 1 eqRhs))
doEqualU i (Container s a s' s'') (Container l a' l' l'') = case (s == l && s' == l' && s'' == l'') of
  False => pure $ Cnstr Bottom
  True => doEqual (cast s) a a'
doEqualU i t u = pure $ Ntrl (EqualU i t u) -- assumption: only logical values reach this far

doCastL (Ntrl a) b t = pure $ Ntrl (CastL a b t)
doCastL Irrel b t = inScope "bug" $ inScope "wrong type for cast" $ fatal "Irrel"
doCastL (Cnstr a) b t = doCastR a b t

doCastR a (Ntrl b) t = pure $ Ntrl (CastR a b t)
doCastR a Irrel t = inScope "bug" $ inScope "wrong type for cast" $ fatal "Irrel"
doCastR a (Cnstr b) t = doCastU a b t

doCastU (Sort s) (Sort s') t = pure t
doCastU (Pi s s' _ a b) (Pi l l' var a' b') t = do
  let x' = Ntrl $ Var var FZ
  let x = doCastL (weaken 1 a') (weaken 1 a) x'
  b <- subst b (add (Logging ann . NormalForm) x (pure . Ntrl . Var "" . FS))
  b' <- subst b' (add (Logging ann . NormalForm) (pure x') (pure . Ntrl . Var "" . FS))
  fx <- doApp (weaken 1 t) !x
  cast <- doCastL b b' fx
  pure $ Cnstr (Lambda var cast)
doCastU (Sigma s s' _ a b) (Sigma l l' var a' b') t = do
  t1 <- doFst t
  t2 <- doSnd t
  t1' <- doCastL a a' t1
  b <- subst1 t1 b
  b' <- subst1 t1' b'
  t2' <- doCastL b b' t2
  pure $ Cnstr (Pair t1' t2')
doCastU (Sum s s' a b) (Sum l l' a' b') t = do
  let x = Ntrl $ Var "x" 0
  castL <- doCastL (weaken 1 a) (weaken 1 a') x
  castR <- doCastL (weaken 1 b) (weaken 1 b') x
  doCase t (Cnstr (Lambda "x" (Cnstr (Left castL)))) (Cnstr (Lambda "x" (Cnstr (Right castR))))
doCastU (Container s a s' s'') (Container l b l' l'') t = do
  t <- expandContainer t
  let a = expandContainerTy s a s' s''
  let b = expandContainerTy l b l' l''
  t <- doCastU a b (Cnstr t)
  t <- contractContainer t
  pure $ Cnstr t
doCastU Top Top t = pure Irrel
doCastU Bottom Bottom t = pure Irrel
doCastU a b t = pure $ Ntrl (CastU a b t)

expandContainer t = do
  tag <- doTag t
  pos <- doPosition t
  next <- doNext t
  pure $ Pair tag (Cnstr $ Pair pos next)

contractContainer t = do
  tag <- doFst t
  t <- doSnd t
  pos <- doFst t
  next <- doSnd t
  pure $ MkContainer tag pos next

-- More utilities --------------------------------------------------------------

export
tagTy : Sort -> NormalForm n -> Sort -> Constructor n
tagTy s a s' = Pi s (suc s') "i" a (cast s')

export
positionTy : Sort -> NormalForm n -> Sort -> NormalForm n -> Sort -> Logging ann (Constructor n)
positionTy s a s' tag s'' = do
  let i = Var "i" 0
  tagI <- doApp (weaken 1 tag) (Ntrl i)
  pure $ Pi s (s' ~> suc s'') "i" a (Cnstr $ Pi s' (suc s'') "t" tagI (cast s''))

export
nextTy : Sort -> NormalForm n -> Sort -> NormalForm n -> Sort -> NormalForm n -> Logging ann (Constructor n)
nextTy s a s' tag s'' pos = do
  let i = Var "i" 0
  tagI <- doApp (weaken 1 tag) (Ntrl i)
  let t = Var "t" 0
  posI <- doApp (weaken 1 pos) (Ntrl i)
  posIT <- doApp (weaken 1 posI) (Ntrl t)
  pure $ Pi s (s' ~> s'' ~> s) "i" a $
  Cnstr $ Pi s' (s'' ~> s) "t" tagI $
  Cnstr $ Pi s'' s "p" posIT (weaken 3 a)