-- A decreasing count of open flags for Dijkstra iterations.
-- 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.DijkstraProgress
import DijkstraStep

namespace Dijkstra


def openCount (flags : List(ℕ)) → Nat
  | List.nil => 0
  | List.cons(head, tail) => chooseNat(equalNat(head, 1), openCount(tail), succ(openCount(tail)))

theorem openCount.closed (flags : List(ℕ))
  shows Nat(openCount(flags))
proof
  induction flags
  | List.nil =>
    have computed : openCount(List.nil) = 0 by settle with openCount
    rw computed
    exact Naturals.constructor.zero
  | List.cons(head, tail) using ih =>
    have headNat : Nat(head) by typecheck
    have condition := Dijkstra.equalNat.boolean at head, 1 with headNat, assumption
    have shape := Dijkstra.boolean_cases with condition
    have nextNat := Naturals.constructor.succ with ih
    cases shape with closed | open
    . have branch := closed
      have computed : openCount(List.cons(head, tail)) = succ(openCount(tail))
        settle with openCount, chooseNat, closed, Grace.Cumulus.Nat.Mod.equalNat.reflexive
      rw computed
      exact nextNat
    . have branch := open
      have computed : openCount(List.cons(head, tail)) = openCount(tail) by settle with openCount, chooseNat, open
      rw computed
      exact ih

theorem openCount.initial (count : Nat)
  shows openCount(Dijkstra.initialFlags(count)) = count
proof
  induction count
  | zero =>
    simp
  | succ(k) using ih =>
    simp [Dijkstra.initialFlags, openCount, equalNat, isZero, chooseNat]

theorem openCount.mark_decreases (index : Nat, flags : List(ℕ))
  requires open : Dijkstra.flagAt(flags, index) ≠ 1
  shows openCount(flags) = succ(openCount(Dijkstra.mark(flags, index)))
proof
  induction index generalizing flags
  | zero =>
    intros flags flagsType open
    have shape := Dijkstra.list_cases with flagsType
    cases shape with empty | cons
    . rw empty at open
      have default := Dijkstra.flagAt.nil with Naturals.constructor.zero
      exfalso
      exact open with default
    . rcases cons with ⟨head, headIn, tail, tailType, equality⟩
      have headNat := nat.omega.backward with headIn
      rw equality at open
      have position : Dijkstra.flagAt(List.cons(head, tail), 0) = head
        settle with Dijkstra.flagAt, Dijkstra.boolHead
      rw position at open
      have condition := Dijkstra.equalNat.false_iff.backward with headNat, assumption, open
      have same := Grace.Cumulus.Nat.Mod.equalNat.reflexive at 1 with assumption
      rw equality
      simp [Dijkstra.mark, Dijkstra.markHead, openCount, chooseNat]
  | succ(k) using ih =>
    intros flags flagsType open
    have shape := Dijkstra.list_cases with flagsType
    cases shape with empty | cons
    . rw empty at open
      have skNat : Nat(succ(k)) by typecheck
      have default := Dijkstra.flagAt.nil with skNat
      exfalso
      exact open with default
    . rcases cons with ⟨head, headIn, tail, tailType, equality⟩
      have headNat := nat.omega.backward with headIn
      rw equality at open
      have position : Dijkstra.flagAt(List.cons(head, tail), succ(k)) = Dijkstra.flagAt(tail, k)
        settle with Dijkstra.flagAt, Dijkstra.boolTail
      rw position at open
      have previous := ih with tail, tailType, open
      have condition := Dijkstra.equalNat.boolean at head, 1 with headNat, assumption
      have alternatives := Dijkstra.boolean_cases with condition
      rw equality
      cases alternatives with closed | openHead
      . simp [Dijkstra.mark, Dijkstra.keepHead, Dijkstra.boolTail, openCount, chooseNat]
      . simp [Dijkstra.mark, Dijkstra.keepHead, Dijkstra.boolTail, openCount, chooseNat]

