Deferred-acceptance states, updates and finite termination measure

import Mathlib.Algebra.BigOperators.Ring.Finset
import Mathlib.Algebra.Order.BigOperators.Group.Finset
import Mathlib.Data.Finset.Basic
import Mathlib.Data.Finset.Max
import Mathlib.Data.Fintype.Basic
import Mathlib.Data.Fintype.Card
import Mathlib.Data.Fintype.Perm
import Mathlib.Data.Real.Basic
import Mathlib.Tactic.Linarith
import Definitions.Def_AMLGS62_AppliedModelingLib_Markets_Matching_Basic

namespace AppliedModelingLib
namespace Matching

variable {M W : Type*} [Fintype M] [Fintype W] [DecidableEq M] [DecidableEq W]

structure DAState (M W : Type*) where
  m_match : M → Option W
  w_match : W → Option M
  m_proposals : M → Finset W
  consistent : ∀ m w, m_match m = some w ↔ w_match w = some m

def initialDAState (M W : Type*) [Fintype W] : DAState M W where
  m_match _ := none
  w_match _ := none
  m_proposals _ := Finset.univ
  consistent m w := by simp

def IsActiveMan (val_m : M → W → ℝ) (s : DAState M W) (m : M) : Prop :=
  s.m_match m = none ∧ ∃ w ∈ s.m_proposals m, 0 ≤ val_m m w

def BestRemainingWoman (val_m : M → W → ℝ) (s : DAState M W) (m : M) (w : W) : Prop :=
  w ∈ s.m_proposals m ∧ 0 ≤ val_m m w ∧
  ∀ w' ∈ s.m_proposals m, 0 ≤ val_m m w' → val_m m w' ≤ val_m m w

/-- Remove the proposal opportunity `m -> w` from a DA state's proposal sets. -/
def removeProposal (s : DAState M W) (m : M) (w : W) : M → Finset W :=
  fun m' => if m' = m then s.m_proposals m \ {w} else s.m_proposals m'

lemma exists_best_woman (val_m : M → W → ℝ) (s : DAState M W) (m : M)
    (hact : IsActiveMan val_m s m) : ∃ w, BestRemainingWoman val_m s m w := by
  rcases hact with ⟨_, w0, hw0_mem, hw0_nonneg⟩
  let eligible : Finset W := (s.m_proposals m).filter fun w => 0 ≤ val_m m w
  have helig_nonempty : eligible.Nonempty := by
    refine ⟨w0, ?_⟩
    simp [eligible, hw0_mem, hw0_nonneg]
  obtain ⟨w, hw_eligible, hw_max⟩ :=
    Finset.exists_max_image eligible (fun w => val_m m w) helig_nonempty
  refine ⟨w, ?_, ?_, ?_⟩
  · exact (Finset.mem_filter.mp hw_eligible).1
  · exact (Finset.mem_filter.mp hw_eligible).2
  · intro w' hw'_mem hw'_nonneg
    exact hw_max w' (by simp [eligible, hw'_mem, hw'_nonneg])

