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 TypeReduceWf (env : Env n) (a, b : Term n) where
  constructor MkTyReduceWf
  tyWf : TypeWf env a
  steps : TypeReduce env a b
  tyWf' : TypeWf env b

public export
record TermReduceWf (env : Env n) (t, u, a : Term n) where
  constructor MkTmReduceWf
  tmWf : TermWf env t a
  steps : TermReduce env t u a
  tmWf' : TermWf env u a

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}
  stepsWf : TypeReduceWf env a a'
  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}
  stepsWf : TypeReduceWf env b b'
  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}
  stepsWf : TermReduceWf env t t' a
  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}
  stepsWf1 : TermReduceWf env t t' a
  stepsWf2 : TermReduceWf env u u' a
  ntrl1 : Neutral t'
  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
  stepsWf : TypeReduceWf env a 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
  stepsWf : TypeReduceWf env b 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}
  stepsWf : TermReduceWf env t t' a
  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}
  stepsWf1 : TermReduceWf env t t' a
  stepsWf2 : TermReduceWf env u u' a
  whnf1 : Whnf t'
  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)}
  stepsWf : TypeReduceWf env a (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)}
  stepsWf : TypeReduceWf env b (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}
  stepsWf : TermReduceWf env t t' a
  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}
  stepsWf1 : TermReduceWf env t t' a
  stepsWf2 : TermReduceWf env u u' a
  whnf1 : Whnf t'
  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 :
  (0 l : Level) ->
  (red : (LogRel eq l).TypeRed env a) ->
  (LogRel eq l).TypeEq env a a red
typeEqRefl l (RedNtrlTy red) =
  EqNtrlTy $ MkNtrlTyEq
    { stepsWf = red.stepsWf
    , ntrl = red.ntrl
    , prf = red.prf
    }
typeEqRefl l (RedSetTy red) =
  EqSetTy $ MkSetTyEq
    { stepsWf = red.stepsWf
    }
typeEqRefl l (RedPiTy red@(MkPiTyRed _ _ domRed codRed _)) =
  EqPiTy $ MkPiTyEq
    { stepsWf = red.stepsWf
    , prf = red.prf
    , domEq = \thinWf => typeEqRefl l (domRed thinWf)
    , codEq = \thinWf, red' => typeEqRefl l (codRed thinWf red')
    }

export
termEqRefl :
  (0 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
    { stepsWf1 = redTm.stepsWf
    , stepsWf2 = redTm.stepsWf
    , ntrl1 = redTm.ntrl
    , ntrl2 = redTm.ntrl
    , prf = redTm.prf
    }
termEqRefl Large (RedSetTy (MkSetTyRed stepsWf LTSmall)) (RedSetTm redTm) =
  EqSetTm $ MkSetTmEq
    { stepsWf1 = redTm.stepsWf
    , stepsWf2 = redTm.stepsWf
    , 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
    { stepsWf1 = redTm.stepsWf
    , stepsWf2 = redTm.stepsWf
    , 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')
    }

-- Escape ----------------------------------------------------------------------

export
escapeTypeRed : (l : Level) -> (red : (LogRel eq l).TypeRed env a) -> TypeWf env a
escapeTypeRed l (RedNtrlTy red) = red.stepsWf.tyWf
escapeTypeRed l (RedSetTy red) = red.stepsWf.tyWf
escapeTypeRed l (RedPiTy red) = red.stepsWf.tyWf

export
escapeTypeEq :
  IsEq eq =>
  (l : Level) ->
  (red : (LogRel eq l).TypeRed env a) ->
  (LogRel eq l).TypeEq env a b red ->
  eq.TypeEq env a b
escapeTypeEq l (RedNtrlTy red) (EqNtrlTy tyEq) =
  expandType
    red.stepsWf.steps tyEq.stepsWf.steps
    (Ntrl red.ntrl) (Ntrl tyEq.ntrl)
    (ntrlEqImpliesTypeEq tyEq.prf)
escapeTypeEq l (RedSetTy red) (EqSetTy tyEq) =
  expandType
    red.stepsWf.steps tyEq.stepsWf.steps
    Set Set
    (typeSet $ typeWfImpliesEnvWf red.stepsWf.tyWf')
escapeTypeEq l (RedPiTy red) (EqPiTy tyEq) =
  expandType
    red.stepsWf.steps tyEq.stepsWf.steps
    Pi Pi
    tyEq.prf

export
escapeTermRed :
  (l : Level) ->
  (red : (LogRel eq l).TypeRed env a) ->
  (LogRel eq l).TermRed env t a red ->
  TermWf env t a
escapeTermRed l (RedNtrlTy red) (RedNtrlTm tmRed) = tmRed.stepsWf.tmWf
escapeTermRed l (RedSetTy red) (RedSetTm tmRed) = tmRed.stepsWf.tmWf
escapeTermRed l (RedPiTy red) (RedPiTm tmRed) = tmRed.stepsWf.tmWf

export
escapeTermEq :
  IsEq eq =>
  (l : Level) ->
  (red : (LogRel eq l).TypeRed env a) ->
  (LogRel eq l).TermEq env t u a red ->
  eq.TermEq env t u a
escapeTermEq l (RedNtrlTy red) (EqNtrlTm tmEq) =
  expandTerm
    tmEq.stepsWf1.steps tmEq.stepsWf2.steps
    (Ntrl tmEq.ntrl1) (Ntrl tmEq.ntrl2)
    (ntrlEqImpliesTermEq tmEq.prf)
escapeTermEq l (RedSetTy red) (EqSetTm tmEq) =
  expandTerm
    tmEq.stepsWf1.steps tmEq.stepsWf2.steps
    tmEq.whnf1 tmEq.whnf2
    tmEq.prf
escapeTermEq l (RedPiTy red) (EqPiTm tmEq) =
  expandTerm
    tmEq.stepsWf1.steps tmEq.stepsWf2.steps
    tmEq.whnf1 tmEq.whnf2
    tmEq.prf