-- File: lean/ShortestPathSpecAeneas.lean
-- Shortest Path (Dijkstra) — Aeneas-compatible specification.
-- Types mirror what Aeneas generates from Vec<(u64,u64,u64)> Rust code.
-- Requires: lake update  (fetches the Aeneas package)

import Aeneas
open Aeneas.Primitives

namespace ShortestPathSpecAeneas

-- Aeneas maps Vec<(u64, u64, u64)> → List (U64 × U64 × U64)
abbrev Graph := List (U64 × U64 × U64)  -- (from, to, weight)

-- Helper: safe list indexing
def listGet (l : List (Option U64)) (i : U64) : Option (Option U64) :=
  l.get? i.val.toNat

-- Update distance at index i if new distance d is smaller
def updateDist (dists : List (Option U64)) (i d : U64) : List (Option U64) :=
  List.zipWith (fun j curr =>
    if j == i.val.toNat then
      match curr with
      | none           => some d
      | some existing  => some (if d.val < existing.val then d else existing)
    else curr)
    (List.range dists.length)
    dists

-- Relax one edge
def relaxEdge (dists : List (Option U64)) (fr to w : U64) : List (Option U64) :=
  match listGet dists fr with
  | some (some df) =>
    match df + w with  -- saturating add avoids overflow for reasonable weights
    | some newD => updateDist dists to newD
    | none      => dists
  | _ => dists

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

-- Bellman-Ford: n-1 relaxation passes
def dijkstra (graph : Graph) (n src : U64) : Result (List (Option U64)) :=
  let init := (List.range n.val.toNat).map (fun i =>
    if i == src.val.toNat then some (0 : U64) else none)
  let result := (List.range (n.val.toNat - 1)).foldl (fun d _ => relaxPass d graph) init
  .ok result

-- Shortest distance src → dst
def shortestDist (graph : Graph) (n src dst : U64) : Result (Option U64) :=
  do let dists ← dijkstra graph n src
     .ok (match listGet dists dst with
          | some (some d) => some d
          | _             => none)

-- Self-distance is zero when src < n
theorem self_distance_zero (graph : Graph) (n src : U64) (h : src.val < n.val) :
    shortestDist graph n src src = .ok (some 0) := by
  simp only [shortestDist, dijkstra, Result.bind]
  simp only [listGet, List.get?]
  simp only [relaxPass, relaxEdge, updateDist]
  -- The init list at index src is some 0; relaxation preserves it.
  sorry -- Full proof requires showing relaxation doesn't worsen src's own distance

end ShortestPathSpecAeneas