path: root/src/Obs/NormalForm/Normalise.idr

module Obs.NormalForm.Normalise

import Data.Bool
import Data.List
import Data.Nat

import Decidable.Equality

import Obs.Logging
import Obs.NormalForm
import Obs.Substitution
import Obs.Universe

import Text.PrettyPrint.Prettyprinter

%default total

-- Aliases ---------------------------------------------------------------------

public export 0
LogConstructor : Type -> Unsorted.Family Relevance
LogConstructor ann ctx = Logging ann (Constructor ctx)

public export 0
LogNormalForm : Type -> Sorted.Family Relevance
LogNormalForm ann b ctx = Logging ann (NormalForm b ctx)

0
LogDeclForm : Type -> Unsorted.Family Relevance
LogDeclForm ann ctx = Logging ann (DeclForm ctx)

0
LogNormalForm' : Type -> Sorted.Family Relevance
LogNormalForm' ann b ctx = Either (Logging ann (NormalForm b ctx)) (Elem b ctx)

-- Copied and specialised from Obs.Substitution
lift : (ctx : List (r ** Maybe String))
  -> Map ctx' (LogNormalForm' ann) ctx''
  -> Map (map DPair.fst ctx ++ ctx') (LogNormalForm' ann) (map DPair.fst ctx ++ ctx'')
lift [] f = f
lift ((s ** y) :: ctx) f = add (LogNormalForm' ann)
  (Left $ pure $ point y Here)
  (\i => bimap (\t => pure (rename !t There)) There (lift ctx f i))

-- Normalisation ---------------------------------------------------------------

subst' : NormalForm ~|> Hom (LogNormalForm' ann) (LogNormalForm ann)

export
subst1 : {s' : _} -> NormalForm s ctx -> NormalForm s' (s :: ctx) -> LogNormalForm ann s' ctx
subst1 t u = subst' u (add (LogNormalForm' ann) (Left $ pure t) Right)

export
map1 : {s' : _} ->
  NormalForm s (s :: ctx) ->
  NormalForm s' (s :: ctx) ->
  LogNormalForm ann s' (s :: ctx)
map1 t u = subst' u (add (LogNormalForm' ann) (Left $ pure t) (Right . There))

export
MkLambda : {rel, domainRel : _} ->
  (var : Maybe String) ->
  (body : NormalForm rel (domainRel :: ctx)) ->
  NormalForm (function domainRel rel) ctx
MkLambda {rel = Irrelevant} var body = Irrel
MkLambda {rel = Relevant} var body = Cnstr $ Lambda {var, body, domainRel}

export
doApp : {domainRel : _} ->
  (fun : NormalForm (function domainRel codomainRel) ctx) ->
  (arg : NormalForm domainRel ctx) ->
  LogNormalForm ann codomainRel ctx

