/-
Gödel-machine specification accompanying
"Recursive Improvement in Schmidhuber's Artificial Intelligence".

Check with Lean 4.24.0: lean GodelMachine.lean
Core Lean only. All application-specific premises are explicit parameters.
-/

import Lean

set_option autoImplicit false
set_option warningAsError true

namespace GodelMachine

/-- Code numbers in the writable machine image. The interpreter is external. -/
structure Image where
  policy : Nat
  searcher : Nat
  theory : Nat
  checker : Nat
  deriving DecidableEq, Repr

/-- An observation history and scratch state can each be encoded as a natural. -/
structure Snapshot where
  tick : Nat
  image : Image
  memory : Nat
  history : Nat
  deriving DecidableEq, Repr

structure Rewrite where
  code : Nat
  activation : Nat
  deriving DecidableEq, Repr

structure Claim where
  snapshot : Snapshot
  rewrite : Rewrite
  deriving DecidableEq, Repr

structure Proposal where
  claim : Claim
  proofCode : Nat
  deriving DecidableEq, Repr

inductive Choice where
  | keepSearching
  | commit (rewrite : Rewrite)
  deriving DecidableEq, Repr

/--
A checkpoint-level interpreter interface. Search and checking dispatch through
the current image. A commit is a macro step that schedules and executes the
rewrite at its designated checkpoint. Ordinary policy execution, environment
updates and accounting are encapsulated in the callbacks.
-/
structure Runtime where
  search : Nat → Snapshot → Option Proposal
  check : Nat → Nat → Claim → Nat → Bool
  schedulable : Snapshot → Rewrite → Bool
  continueSearch : Snapshot → Snapshot
  executeRewrite : Snapshot → Rewrite → Snapshot

/--
Branch values describe complete counterfactual continuations from one snapshot.
The keep-searching value includes its costs and any later rewrites.
-/
def Target {U : Type} [LT U] (value : Snapshot → Choice → U)
    (s : Snapshot) (r : Rewrite) : Prop :=
  value s .keepSearching < value s (.commit r)

/-- Semantic soundness of the live checker and theory on claims at this snapshot. -/
def SoundAt {U : Type} [LT U] (rt : Runtime) (value : Snapshot → Choice → U)
    (s : Snapshot) : Prop :=
  ∀ (r : Rewrite) (proofCode : Nat),
    rt.check s.image.checker s.image.theory ⟨s, r⟩ proofCode = true →
      Target value s r

def Accepted (rt : Runtime) (s : Snapshot) (p : Proposal) : Prop :=
  p.claim.snapshot = s ∧
  s.tick < p.claim.rewrite.activation ∧
  rt.schedulable s p.claim.rewrite = true ∧
  rt.check s.image.checker s.image.theory p.claim p.proofCode = true

instance (rt : Runtime) (s : Snapshot) (p : Proposal) :
    Decidable (Accepted rt s p) := by
  unfold Accepted
  infer_instance

def choose (rt : Runtime) (s : Snapshot) : Choice :=
  match rt.search s.image.searcher s with
  | none => .keepSearching
  | some p =>
    if Accepted rt s p then .commit p.claim.rewrite else .keepSearching

def step (rt : Runtime) (s : Snapshot) : Snapshot :=
  match choose rt s with
  | .keepSearching => rt.continueSearch s
  | .commit r => rt.executeRewrite s r

def run (rt : Runtime) (initial : Snapshot) : Nat → Snapshot
  | 0 => initial
  | n + 1 => step rt (run rt initial n)

theorem commit_has_certificate (rt : Runtime) (s : Snapshot) (r : Rewrite)
    (h : choose rt s = .commit r) :
    ∃ p, rt.search s.image.searcher s = some p ∧
      p.claim.rewrite = r ∧ Accepted rt s p := by
  unfold choose at h
  cases hs : rt.search s.image.searcher s with
  | none =>
    rw [hs] at h
    cases h
  | some p =>
    rw [hs] at h
    change (if Accepted rt s p then Choice.commit p.claim.rewrite
      else Choice.keepSearching) = Choice.commit r at h
    by_cases ha : Accepted rt s p
    · rw [if_pos ha] at h
      cases h
      exact ⟨p, rfl, rfl, ha⟩
    · rw [if_neg ha] at h
      cases h

theorem choose_of_accepted (rt : Runtime) (s : Snapshot) (p : Proposal)
    (found : rt.search s.image.searcher s = some p) (valid : Accepted rt s p) :
    choose rt s = .commit p.claim.rewrite := by
  simp [choose, found, valid]

theorem stale_offer_waits (rt : Runtime) (s : Snapshot) (p : Proposal)
    (found : rt.search s.image.searcher s = some p)
    (stale : p.claim.snapshot ≠ s) :
    choose rt s = .keepSearching := by
  simp [choose, found, Accepted, stale]

theorem expired_offer_waits (rt : Runtime) (s : Snapshot) (p : Proposal)
    (found : rt.search s.image.searcher s = some p)
    (expired : p.claim.rewrite.activation ≤ s.tick) :
    choose rt s = .keepSearching := by
  have late : ¬ s.tick < p.claim.rewrite.activation := Nat.not_lt.mpr expired
  simp [choose, found, Accepted, late]

