-- Lentil/Rules/WF.lean

import Lentil.Rules.Basic

import Lentil.Tactics.Structural

import Lentil.Gadgets.TheoremDeriving

open Classical

namespace TLA

section wf

variable {σ : Type u}

section wf_def

variable {a : action σ}

end wf_def

-- Repository definitions used:
content: syntax:15 tlafml:16 " → " tlafml:15 : tlafml

syntax:35 tlafml:36 " ∧ " tlafml:35 : tlafml

syntax:30 tlafml:31 " ∨ " tlafml:30 : tlafml

syntax:20 tlafml:21 " ↝ " tlafml:20 : tlafml

syntax:25 tlafml:26 " 𝑈 " tlafml:25 : tlafml

syntax:17 tlafml:18 " ⇒ " tlafml:17 : tlafml

syntax:max tlafml:max " |-tla- " tlafml:max : term

syntax:max tlafml:max " =tla= " tlafml:max : term

syntax "∀ " extBinder ", " tlafml:51 : tlafml

syntax tlafml_heading_op := "¬" <|> "□" <|> "◇" <|> "◯"

syntax:max tlafml_heading_op tlafml:40 : tlafml

syntax "[tlafml|" tlafml "]" : term

syntax "∃ " extBinder ", " tlafml:51 : tlafml

syntax:max "Enabled" term:40 : tlafml

syntax "⟨ " term " ⟩" : tlafml

content: macro "tla_unfold_simp" : tactic => `(tactic| (simp [tlasimp_def] at *))

syntax:max "|-tla- " tlafml:max : term

syntax:arg "𝒲ℱ" term:max : tlafml

syntax (name := tlaDerive) "tla_derive" : attr

content: macro "tla_unfold_simp'" : tactic => `(tactic| (tla_unfold_simp ; (try simp only [execsimp] at *)))

