-- Juvix/Core/Main/Semantics/Approx.lean

import Juvix.Core.Main.Semantics.Approx.Indexed

namespace Juvix.Core.Main

def Value.Approx (v v' : Value) : Prop :=
  ∀ n, v ≲ᵥ(n) v'

notation:40 v:40 " ≲ᵥ " v':40 => Value.Approx v v'

notation:40 args₁:40 " ≲ₐ " args₂:40 => List.Forall₂ Value.Approx args₁ args₂

notation:40 e:40 " ≲⟨" env:40 ", " env':40 "⟩ " e':40 => Expr.Approx.Param env env' e e'

notation:40 v:40 " ≲ₒ " v':40 => Object.Approx v v'

notation:40 env₁:40 " ≲ₑ " env₂:40 => Env.Approx env₁ env₂

notation:40 e₁:40 " ≲ " e₂:40 => Expr.Approx e₁ e₂

@[aesop safe cases]
inductive Value.Approx.Inversion : Value -> Value -> Prop where
  | unit : Value.Approx.Inversion Value.unit Value.unit
  | const {c} : Value.Approx.Inversion (Value.const c) (Value.const c)
  | constr_app {ctr_name args_rev args_rev'} :
    args_rev ≲ₐ args_rev' →
    Value.Approx.Inversion (Value.constr_app ctr_name args_rev) (Value.constr_app ctr_name args_rev')
  | closure {env₁ body₁ env₂ body₂} :
    (∀ a₁ a₂, a₁ ≲ᵥ a₂ → body₁ ≲⟨a₁ ∷ env₁, a₂ ∷ env₂⟩ body₂) →
    Value.Approx.Inversion (Value.closure env₁ body₁) (Value.closure env₂ body₂)

-- Repository definitions used:
content: notation:40 v:40 " ≲ᵥ(" n:40 ") " v':40 => Value.Approx.Indexed n v v'

content: notation:40 v:40 " ≲ᵥ " v':40 => Value.Approx v v'

content: notation:40 e₁:40 " ≲ " e₂:40 => Expr.Approx e₁ e₂

content: notation:40 env:40 " ⊢ " e:40 " ↦ " v:40 => Eval env e v

content: notation:40 env:40 " ⊢ " e:40 " ↓" => Eval.Defined env e

content: notation:40 env:40 " ⊢ " e:40 " ↦(" n ") " v:40 => Eval.Indexed n env e v

content: notation:40 e:40 " ≲⟨" env:40 ", " env':40 "⟩ " e':40 => Expr.Approx.Param env env' e e'

content: notation:40 args₁:40 " ≲ₐ " args₂:40 => List.Forall₂ Value.Approx args₁ args₂

content: macro "invert" h:term : tactic => `(tactic| (cases ($h).invert <;> try clear $h))

content: notation:40 args₁:40 " ≲ₐ(" n:40 ") " args₂:40 => List.Forall₂ (Value.Approx.Indexed n) args₁ args₂

syntax "case " expr " of " cases " end" : expr

content: macro_rules
  | `(⟪$s:ident ♯ $i:num⟫) => `(Expr.var $(Lean.Syntax.mkStrLit s.getId.toString) $i)
  | `(⟪$num:num⟫) => `(Expr.const (Constant.int $num))
  | `(⟪$s:str⟫) => `(Expr.const (Constant.string $s))
  | `(⟪υ⟫) => `(Expr.unit)
  | `(⟪$e₁:expr $e₂:expr⟫) => `(Expr.app ⟪$e₁⟫ ⟪$e₂⟫)
  | `(⟪λ $s:ident $e:expr⟫) => `(Expr.lambda $(Lean.Syntax.mkStrLit s.getId.toString) ⟪$e⟫)
  | `(⟪λ $ss:ident* . $e:expr⟫) => mkLambdas ss e
  | `(⟪ $s:ident ⟫) => `(Expr.constr $(Lean.Syntax.mkStrLit s.getId.toString))
  | `(⟪ $s:ident $es:expr* ⟫) => mkConstrApp s es
  | `(⟪$e₁ + $e₂⟫) => `(Expr.binop BinaryOp.add_int ⟪$e₁⟫ ⟪$e₂⟫)
  | `(⟪$e₁ - $e₂⟫) => `(Expr.binop BinaryOp.sub_int ⟪$e₁⟫ ⟪$e₂⟫)
  | `(⟪$e₁ * $e₂⟫) => `(Expr.binop BinaryOp.mul_int ⟪$e₁⟫ ⟪$e₂⟫)
  | `(⟪$e₁ / $e₂⟫) => `(Expr.binop BinaryOp.div_int ⟪$e₁⟫ ⟪$e₂⟫)
  | `(⟪let $s:ident := $e₁:expr in $e₂:expr⟫) => `(Expr.save $(Lean.Syntax.mkStrLit s.getId.toString) ⟪$e₁⟫ ⟪$e₂⟫)
  | `(⟪letrec $s:ident := $e₁:expr in $e₂:expr⟫) =>
    let name := Lean.Syntax.mkStrLit s.getId.toString
    `(Expr.save $name (Expr.recur $name ⟪$e₁⟫) ⟪$e₂⟫)
  | `(⟪rec $s:ident $e:expr⟫) =>
    let name := Lean.Syntax.mkStrLit s.getId.toString
    `(Expr.recur $name ⟪$e⟫)
  | `(⟪⊥⟫) => `(Expr.fail)
  | `(⟪cases| | $s:ident => $e:expr $cs:cases ⟫) =>
    let name := Lean.Syntax.mkStrLit s.getId.toString
    `(Expr.branch $name [] ⟪$e⟫ ⟪cases|$cs⟫)
  | `(⟪cases| | $s:ident $ss:ident* => $e:expr $cs:cases ⟫) =>
    let name := Lean.Syntax.mkStrLit s.getId.toString
    let vnames := ss.toList.map fun s => s.getId.toString
    `(Expr.branch $name $(quote vnames) ⟪$e⟫ ⟪cases|$cs⟫)
  | `(⟪cases| | $s:ident => $e:expr ⟫) =>
    let name := Lean.Syntax.mkStrLit s.getId.toString
    `(Expr.branch $name [] ⟪$e⟫ Expr.fail)
  | `(⟪cases| | $s:ident $ss:ident* => $e:expr ⟫) =>
    let name := Lean.Syntax.mkStrLit s.getId.toString
    let vnames := ss.toList.map fun s => s.getId.toString
    `(Expr.branch $name $(quote vnames) ⟪$e⟫ Expr.fail)
  | `(⟪cases| | _ => $e:expr ⟫) =>
    `(⟪$e⟫)
  | `(⟪case $e:expr of $cs:cases end⟫) => do
    `(Expr.save "_case_" ⟪$e⟫ ⟪cases|$cs⟫)
  | `(⟪($e)⟫) => `(⟪$e⟫)

content: inductive BinaryOp : Type where
  | add_int : BinaryOp
  | sub_int : BinaryOp
  | mul_int : BinaryOp
  | div_int : BinaryOp
  deriving Inhabited, BEq, DecidableEq

content: def Value.Approx.Indexed (n : Nat) (v₁ v₂ : Value) : Prop :=
  (v₁ = Value.unit ∧ v₂ = Value.unit) ∨
  (∃ c, v₁ = Value.const c ∧ v₂ = Value.const c) ∨
  (∃ ctr_name args_rev args_rev',
    v₁ = Value.constr_app ctr_name args_rev ∧
    v₂ = Value.constr_app ctr_name args_rev' ∧
    (∀ k < n, List.Forall₂ (Value.Approx.Indexed k) args_rev args_rev')) ∨
  (∃ env₁ body₁ env₂ body₂,
    v₁ = Value.closure env₁ body₁ ∧
    v₂ = Value.closure env₂ body₂ ∧
    (∀ n₁ n₂, n₁ + n₂ < n →
      ∀ a₁ a₂ r₁,
        Value.Approx.Indexed (n₁ + n₂) a₁ a₂ →
        a₁ ∷ env₁ ⊢ body₁ ↦(n₁) r₁ →
        ∃ r₂,
          a₂ ∷ env₂ ⊢ body₂ ↦ r₂ ∧
          Value.Approx.Indexed n₂ r₁ r₂))

  inductive Value : Type where
    | unit : Value
    | const : Constant → Value
    | constr_app : (constr : Name) → (args_rev : List Value) → Value
    | closure : (env : List Object) → (value : Expr) → Value
    deriving Inhabited

content: abbrev cons_value (v : Value) (env : Env) : Env := Object.value v :: env

  inductive Object : Type where
    | value : Value → Object
    | delayed : (env : List Object) → Expr → Object
    deriving Inhabited

content: inductive Expr : Type where
  | var : (name : String) → (index : Nat) → Expr
  | unit : Expr
  | const : Constant → Expr
  | constr : Name → Expr
  | app : Expr → Expr → Expr
  | constr_app : Expr → Expr → Expr
  | binop : (oper : BinaryOp) → (arg₁ arg₂ : Expr) → Expr
  | lambda : (var_name : String) → (body : Expr) → Expr
  | save : (var_name : String) → (value : Expr) → (body : Expr) → Expr
  | branch : (constr : Name) → (var_names : List Name) → (body : Expr) → (next : Expr) → Expr
  | recur : (var_name : Name) → (body : Expr) → Expr
  | fail : Expr
  deriving Inhabited, BEq, DecidableEq

content: inductive Constant : Type where
  | int : Int → Constant
  | string : String → Constant
  deriving Inhabited, BEq, DecidableEq

content: abbrev Name : Type := String

content: abbrev Env : Type := List Object

content: @[aesop safe cases]
inductive Value.Approx.Indexed.Inversion (n : Nat) : Value → Value → Prop where
  | unit : Value.Approx.Indexed.Inversion n Value.unit Value.unit
  | const {c} : Value.Approx.Indexed.Inversion n (Value.const c) (Value.const c)
  | constr_app {ctr_name args_rev args_rev'} :
    (∀ k < n, args_rev ≲ₐ(k) args_rev') →
    Value.Approx.Indexed.Inversion n (Value.constr_app ctr_name args_rev) (Value.constr_app ctr_name args_rev')
  | closure {env₁ body₁ env₂ body₂} :
    (∀ n₁ n₂, n₁ + n₂ < n →
      ∀ a₁ a₂ v₁,
        a₁ ≲ᵥ(n₁ + n₂) a₂ →
        a₁ ∷ env₁ ⊢ body₁ ↦(n₁) v₁ →
        ∃ v₂, a₂ ∷ env₂ ⊢ body₂ ↦ v₂ ∧ v₁ ≲ᵥ(n₂) v₂) →
    Value.Approx.Indexed.Inversion n (Value.closure env₁ body₁) (Value.closure env₂ body₂)

content: infixr:50 " ∷ " => cons_value

content: notation:40 v:40 " ≲ᵥ(" n:40 ") " v':40 => Value.Approx.Indexed n v v'

content: infixl:100 " @@ " => Expr.app

-- Repository lemmas used:
content: @[aesop unsafe destruct]
lemma Value.Approx.Indexed.invert {n v v'} :
  v ≲ᵥ(n) v' →
  Value.Approx.Indexed.Inversion n v v'

content: @[aesop unsafe apply]
lemma Value.Approx.Indexed.constr_app {n ctr_name args_rev args_rev'} :
    (∀ k < n, args_rev ≲ₐ(k) args_rev') →
    Value.constr_app ctr_name args_rev ≲ᵥ(n) Value.constr_app ctr_name args_rev'

content: @[aesop unsafe apply]
lemma Value.Approx.Indexed.unit {n} : Value.unit ≲ᵥ(n) Value.unit

content: @[aesop unsafe apply]
lemma Value.Approx.Indexed.const {n c} : Value.const c ≲ᵥ(n) Value.const c

content: @[aesop unsafe apply]
lemma Value.Approx.Indexed.closure {n env₁ body₁ env₂ body₂} :
    (∀ n₁ n₂, n₁ + n₂ < n →
      ∀ a₁ a₂ v₁,
        a₁ ≲ᵥ(n₁ + n₂) a₂ →
        a₁ ∷ env₁ ⊢ body₁ ↦(n₁) v₁ →
        ∃ v₂, a₂ ∷ env₂ ⊢ body₂ ↦ v₂ ∧ v₁ ≲ᵥ(n₂) v₂) →
    Value.closure env₁ body₁ ≲ᵥ(n) Value.closure env₂ body₂

content: theorem forall₂_trans' {α} {P Q R : α → α → Prop} {l₁ l₂ l₃}
  (h : ∀ x y z, P x y → Q y z → R x z)
  (h₁ : List.Forall₂ P l₁ l₂)
  (h₂ : List.Forall₂ Q l₂ l₃)
  : List.Forall₂ R l₁ l₃

-- Local definitions used:
content: def Value.Approx (v v' : Value) : Prop :=
  ∀ n, v ≲ᵥ(n) v'

content: @[aesop safe cases]
inductive Value.Approx.Inversion : Value -> Value -> Prop where
  | unit : Value.Approx.Inversion Value.unit Value.unit
  | const {c} : Value.Approx.Inversion (Value.const c) (Value.const c)
  | constr_app {ctr_name args_rev args_rev'} :
    args_rev ≲ₐ args_rev' →
    Value.Approx.Inversion (Value.constr_app ctr_name args_rev) (Value.constr_app ctr_name args_rev')
  | closure {env₁ body₁ env₂ body₂} :
    (∀ a₁ a₂, a₁ ≲ᵥ a₂ → body₁ ≲⟨a₁ ∷ env₁, a₂ ∷ env₂⟩ body₂) →
    Value.Approx.Inversion (Value.closure env₁ body₁) (Value.closure env₂ body₂)

-- Local lemmas used:
content: @[aesop unsafe destruct]
lemma Value.Approx.invert {v v'} :
  v ≲ᵥ v' →
  Value.Approx.Inversion v v'

-- Target theorem and proof:
@[trans]
lemma Value.Approx.Indexed.trans {n v₁ v₂ v₃} : v₁ ≲ᵥ(n) v₂ → v₂ ≲ᵥ v₃ → v₁ ≲ᵥ(n) v₃
:= by
  revert n
  suffices ∀ n, ∀ k ≤ n, v₁ ≲ᵥ(k) v₂ → v₂ ≲ᵥ v₃ → v₁ ≲ᵥ(k) v₃ by
    intro k
    exact this k k k.le_refl
  intros n
  induction n generalizing v₁ v₂ v₃ with
  | zero =>
    intros k hk h₁ h₂
    invert h₁
    case unit =>
      invert h₂
      case unit =>
        exact Value.Approx.Indexed.unit
    case const =>
      invert h₂
      case const =>
        exact Value.Approx.Indexed.const
    case constr_app ctr_name args_rev₁ args_rev₁' ch₁ =>
      cases h₂.invert
      case constr_app args_rev₂ ch₂ =>
        apply Value.Approx.Indexed.constr_app;
        intros
        simp_all only [nonpos_iff_eq_zero, not_lt_zero', IsEmpty.forall_iff, implies_true]
    case closure env₁ body₁ env₁' body₁' ch₁ =>
      cases h₂.invert
      case closure env₂ body₂ ch₂ =>
        apply Value.Approx.Indexed.closure
        · intro k' hk' v v₁ h
          have : k = 0 := by linarith
          subst k
          contradiction
  | succ n ih =>
    intros k hk h₁ h₂
    invert h₁
    case unit =>
      invert h₂
      case unit =>
        exact Value.Approx.Indexed.unit
    case const =>
      invert h₂
      case const =>
        exact Value.Approx.Indexed.const
    case constr_app ctr_name args_rev args_rev' ch₁ =>
      invert h₂
      case constr_app args_rev'' ch₂ =>
        apply Value.Approx.Indexed.constr_app
        · intro k' hk'
          have hk' : k' ≤ n := by linarith
          apply Utils.forall₂_trans' (P := fun v₁ v₂ => v₁ ≲ᵥ(k') v₂) (Q := fun v₁ v₂ => v₁ ≲ᵥ v₂)
          · intros v₁ v₂ v₃ h₁ h₂
            apply ih <;> assumption
          · apply ch₁
            simp_all only
          · apply ch₂
    case closure env₁ body₁ env₂ body₂ ch₁ =>
      invert h₂
      case closure env₃ body₃ ch₂ =>
        apply Value.Approx.Indexed.closure
        · intro n₁ n₂ hn a₁ a₃ v₁ happrox heval₁
          have ah₁ : ∃ v₂, (a₃ ∷ env₂) ⊢ body₂ ↦ v₂ ∧ v₁ ≲ᵥ(n₂) v₂ := by
            apply ch₁
            · assumption
            · assumption
            · assumption
          obtain ⟨v₂, heval₂, h₂⟩ := ah₁
          have ah₂ : ∃ v₃, (a₃ ∷ env₃) ⊢ body₃ ↦ v₃ ∧ v₂ ≲ᵥ v₃ := by
            apply ch₂
            · rfl
            · assumption
          obtain ⟨v₃, _, h₃⟩ := ah₂
          have : n₂ ≤ n := by linarith
          exists v₃
          aesop