lemma acceptMatch_consistent (s : DAState M W) {m : M} {w : W}
    (hm_unmatched : s.m_match m = none) :
    ∀ m0 w0,
      (if m0 = m then some w
        else if s.w_match w = some m0 then none
        else s.m_match m0) = some w0 ↔
      Function.update s.w_match w (some m) w0 = some m0 := by
  intro m0 w0
  by_cases hm : m0 = m
  · subst m0
    by_cases hw : w0 = w
    · subst w0
      simp
    · have hw' : w ≠ w0 := fun h => hw h.symm
      have hnot_old : s.w_match w0 ≠ some m := by
        intro hwm
        have hmmatch : s.m_match m = some w0 := (s.consistent m w0).2 hwm
        rw [hm_unmatched] at hmmatch
        cases hmmatch
      simp [hw, hw', hnot_old]
  · by_cases hw : w0 = w
    · subst w0
      have hm' : m ≠ m0 := fun h => hm h.symm
      by_cases hcur : s.w_match w = some m0
      · simp [hm, hm', hcur]
      · have hnot_old : s.m_match m0 ≠ some w := by
          intro hmmatch
          exact hcur ((s.consistent m0 w).1 hmmatch)
        simp [hm, hm', hcur, hnot_old]
    · by_cases hcur : s.w_match w = some m0
      · have hnot_old : s.w_match w0 ≠ some m0 := by
          intro hwm0
          have hm_w : s.m_match m0 = some w := (s.consistent m0 w).2 hcur
          have hm_w0 : s.m_match m0 = some w0 := (s.consistent m0 w0).2 hwm0
          have : w0 = w := Option.some.inj (hm_w0.symm.trans hm_w)
          exact hw this
        simp [hm, hw, hcur, hnot_old]
      · simpa [hm, hw, hcur] using (s.consistent m0 w0)

noncomputable def daStep (val_m : M → W → ℝ) (val_w : W → M → ℝ)
    (s : DAState M W) : DAState M W :=
  have _ : Decidable (∃ m, IsActiveMan val_m s m) := Classical.propDecidable _
  if h : ∃ m, IsActiveMan val_m s m then
    let m := Classical.choose h
    let hact := Classical.choose_spec h
    let w_exists := exists_best_woman val_m s m hact
    let w := Classical.choose w_exists
    let hw_best := Classical.choose_spec w_exists
    let new_proposals := removeProposal s m w
    let m_current := s.w_match w
    let accepts :=
      match m_current with
      | none => 0 ≤ val_w w m
      | some m' => val_w w m' < val_w w m
    have _ : Decidable accepts := Classical.propDecidable _
    if hacc : accepts then
      let new_w_match := Function.update s.w_match w (some m)
      let new_m_match := fun m'' =>
        if m'' = m then some w
        else if m_current = some m'' then none
        else s.m_match m''
      { m_match := new_m_match
        w_match := new_w_match
        m_proposals := new_proposals
        consistent := by
          simpa [new_m_match, new_w_match, m_current] using
            acceptMatch_consistent s hact.1 }
    else
      { s with m_proposals := new_proposals }
  else
    s

set_option linter.unusedSimpArgs false

set_option linter.unusedSimpArgs true

-- DA algorithm fold
noncomputable def deferredAcceptanceState (val_m : M → W → ℝ) (val_w : W → M → ℝ) : DAState M W :=
  let max_steps := Fintype.card M * Fintype.card W
  (List.range max_steps).foldl (fun s _ => daStep val_m val_w s) (initialDAState M W)

noncomputable def deferredAcceptance (val_m : M → W → ℝ) (val_w : W → M → ℝ) : Assignment M W where
  m_match := (deferredAcceptanceState val_m val_w).m_match
  w_match := (deferredAcceptanceState val_m val_w).w_match
  consistent_m := (deferredAcceptanceState val_m val_w).consistent

/-- Number of proposal opportunities that have not yet been used. -/
def remainingProposalCount (s : DAState M W) : ℕ :=
  ∑ m : M, (s.m_proposals m).card

set_option linter.unusedSimpArgs false

set_option linter.unusedSimpArgs true

def ManIRInvariant (val_m : M → W → ℝ) (s : DAState M W) : Prop :=
  ∀ m w, s.m_match m = some w → 0 ≤ val_m m w

def WomanIRInvariant (val_w : W → M → ℝ) (s : DAState M W) : Prop :=
  ∀ w m, s.w_match w = some m → 0 ≤ val_w w m

def MatchedProposedInvariant (s : DAState M W) : Prop :=
  ∀ m w, s.m_match m = some w → w ∉ s.m_proposals m

def WomanRejectionInvariant (val_w : W → M → ℝ) (s : DAState M W) : Prop :=
  ∀ w m', w ∉ s.m_proposals m' →
    s.m_match m' ≠ some w →
    val_w w m' < 0 ∨ (∃ m, s.w_match w = some m ∧ val_w w m' ≤ val_w w m)

def ManProposalOrderInvariant (val_m : M → W → ℝ) (s : DAState M W) : Prop :=
  ∀ m w w', w ∉ s.m_proposals m → w' ∈ s.m_proposals m → 0 ≤ val_m m w' →
    val_m m w' ≤ val_m m w

def DAInvariants (val_m : M → W → ℝ) (val_w : W → M → ℝ) (s : DAState M W) : Prop :=
  ManIRInvariant val_m s ∧
  WomanIRInvariant val_w s ∧
  MatchedProposedInvariant s ∧
  WomanRejectionInvariant val_w s ∧
  ManProposalOrderInvariant val_m s

def DAStepPreservesInvariantsCertificate (val_m : M → W → ℝ) (val_w : W → M → ℝ) : Prop :=
  ∀ s, DAInvariants val_m val_w s → DAInvariants val_m val_w (daStep val_m val_w s)

def DAStateInvariantCertificate (val_m : M → W → ℝ) (val_w : W → M → ℝ) : Prop :=
  DAInvariants val_m val_w (deferredAcceptanceState val_m val_w)

def DATerminationCertificate (val_m : M → W → ℝ) (val_w : W → M → ℝ) : Prop :=
  ¬ ∃ m, IsActiveMan val_m (deferredAcceptanceState val_m val_w) m

def DaProducesStableMatchingCertificate (val_m : M → W → ℝ) (val_w : W → M → ℝ) : Prop :=
  DAStepPreservesInvariantsCertificate val_m val_w

/-- Men have strict preferences over women. -/
def MenStrictPreferenceProfile (val_m : M → W → ℝ) : Prop :=
  ∀ m w w', val_m m w = val_m m w' → w = w'

/-- Women have strict preferences over men. -/
def WomenStrictPreferenceProfile (val_w : W → M → ℝ) : Prop :=
  ∀ w m m', val_w w m = val_w w m' → m = m'

/--
All potential pairs are acceptable to both sides. This is the marriage-problem
specialization where every agent must be matched.
-/
def AllPairsAcceptable (val_m : M → W → ℝ) (val_w : W → M → ℝ) : Prop :=
  (∀ m w, 0 < val_m m w) ∧ (∀ w m, 0 < val_w w m)

end Matching
end AppliedModelingLib