-- Clean/Utils/OfflineMemory.lean

import Clean.Circuit.Basic

import Clean.Circuit.Provable

import Clean.Gadgets.Equality

import Clean.Utils.Field

import Clean.Utils.Primes

import Clean.Utils.Tactics

import Mathlib.Data.List.Sort

def MemoryAccess := ℕ × ℕ × ℕ × ℕ 

def MemoryAccessList := List MemoryAccess

abbrev timestamp_ordering (x y : MemoryAccess) := match x, y with
| (t2, _a2, _r2, _w2), (t1, _a1, _r1, _w1) => t1 < t2

def MemoryAccessList.isTimestampSorted (accesses : MemoryAccessList) : Prop :=
  accesses.Sorted timestamp_ordering

def MemoryAccessList.timestamps_neq (x y: MemoryAccess) : Prop :=
  match x, y with
  | (t_x, _a_x, _r_x, _w_x), (t_y, _a_y, _r_y, _w_y) => t_x ≠ t_y

def MemoryAccessList.Notimestampdup (accesses : MemoryAccessList) : Prop :=
  List.Pairwise timestamps_neq accesses

abbrev address_timestamp_ordering (x y : MemoryAccess) := match x, y with
| (t2, a2, _, _), (t1, a1, _, _) => if a1 = a2 then t1 ≤ t2 else a1 < a2

abbrev address_strict_timestamp_ordering (x y : MemoryAccess) := match x, y with
| (t2, a2, _, _), (t1, a1, _, _) => if a1 = a2 then t1 < t2 else a1 < a2

@[reducible]
def MemoryAccessList.isAddressTimestampSorted (accesses : MemoryAccessList) : Prop :=
  accesses.Sorted address_timestamp_ordering

@[reducible]
def MemoryAccessList.isAddressStrictTimestampSorted (accesses : MemoryAccessList) : Prop :=
  accesses.Sorted address_strict_timestamp_ordering

