module Quox.Typing

import public Quox.Syntax
import public Quox.Context

import Data.Nat
import Data.SortedMap
import Control.Monad.Reader
import Control.Monad.State


%default total

public export
data DContext : Nat -> Type where
  DNil  : DContext 0
  DBind : DContext d -> DContext (S d)
  DEq   : Dim d -> Dim d -> DContext d -> DContext d

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

public export
QContext : Nat -> Type
QContext = Triangle' Qty

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


public export
record Global where
  constructor MkGlobal
  type, term : forall d, n. Term d n

public export
Globals : Type
Globals = SortedMap Name Global


public export
record TyContext (d, n : Nat) where
  constructor MkTyContext
  globals : Globals
  dctx : DContext d
  tctx : TContext d n
  qctx : QContext n

%name TyContext ctx


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


namespace TyContext
  export
  extendTy : Term d n -> QOutput n -> TyContext d n -> TyContext d (S n)
  extendTy s rhos = {tctx $= (:< s), qctx $= (:< rhos)}

  export
  extendDim : TyContext d n -> TyContext (S d) n
  extendDim = {dctx $= DBind, tctx $= pushD}

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



namespace QOutput
  export
  (+) : QOutput n -> QOutput n -> QOutput n
  (+) = zipWith (+)

  export
  (*) : Qty -> QOutput n -> QOutput n
  (*) pi = map (pi *)


namespace Zero
  public export
  data IsZero : QOutput n -> Type where
    LinZ  : IsZero [<]
    SnocZ : IsZero qctx -> IsZero (qctx :< Zero)

  export
  isZeroIrrel : {qctx : _} -> (0 _ : IsZero qctx) -> IsZero qctx
  isZeroIrrel {qctx = [<]}    LinZ      = LinZ
  isZeroIrrel {qctx = _ :< _} (SnocZ x) = SnocZ $ isZeroIrrel x

  export
  zero : Context _ n -> Subset (QOutput n) IsZero
  zero [<]        = Element [<] LinZ
  zero (ctx :< _) = let zeroN = zero ctx in
                    Element (zeroN.fst :< Zero) (SnocZ zeroN.snd)


namespace Only
  public export
  data IsOnly : QOutput n -> Var n -> Type where
    Here  : IsZero qctx   -> IsOnly (qctx :< One) VZ
    There : IsOnly qctx i -> IsOnly (qctx :< Zero) (VS i)

  export
  isOnlyIrrel : {qctx, i : _} -> (0 _ : IsOnly qctx i) -> IsOnly qctx i
  isOnlyIrrel {i = VZ}     (Here  z) = Here  $ isZeroIrrel z
  isOnlyIrrel {i = (VS i)} (There o) = There $ isOnlyIrrel o

  export
  only : Context _ n -> (i : Var n) ->
         Subset (QOutput n) (\qctx => IsOnly qctx i)
  only (ctx :< _) VZ     =
    let zeroN = zero ctx in
    Element (zeroN.fst :< One) (Here zeroN.snd)
  only (ctx :< _) (VS i) =
    let onlyN = only ctx i in
    Element (onlyN.fst :< Zero) (There onlyN.snd)


namespace Universe
  public export
  data LT : Universe -> Universe -> Type where
    Fin : k `LT` l -> U k `LT` U l
    Any : U _ `LT` UAny

  export
  Uninhabited (UAny `LT` j) where
    uninhabited _ impossible

  export
  isLT : (i, j : Universe) -> Dec (i `LT` j)
  isLT (U i) (U j) with (i `isLT` j)
    _ | Yes prf   = Yes $ Fin prf
    _ | No contra = No  $ \(Fin prf) => contra prf
  isLT (U _) UAny = Yes Any
  isLT UAny  _    = No uninhabited