theorem advance.openCount_decreases (rows : Set, values : List(Dijkstra.DistanceValues), flags : List(ℕ), index distance : Nat)
  requires chosen : Dijkstra.minimum(values, flags, 0) = Option.some(Dijkstra.Candidate.value(index, distance))
  shows openCount(flags) = succ(openCount(Dijkstra.settled(Dijkstra.advance(rows, Dijkstra.State.current(values, flags)))))
proof
  have selected := Dijkstra.minimum.selected_entry at Dijkstra.DistanceValues, values, flags, index, distance with *{4}, Dijkstra.DistanceValues.cells, chosen
  rcases selected with ⟨bound, entryValue, open⟩
  have decreases := openCount.mark_decreases at index, flags with *
  have computed : Dijkstra.settled(Dijkstra.advance(rows, Dijkstra.State.current(values, flags))) = Dijkstra.mark(flags, index)
    settle with Dijkstra.advance, Dijkstra.advanceChosen, Dijkstra.settled, chosen
  rw computed
  exact decreases

theorem openCount.zero_all_closed (flags : List(ℕ), index : Nat)
  requires zeroCount : openCount(flags) = 0
  shows Dijkstra.flagAt(flags, index) = 1
proof
  induction flags generalizing index
  | List.nil =>
    intros index indexNat zeroCount
    exact Dijkstra.flagAt.nil with indexNat
  | List.cons(head, tail) using ih =>
    intros index indexNat zeroCount
    have headNat : Nat(head) by typecheck
    have condition := Dijkstra.equalNat.boolean at head, 1 with headNat, assumption
    have alternatives := Dijkstra.boolean_cases with condition
    cases alternatives with open | closed
    . have impossible := Naturals.disjoint.zero.succ at openCount(tail)
      have computed : openCount(List.cons(head, tail)) = succ(openCount(tail))
        settle with openCount, chooseNat, open
      rw computed at zeroCount
      exfalso
      apply impossible
      symm
      exact zeroCount
    . have headOne := Grace.Cumulus.Nat.Mod.equalNat.true_iff.forward at head, 1 with headNat, assumption, closed
      have computed : openCount(List.cons(head, tail)) = openCount(tail)
        settle with openCount, chooseNat, closed, Grace.Cumulus.Nat.Mod.equalNat.reflexive
      rw computed at zeroCount
      have indexShape := natural.zero_or_successor with indexNat
      cases indexShape with zeroIndex | nextIndex
      . rw zeroIndex
        simp [Dijkstra.flagAt, Dijkstra.boolHead]
      . rcases nextIndex with ⟨j, jNat, indexValue⟩
        have previous := ih with j, jNat, zeroCount
        rw indexValue
        simp [Dijkstra.flagAt, Dijkstra.boolTail]

predicate TypedRows (rows : Set) :=
  ∀ index ∈ ℕ, List(Dijkstra.DistanceValues, Dijkstra.rowAt(rows, index))

theorem iterate.exhausts_finite_frontier (fuel : Nat, rows : Set, values : List(Dijkstra.DistanceValues), flags : List(ℕ))
  requires rowsTyped : TypedRows(rows)
  requires budget : openCount(flags) ≤ fuel
  shows Dijkstra.minimum(Dijkstra.distances(Dijkstra.iterate(fuel, rows, Dijkstra.State.current(values, flags))),
    Dijkstra.settled(Dijkstra.iterate(fuel, rows, Dijkstra.State.current(values, flags))), 0) = Option.none