content: macro_rules
  | `([tlafml| ( $f:tlafml ) ]) => `([tlafml| $f ])
  | `([tlafml| ⌜ $t:term ⌝ ]) => `(TLA.state_pred $t)
  | `([tlafml| ⌞ $t:term ⌟ ]) => `(TLA.pure_pred $t)
  | `([tlafml| ⟨ $t:term ⟩ ]) => `(TLA.action_pred $t)
  | `([tlafml| ⊤ ]) => `(TLA.tla_true)
  | `([tlafml| ⊥ ]) => `(TLA.tla_false)
  | `([tlafml| $op:tlafml_heading_op $f:tlafml ]) => do
    let opterm ← match op with
      | `(tlafml_heading_op|¬) => `(TLA.tla_not)
      | `(tlafml_heading_op|□) => `(TLA.always)
      | `(tlafml_heading_op|◇) => `(TLA.eventually)
      | `(tlafml_heading_op|◯) => `(TLA.later)
      | _ => Macro.throwUnsupported
    `($opterm [tlafml| $f ])
  | `([tlafml| Enabled $t:term ]) => `(TLA.tla_enabled $t)
  | `([tlafml| $f1:tlafml → $f2:tlafml ]) => `(TLA.tla_implies [tlafml| $f1 ] [tlafml| $f2 ])
  | `([tlafml| $f1:tlafml ∧ $f2:tlafml ]) => `(TLA.tla_and [tlafml| $f1 ] [tlafml| $f2 ])
  | `([tlafml| $f1:tlafml ∨ $f2:tlafml ]) => `(TLA.tla_or [tlafml| $f1 ] [tlafml| $f2 ])
  | `([tlafml| $f1:tlafml 𝑈 $f2:tlafml ]) => `(TLA.tla_until [tlafml| $f1 ] [tlafml| $f2 ])
  | `([tlafml| ∀ $x:ident, $f:tlafml]) => `(TLA.tla_forall fun $x:ident => [tlafml| $f ])
  | `([tlafml| ∀ $x:ident : $t, $f:tlafml]) => `(TLA.tla_forall fun $x:ident : $t => [tlafml| $f ])
  | `([tlafml| ∃ $x:ident, $f:tlafml]) => `(TLA.tla_exists fun $x:ident => [tlafml| $f ])
  | `([tlafml| ∃ $x:ident : $t, $f:tlafml]) => `(TLA.tla_exists fun $x:ident : $t => [tlafml| $f ])
  | `([tlafml| $op:tlafml_bigop $x:binderIdent ∈ $l:term, $f:tlafml]) =>
    
    match op with
    | `(tlafml_bigop|⋀ ) => do `(TLA.tla_bigwedge (fun $(← binderIdentToFunBinder x) => [tlafml| $f ]) $l)
    | `(tlafml_bigop|⋁ ) => do `(TLA.tla_bigvee (fun $(← binderIdentToFunBinder x) => [tlafml| $f ]) $l)
    | _ => Macro.throwUnsupported
  | `([tlafml| $t:term ]) => `($t)

content: macro_rules
  | `([tlafml| $f1:tlafml ↝ $f2:tlafml ]) => `(TLA.leads_to [tlafml| $f1 ] [tlafml| $f2 ])
  | `([tlafml| $f1:tlafml ⇒ $f2:tlafml ]) => `(TLA.always_implies [tlafml| $f1 ] [tlafml| $f2 ])
  | `([tlafml| 𝒲ℱ $t:term ]) => `(TLA.weak_fairness $t)

content: macro_rules
  | `($f1:tlafml |-tla- $f2:tlafml) => `(TLA.pred_implies [tlafml| $f1 ] [tlafml| $f2 ])
  | `(|-tla- $f1:tlafml) => `(TLA.valid [tlafml| $f1 ])
  | `($f1:tlafml =tla= $f2:tlafml) => `([tlafml| $f1 ] = [tlafml| $f2 ])
  | `($e:term |=tla= $f:tlafml) => `(TLA.exec.satisfies [tlafml| $f ] $e)

content: macro_rules
  | `(tactic| try_unfold_at_all $idt:ident ) => `(tactic| (try unfold $idt at *) )
  | `(tactic| try_unfold_at_all $idt:ident $idts:ident* ) => `(tactic| (try unfold $idt at *) ; try_unfold_at_all $idts* )

content: def exec.drop {α : Type u} (k : Nat) (σ : exec α) : exec α := λ n => σ (n + k)

content: def tla_true {α : Type u} : pred α := pure_pred True

content: def pred (σ : Type u) := exec σ → Prop

content: def valid {α : Type u} (p : pred α) : Prop := ∀ (σ : exec α), σ.satisfies p

content: def tla_until {α : Type u} (p q : pred α) : pred α := λ σ => ∃ i, (q <| σ.drop i) ∧ ∀ j < i, (p <| σ.drop j)

content: def tla_false {α : Type u} : pred α := pure_pred False

content: def pure_pred {α : Type u} (p : Prop) : pred α := state_pred (fun _ => p)

content: def state_pred {σ : Type u} (f : σ → Prop) : pred σ :=
  fun e => f (e 0)

content: def tla_and {α : Type u} (p q : pred α) : pred α := fun σ => p σ ∧ q σ

content: def tla_exists {α : Sort u} {β : Type v} (p : α → pred β) : pred β := fun σ => ∃ x, p x σ

content: def tla_forall {α : Sort u} {β : Type v} (p : α → pred β) : pred β := fun σ => ∀ x, p x σ

content: def TLA.always_implies {α : Type u} (p q : TLA.pred α) : TLA.pred α := [tlafml| □ (p → q) ]

content: def tla_bigwedge {α : Type u} {β : Type v} {c} [Foldable c] (f : β → pred α) (s : c β) : pred α :=
  Foldable.fold tla_and tla_true f s

content: def TLA.weak_fairness {α : Type u} (a : action α) : pred α := [tlafml| □ ((□ (Enabled a)) → ◇ ⟨a⟩)]

content: def action (σ : Type u) := σ → σ → Prop

content: def eventually {α : Type u} (p : pred α) : pred α := λ σ => ∃ k, p <| σ.drop k

content: def tla_not {α : Type u} (p : pred α) : pred α := fun σ => ¬ p σ

content: def tla_or {α : Type u} (p q : pred α) : pred α := fun σ => p σ ∨ q σ

content: def pred_implies {α : Type u} (p q : pred α) : Prop := ∀ (σ : exec α), σ.satisfies p → σ.satisfies q

content: def exec (σ : Type u) := Nat → σ

content: def exec.satisfies {α : Type u} (p : pred α) (σ : exec α) : Prop := p σ

content: def tla_enabled {α : Type u} (a : action α) : pred α := state_pred (enabled a)

content: def enabled {α : Type u} (a : action α) (s : α) : Prop := ∃ s', a s s'

content: def TLA.leads_to {α : Type u} (p q : TLA.pred α) : TLA.pred α := [tlafml| □ (p → ◇ q) ]

content: def action_pred {σ : Type u} (a : action σ) : pred σ :=
  fun e => a (e 0) (e 1)

content: def tla_implies {α : Type u} (p q : pred α) : pred α := fun σ => p σ → q σ

content: def later {α : Type u} (p : pred α) : pred α := λ σ => p <| σ.drop 1

content: def always {α : Type u} (p : pred α) : pred α := λ σ => ∀ k, p <| σ.drop k

content: def tla_bigvee {α : Type u} {β : Type v} {c} [Foldable c] (f : β → pred α) (s : c β) : pred α :=
  Foldable.fold tla_or tla_false f s

content: class Foldable (c : Type u → Type v) where
  fold {α : Type u} {β : Type w} (op : β → β → β) [Std.Commutative op] [Std.Associative op]
    (b : β) (f : α → β) (s : c α) : β

-- Repository lemmas used:
content: @[tladual]
theorem always_eventually_or_distrib : (□ ◇ (p ∨ q)) =tla= (□ ◇ p ∨ □ ◇ q)

content: theorem dual_lemma (p q : pred σ) : ¬ p =tla= ¬ q → (p) =tla= (q)

content: theorem eventually_always_and_distrib : (◇ □ (p ∧ q)) =tla= (◇ □ p ∧ ◇ □ q)

content: theorem always_and : (□ (p ∧ q)) =tla= (□ p ∧ □ q)

content: theorem pred_eq_iff_iff (p q : pred σ) : (p) =tla= (q) ↔ (p) |-tla- (q) ∧ (q) |-tla- (p)

content: @[tladual]
theorem eventually_and_split : (◇ (p ∧ q)) |-tla- (◇ p ∧ ◇ q)

content: @[tladual]
theorem eventually_or : (◇ (p ∨ q)) =tla= (◇ p ∨ ◇ q)

content: theorem wf1' (p q init inv : σ → Prop) (next a : action σ)
  (hpuntilq : ∀ s s', p s → next s s' → p s' ∨ q s')
  (haq : ∀ s s', p s → next s s' → a s s' → q s')
  (henable : ∀ s, inv s → p s → enabled a s ∨ q s)
  (hinit_inv : ∀ s, init s → inv s)
  (hnext_inv : ∀ s s', next s s' → inv s → inv s') :
  (⌜ init ⌝ ∧ □ ⟨next⟩ ∧ 𝒲ℱ a) |-tla- (⌜ p ⌝ ↝ ⌜ q ⌝)

-- Local definitions used:

-- Local lemmas used:
content: theorem wf_alt1 : (𝒲ℱ a) =tla= □ ◇ ((¬ Enabled a) ∨ □ ◇ ⟨a⟩)

content: theorem wf_alt1' : (𝒲ℱ a) =tla= □ ◇ ((¬ Enabled a) ∨ ⟨a⟩)

-- Target theorem and proof:
@[tla_derive]
theorem wf1 (p q : pred σ) (next a : action σ) :
  ((p ∧ ⟨next⟩ ⇒ ◯ p ∨ ◯ q) ∧
   (p ∧ ⟨next⟩ ∧ ⟨a⟩ ⇒ ◯ q) ∧
   (p ⇒ Enabled a ∨ q) ∧
   (□ ⟨next⟩ ∧ 𝒲ℱ a)) |-tla- (p ↝ q)
:= by
  rw [wf_alt1']
  intro e ⟨hpuntilq, haq, henable, hnext, hwf_alt⟩ k hp
  specialize hwf_alt k ; rcases hwf_alt with ⟨k1, hwf_alt⟩
  -- know that: either `q` holds between `k` and `k + k1`, or `p` holds at `k1`
  -- use `henable` to know that if it is the latter case, then `q` must hold in the next step
  have htmp : (∃ k' ≤ k1, q <| e.drop (k + k')) ∨ (p <| e.drop (k + k1)) := by
    clear hwf_alt
    induction k1 with
    | zero => right ; assumption
    | succ n ih => {
      rw [← Nat.add_assoc]
      rcases ih with ⟨k', hle, ih⟩ | ih
      · left ; exists k' ; constructor ; omega ; apply ih
      · specialize hpuntilq _ ⟨ih, (hnext _)⟩
        rcases hpuntilq with hq | hq <;> tla_unfold_simp'
        · right ; apply hq
        · left ; exists (n + 1) ; aesop
    }
  rcases htmp with ⟨k', _, hq⟩ | hq <;> tla_unfold_simp'
  · aesop
  · rcases hwf_alt with hq2 | hq2
    · specialize henable _ hq ; aesop
    · exists (k1 + 1)
      specialize haq (k + k1) hq ; rw [← Nat.add_assoc] ; apply haq <;> aesop