module Core.LogRel

import Core.Declarative
import Core.Environment
import Core.Generic
import Core.Reduction
import Core.Term
import Core.Term.Substitution
import Core.Term.Thinned
import Core.Term.NormalForm
import Core.Thinning

%prefix_record_projections off

-- Levels ----------------------------------------------------------------------

data Level = Small | Large

%name Level l

public export
data LT : Level -> Level -> Type where
  LTSmall : Small `LT` Large

Uninhabited (l `LT` Small) where uninhabited prf impossible
Uninhabited (Large `LT` l) where uninhabited prf impossible

-- Logical Relation ------------------------------------------------------------

public export
record LogicalRelation (eq : Equality) where
  constructor MkLogRel
  0 TypeRed : forall n. Env n -> Term n -> Type
  0 TypeEq : forall n. (env : Env n) -> (a, b : Term n) -> TypeRed {n} env a -> Type
  0 TermRed : forall n. (env : Env n) -> (t, a : Term n) -> TypeRed {n} env a -> Type
  0 TermEq : forall n. (env : Env n) -> (t, u, a : Term n) -> TypeRed {n} env a -> Type

public export
data TypeRed :
  (eq : Equality) ->
  (l : Level) ->
  (rec : (l' : Level) -> l' `LT` l -> LogicalRelation eq) ->
  (env : Env n) ->
  (a : Term n) ->
  Type
public export
data TypeEq :
  (eq : Equality) ->
  (l : Level) ->
  (rec : (l' : Level) -> l' `LT` l -> LogicalRelation eq) ->
  (env : Env n) ->
  (a, b : Term n) ->
  TypeRed eq l rec {n} env a ->
  Type
public export
data TermRed :
  (eq : Equality) ->
  (l : Level) ->
  (rec : (l' : Level) -> l' `LT` l -> LogicalRelation eq) ->
  (env : Env n) ->
  (t, a : Term n) ->
  TypeRed eq l rec {n} env a ->
  Type
public export
data TermEq :
  (eq : Equality) ->
  (l : Level) ->
  (rec : (l' : Level) -> l' `LT` l -> LogicalRelation eq) ->
  (env : Env n) ->
  (t, u, a : Term n) ->
  TypeRed eq l rec {n} env a ->
  Type

-- Neutrals

public export
record NeutralTyRed
  (eq : Equality)
  (l : Level)
  (rec : (l' : Level) -> l' `LT` l -> LogicalRelation eq)
  (env : Env n)
  (a : Term n)
  where
  constructor MkNtrlTyRed
  {0 a' : Term n}
  tyWf : TypeWf env a
  steps : TypeReduce env a a'
  tyWf' : TypeWf env a'
  0 ntrl : Neutral a'
  prf : eq.NtrlEq env (Element a' ntrl) (Element a' ntrl) Set

public export
record NeutralTyEq
  (eq : Equality)
  (l : Level)
  (rec : (l' : Level) -> l' `LT` l -> LogicalRelation eq)
  (env : Env n)
  (a, b : Term n)
  (red : NeutralTyRed eq l rec {n} env a)
  where
  constructor MkNtrlTyEq
  {0 b' : Term n}
  tyWf : TypeWf env b
  steps : TypeReduce env b b'
  tyWf' : TypeWf env b'
  0 ntrl : Neutral b'
  prf : eq.NtrlEq env (Element red.a' red.ntrl) (Element b' ntrl) Set

public export
record NeutralTmRed
  (eq : Equality)
  (l : Level)
  (rec : (l' : Level) -> l' `LT` l -> LogicalRelation eq)
  (env : Env n)
  (t, a : Term n)
  (red : NeutralTyRed eq l rec {n} env a)
  where
  constructor MkNtrlTmRed
  {0 t' : Term n}
  tmWf : TermWf env t a
  steps : TermReduce env t t' a
  tmWf' : TermWf env t' a
  0 ntrl : Neutral t'
  prf : eq.NtrlEq env (Element t' ntrl) (Element t' ntrl) a

public export
record NeutralTmEq
  (eq : Equality)
  (l : Level)
  (rec : (l' : Level) -> l' `LT` l -> LogicalRelation eq)
  (env : Env n)
  (t, u, a : Term n)
  (red : NeutralTyRed eq l rec {n} env a)
  where
  constructor MkNtrlTmEq
  {0 t', u' : Term n}
  tmWf1 : TermWf env t a
  tmWf2 : TermWf env u a
  steps1 : TermReduce env t t' a
  steps2 : TermReduce env u u' a
  tmWf1' : TermWf env t' a
  tmWf2' : TermWf env u' a
  0 ntrl1 : Neutral t'
  0 ntrl2 : Neutral u'
  prf : eq.NtrlEq env (Element t' ntrl1) (Element u' ntrl2) a

-- Set

public export
record SetTyRed
  (eq : Equality)
  (l : Level)
  (rec : (l' : Level) -> l' `LT` l -> LogicalRelation eq)
  (env : Env n)
  (a : Term n)
  where
  constructor MkSetTyRed
  tyWf : TypeWf env a
  steps : TypeReduce env a Set
  tyWf' : TypeWf env Set
  prf : Small `LT` l

public export
record SetTyEq
  (eq : Equality)
  (l : Level)
  (rec : (l' : Level) -> l' `LT` l -> LogicalRelation eq)
  (env : Env n)
  (a, b : Term n)
  (red : SetTyRed eq l rec {n} env a)
  where
  constructor MkSetTyEq
  tyWf : TypeWf env b
  steps : TypeReduce env b Set
  tyWf' : TypeWf env Set

public export
record SetTmRed
  (eq : Equality)
  (l : Level)
  (rec : (l' : Level) -> l' `LT` l -> LogicalRelation eq)
  (env : Env n)
  (t, a : Term n)
  (red : SetTyRed eq l rec {n} env a)
  where
  constructor MkSetTmRed
  {0 t' : Term n}
  tmWf : TermWf env t a
  steps : TypeReduce env t t'
  tmWf' : TermWf env t' a
  0 whnf : Whnf t'
  prf : eq.TermEq env t' t' a
  tyRed : (rec Small red.prf).TypeRed env t

public export
record SetTmEq
  (eq : Equality)
  (l : Level)
  (rec : (l' : Level) -> l' `LT` l -> LogicalRelation eq)
  (env : Env n)
  (t, u, a : Term n)
  (red : SetTyRed eq l rec {n} env a)
  where
  constructor MkSetTmEq
  {0 t', u' : Term n}
  tmWf1 : TermWf env t a
  tmWf2 : TermWf env u a
  steps1 : TypeReduce env t t'
  steps2 : TypeReduce env u u'
  tmWf1' : TermWf env t' a
  tmWf2' : TermWf env u' a
  0 whnf1 : Whnf t'
  0 whnf2 : Whnf u'
  prf : eq.TermEq env t' u' a
  tyRed1 : (rec Small red.prf).TypeRed env t
  tyRed2 : (rec Small red.prf).TypeRed env u
  tyEq : (rec Small red.prf).TypeEq env t u tyRed1

-- Pi

public export
record PiTyRed
  (eq : Equality)
  (l : Level)
  (rec : (l' : Level) -> l' `LT` l -> LogicalRelation eq)
  (env : Env n)
  (a : Term n)
  where
  constructor MkPiTyRed
  {0 f : Term n}
  {0 g : Term (S n)}
  tyWf : TypeWf env a
  steps : TypeReduce env a (Pi f g)
  tyWf' : TypeWf env (Pi f g)
  prf : eq.TypeEq env (Pi f g) (Pi f g)
  domRed :
    forall m.
    {0 thin : n `Thins` m} ->
    forall env'.
    ExtendsWf thin env' env ->
    TypeRed eq l rec env' (wkn f thin)
  codRed :
    forall m.
    {0 thin : n `Thins` m} ->
    forall env'.
    (thinWf : ExtendsWf thin env' env) ->
    forall t.
    TermRed eq l rec env' t (wkn f thin) (domRed thinWf) ->
    TypeRed eq l rec env' (subst g (Wkn thin :< pure t))
  codEq :
    forall m.
    {0 thin : n `Thins` m} ->
    forall env'.
    (thinWf : ExtendsWf thin env' env) ->
    forall t, u.
    (red : TermRed eq l rec env' t (wkn f thin) (domRed thinWf)) ->
    TermRed eq l rec env' u (wkn f thin) (domRed thinWf) ->
    TermEq eq l rec env' t u (wkn f thin) (domRed thinWf) ->
    TypeEq eq l rec env'
      (subst g (Wkn thin :< pure t))
      (subst g (Wkn thin :< pure u))
      (codRed thinWf red)

public export
record PiTyEq
  (eq : Equality)
  (l : Level)
  (rec : (l' : Level) -> l' `LT` l -> LogicalRelation eq)
  (env : Env n)
  (a, b : Term n)
  (red : PiTyRed eq l rec env a)
  where
  constructor MkPiTyEq
  {0 f : Term n}
  {0 g : Term (S n)}
  tyWf : TypeWf env b
  steps : TypeReduce env b (Pi f g)
  tyWf' : TypeWf env (Pi f g)
  prf : eq.TypeEq env (Pi red.f red.g) (Pi f g)
  domEq :
    forall m.
    {0 thin : n `Thins` m} ->
    forall env'.
    (thinWf : ExtendsWf thin env' env) ->
    TypeEq eq l rec env' (wkn red.f thin) (wkn f thin) (red.domRed thinWf)
  codEq :
    forall m.
    {0 thin : n `Thins` m} ->
    forall env'.
    (thinWf : ExtendsWf thin env' env) ->
    forall t.
    (red' : TermRed eq l rec env' t (wkn red.f thin) (red.domRed thinWf)) ->
    TypeEq eq l rec env'
      (subst red.g (Wkn thin :< pure t))
      (subst g (Wkn thin :< pure t))
      (red.codRed thinWf red')

public export
record PiTmRed
  (eq : Equality)
  (l : Level)
  (rec : (l' : Level) -> l' `LT` l -> LogicalRelation eq)
  (env : Env n)
  (t, a : Term n)
  (red : PiTyRed eq l rec env a)
  where
  constructor MkPiTmRed
  {0 t' : Term n}
  tmWf : TermWf env t a
  steps : TermReduce env t t' a
  tmWf' : TermWf env t' a
  0 whnf : Whnf t'
  prf : eq.TermEq env t' t' a
  codRed :
    forall m.
    {0 thin : n `Thins` m} ->
    forall env'.
    (thinWf : ExtendsWf thin env' env) ->
    forall u.
    (red' : TermRed eq l rec env' u (wkn red.f thin) (red.domRed thinWf)) ->
    TermRed eq l rec env'
      (App (wkn t' thin) u)
      (subst red.g (Wkn thin :< pure u))
      (red.codRed thinWf red')
  codEq :
    forall m.
    {0 thin : n `Thins` m} ->
    forall env'.
    (thinWf : ExtendsWf thin env' env) ->
    forall u, v.
    (red' : TermRed eq l rec env' u (wkn red.f thin) (red.domRed thinWf)) ->
    TermRed eq l rec env' v (wkn red.f thin) (red.domRed thinWf) ->
    TermEq eq l rec env' u v (wkn red.f thin) (red.domRed thinWf) ->
    TermEq eq l rec env'
      (App (wkn t' thin) u)
      (App (wkn t' thin) v)
      (subst red.g (Wkn thin :< pure u))
      (red.codRed thinWf red')

public export
record PiTmEq
  (eq : Equality)
  (l : Level)
  (rec : (l' : Level) -> l' `LT` l -> LogicalRelation eq)
  (env : Env n)
  (t, u, a : Term n)
  (red : PiTyRed eq l rec env a)
  where
  constructor MkPiTmEq
  {0 t', u' : Term n}
  tmWf1 : TermWf env t a
  tmWf2 : TermWf env u a
  steps1 : TermReduce env t t' a
  steps2 : TermReduce env u u' a
  tmWf1' : TermWf env t' a
  tmWf2' : TermWf env u' a
  0 whnf1 : Whnf t'
  0 whnf2 : Whnf u'
  prf : eq.TermEq env t' u' a
  red1 : PiTmRed eq l rec env t a red
  red2 : PiTmRed eq l rec env u a red
  codEq :
    forall m.
    {0 thin : n `Thins` m} ->
    forall env'.
    (thinWf : ExtendsWf thin env' env) ->
    forall v.
    (red' : TermRed eq l rec env' v (wkn red.f thin) (red.domRed thinWf)) ->
    TermEq eq l rec env'
      (App (wkn t' thin) v)
      (App (wkn u' thin) v)
      (subst red.g (Wkn thin :< pure v))
      (red.codRed thinWf red')

-- Putting it all together

data TypeRed where
  RedNtrlTy : NeutralTyRed eq l rec env a -> TypeRed eq l rec env a
  RedSetTy : SetTyRed eq l rec env a -> TypeRed eq l rec env a
  RedPiTy : PiTyRed eq l rec env a -> TypeRed eq l rec env a

data TypeEq where
  EqNtrlTy : NeutralTyEq eq l rec env a b red -> TypeEq eq l rec env a b (RedNtrlTy red)
  EqSetTy : SetTyEq eq l rec env a b red -> TypeEq eq l rec env a b (RedSetTy red)
  EqPiTy : PiTyEq eq l rec env a b red -> TypeEq eq l rec env a b (RedPiTy red)

data TermRed where
  RedNtrlTm : NeutralTmRed eq l rec env t a red -> TermRed eq l rec env t a (RedNtrlTy red)
  RedSetTm : SetTmRed eq l rec env t a red -> TermRed eq l rec env t a (RedSetTy red)
  RedPiTm : PiTmRed eq l rec env t a red -> TermRed eq l rec env t a (RedPiTy red)

data TermEq where
  EqNtrlTm : NeutralTmEq eq l rec env t u a red -> TermEq eq l rec env t u a (RedNtrlTy red)
  EqSetTm : SetTmEq eq l rec env t u a red -> TermEq eq l rec env t u a (RedSetTy red)
  EqPiTm : PiTmEq eq l rec env t u a red -> TermEq eq l rec env t u a (RedPiTy red)

public export
LogRelKit :
  (eq : Equality) ->
  (l : Level) ->
  (rec : (l' : Level) -> l' `LT` l -> LogicalRelation eq) ->
  LogicalRelation eq
LogRelKit eq l rec =
  MkLogRel
    (TypeRed eq l rec)
    (TypeEq eq l rec)
    (TermRed eq l rec)
    (TermEq eq l rec)

-- Induction -------------------------------------------------------------------

public export
LogRelRec : (eq : Equality) -> (l, l' : Level) -> l' `LT` l -> LogicalRelation eq
LogRelRec eq Small _ prf = absurd prf
LogRelRec eq Large Small LTSmall = LogRelKit eq Small (\_, prf => absurd prf)

public export
LogRel : (eq : Equality) -> Level -> LogicalRelation eq
LogRel eq l = LogRelKit eq l (LogRelRec eq l)

-- Reflexivity -----------------------------------------------------------------

export
typeEqRefl :
  (l : Level) ->
  (red : (LogRel eq l).TypeRed env a) ->
  (LogRel eq l).TypeEq env a a red
typeEqRefl l (RedNtrlTy red) =
  EqNtrlTy $ MkNtrlTyEq
    { tyWf = red.tyWf
    , steps = red.steps
    , tyWf' = red.tyWf'
    , ntrl = red.ntrl
    , prf = red.prf
    }
typeEqRefl l (RedSetTy red) =
  EqSetTy $ MkSetTyEq
    { tyWf = red.tyWf
    , steps = red.steps
    , tyWf' = red.tyWf'
    }
typeEqRefl l (RedPiTy (MkPiTyRed tyWf steps tyWf' prf domRed codRed codEq)) =
  EqPiTy $ MkPiTyEq
    { tyWf = tyWf
    , steps = steps
    , tyWf' = tyWf'
    , prf = prf
    , domEq = \thinWf => typeEqRefl l (domRed thinWf)
    , codEq = \thinWf, red' => typeEqRefl l (codRed thinWf red')
    }

export
termEqRefl :
  (l : Level) ->
  (red : (LogRel eq l).TypeRed env a) ->
  (LogRel eq l).TermRed env t a red ->
  (LogRel eq l).TermEq env t t a red
termEqRefl l (RedNtrlTy red) (RedNtrlTm redTm) =
  EqNtrlTm $ MkNtrlTmEq
    { tmWf1 = redTm.tmWf
    , tmWf2 = redTm.tmWf
    , steps1 = redTm.steps
    , steps2 = redTm.steps
    , tmWf1' = redTm.tmWf'
    , tmWf2' = redTm.tmWf'
    , ntrl1 = redTm.ntrl
    , ntrl2 = redTm.ntrl
    , prf = redTm.prf
    }
termEqRefl Small (RedSetTy (MkSetTyRed tyWf steps tyWf' prf)) (RedSetTm redTm) = absurd prf
termEqRefl Large (RedSetTy (MkSetTyRed tyWf steps tyWf' LTSmall)) (RedSetTm redTm) =
  EqSetTm $ MkSetTmEq
    { tmWf1 = redTm.tmWf
    , tmWf2 = redTm.tmWf
    , steps1 = redTm.steps
    , steps2 = redTm.steps
    , tmWf1' = redTm.tmWf'
    , tmWf2' = redTm.tmWf'
    , whnf1 = redTm.whnf
    , whnf2 = redTm.whnf
    , prf = redTm.prf
    , tyRed1 = redTm.tyRed
    , tyRed2 = redTm.tyRed
    , tyEq = typeEqRefl {eq} Small redTm.tyRed
    }
termEqRefl l (RedPiTy red@(MkPiTyRed _ _ _ _ domRed _ _)) (RedPiTm redTm) =
  EqPiTm $ MkPiTmEq
    { tmWf1 = redTm.tmWf
    , tmWf2 = redTm.tmWf
    , steps1 = redTm.steps
    , steps2 = redTm.steps
    , tmWf1' = redTm.tmWf'
    , tmWf2' = redTm.tmWf'
    , whnf1 = redTm.whnf
    , whnf2 = redTm.whnf
    , prf = redTm.prf
    , red1 = redTm
    , red2 = redTm
    , codEq = \thinWf, red' => redTm.codEq thinWf red' red' (termEqRefl l (domRed thinWf) red')
    }