module Quox.Typing.Context

import Quox.Syntax
import Quox.Context

%default total


public export
TContext : Type -> Nat -> Nat -> Type
TContext q d = Context (Term q d)

public export
QOutput : Type -> Nat -> Type
QOutput = Context'

public export
NContext : Nat -> Type
NContext = Context' BaseName

public export
DimAssign : Nat -> Type
DimAssign = Context' DimConst


public export
record TyContext q d n where
  constructor MkTyContext
  dctx   : DimEq d
  dnames : NContext d
  tctx   : TContext q d n
  tnames : NContext n
%name TyContext ctx


public export
record EqContext q n where
  constructor MkEqContext
  -- (only used for errors; not needed by the actual equality test)
  dassign : DimAssign dimLen
  dnames  : NContext dimLen
  tctx    : TContext q 0 n
  tnames  : NContext n
%name EqContext ctx


namespace TContext
  export %inline
  pushD : TContext q d n -> TContext q (S d) n
  pushD tel = map (// shift 1) tel

  export %inline
  zeroFor : IsQty q => Context tm n -> QOutput q n
  zeroFor ctx = zero <$ ctx

namespace TyContext
  public export %inline
  empty : TyContext q 0 0
  empty = MkTyContext {dctx = new, dnames = [<], tctx = [<], tnames = [<]}

  export %inline
  extendTyN : Telescope (\n => (BaseName, Term q d n)) from to ->
              TyContext q d from -> TyContext q d to
  extendTyN xss ctx =
    let (xs, ss) = unzip xss in {tctx $= (. ss), tnames $= (. xs)} ctx

  export %inline
  extendTy : BaseName -> Term q d n -> TyContext q d n -> TyContext q d (S n)
  extendTy x s = extendTyN [< (x, s)]

  export %inline
  extendDim : BaseName -> TyContext q d n -> TyContext q (S d) n
  extendDim x = {dctx $= (:<? Nothing), dnames $= (:< x), tctx $= pushD}

  export %inline
  eqDim : Dim d -> Dim d -> TyContext q d n -> TyContext q d n
  eqDim p q = {dctx $= set p q}


namespace QOutput
  parameters {auto _ : IsQty q}
    export %inline
    (+) : QOutput q n -> QOutput q n -> QOutput q n
    (+) = zipWith (+)

    export %inline
    (*) : q -> QOutput q n -> QOutput q n
    (*) pi = map (pi *)

    export %inline
    zeroFor : TyContext q _ n -> QOutput q n
    zeroFor ctx = zeroFor ctx.tctx


export
makeDAssign : DSubst d 0 -> DimAssign d
makeDAssign (Shift SZ)   = [<]
makeDAssign (K e ::: th) = makeDAssign th :< e

export
makeEqContext : TyContext q d n -> DSubst d 0 -> EqContext q n
makeEqContext ctx th = MkEqContext {
  dassign = makeDAssign th,
  dnames  = ctx.dnames,
  tctx    = map (// th) ctx.tctx,
  tnames  = ctx.tnames
}

namespace EqContext
  public export %inline
  empty : EqContext q 0
  empty = MkEqContext [<] [<] [<] [<]

  export %inline
  extendTyN : Telescope (\n => (BaseName, Term q 0 n)) from to ->
              EqContext q from -> EqContext q to
  extendTyN tel ctx =
    let (xs, ss) = unzip tel in {tctx $= (. ss), tnames $= (. xs)} ctx

  export %inline
  extendTy : BaseName -> Term q 0 n -> EqContext q n -> EqContext q (S n)
  extendTy x s = extendTyN [< (x, s)]

  export %inline
  extendDim : BaseName -> DimConst -> EqContext q n -> EqContext q n
  extendDim x e ctx = {dassign $= (:< e), dnames $= (:< x)} ctx

  export
  toTyContext : (e : EqContext q n) -> (d ** TyContext q d n)
  toTyContext (MkEqContext {dassign, dnames, tctx, tnames})
      with (lengthPrf0 dnames)
    _ | Element d eq =
      (d ** MkTyContext {
        dctx   = rewrite eq in fromGround dassign,
        dnames = rewrite eq in dnames,
        tctx   = map (// shift0 d) tctx,
        tnames
      })