doApp (Ntrl fun) arg = pure $ Ntrl $ App {argRel = _, fun, arg}
doApp (Cnstr (Lambda {domainRel = domainRel', var, body})) arg = do
  let Yes Refl = decEq domainRel domainRel'
    | No _ => fatal "internal relevance mismatch"

  subst1 arg body
doApp (Cnstr t) arg = inScope "wrong constructor for apply" $ fatal t
doApp Irrel arg = pure Irrel

export
MkPair : {indexRel, elementRel : Relevance}
  -> (first : NormalForm indexRel ctx)
  -> (second : NormalForm elementRel ctx)
  -> NormalForm (pair indexRel elementRel) ctx
MkPair {indexRel = Irrelevant, elementRel = Irrelevant, first, second} = Irrel
MkPair {indexRel = Irrelevant, elementRel = Relevant, first, second} =
  Cnstr $ Pair {indexRel = Irrelevant, elementRel = Relevant, prf = Relevant, first, second}
MkPair {indexRel = Relevant, elementRel, first, second} =
  Cnstr $ Pair {indexRel = Relevant, elementRel, prf = Relevant, first, second}

export
doFst : (firstRel, secondRel : _) ->
  (arg : NormalForm (pair firstRel secondRel) ctx) ->
  LogNormalForm ann firstRel ctx

doFst Irrelevant secondRel arg = pure Irrel
doFst Relevant secondRel (Ntrl arg) = pure $ Ntrl $ First {secondRel, arg}
doFst Relevant secondRel (Cnstr (Pair {indexRel = Relevant, elementRel, prf, first, second})) =
  pure first
doFst Relevant secondRel (Cnstr t) = inScope "wrong constructor for fst" $ fatal t

export
doSnd : (firstRel, secondRel : _) ->
  (arg : NormalForm (pair firstRel secondRel) ctx) ->
  LogNormalForm ann secondRel ctx

doSnd firstRel Irrelevant arg = pure Irrel
doSnd firstRel Relevant arg =
  let arg' : NormalForm Relevant ctx
      arg' = rewrite sym $ pairRelevantRight firstRel in arg
  in case arg' of
  Ntrl arg => pure $ Ntrl $ Second {firstRel, arg}
  Cnstr (Pair {indexRel, elementRel = Relevant, prf, first, second}) => pure second
  Cnstr t => inScope "wrong constructor for snd" $ fatal t

export
doContainer : ContainerTy ctx -> TypeNormalForm ctx
doContainer container =
  Cnstr $ Pi
    { domainSort = container.inputSort
    , codomainSort =
        max (succ container.shapeSort)
            (container.shapeSort ~>
             max (succ container.positionSort) (container.positionSort ~> container.outputSort))
    , domain = MkDecl Nothing container.input
    , codomain = Cnstr $ Sigma
      { indexSort = succ container.shapeSort
      , elementSort = container.shapeSort ~>
                      max (succ container.positionSort)
                          (container.positionSort ~> container.outputSort)
      , index = MkDecl (Just "Shape") (cast container.shapeSort)
      , element = Cnstr $ Pi
        { domainSort = container.shapeSort
        , codomainSort = max (succ container.positionSort)
                             (container.positionSort ~> container.outputSort)
        , domain = MkDecl Nothing (Ntrl $ Var "Shape" Here)
        , codomain = Cnstr $ Sigma
          { indexSort = succ container.positionSort
          , elementSort = container.positionSort ~> container.outputSort
          , index = MkDecl (Just "Position") (cast container.positionSort)
          , element = Cnstr $ Pi
            { domainSort = container.positionSort
            , codomainSort = container.outputSort
            , domain = MkDecl Nothing (Ntrl $ Var "Position" Here)
            , codomain = weaken [_, _, _, _, _] container.output
            }
          }
        }
      }
    }

export
MkContainer : (inputRel, shapeRel : Relevance) ->
  {outputRel : Relevance} ->
  (shape : TypeNormalForm ctx) ->
  (position : TypeNormalForm ctx) ->
  (next : NormalForm outputRel ctx) ->
  LogNormalForm ann Relevant ctx
MkContainer {inputRel, shapeRel, outputRel, shape, position, next} = do
  shape <- doApp (Sorted.weaken [inputRel] shape) (point Nothing Here)
  position <- doApp (weaken [shapeRel, inputRel] position) (point Nothing (There Here))
  position <- doApp position (point Nothing Here)
  next <- doApp (weaken [shapeRel, inputRel] next) (point Nothing (There Here))
  next <- doApp next (point Nothing Here)
  pure $ MkLambda
    { var = Nothing
    , body = MkPair
      { first = shape
      , second = MkLambda
        { var = Nothing
        , body = MkPair
          { first = position
          , second = next
          }
        }
      }
    }

export
doShape : (inputRel : Relevance) ->
  (arg : NormalForm Relevant ctx) ->
  LogNormalForm ann Relevant ctx
doShape inputRel arg =
  let inputIndex : NormalForm inputRel (inputRel :: ctx)
      inputIndex = point Nothing Here
  in do
  inputIndexed <- doApp (weaken [_] arg) inputIndex
  body <- doFst Relevant Relevant inputIndexed

  pure $ MkLambda {var = Nothing, body}

export
doPosition : (inputRel, shapeRel, outputRel : Relevance) ->
  (arg : NormalForm Relevant ctx) ->
  LogNormalForm ann Relevant ctx
doPosition {inputRel, shapeRel, outputRel, arg} =
  let inputIndex : NormalForm inputRel (inputRel :: ctx)
      inputIndex = point Nothing Here
  in
  let shapeIndex : NormalForm shapeRel (shapeRel :: inputRel :: ctx)
      shapeIndex = point Nothing Here
  in do
  inputIndexed <- doApp (weaken [_] arg) inputIndex

  positionNextPair <- doSnd Relevant Relevant inputIndexed
  let positionNextPair = weaken [_] positionNextPair

  shapeIndexed <- doApp positionNextPair shapeIndex
  body <- doFst Relevant outputRel shapeIndexed

  pure $ MkLambda
    { var = Nothing
    , body = MkLambda {var = Nothing, body}
    }

export
doNext : (inputRel, shapeRel, outputRel : Relevance) ->
  (arg : NormalForm Relevant ctx) ->
  LogNormalForm ann outputRel ctx
doNext {inputRel, shapeRel, outputRel, arg} =
  let inputIndex : NormalForm inputRel (inputRel :: ctx)
      inputIndex = point Nothing Here
  in
  let shapeIndex : NormalForm shapeRel (shapeRel :: inputRel :: ctx)
      shapeIndex = point Nothing Here
  in do
  inputIndexed <- doApp (weaken [_] arg) inputIndex

  positionNextPair <- doSnd Relevant Relevant inputIndexed
  let positionNextPair = weaken [_] positionNextPair

  shapeIndexed <- doApp positionNextPair shapeIndex
  body <- doSnd Relevant outputRel shapeIndexed

  pure $ MkLambda
    { var = Nothing
    , body = MkLambda {var = Nothing, body}
    }

export
doSem : (container : ContainerTy ctx) ->
  (predSort : Universe) ->
  (pred : TypeNormalForm (relevance container.outputSort :: ctx)) ->
  (arg : NormalForm Relevant ctx) ->
  LogNormalForm ann Relevant ctx
doSem {container, predSort, pred, arg} = do
  let inputVar = Nothing
  let shapeVar = Nothing
  let positionVar = Nothing

  shape <- doShape {inputRel = relevance container.inputSort, arg}
  let shape = Sorted.weaken [_] shape

  shapeType <- doApp shape (point inputVar Here)
  let shapeIndex = MkDecl shapeVar shapeType

  position <- doPosition
    { inputRel = relevance container.inputSort
    , shapeRel = relevance container.shapeSort
    , outputRel = relevance container.outputSort
    , arg
    }
  let position = Sorted.weaken [_,_] position

  positionType <- doApp position (point inputVar (There Here))
  positionType <- doApp positionType (point shapeVar Here)
  let positionDomain = MkDecl positionVar positionType

  next <- doNext
    { inputRel = relevance container.inputSort
    , shapeRel = relevance container.shapeSort
    , outputRel = relevance container.outputSort
    , arg
    }
  let next = Sorted.weaken [_,_,_] next

  next <- doApp next (point inputVar (There (There Here)))
  next <- doApp next (point shapeVar (There Here))
  next <- doApp next (point positionVar Here)

  let f = add (LogNormalForm' ann) (Left $ pure next) (Right . There . There . There)
  codomain <- subst' pred f

  pure $ MkLambda
    { var = inputVar
    , domainRel = relevance container.inputSort
    , body = Cnstr $ Sigma
      { indexSort = container.shapeSort
      , elementSort = container.positionSort ~> predSort
      , index = shapeIndex
      , element = Cnstr $ Pi
        { domainSort = container.positionSort
        , codomainSort = predSort
        , domain = positionDomain
        , codomain
        }
      }
    }

export
doIf : {rel : _} ->
  (discriminant : NormalForm Relevant ctx) ->
  (true, false : NormalForm rel ctx) ->
  LogNormalForm ann rel ctx
doIf {rel = Irrelevant} discriminant true false = pure Irrel
doIf {rel = Relevant} (Ntrl discriminant) true false = pure $ Ntrl $ If {discriminant, true, false}
doIf {rel = Relevant} (Cnstr True) true false = pure true
doIf {rel = Relevant} (Cnstr False) true false = pure false
doIf {rel = Relevant} (Cnstr t) true false = inScope "wrong constructor for if" $ fatal t

export
doAbsurd : (rel : _) -> NormalForm rel ctx
doAbsurd Irrelevant = Irrel
doAbsurd Relevant = Ntrl Absurd

export
doCast : (rel : _) ->
  (oldType, newType : TypeNormalForm ctx) ->
  (value : NormalForm rel ctx) ->
  LogNormalForm ann rel ctx

doCastHelper : (oldType, newType : Constructor ctx) ->
  (value : NormalForm Relevant ctx) ->
  LogNormalForm ann Relevant ctx
doCastHelper (Universe {s}) (Universe {s = s'}) value = pure value
doCastHelper
  oldType@(Pi {domainSort, codomainSort, domain, codomain})
  newType@(Pi
    { domainSort = domainSort'
    , codomainSort = codomainSort'
    , domain = domain'
    , codomain = codomain'
    })
  value =
  let y : NormalForm (relevance domainSort) (relevance domainSort :: ctx)
      y = point domain.var Here
  in do

  let Yes Refl = decEq (domainSort, codomainSort) (domainSort', codomainSort')
    | No _ => pure $ Ntrl $ CastStuck {oldType, newType, value}

  let domainType = weaken [relevance domainSort] domain.type
  let domainType' = weaken [relevance domainSort] domain'.type

  x <- doCast
    { rel = relevance domainSort
    , oldType = assert_smaller oldType domainType'
    , newType = assert_smaller newType domainType
    , value = y
    }

  codomainType <- map1 x codomain
  codomainType' <- map1 y codomain'

  call <- doApp (weaken [relevance domainSort] value) x
  body <- doCast
    { rel = Relevant
    , oldType = assert_smaller oldType codomainType
    , newType = assert_smaller newType codomainType'
    , value = call
    }

  pure $ Cnstr $ Lambda {domainRel = relevance domainSort, var = Nothing, body}
doCastHelper
  oldType@(Sigma {indexSort = indexSort@(Set k), elementSort, index, element})
  newType@(Sigma
    { indexSort = indexSort'
    , elementSort = elementSort'
    , index = index'
    , element = element'
    })
  value = do
  let Yes Refl = decEq (indexSort, elementSort) (indexSort', elementSort')
    | No _ => pure $ Ntrl $ CastStuck {oldType, newType, value}

  first <- doFst Relevant (relevance elementSort) value
  second <- doSnd Relevant (relevance elementSort) value

  first' <- doCast
    { rel = Relevant
    , oldType = assert_smaller oldType index.type
    , newType = assert_smaller newType index'.type
    , value = first
    }

  elementType <- subst1 first element
  elementType' <- subst1 first' element'

  second' <- doCast
    { rel = relevance elementSort
    , oldType = assert_smaller oldType elementType
    , newType = assert_smaller newType elementType'
    , value = second
    }

  pure $ Cnstr $ Pair
    { indexRel = Relevant
    , elementRel = relevance elementSort
    , prf = Relevant
    , first = first'
    , second = second'
    }
doCastHelper
  oldType@(Sigma {indexSort = Prop, elementSort = elementSort@(Set k), index, element})
  newType@(Sigma
    { indexSort = Prop
    , elementSort = elementSort'
    , index = index'
    , element = element'
    })
  value = do
  let Yes Refl = decEq elementSort elementSort'
    | No _ => pure $ Ntrl $ CastStuck {oldType, newType, value}

  first <- doFst Irrelevant Relevant value
  second <- doSnd Irrelevant Relevant value

  first' <- doCast
    { rel = Irrelevant
    , oldType = assert_smaller oldType index.type
    , newType = assert_smaller newType index'.type
    , value = first
    }

  elementType <- subst1 first element
  elementType' <- subst1 first' element'

  second' <- doCast
    { rel = Relevant
    , oldType = assert_smaller oldType elementType
    , newType = assert_smaller newType elementType'
    , value = second
    }

  pure $ Cnstr $ Pair
    { indexRel = Irrelevant
    , elementRel = Relevant
    , prf = Relevant
    , first = first'
    , second = second'
    }
doCastHelper Bool Bool value = pure value
doCastHelper oldType newType value = pure $ Ntrl $ CastStuck {oldType, newType, value}

doCast Irrelevant oldType newType value = pure Irrel
doCast Relevant (Ntrl oldType) newType value = pure $ Ntrl $ CastL {oldType, newType, value}
doCast Relevant (Cnstr oldType) (Ntrl newType) value =
  pure $ Ntrl $ CastR {oldType, newType, value}
doCast Relevant (Cnstr oldType) (Cnstr newType) value = doCastHelper oldType newType value

export
doEqual : (rel : _) ->
  (type : TypeNormalForm ctx) ->
  (left, right : NormalForm rel ctx) ->
  LogNormalForm ann Relevant ctx

equalHelper : (left, right : Constructor ctx) -> LogNormalForm ann Relevant ctx
equalHelper (Universe {s}) (Universe {s = s'}) = pure (Cnstr Top)
equalHelper
  left@(Pi {domainSort, codomainSort, domain, codomain})
  right@(Pi
    { domainSort = domainSort'
    , codomainSort = codomainSort'
    , domain = domain'
    , codomain = codomain'
    }) =
  let y : NormalForm (relevance domainSort) (relevance domainSort :: Irrelevant :: ctx)
      y = point domain.var Here
  in do

  let Yes Refl = decEq (domainSort, codomainSort) (domainSort', codomainSort')
    | No _ => pure $ Ntrl $ EqualStuck {left, right}

  domainEqual <- doEqual
    { rel = Relevant
    , type = cast domainSort
    , left = assert_smaller right domain'.type
    , right = assert_smaller left domain.type
    }

  let domainType = Sorted.weaken [relevance domainSort, Irrelevant] domain.type
  let domainType' = Sorted.weaken [relevance domainSort, Irrelevant] domain'.type

  x <- doCast
    { rel = relevance domainSort
    , oldType = domainType'
    , newType = domainType
    , value = y
    }

  codomainType <- map1 x (rename codomain (add Elem Here (There . There)))
  codomainType' <- map1 y (rename codomain' (add Elem Here (There . There)))

  codomainEqual <- doEqual
    { rel = Relevant
    , type = cast codomainSort
    , left = assert_smaller left codomainType
    , right = assert_smaller right codomainType'
    }

  let returnElement = Cnstr $ Pi
       { domainSort = domainSort
       , codomainSort = Prop
       , domain = MkDecl Nothing (Sorted.weaken [Irrelevant] domain.type)
       , codomain = codomainEqual
       }

  pure $ Cnstr $ Sigma
    { indexSort = Prop
    , elementSort = Prop
    , index = MkDecl Nothing domainEqual
    , element = returnElement
    }
equalHelper
  left@(Sigma {indexSort, elementSort, index, element})
  right@(Sigma
    { indexSort = indexSort'
    , elementSort = elementSort'
    , index = index'
    , element = element'
    }) =
  let x : NormalForm (relevance indexSort) (relevance indexSort :: Irrelevant :: ctx)
      x = point index.var Here
  in do

  let Yes Refl = decEq (indexSort, elementSort) (indexSort', elementSort')
    | No _ => pure $ Ntrl $ EqualStuck {left, right}

  indexEqual <- doEqual
    { rel = Relevant
    , type = cast indexSort
    , left = assert_smaller left index.type
    , right = assert_smaller right index'.type
    }

  let indexType = Sorted.weaken [relevance indexSort, Irrelevant] index.type
  let indexType' = Sorted.weaken [relevance indexSort, Irrelevant] index'.type

  y <- doCast
    { rel = relevance indexSort
    , oldType = indexType
    , newType = indexType'
    , value = x
    }

  elementType <- map1 x (rename element (add Elem Here (There . There)))
  elementType' <- map1 y (rename element' (add Elem Here (There . There)))

  elementEqual <- doEqual
    { rel = Relevant
    , type = cast elementSort
    , left = assert_smaller left elementType
    , right = assert_smaller right elementType'
    }

  let returnElement = Cnstr $ Pi
       { domainSort = indexSort
       , codomainSort = Prop
       , domain = MkDecl Nothing (Sorted.weaken [Irrelevant] index.type)
       , codomain = elementEqual
       }

  pure $ Cnstr $ Sigma
    { indexSort = Prop
    , elementSort = Prop
    , index = MkDecl Nothing indexEqual
    , element = returnElement
    }
equalHelper Bool Bool = pure (Cnstr Top)
equalHelper left right = pure $ Ntrl $ EqualStuck {left, right}

doEqualType : (left, right : TypeNormalForm ctx) -> LogNormalForm ann Relevant ctx
doEqualType (Ntrl left) right = pure $ Ntrl $ EqualL {left, right}
doEqualType (Cnstr left) (Ntrl right) = pure $ Ntrl $ EqualR {left, right}
doEqualType (Cnstr left) (Cnstr right) = equalHelper left right

doEqual Irrelevant type left right = pure $ Cnstr Top
doEqual Relevant (Ntrl type) left right = pure $ Ntrl $ Equal {type, left, right}
doEqual Relevant (Cnstr (Universe {s = Prop})) left right = do
  let leftToRight = Cnstr $ Pi
        { domainSort = Prop
        , codomainSort = Prop
        , domain = MkDecl Nothing left
        , codomain = weaken [Irrelevant] right
        }
  let rightToLeft = Cnstr $ Pi
        { domainSort = Prop
        , codomainSort = Prop
        , domain = MkDecl Nothing right
        , codomain = weaken [Irrelevant] left
        }
  pure $ Cnstr $ Sigma
    { indexSort = Prop
    , elementSort = Prop
    , index = MkDecl Nothing leftToRight
    , element = Sorted.weaken [Irrelevant] rightToLeft
    }
doEqual Relevant (Cnstr (Universe {s = Set _})) left right = doEqualType left right
doEqual Relevant (Cnstr (Pi {domainSort, codomainSort, domain, codomain})) left right =
  let var : NormalForm (relevance domainSort) (relevance domainSort :: ctx)
      var = point domain.var Here
  in do

  leftApp <- doApp (weaken [relevance domainSort] left) var
  rightApp <- doApp (weaken [relevance domainSort] right) var

  equality <- doEqual
    { rel = Relevant
    , type = codomain
    , left = assert_smaller left leftApp
    , right = assert_smaller right rightApp
    }

  pure $ Cnstr $ Pi {domainSort, codomainSort = Prop, domain, codomain = equality}
doEqual
  Relevant
  (Cnstr (Sigma {indexSort = indexSort@(Set _), elementSort, index, element}))
  left
  right = do
  leftFirst <- doFst Relevant (relevance elementSort) left
  rightFirst <- doFst Relevant (relevance elementSort) right
  leftSecond <- doSnd Relevant (relevance elementSort) left
  rightSecond <- doSnd Relevant (relevance elementSort) right

  leftEquality <- doEqual
    { rel = Relevant
    , type = index.type
    , left = assert_smaller left leftFirst
    , right = assert_smaller right rightFirst
    }

  leftElementType <- subst1 leftFirst element
  rightElementType <- subst1 rightFirst element

  leftSecond <- doCast (relevance elementSort) leftElementType rightElementType leftSecond

  rightEquality <- doEqual
    { rel = relevance elementSort
    , type = rightElementType
    , left = assert_smaller left leftSecond
    , right = assert_smaller right rightSecond
    }

  pure $ Cnstr $ Sigma
    { indexSort = Prop
    , elementSort = Prop
    , index = MkDecl Nothing leftEquality
    , element = Sorted.weaken [Irrelevant] rightEquality
    }
doEqual
  Relevant
  (Cnstr (Sigma {indexSort = Prop, elementSort = elementSort@(Set _), index, element}))
  left
  right = do
  leftFirst <- doFst Irrelevant Relevant left
  rightFirst <- doFst Irrelevant Relevant right
  leftSecond <- doSnd Irrelevant Relevant left
  rightSecond <- doSnd Irrelevant Relevant right

  leftEquality <- doEqual
    { rel = Irrelevant
    , type = index.type
    , left = assert_smaller left leftFirst
    , right = assert_smaller right rightFirst
    }

  leftElementType <- subst1 leftFirst element
  rightElementType <- subst1 rightFirst element

  leftSecond <- doCast Relevant leftElementType rightElementType leftSecond

  rightEquality <- doEqual
    { rel = Relevant
    , type = rightElementType
    , left = assert_smaller left leftSecond
    , right = assert_smaller right rightSecond
    }

  pure $ Cnstr $ Sigma
    { indexSort = Prop
    , elementSort = Prop
    , index = MkDecl Nothing leftEquality
    , element = Sorted.weaken [Irrelevant] rightEquality
    }
doEqual Relevant (Cnstr Bool) left right = do
  whenLeftTrue <- doIf right (Cnstr Top) (Cnstr Bottom)
  whenLeftFalse <- doIf right (Cnstr Bottom) (Cnstr Top)
  doIf left whenLeftTrue whenLeftFalse
doEqual Relevant (Cnstr t) left right = inScope "wrong constructor for equal" $ fatal t

substDecl : DeclForm ~|> Hom (LogNormalForm' ann) (LogDeclForm ann)
substDecl (MkDecl var type) f = pure (MkDecl var !(subst' type f))

substCnstr : Constructor ~|> Hom (LogNormalForm' ann) (LogConstructor ann)
substCnstr (Universe {s}) f = pure (Universe {s})
substCnstr (Pi {domainSort, codomainSort, domain, codomain}) f = do
  domain <- substDecl domain f
  codomain <- subst' codomain (lift [(_ ** domain.var)] f)
  pure (Pi {domainSort, codomainSort, domain, codomain})
substCnstr (Lambda {domainRel, var, body}) f = do
  body <- subst' body (lift [(_ ** var)] f)
  pure (Lambda {domainRel, var, body})
substCnstr (Sigma {indexSort, elementSort, index, element}) f = do
  index <- substDecl index f
  element <- subst' element (lift [(_ ** index.var)] f)
  pure (Sigma {indexSort, elementSort, index, element})
substCnstr (Pair {indexRel, elementRel, prf, first, second}) f = do
  first <- subst' first f
  second <- subst' second f
  pure (Pair {indexRel, elementRel, prf, first, second})
substCnstr Bool f = pure Bool
substCnstr True f = pure True
substCnstr False f = pure False
substCnstr Top f = pure Top
substCnstr Bottom f = pure Bottom

substNtrl : Neutral ~|> Hom (LogNormalForm' ann) (LogNormalForm ann Relevant)
substNtrl (Var {var, i}) f = case f i of
  Left t => t
  Right i => pure $ Ntrl $ Var {var, i}
substNtrl (App {argRel, fun, arg}) f = do
  fun <- substNtrl fun f
  arg <- subst' arg f
  assert_total (doApp fun arg)
substNtrl (First {secondRel, arg}) f = do
  arg <- substNtrl arg f
  doFst Relevant secondRel arg
substNtrl (Second {firstRel, arg}) f = do
  arg <- substNtrl arg f
  let arg = rewrite pairRelevantRight firstRel in arg
  doSnd firstRel Relevant arg
substNtrl (If {discriminant, true, false}) f = do
  discriminant <- substNtrl discriminant f
  true <- subst' true f
  false <- subst' false f
  doIf discriminant true false
substNtrl Absurd f = pure (doAbsurd Relevant)
substNtrl (Equal {type, left, right}) f = do
  type <- substNtrl type f
  left <- subst' left f
  right <- subst' right f
  assert_total (doEqual Relevant type left right)
substNtrl (EqualL {left, right}) f = do
  left <- substNtrl left f
  right <- subst' right f
  assert_total (doEqualType left right)
substNtrl (EqualR {left, right}) f = do
  left <- substCnstr left f
  right <- substNtrl right f
  assert_total (doEqualType (Cnstr left) right)
substNtrl (EqualStuck {left, right}) f = do
  left <- substCnstr left f
  right <- substCnstr right f
  assert_total (doEqualType (Cnstr left) (Cnstr right))
substNtrl (CastL {oldType, newType, value}) f = do
  oldType <- substNtrl oldType f
  newType <- subst' newType f
  value <- subst' value f
  assert_total (doCast Relevant oldType newType value)
substNtrl (CastR {oldType, newType, value}) f = do
  oldType <- substCnstr oldType f
  newType <- substNtrl newType f
  value <- subst' value f
  assert_total (doCast Relevant (Cnstr oldType) newType value)
substNtrl (CastStuck {oldType, newType, value}) f = do
  oldType <- substCnstr oldType f
  newType <- substCnstr newType f
  value <- subst' value f
  assert_total (doCast Relevant (Cnstr oldType) (Cnstr newType) value)

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

export
subst : NormalForm ~|> Hom (LogNormalForm ann) (LogNormalForm ann)
subst t f = subst' t (Left . f)

export
strengthen : (ctx' : List Relevance) ->
  Map ctx' (LogNormalForm' ann) ctx ->
  Map (ctx' ++ ctx) (LogNormalForm' ann) ctx
strengthen [] f = Right
strengthen (rel :: ctx) f = add (LogNormalForm' ann) (f Here) (strengthen ctx (f . There))

-- Container Utilities ---------------------------------------------------------

public export
containerSort : (container : ContainerTy ctx) -> Universe
containerSort container =
  container.inputSort ~>
  max (succ container.shapeSort)
      (container.shapeSort ~>
       max (succ container.positionSort) (container.positionSort ~> container.outputSort))

export
matchContainer : (type : TypeNormalForm ctx) -> Logging ann (ContainerTy ctx)
matchContainer type@(Cnstr (Pi
  { domainSort = inputSort
  , codomainSort = inputCodomainSort
  , domain = MkDecl _ input
  , codomain = Cnstr (Sigma
    { indexSort = succShapeSort@(Set _)
    , elementSort = shapeElementSort
    , index = MkDecl _ (Cnstr (Universe shapeSort))
    , element = Cnstr (Pi
      { domainSort = shapeSort'
      , codomainSort = shapeCodomainSort
      , domain = MkDecl _ (Ntrl (Var _ Here))
      , codomain = Cnstr (Sigma
        { indexSort = succPositionSort@(Set _)
        , elementSort = positionElementSort
        , index = MkDecl _ (Cnstr (Universe positionSort))
        , element = Cnstr (Pi
          { domainSort = positionSort'
          , codomainSort = outputSort
          , domain = MkDecl _ (Ntrl (Var _ Here))
          , codomain = output'
          })
        })
      })
    })
  })) = do
  let Yes Refl = decEq
                   ( inputCodomainSort
                   , succShapeSort
                   , shapeElementSort
                   , shapeSort'
                   , shapeCodomainSort
                   , succPositionSort
                   , positionElementSort
                   , positionSort')
                   ( max (succ shapeSort)
                         (shapeSort ~> max (succ positionSort) (positionSort ~> outputSort))
                   , succ shapeSort
                   , shapeSort ~> max (succ positionSort) (positionSort ~> outputSort)
                   , shapeSort
                   , max (succ positionSort) (positionSort ~> outputSort)
                   , succ positionSort
                   , positionSort ~> outputSort
                   , positionSort)
    | No _ => typeMismatch (pretty "container") (pretty type)

  output <- subst' output' $
    strengthen [_, _, _, _, _] $
    const $ Left $ typeMismatch (pretty "container") (pretty type)
  pure (MkContainerTy {inputSort, shapeSort, positionSort, outputSort, input, output})
matchContainer type = typeMismatch (pretty "container") (pretty type)

export
containerShapeType : ContainerTy ctx -> TypeNormalForm ctx
containerShapeType container =
  Cnstr $ Pi
    { domainSort = container.inputSort
    , codomainSort = succ container.shapeSort
    , domain = MkDecl Nothing container.input
    , codomain = cast container.shapeSort
    }

export
containerPositionType : (container : ContainerTy ctx) ->
  (shape : TypeNormalForm (relevance container.inputSort :: ctx)) ->
  TypeNormalForm ctx
containerPositionType container shape =
  Cnstr $ Pi
    { domainSort = container.inputSort
    , codomainSort = container.shapeSort ~> succ container.positionSort
    , domain = MkDecl Nothing container.input
    , codomain = Cnstr $ Pi
      { domainSort = container.shapeSort
      , codomainSort = succ container.positionSort
      , domain = MkDecl Nothing shape
      , codomain = cast container.positionSort
      }
    }

export
containerNextType : (container : ContainerTy ctx) ->
  (shape : TypeNormalForm (relevance container.inputSort :: ctx)) ->
  (position :
    TypeNormalForm (relevance container.shapeSort :: relevance container.inputSort :: ctx)) ->
  TypeNormalForm ctx
containerNextType container shape position =
  Cnstr $ Pi
    { domainSort = container.inputSort
    , codomainSort = container.shapeSort ~> container.positionSort ~> container.outputSort
    , domain = MkDecl Nothing container.input
    , codomain = Cnstr $ Pi
      { domainSort = container.shapeSort
      , codomainSort = container.positionSort ~> container.outputSort
      , domain = MkDecl Nothing shape
      , codomain = Cnstr $ Pi
        { domainSort = container.positionSort
        , codomainSort = container.outputSort
        , domain = MkDecl Nothing position
        , codomain = weaken [_, _, _] container.output
        }
      }
    }

export
containerSemType : (container : ContainerTy ctx) -> (predSort : Universe) -> TypeNormalForm ctx
containerSemType container predSort =
  Cnstr $ Pi
    { domainSort = container.inputSort
    , codomainSort = succ (max container.shapeSort (container.positionSort ~> predSort))
    , domain = MkDecl Nothing container.input
    , codomain = cast $ max container.shapeSort (container.positionSort ~> predSort)
    }