def MemoryAccessList.filterAddress (accesses : MemoryAccessList) (addr : ℕ) : MemoryAccessList :=
  accesses.filter (fun (_timestamp, addr', _readValue, _writeValue) => addr' = addr)

def MemoryAccessList.isConsistentSingleAddress (accesses : MemoryAccessList) (h_sorted : accesses.isTimestampSorted) : Prop := match accesses with
  
  | [] => True
  
  | (_timestamp, _addr, readValue, _writeValue) :: [] => readValue = 0
  
  | (_t2, _addr2, readValue2, _writeValue2) :: (t1, addr1, readValue1, writeValue1) :: rest =>
    readValue2 = writeValue1 ∧
    MemoryAccessList.isConsistentSingleAddress ((t1, addr1, readValue1, writeValue1) :: rest) (List.Sorted.of_cons h_sorted)

def MemoryAccessList.isConsistentOffline (accesses : MemoryAccessList) (h_sorted : accesses.isAddressTimestampSorted) : Prop := match accesses with
  | [] => True 
  | (_timestamp, _addr, readValue, _writeValue) :: [] => readValue = 0
  | (_t2, addr2, readValue2, _writeValue2) :: (t1, addr1, readValue1, writeValue1) :: rest =>
    (if addr1 = addr2 then readValue2 = writeValue1 else readValue2 = 0) ∧
    MemoryAccessList.isConsistentOffline ((t1, addr1, readValue1, writeValue1) :: rest) (List.Sorted.of_cons h_sorted)

-- Repository definitions used:
content: ...

-- Repository lemmas used:
content: theorem List.getElem_append_left{α : Type u_1} {l₁ l₂ : List α} {i : Nat} (hn : i < l₁.length) :\n(l₁ ++ l₂)[i] = l₁[i]

-- Local definitions used:
content: def MemoryAccess := ℕ × ℕ × ℕ × ℕ

content: def MemoryAccessList := List MemoryAccess

content: abbrev timestamp_ordering (x y : MemoryAccess) := match x, y with
| (t2, _a2, _r2, _w2), (t1, _a1, _r1, _w1) => t1 < t2

content: def MemoryAccessList.isTimestampSorted (accesses : MemoryAccessList) : Prop :=
  accesses.Sorted timestamp_ordering

content: def MemoryAccessList.timestamps_neq (x y: MemoryAccess) : Prop :=
  match x, y with
  | (t_x, _a_x, _r_x, _w_x), (t_y, _a_y, _r_y, _w_y) => t_x ≠ t_y

content: def MemoryAccessList.Notimestampdup (accesses : MemoryAccessList) : Prop :=
  List.Pairwise timestamps_neq accesses

content: abbrev address_timestamp_ordering (x y : MemoryAccess) := match x, y with
| (t2, a2, _, _), (t1, a1, _, _) => if a1 = a2 then t1 ≤ t2 else a1 < a2

content: abbrev address_strict_timestamp_ordering (x y : MemoryAccess) := match x, y with
| (t2, a2, _, _), (t1, a1, _, _) => if a1 = a2 then t1 < t2 else a1 < a2

content: @[reducible]
def MemoryAccessList.isAddressTimestampSorted (accesses : MemoryAccessList) : Prop :=
  accesses.Sorted address_timestamp_ordering

content: @[reducible]
def MemoryAccessList.isAddressStrictTimestampSorted (accesses : MemoryAccessList) : Prop :=
  accesses.Sorted address_strict_timestamp_ordering

content: def MemoryAccessList.filterAddress (accesses : MemoryAccessList) (addr : ℕ) : MemoryAccessList :=
  accesses.filter (fun (_timestamp, addr', _readValue, _writeValue) => addr' = addr)

content: def MemoryAccessList.isConsistentSingleAddress (accesses : MemoryAccessList) (h_sorted : accesses.isTimestampSorted) : Prop := match accesses with
  
  | [] => True
  
  | (_timestamp, _addr, readValue, _writeValue) :: [] => readValue = 0
  
  | (_t2, _addr2, readValue2, _writeValue2) :: (t1, addr1, readValue1, writeValue1) :: rest =>
    readValue2 = writeValue1 ∧
    MemoryAccessList.isConsistentSingleAddress ((t1, addr1, readValue1, writeValue1) :: rest) (List.Sorted.of_cons h_sorted)

content: def MemoryAccessList.isConsistentOffline (accesses : MemoryAccessList) (h_sorted : accesses.isAddressTimestampSorted) : Prop := match accesses with
  | [] => True 
  | (_timestamp, _addr, readValue, _writeValue) :: [] => readValue = 0
  | (_t2, addr2, readValue2, _writeValue2) :: (t1, addr1, readValue1, writeValue1) :: rest =>
    (if addr1 = addr2 then readValue2 = writeValue1 else readValue2 = 0) ∧
    MemoryAccessList.isConsistentOffline ((t1, addr1, readValue1, writeValue1) :: rest) (List.Sorted.of_cons h_sorted)

-- Local lemmas used:
content: theorem MemoryAccessList.addressStrictTimestampSorted_of_AddressTimestampSorted_noTimestampDup
    (accesses : MemoryAccessList) (h_sorted : accesses.isAddressTimestampSorted)
    (h_no_timestamp_dup : accesses.Notimestampdup) :
    accesses.isAddressStrictTimestampSorted

content: theorem MemoryAccessList.noTimestampDup_of_cons (head : MemoryAccess) (tail : MemoryAccessList)
    (h : Notimestampdup (head :: tail)) :
    Notimestampdup tail

content: theorem MemoryAccessList.isAddressTimestampSorted_of_cons (head : MemoryAccess) (tail : MemoryAccessList)
    (h : isAddressTimestampSorted (head :: tail)) :
    isAddressTimestampSorted tail

content: theorem MemoryAccessList.filterAddress_sorted_from_addressTimestampSorted
    (accesses : MemoryAccessList)
    (h_sorted : accesses.isAddressTimestampSorted)
    (h_nodup : accesses.Notimestampdup)
    (addr : ℕ) :
    (accesses.filterAddress addr).isTimestampSorted

content: theorem MemoryAccessList.filterAddress_empty_when_address_changes
    (head : MemoryAccess) (second : MemoryAccess) (tail : MemoryAccessList)
    (h_sorted : isAddressTimestampSorted (head :: second :: tail))
    (h_addr_ne : head.2.1 ≠ second.2.1) :
    filterAddress (second :: tail) head.2.1 = []

content: theorem MemoryAccessList.isConsistentOffline_of_cons
    (head : MemoryAccess) (tail : MemoryAccessList)
    (h_sorted : isAddressTimestampSorted (head :: tail))
    (h_offline : isConsistentOffline (head :: tail) h_sorted) :
    isConsistentOffline tail (isAddressTimestampSorted_of_cons head tail h_sorted)

-- Target theorem and proof:
theorem MemoryAccessList.isConsistentOffline_implies_single_address
    (accesses : MemoryAccessList)
    (h_sorted : accesses.isAddressTimestampSorted)
    (h_nodup : accesses.Notimestampdup)
    (h_offline : accesses.isConsistentOffline h_sorted)
    (addr : ℕ) :
    (accesses.filterAddress addr).isConsistentSingleAddress
      (filterAddress_sorted_from_addressTimestampSorted accesses h_sorted h_nodup addr)
:= by
  induction accesses with
  | nil =>
    simp [filterAddress, isConsistentSingleAddress]
  | cons head tail ih =>
    obtain ⟨t_head, a_head, r_head, w_head⟩ := head
    simp only [filterAddress, List.filter_cons]
    split_ifs with h_addr
    · -- head has the target address
      rw [decide_eq_true_eq] at h_addr
      subst h_addr
      -- Now we need to show consistency for (head :: filterAddress tail addr)
      cases tail with
      | nil =>
        simp [isConsistentSingleAddress]
        -- For a singleton list, need to show r_head = 0
        exact h_offline
      | cons second tail_rest =>
        obtain ⟨t_second, a_second, r_second, w_second⟩ := second
        -- Filter the tail
        simp only [List.filter_cons]
        split_ifs with h_second_addr
        · -- second also has address a_head
          rw [decide_eq_true_eq] at h_second_addr
          subst h_second_addr
          -- Now we show consistency for (head :: second :: filterAddress tail_rest a_head)
          simp only [isConsistentSingleAddress]
          constructor
          · -- Show r_second = w_head from offline consistency
            simp only [isConsistentOffline] at h_offline
            simp at h_offline
            exact h_offline.1
          · -- Apply IH to (second :: tail_rest)
            have ih_applied := ih
              (isAddressTimestampSorted_of_cons (t_head, a_second, r_head, w_head) ((t_second, a_second, r_second, w_second) :: tail_rest) h_sorted)
              (noTimestampDup_of_cons (t_head, a_second, r_head, w_head) ((t_second, a_second, r_second, w_second) :: tail_rest) h_nodup)
              (isConsistentOffline_of_cons (t_head, a_second, r_head, w_head) ((t_second, a_second, r_second, w_second) :: tail_rest) h_sorted h_offline)
            simp only [filterAddress, List.filter_cons, decide_true] at ih_applied
            exact ih_applied
        · -- second has a different address
          -- Need to show isConsistentSingleAddress for (head :: filterAddress tail_rest a_head)
          -- Key insight: since a_head ≠ a_second and list is address-sorted,
          -- all elements in (second :: tail_rest) have address ≠ a_head
          have h_addr_ne : ¬(a_second = a_head) := by
            simp only [decide_eq_true_eq] at h_second_addr
            exact h_second_addr
          -- Use filterAddress_empty_when_address_changes to show filterAddress (second :: tail_rest) a_head = []
          have h_empty := filterAddress_empty_when_address_changes
            (t_head, a_head, r_head, w_head) (t_second, a_second, r_second, w_second) tail_rest h_sorted
            (by simp; intro h_eq; exact h_addr_ne h_eq.symm)
          simp only [filterAddress, List.filter_cons] at h_empty
          have h_second_ne : decide (a_second = a_head) = false := by
            simp [h_addr_ne]
          simp only [h_second_ne] at h_empty
          have h_empty_simplified : List.filter (fun x => match x with | (_, addr', _, _) => decide (addr' = a_head)) tail_rest = [] := by
            simp at h_empty
            exact h_empty
          simp only [h_empty_simplified, isConsistentSingleAddress]
          simp only [isConsistentOffline] at h_offline
          simp [h_addr_ne] at h_offline
          exact h_offline.1
    · -- head doesn't have the target address
      -- Apply IH to tail
      apply ih (isAddressTimestampSorted_of_cons (t_head, a_head, r_head, w_head) tail h_sorted)
        (noTimestampDup_of_cons (t_head, a_head, r_head, w_head) tail h_nodup)
      exact isConsistentOffline_of_cons (t_head, a_head, r_head, w_head) tail h_sorted h_offline