-- File: lean/ShortestPathSpec.lean
-- Shortest Path (Dijkstra/Bellman-Ford) Specification
-- Agents migrate a C++ Dijkstra to Rust.
-- Graph: List (from, to, weight) with non-negative Nat weights.

namespace ShortestPathSpec

abbrev Graph := List (Nat × Nat × Nat)   -- (from, to, weight)

-- Safe index into a list (returns none if out of bounds)
def listGet (l : List α) (i : Nat) : Option α :=
  (l.drop i).head?

-- Update distance at index i in the distance list
def updateDist (dists : List (Option Nat)) (i d : Nat) : List (Option Nat) :=
  List.zipWith
    (fun j curr =>
      if j == i then
        match curr with
        | none => some d
        | some existing => some (min existing d)
      else curr)
    (List.range dists.length)
    dists

-- Relax one edge (fr → to with weight w)
def relaxEdge (dists : List (Option Nat)) (fr to w : Nat) : List (Option Nat) :=
  match listGet dists fr with
  | some (some df) => updateDist dists to (df + w)
  | _ => dists

-- One Bellman-Ford pass over all edges
def relaxPass (dists : List (Option Nat)) (edges : Graph) : List (Option Nat) :=
  edges.foldl (fun d e => relaxEdge d e.1 e.2.1 e.2.2) dists

-- Dijkstra via n-1 Bellman-Ford iterations (correct for non-negative weights)
def dijkstra (graph : Graph) (n src : Nat) : List (Option Nat) :=
  let init := (List.range n).map (fun i => if i == src then some 0 else none)
  (List.range (n - 1)).foldl (fun d _ => relaxPass d graph) init

-- Shortest distance between two nodes
def shortestDist (graph : Graph) (n src dst : Nat) : Option Nat :=
  match listGet (dijkstra graph n src) dst with
  | some (some d) => some d
  | _ => none

-- listGet via getElem?
theorem listGet_eq (l : List α) (i : Nat) :
    listGet l i = l[i]? := List.head?_drop

-- (List.range n)[src]? = some src when src < n
theorem range_getElem? (src n : Nat) (h : src < n) :
    (List.range n)[src]? = some src := by
  induction n generalizing src with
  | zero => omega
  | succ n ih =>
    rw [List.range_succ]
    by_cases hlt : src < n
    · rw [List.getElem?_append_left (by simp [List.length_range, hlt])]
      exact ih src hlt
    · have heq : src = n := Nat.le_antisymm (Nat.lt_succ_iff.mp h) (by omega)
      subst heq
      rw [List.getElem?_append_right (by simp [List.length_range])]
      simp [Nat.sub_self]

-- Initial dist list has some (some 0) at src
theorem listGet_init (n src : Nat) (h : src < n) :
    listGet ((List.range n).map (fun i => if i == src then some 0 else none)) src =
    some (some 0) := by
  rw [listGet_eq, List.getElem?_map, range_getElem? src n h]
  simp

-- updateDist preserves some (some 0) at position src
theorem updateDist_preserves_zero (dists : List (Option Nat)) (i d src : Nat)
    (h : listGet dists src = some (some 0)) :
    listGet (updateDist dists i d) src = some (some 0) := by
  rw [listGet_eq] at h ⊢
  have hlen : src < dists.length := by
    rcases Nat.lt_or_ge src dists.length with lt | ge
    · exact lt
    · exact absurd h (by simp [List.getElem?_eq_none ge])
  unfold updateDist
  rw [List.getElem?_zipWith, range_getElem? src dists.length hlen, h]
  -- Goal: some ((fun j curr => if j==i then … else curr) src (some 0)) = some (some 0)
  simp only  -- beta-reduce
  by_cases h_eq : i = src
  · subst h_eq; simp
  · simp [beq_eq_false_iff_ne.mpr (Ne.symm h_eq)]

-- relaxEdge preserves some (some 0) at src
theorem relaxEdge_preserves_zero (dists : List (Option Nat)) (fr to w src : Nat)
    (h : listGet dists src = some (some 0)) :
    listGet (relaxEdge dists fr to w) src = some (some 0) := by
  unfold relaxEdge
  split
  · exact updateDist_preserves_zero dists to _ src h
  · exact h

-- relaxPass preserves some (some 0) at src
theorem relaxPass_preserves_zero (dists : List (Option Nat)) (edges : Graph) (src : Nat)
    (h : listGet dists src = some (some 0)) :
    listGet (relaxPass dists edges) src = some (some 0) := by
  unfold relaxPass
  induction edges generalizing dists with
  | nil => simpa
  | cons e rest ih =>
    simp only [List.foldl_cons]
    exact ih _ (relaxEdge_preserves_zero dists e.1 e.2.1 e.2.2 src h)

-- foldl of relaxPass preserves the invariant (generalizing over init)
theorem foldl_relaxPass_preserves_zero
    (rounds : Nat) (graph : Graph) (src : Nat) (init : List (Option Nat))
    (h : listGet init src = some (some 0)) :
    listGet ((List.range rounds).foldl (fun d _ => relaxPass d graph) init) src =
    some (some 0) := by
  induction rounds generalizing init with
  | zero => simpa
  | succ k ih =>
    rw [List.range_succ, List.foldl_append]
    simp only [List.foldl_cons, List.foldl_nil]
    exact relaxPass_preserves_zero _ graph src (ih _ h)

-- Self-distance is always 0 for valid source nodes
theorem self_distance_zero (graph : Graph) (n src : Nat) (h : src < n) :
    shortestDist graph n src src = some 0 := by
  let init := (List.range n).map (fun i => if i == src then some 0 else none)
  unfold shortestDist
  rw [show dijkstra graph n src =
      (List.range (n - 1)).foldl (fun d _ => relaxPass d graph) init from rfl,
    foldl_relaxPass_preserves_zero (n - 1) graph src init (listGet_init n src h)]

-- DRT checks
example : shortestDist [(0, 1, 5), (1, 2, 3)] 3 0 2 = some 8 := by native_decide
example : shortestDist [(0, 1, 5)] 3 0 2 = none := by native_decide
example : shortestDist [] 3 0 0 = some 0 := by native_decide
example : shortestDist [(0, 1, 1), (0, 2, 4), (1, 2, 2)] 3 0 2 = some 3 := by native_decide
example : dijkstra [(0, 1, 5), (1, 2, 3)] 3 0 = [some 0, some 5, some 8] := by native_decide

end ShortestPathSpec