-- Initialization of Dijkstra's executable distance and flag vectors.
-- Part of Cumulus, the Grace17 mathematics library.
-- Copyright (c) 2026 Grant Passmore.
-- Released under the Apache License 2.0. See LICENSE in the repository root.
-- Authors: Grant Passmore

module Grace.Cumulus.GraphTheory.DijkstraInitial
import DijkstraFlags

namespace Dijkstra


theorem initialCell.closed (index source : Nat)
  shows Option(ℕ, Dijkstra.initialCell(index, source))
proof
  have condition := Dijkstra.equalNat.boolean at index, source with *
  have choices := Dijkstra.chooseDistance.selects at Option.some(0), Option.none, equalNat(index, source) with condition
  have computed : Dijkstra.initialCell(index, source) = Dijkstra.chooseDistance(equalNat(index, source), Option.some(0), Option.none)
    settle with Dijkstra.initialCell
  rw computed
  cases choices with finite | absent
  . rw finite
    have zeroIn := nat.omega.forward with Naturals.constructor.zero
    exact Option.inductive.constructor.some with zeroIn
  . rw absent
    exact Option.inductive.constructor.none

theorem initialDistances.closed (count index source : Nat)
  shows List(Dijkstra.DistanceValues, Dijkstra.initialDistances(count, index, source))
proof
  induction count generalizing index
  | zero =>
    intros index indexNat
    have computed : Dijkstra.initialDistances(0, index, source) = List.nil by settle with Dijkstra.initialDistances
    rw computed
    exact List.inductive.constructor.nil
  | succ(k) using ih =>
    intros index indexNat
    have cellType := initialCell.closed at index, source with indexNat, assumption
    have cellIn := Dijkstra.DistanceValues.membership.backward with cellType
    have nextNat := Naturals.constructor.succ with indexNat
    have tailType := ih with succ(index), nextNat
    have constructed := List.inductive.constructor.cons at Dijkstra.DistanceValues,
      Dijkstra.initialCell(index, source), Dijkstra.initialDistances(k, succ(index), source) with cellIn, tailType
    have computed : Dijkstra.initialDistances(succ(k), index, source) = List.cons(Dijkstra.initialCell(index, source), Dijkstra.initialDistances(k, succ(index), source))
      settle with Dijkstra.initialDistances
    rw computed
    exact constructed

theorem initialDistances.length (count index source : Nat)
  shows Dijkstra.vectorLength(Dijkstra.initialDistances(count, index, source)) = count
proof
  induction count generalizing index
  | zero =>
    intros index indexNat
    simp [Dijkstra.initialDistances, Dijkstra.vectorLength, ListOps.length]
  | succ(k) using ih =>
    intros index indexNat
    have nextNat := Naturals.constructor.succ with indexNat
    have previous := ih with succ(index), nextNat
    unfold Dijkstra.vectorLength at previous
    simp [Dijkstra.initialDistances, Dijkstra.vectorLength, ListOps.length]

theorem initialFlags.closed (count : Nat)
  shows List(ℕ, Dijkstra.initialFlags(count))
proof
  induction count
  | zero =>
    have computed : Dijkstra.initialFlags(0) = List.nil by simp [Dijkstra.initialFlags]
    rw computed
    exact List.inductive.constructor.nil
  | succ(k) using ih =>
    have zeroIn := nat.omega.forward with Naturals.constructor.zero
    have constructed := List.inductive.constructor.cons with zeroIn, ih
    have computed : Dijkstra.initialFlags(succ(k)) = List.cons(0, Dijkstra.initialFlags(k))
      settle with Dijkstra.initialFlags
    rwa computed

theorem initialFlags.length (count : Nat)
  shows ListOps.length{A := ℕ}(Dijkstra.initialFlags(count)) = count
proof
  induction count
  | zero =>
    simp
  | succ(k) using ih =>
    simp [Dijkstra.initialFlags, ListOps.length]

-- Indexed semantics of initialization, for arbitrary vector lengths.
theorem initialFlags.entry (count target : Nat)
  requires bound : Less(target, count)
  shows Dijkstra.flagAt(Dijkstra.initialFlags(count), target) = 0
proof
  induction count generalizing target
  | zero =>
    intros target targetNat bound
    have impossible := Dijkstra.no_less_zero with targetNat
    exfalso
    exact impossible with bound
  | succ(k) using ih =>
    intros target targetNat bound
    have shape := natural.zero_or_successor with targetNat
    cases shape with atZero | atNext
    . rw atZero
      settle with Dijkstra.initialFlags, Dijkstra.flagAt, Dijkstra.boolHead
    . rcases atNext with ⟨j, jNat, position⟩
      rw position at bound
      have smaller := Dijkstra.successor_less_iff.forward with jNat, bound
      have previous := ih with j, jNat, smaller
      rw position
      simp [Dijkstra.initialFlags, Dijkstra.flagAt, Dijkstra.boolTail]

theorem initialDistances.entry (count index source target : Nat)
  requires bound : Less(target, count)
  shows Dijkstra.entry(Dijkstra.initialDistances(count, index, source), target) = Dijkstra.initialCell(add(target, index), source)
proof
  induction count generalizing index target
  | zero =>
    intros index target indexNat targetNat bound
    have impossible := Dijkstra.no_less_zero with targetNat
    exfalso
    exact impossible with bound
  | succ(k) using ih =>
    intros index target indexNat targetNat bound
    have shape := natural.zero_or_successor with targetNat
    cases shape with atZero | atNext
    . rw atZero
      settle with Dijkstra.initialDistances, Dijkstra.entry, Dijkstra.optionHead, add
    . rcases atNext with ⟨j, jNat, position⟩
      rw position at bound
      have smaller := Dijkstra.successor_less_iff.forward with jNat, bound
      have previous := ih with succ(index), j, jNat, smaller
      have shifted : add(j, succ(index)) = add(succ(j), index) by simp [add, add.succ_right]
      rw shifted at previous
      rw position
      simp [Dijkstra.initialDistances, Dijkstra.entry, Dijkstra.optionTail]

theorem initialCell.at_source (source : Nat)
  shows Dijkstra.initialCell(source, source) = Option.some(0)
proof
  have same := Grace.Cumulus.Nat.Mod.equalNat.reflexive at source
  simp [Dijkstra.initialCell, Dijkstra.chooseDistance]

theorem initialCell.witnessed (V E weights : Set, source target : Nat)
  shows Dijkstra.WitnessedDistance(V, E, weights, source, target, Dijkstra.initialCell(target, source))
proof
  have condition := Dijkstra.equalNat.boolean at target, source with *
  have choices := Dijkstra.boolean_cases with condition
  cases choices with no | yes
  . have computed : Dijkstra.initialCell(target, source) = Option.none
      settle with Dijkstra.initialCell, Dijkstra.chooseDistance
    rw computed
    exact Dijkstra.WitnessedDistance.none
  . have same := Grace.Cumulus.Nat.Mod.equalNat.true_iff.forward at target, source with yes
    rw [same, initialCell.at_source]
    have path := HasPathCost.source at V, E, weights, source
    exact Dijkstra.WitnessedDistance.some with Naturals.constructor.zero, path

theorem initialDistances.lookup (count source target : Nat)
  requires bound : Less(target, count)
  shows Dijkstra.entry(Dijkstra.initialDistances(count, 0, source), target) = Dijkstra.initialCell(target, source)
proof
  have indexed := initialDistances.entry at count, 0, source, target with Naturals.constructor.zero, bound
  simpa [add.right_unit] using indexed

end Dijkstra