proof
  induction fuel generalizing values flags
  | zero =>
    intros values flags valuesType flagsType budget
    have countNat := openCount.closed with flagsType
    have nonnegative := Grace.Core.Nat.Order.LessEqual.zero with countNat
    have zeroCount := LessEqual.antisymmetric at openCount(flags), 0 with countNat, Naturals.constructor.zero, budget, nonnegative
    have noOpen : ∀ target ∈ ℕ, ∀ distance ∈ ℕ, Less(target, ListOps.length{A := Dijkstra.DistanceValues}(values)) → Dijkstra.entry(values, target) = Option.some(distance) → Dijkstra.flagAt(flags, target) = 1
      intros target targetIn distance distanceIn bound value
      have targetNat := nat.omega.backward with targetIn
      exact openCount.zero_all_closed with flagsType, targetNat, zeroCount
    have absent := Dijkstra.minimum.none_iff_no_open_finite.backward at Dijkstra.DistanceValues, values, flags with valuesType, flagsType, Dijkstra.DistanceValues.cells, noOpen
    simp [Dijkstra.iterate, Dijkstra.distances, Dijkstra.settled]
  | succ(k) using ih =>
    intros values flags valuesType flagsType budget
    have correct := Dijkstra.minimum.correct at Dijkstra.DistanceValues, values, flags, 0 with valuesType, flagsType, Naturals.constructor.zero, Dijkstra.DistanceValues.cells
    rcases correct with ⟨valid, member, least⟩
    have skNat : Nat(succ(k)) by typecheck
    cases valid with absent | finite
    . have fixed := Dijkstra.iterate.none at succ(k), rows, values, flags with skNat, absent
      rw fixed
      simp [Dijkstra.distances, Dijkstra.settled]
    . rcases finite with ⟨index, indexIn, distance, distanceIn, chosen⟩
      have indexNat := nat.omega.backward with indexIn
      have distanceNat := nat.omega.backward with distanceIn
      have selected := Dijkstra.minimum.selected_entry at Dijkstra.DistanceValues, values, flags, index, distance with valuesType, flagsType, indexNat, distanceNat, Dijkstra.DistanceValues.cells, chosen
      rcases selected with ⟨bound, value, open⟩
      have decreases := openCount.mark_decreases with indexNat, flagsType, open
      have markedType := Dijkstra.mark.closed with indexNat, flagsType
      have markedCountNat := openCount.closed with markedType
      have kNat : Nat(k) by typecheck
      rw decreases at budget
      have smaller : openCount(Dijkstra.mark(flags, index)) ≤ k
        apply LessEqual.add_right_cancel at openCount(Dijkstra.mark(flags, index)), k, 1
        . have leftStep := add.right_one at openCount(Dijkstra.mark(flags, index)) with markedCountNat
          have rightStep := add.right_one with kNat
          rw leftStep
          rw rightStep
          exact budget
      have rowType := rowsTyped with index, indexIn
      have sourceType := Option.inductive.constructor.some at ℕ, distance with distanceIn
      have relaxedType := Dijkstra.relax.closed with valuesType, rowType, markedType, sourceType
      have previous := ih with Dijkstra.relax(Dijkstra.rowAt(rows, index), values, Dijkstra.mark(flags, index), Option.some(distance)), Dijkstra.mark(flags, index), relaxedType, markedType, smaller
      have advanced : Dijkstra.advance(rows, Dijkstra.State.current(values, flags)) = Dijkstra.State.current(Dijkstra.relax(Dijkstra.rowAt(rows, index), values, Dijkstra.mark(flags, index), Option.some(distance)), Dijkstra.mark(flags, index))
        settle with Dijkstra.advance, Dijkstra.advanceChosen, chosen
      simp [Dijkstra.iterate]

theorem solve.exhausts_finite_frontier (count source : Nat, rows : Set)
  requires rowsTyped : TypedRows(rows)
  shows Dijkstra.minimum(Dijkstra.distances(Dijkstra.solve(count, source, rows)), Dijkstra.settled(Dijkstra.solve(count, source, rows)), 0) = Option.none
proof
  have valuesType := Dijkstra.initialDistances.closed at count, 0, source with Naturals.constructor.zero
  have countValue := openCount.initial at count
  have budget : openCount(Dijkstra.initialFlags(count)) ≤ count
    rw countValue
    exact LessEqual.reflexive with assumption
  have terminated := iterate.exhausts_finite_frontier with rowsTyped, budget
  unfold Dijkstra.solve
  exact terminated

end Dijkstra