-- Pure strategy profiles: choosing one pure strategy per player gives a
-- point of the product of simplices, the anchor Nash's theorem needs
-- 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.GameTheory.PureProfile
import FixedPoint.BlockSimplex

-- The sequence that is one at r and zero elsewhere.
defgraph pointMass (r : Set) (i ∈ ℕ) : ℝ := (realNat(indicator({r}, i)):R)
proof
  have natural := indicator.natural at {r}, i
  exact realNat.real at indicator({r}, i)

theorem pointMass.sequence (r : Set)
  shows RealSequence(pointMass(r))
proof
  exact pointMass.function at r

theorem realPartialSum.pointMass_outside (r : Set) (m : Nat)
  shows ¬r ∈ m → realPartialSum(pointMass(r), m) = (0:R)
proof
  induction m
  | zero =>
    intro outside
    exact realPartialSum.zero at pointMass(r)
  | succ(j) using ih =>
    intro outside
    have jNat : Nat(j) by typecheck
    have jIn := nat.omega.forward with jNat
    have notInJ : ¬r ∈ j
      intro inJ
      apply outside
      apply bootstrap.succ.membership.backward
      left
      exact inJ
    have notEqual : r ≠ j
      intro equal
      apply outside
      rw equal
      exact bootstrap.succ.self at j
    have previous := ih with notInJ
    have massValue := pointMass.value at r, j
    have other := indicator.singleton_other at r, j with notEqual
    have step := realPartialSum.succ at pointMass(r), j
    have sequence := pointMass.sequence at r
    have termReal := RealSequence.value at pointMass(r), j with sequence
    have termZero : pointMass(r)(j) = (0:R)
      rw [massValue, other]
      exact realNat.zero
    rw step
    linarith

theorem realPartialSum.pointMass_below (r : Set) (m : Nat)
  shows r ∈ m → realPartialSum(pointMass(r), m) = (1:R)
proof
  induction m
  | zero =>
    intro below
    exfalso
    rw nat.zero at below
    simpa using below
  | succ(j) using ih =>
    intro below
    have jNat : Nat(j) by typecheck
    have jIn := nat.omega.forward with jNat
    have massValue := pointMass.value at r, j
    have step := realPartialSum.succ at pointMass(r), j
    have sequence := pointMass.sequence at r
    have sumReal := realPartialSum.real at pointMass(r), j with sequence
    have termReal := RealSequence.value at pointMass(r), j with sequence
    have split := bootstrap.succ.membership.forward with below
    cases split with inJ | same
    . have previous := ih with inJ
      have notEqual : r ≠ j
        intro equal
        rw equal at inJ
        have irreflexive := membership.irreflexive at j
        exact irreflexive with inJ
      have other := indicator.singleton_other at r, j with notEqual
      have termZero : pointMass(r)(j) = (0:R)
        rw [massValue, other]
        exact realNat.zero
      rw step
      linarith
    . have notInJ : ¬r ∈ j
        intro inJ
        rw same at inJ
        have irreflexive := membership.irreflexive at j
        exact irreflexive with inJ
      have outside := realPartialSum.pointMass_outside at r, j
      have previous := outside with notInJ
      have termOne : pointMass(r)(j) = (1:R)
        rw [massValue, same]
        have self := indicator.singleton_self at j
        rw self
        exact realNat.one
      rw step
      linarith

-- The profile that plays the strategies in R with probability one.
defgraph pureProfile (n R : Set) (i ∈ n) : ℝ := (realNat(indicator(R, i)):R)
requires nNat : Nat(n)
proof
  have natural := indicator.natural at R, i
  exact realNat.real at indicator(R, i)

theorem pureProfile.in_cube (n : Nat) (R : Set)
  shows pureProfile(n, R) ∈ unitCube(n)
proof
  have function := pureProfile.function at n, R with *
  have zeroIn : (0:R) ∈ ℝ by typecheck
  have oneIn : (1:R) ∈ ℝ by typecheck
  apply unitCube.membership.backward
  constructor
  . exact functionSpace.mem_iff.backward with function
  . intros i hi
    have value := pureProfile.value at n, R, i with *
    apply closedInterval.membership.backward
    rw value
    have alternatives : i ∈ R ∨ ¬i ∈ R by tauto
    cases alternatives with inR | outside
    . have one := indicator.member with inR
      rw [one, realNat.one]
      have oneLe : (1:R) ≤ (1:R) by linarith
      have zeroLe : (0:R) ≤ (1:R) by linarith
      exact ⟨oneIn, ⟨zeroLe, oneLe⟩⟩
    . have zero := indicator.absent with outside
      rw [zero, realNat.zero]
      have zeroLe : (0:R) ≤ (0:R) by linarith
      have zeroLeOne : (0:R) ≤ (1:R) by linarith
      exact ⟨zeroIn, ⟨zeroLe, zeroLeOne⟩⟩