namespace Lookup
  public export
  data IsLookup :
      (tctx : TContext d n) ->
      (i    : Var n) ->
      (ty   : Term d from) ->
      (by   : Shift from n) -> Type
    where
      Here  : IsLookup {tctx=(tctx :< ty), i=VZ, ty, by=(SS SZ)}
      There : IsLookup {tctx, i, ty, by} ->
              IsLookup {tctx=(tctx :< ty'), i=(VS i), ty, by=(SS by)}

  public export
  record Lookup'
    {0 d, n : Nat}
    (0 tctx : TContext d n)
    (0 qctx : QContext n)
    (0 i : Var n)
  where
    constructor MkLookup
    tmout, tyout : QOutput n
    type : Term d from
    by   : Shift from n
    0 only : IsOnly tmout i
    0 look : IsLookup tctx i type by

  public export
  lookup' : (tctx : TContext d n) -> (qctx : QContext n) ->
            (i : Var n) -> Lookup' tctx qctx i
  lookup' tctx@(_ :< _) (_ :< tyout) VZ =
    MkLookup {only = (only tctx VZ).snd, look = Here,
              tyout = tyout :< Zero, _}
  lookup' (tctx :< _) (qctx :< _) (VS i) =
    let inner = lookup' tctx qctx i in
    MkLookup {only = There inner.only, look = There inner.look,
              tyout = inner.tyout :< Zero, _}


  public export %inline
  Lookup : TyContext d n -> Var n -> Type
  Lookup ctx i = Lookup' ctx.tctx ctx.qctx i

  public export %inline
  lookup : (ctx : TyContext d n) -> (i : Var n) -> Lookup ctx i
  lookup ctx = lookup' ctx.tctx ctx.qctx



mutual
  public export
  data HasTypeT :
      {0 d, n : Nat} ->
      (ctx  : TyContext d n) ->
      (subj : Term d n) ->
      (ty   : Term d n) ->
      (tmout, tyout : QOutput n) ->
      Type
    where
      TYPE :
        (0 wf : IsWf ctx) ->
        (lt : k `LT` l) ->
        (0 ztm : IsZero tmout) -> (0 zty : IsZero tyout) ->
        HasTypeT {ctx, subj = TYPE k, ty = TYPE l, tmout, tyout}
      Pi :
        (tyA : HasTypeT {ctx, subj = a, ty = TYPE l, tmout = tmoutA, tyout}) ->
        (tyB : HasTypeT {ctx = extendTy a tmoutA ctx, subj = b, ty = TYPE l,
                         tmout = tmoutB :< qty, tyout = tyout :< Zero}) ->
        (0 zty : IsZero tyout) ->
        HasTypeT {ctx, subj = Pi qtm qty x a b, ty = TYPE l,
                  tmout = tmoutA + tmoutB, tyout}
      PushT :
        (0 eq : ty' = pushSubstsT' ty) ->
        (j : HasTypeT {ctx, subj, ty = ty', tmout, tyout}) ->
        HasTypeT {ctx, subj, ty, tmout, tyout}

  public export
  data HasTypeE :
      {0 d, n : Nat} ->
      (ctx  : TyContext d n) ->
      (subj : Elim d n) ->
      (ty   : Term d n) ->
      (tmout, tyout : QOutput n) ->
      Type
    where
      F :
        (0 wf : IsWf ctx) ->
        (g : Global) ->
        (0 eq : lookup x ctx.globals = Just g) ->
        (0 ztm : IsZero tmout) -> (0 zty : IsZero tyout) ->
        HasTypeE {ctx, subj = F x, ty = g.type, tmout, tyout}
      B :
        {0 ctx : TyContext d n} ->
        (0 wf : IsWf ctx) ->
        (look : Lookup ctx i) ->
        (0 eq : ty = look.type // look.by) ->
        HasTypeE {ctx, subj = B i, ty, tmout = look.tmout, tyout = look.tyout}
      Ann :
        HasTypeT {ctx, subj = tm,       ty, tmout, tyout} ->
        HasTypeE {ctx, subj = tm :# ty, ty, tmout, tyout}
      PushE :
        (0 eq : ty' = pushSubstsT' ty) ->
        HasTypeE {ctx, subj, ty = ty', tmout, tyout} ->
        HasTypeE {ctx, subj, ty, tmout, tyout}

  public export
  data IsWf' :
      (globals : Globals) ->
      (dctx : DContext d) ->
      (tctx : TContext d n) ->
      (qctx : QContext n) -> Type
    where
      NIL  : IsWf' globals dctx [<] [<]
      BIND :
        IsWf' {globals, dctx, tctx, qctx} ->
        HasTypeT {ctx = MkTyContext {globals, dctx, tctx, qctx},
                  subj, ty = TYPE l, tmout, tyout} ->
        IsWf' {globals, dctx, tctx = tctx :< subj, qctx = qctx :< tmout}

  public export %inline
  IsWf : TyContext d n -> Type
  IsWf (MkTyContext {globals, dctx, tctx, qctx}) =
    IsWf' {globals, dctx, tctx, qctx}

public export
record WfContext (d, n : Nat) where
  constructor MkWfContext
  ctx  : TyContext d n
  0 wf : IsWf ctx


public export
record TypeTerm {0 d, n : Nat} (ctx : TyContext d n) where
  constructor MkTypeTerm
  term : Term d n
  univ : Universe
  tmout, tyout : QOutput n
  isType : HasTypeT {ctx, subj = term, ty = TYPE univ, tmout, tyout}