/-- Strict dominance over the specified continuation, at this checkpoint. -/
theorem accepted_commit_better {U : Type} [LT U]
    (rt : Runtime) (value : Snapshot → Choice → U) (s : Snapshot) (r : Rewrite)
    (sound : SoundAt rt value s) (commits : choose rt s = .commit r) :
    Target value s r := by
  obtain ⟨p, _, rewrite_eq, snapshot_eq, _, _, check_eq⟩ :=
    commit_has_certificate rt s r commits
  have claim_eq : p.claim = ⟨s, r⟩ := by
    calc
      p.claim = ⟨p.claim.snapshot, p.claim.rewrite⟩ := rfl
      _ = ⟨s, r⟩ := by rw [snapshot_eq, rewrite_eq]
  rw [claim_eq] at check_eq
  exact sound r p.proofCode check_eq

theorem invariant_at_run (rt : Runtime) (I : Snapshot → Prop) (initial : Snapshot)
    (boot : I initial) (closed : ∀ s, I s → I (step rt s)) (n : Nat) :
    I (run rt initial n) := by
  induction n with
  | zero => exact boot
  | succ n ih => exact closed (run rt initial n) ih

/-- Every commit on an invariant-preserving run passes its own local comparison. -/
theorem every_commit_better {U : Type} [LT U] (rt : Runtime)
    (value : Snapshot → Choice → U) (I : Snapshot → Prop) (initial : Snapshot)
    (boot : I initial) (closed : ∀ s, I s → I (step rt s))
    (sound : ∀ s, I s → SoundAt rt value s)
    (n : Nat) (r : Rewrite) (commits : choose rt (run rt initial n) = .commit r) :
    Target value (run rt initial n) r :=
  accepted_commit_better rt value (run rt initial n) r
    (sound _ (invariant_at_run rt I initial boot closed n)) commits

/-- The largest step-closed predicate contained in an audit predicate. -/
def InvariantCore (rt : Runtime) (audit : Snapshot → Prop) (s : Snapshot) : Prop :=
  ∀ n, audit (run rt s n)

theorem core_subset (rt : Runtime) (audit : Snapshot → Prop) (s : Snapshot)
    (h : InvariantCore rt audit s) : audit s :=
  h 0

theorem run_after_step (rt : Runtime) (s : Snapshot) (n : Nat) :
    run rt (step rt s) n = run rt s (n + 1) := by
  induction n with
  | zero => rfl
  | succ n ih => simp only [run, ih]

theorem core_closed (rt : Runtime) (audit : Snapshot → Prop) (s : Snapshot)
    (h : InvariantCore rt audit s) : InvariantCore rt audit (step rt s) := by
  intro n
  rw [run_after_step]
  exact h (n + 1)

theorem core_greatest (rt : Runtime) (audit I : Snapshot → Prop)
    (contained : ∀ s, I s → audit s)
    (closed : ∀ s, I s → I (step rt s)) (s : Snapshot) (h : I s) :
    InvariantCore rt audit s :=
  fun n => contained _ (invariant_at_run rt I s h closed n)

namespace Example

def initial : Snapshot := ⟨0, ⟨10, 20, 30, 40⟩, 0, 0⟩
def first : Rewrite := ⟨1, 2⟩
def second : Rewrite := ⟨2, 3⟩

def value (s : Snapshot) (branch : Choice) : Int :=
  if s.image.checker = 40 then
    match branch with
    | .keepSearching => 6
    | .commit r => if r.code = 1 then 8 else 0
  else
    match branch with
    | .keepSearching => 5
    | .commit _ => 2

/-- Checker 40 verifies a deliberately tiny model. Checker 41 accepts every claim. -/
def runtime : Runtime where
  search id s :=
    if id = 20 then some ⟨⟨s, first⟩, 7⟩
    else if id = 21 then some ⟨⟨s, second⟩, 0⟩
    else none
  check id theory claim proofCode :=
    if id = 40 then decide (theory = 30 ∧ claim.rewrite.code = 1 ∧ proofCode = 7)
    else id == 41
  schedulable _ _ := true
  continueSearch s := { s with tick := s.tick + 1 }
  executeRewrite s r :=
    if r.code = 1 then { s with tick := r.activation, image := ⟨11, 21, 31, 41⟩ }
    else { s with tick := r.activation, image := ⟨12, 22, 32, 42⟩ }

theorem initial_sound : SoundAt runtime value initial := by
  intro r proofCode ok
  have accepted : r.code = 1 ∧ proofCode = 7 := by
    simpa [runtime, initial] using ok
  simp [Target, value, initial, accepted.1]

theorem first_commits : choose runtime initial = .commit first := by decide

theorem first_is_better : Target value initial first :=
  accepted_commit_better runtime value initial first initial_sound first_commits

theorem second_commits :
    choose runtime (run runtime initial 1) = .commit second := by decide

theorem second_is_worse : ¬ Target value (run runtime initial 1) second := by
  unfold Target
  decide

theorem no_preserved_soundness :
    ¬ InvariantCore runtime (SoundAt runtime value) initial := by
  intro h
  exact second_is_worse
    (accepted_commit_better runtime value _ second (h 1) second_commits)

#eval choose runtime initial
#eval (run runtime initial 1).image
#eval choose runtime (run runtime initial 1)

end Example

#print axioms accepted_commit_better
#print axioms every_commit_better
#print axioms core_greatest
#print axioms Example.no_preserved_soundness

end GodelMachine