-- With one representative strategy per player, the masked terms of a block
-- form the point mass at that player's representative.
theorem pureProfile.masked (n : Nat) (m blk rep p i : Set)
  requires blkFunction : blk : n → m
  requires repFunction : rep : m → n
  requires section : ∀ q ∈ m, blk(rep(q)) = q
  requires hp : p ∈ m
  requires hi : i ∈ n
  shows maskedPoint(n, blockSet(n, blk, p), pureProfile(n, functionImage(rep, m, n)))(i) = pointMass(rep(p))(i)
proof
  have aFunction := pureProfile.function at n, functionImage(rep, m, n) with *
  have iNat := natural.member at n with hi
  have iIn := nat.omega.forward with iNat
  have massValue := pointMass.value at rep(p), i
  have alternatives : i = rep(p) ∨ i ≠ rep(p) by tauto
  cases alternatives with same | different
  . have inBlock : i ∈ blockSet(n, blk, p)
      apply blockSet.membership.backward
      constructor
      . exact hi
      . rw same
        exact section with hp
    have onBlock := maskedPoint.on_block with aFunction, hi, inBlock
    have value := pureProfile.value at n, functionImage(rep, m, n), i with *
    have inImage : i ∈ functionImage(rep, m, n)
      apply functionImage.membership.backward
      constructor
      . exact hi
      . witness p
        exact ⟨hp, same⟩
    have one := indicator.member with inImage
    rw [onBlock, value, one, massValue, same]
    have self := indicator.singleton_self at rep(p)
    rw self
    rfl
  . have differentReversed : rep(p) ≠ i
      intro equal
      apply different
      rw equal
      rfl
    have other := indicator.singleton_other with differentReversed
    have massZero : pointMass(rep(p))(i) = (0:R)
      rw [massValue, other]
      exact realNat.zero
    rw massZero
    have blockAlternatives : i ∈ blockSet(n, blk, p) ∨ ¬i ∈ blockSet(n, blk, p) by tauto
    cases blockAlternatives with inBlock | outside
    . have onBlock := maskedPoint.on_block with aFunction, hi, inBlock
      have value := pureProfile.value at n, functionImage(rep, m, n), i with *
      have notImage : ¬i ∈ functionImage(rep, m, n)
        intro inImage
        have parts := functionImage.membership.forward with inImage
        obtain ⟨q, hq, equal⟩ := parts.right
        have blockParts := blockSet.membership.forward with inBlock
        have blockI := blockParts.right
        have sectionQ := section with hq
        rw [equal, sectionQ] at blockI
        apply different
        rw [equal, blockI]
        rfl
      have zero := indicator.absent with notImage
      rw [onBlock, value, zero]
      exact realNat.zero
    . exact maskedPoint.off_block with aFunction, hi, outside

theorem pureProfile.in_product (n : Nat) (m blk rep : Set)
  requires blkFunction : blk : n → m
  requires repFunction : rep : m → n
  requires section : ∀ q ∈ m, blk(rep(q)) = q
  shows pureProfile(n, functionImage(rep, m, n)) ∈ simplexProduct(n, m, blk)
proof
  have aCube := pureProfile.in_cube at n, functionImage(rep, m, n)
  apply simplexProduct.membership.backward
  constructor
  . exact aCube
  . intros p hp
    have aFunction := unitCube.point with aCube
    have sequence := maskedPoint.sequence at n, blockSet(n, blk, p), pureProfile(n, functionImage(rep, m, n)) with aFunction
    have mass := pointMass.sequence at rep(p)
    have agree : ∀ k ∈ ℕ, k < n → maskedPoint(n, blockSet(n, blk, p), pureProfile(n, functionImage(rep, m, n)))(k) = pointMass(rep(p))(k)
      intros k hk hkn
      have hkIn : k ∈ n by exact membership.of_Less with hkn
      exact pureProfile.masked with blkFunction, repFunction, section, hp, hkIn
    have equal := realPartialSum.eq_prefix with sequence, mass, agree
    have repIn := FunctionBetween.apply_closed with repFunction, hp
    have one := realPartialSum.pointMass_below at rep(p), n
    have blockValue := blockSum.definition at n, blk, p
    have cubeValue := sumCube.value at n, blockSet(n, blk, p), pureProfile(n, functionImage(rep, m, n)) with aCube
    rw blockValue at cubeValue
    rw [cubeValue, equal]
    exact one with repIn