-- Lifting a Kuhn simplex of the lower cube to one whose far-wall face
-- restricts back to it: the base gains the value p − 1 at the last
-- coordinate and the permutation sends 0 there first
-- 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.FixedPoint.KuhnLift
import KuhnWallSimplex

def liftBase (m q b : Set) : Set := b ∪ constantFunction({m}, ℕ, q)

def liftPerm (m u : Set) : Set := unrotated(succ(m), u ∪ constantFunction({m}, {m}, m))

theorem successor.as_union (m : Nat)
  shows m ∪ {m} = succ(m)
proof
  have asInsert := natural.successor_as_insert at m
  have unfolded : insert(m, m) = {m} ∪ m
    unfold insert
    rfl
  rw [asInsert, unfolded]
  exact union2.commutative

theorem liftBase.function (m q : Nat) (b : Set)
  requires baseFunction : b : m → ℕ
  shows liftBase(m, q, b) : succ(m) → ℕ
proof
  have qNat : Nat(q) by typecheck
  have qIn := nat.omega.forward with qNat
  have constant := constantFunction.function at {m}, ℕ, q with qIn
  have disjoint := cubeExtend.disjoint at m
  have union := FunctionBetween.union_disjoint with baseFunction, constant, disjoint
  have domain := successor.as_union at m
  have codomain := union2.idempotent at ℕ
  rw [domain, codomain] at union
  have same : liftBase(m, q, b) = b ∪ constantFunction({m}, ℕ, q)
    unfold liftBase
    rfl
  rwa same

theorem liftBase.apply_below (m q : Nat) (b j : Set)
  requires baseFunction : b : m → ℕ
  requires hj : j ∈ m
  shows liftBase(m, q, b)(j) = b(j)
proof
  have qNat : Nat(q) by typecheck
  have qIn := nat.omega.forward with qNat
  have constant := constantFunction.function at {m}, ℕ, q with qIn
  have disjoint := cubeExtend.disjoint at m
  have same : liftBase(m, q, b) = b ∪ constantFunction({m}, ℕ, q)
    unfold liftBase
    rfl
  rw same
  exact FunctionBetween.union_apply_left with baseFunction, constant, disjoint, hj

theorem liftBase.apply_top (m q : Nat) (b : Set)
  requires baseFunction : b : m → ℕ
  shows liftBase(m, q, b)(m) = q
proof
  have qNat : Nat(q) by typecheck
  have qIn := nat.omega.forward with qNat
  have constant := constantFunction.function at {m}, ℕ, q with qIn
  have disjoint := cubeExtend.disjoint at m
  have mIn : m ∈ {m}
    apply singleton.membership.backward
    rfl
  have same : liftBase(m, q, b) = b ∪ constantFunction({m}, ℕ, q)
    unfold liftBase
    rfl
  rw same
  have value := FunctionBetween.union_apply_right with baseFunction, constant, disjoint, mIn
  rw value
  exact constantFunction.apply with mIn, qIn

theorem liftBase.restrict (m q : Nat) (b : Set)
  requires baseFunction : b : m → ℕ
  shows restrict(liftBase(m, q, b), m) = b
proof
  have lifted := liftBase.function at m, q, b with baseFunction
  have mSubset := successor.subset at m
  have restricted := restrict.domain_function with lifted, mSubset
  apply FunctionBetween.ext with restricted, baseFunction
  intros j hj
  have value := restrict.domain_apply with lifted, mSubset, hj
  rw value
  exact liftBase.apply_below with baseFunction, hj

theorem constantFunction.bijective_singleton (a : Set)
  shows BijectiveBetween(constantFunction({a}, {a}, a), {a}, {a})
proof
  have aIn : a ∈ {a}
    apply singleton.membership.backward
    rfl
  have function := constantFunction.function at {a}, {a}, a with aIn
  constructor
  . apply Injective.of_cancel with function
    intros x hx y hy same
    have xEq := singleton.membership.forward with hx
    have yEq := singleton.membership.forward with hy
    rw [xEq, yEq]
    rfl
  . constructor
    . exact function
    . intros y hy
      have yEq := singleton.membership.forward with hy
      witness a
      constructor
      . exact aIn
      . have edge := FunctionBetween.relates_apply with function, aIn
        have value := constantFunction.apply with aIn, aIn
        rw value at edge
        rwa yEq

theorem liftPerm.union_bijective (m : Nat) (u : Set)
  requires updBijective : BijectiveBetween(u, m, m)
  shows BijectiveBetween(u ∪ constantFunction({m}, {m}, m), succ(m), succ(m))
proof
  have single := constantFunction.bijective_singleton at m
  have disjoint := cubeExtend.disjoint at m
  have union := Bijective.union_disjoint with updBijective, single, disjoint, disjoint
  have domain := successor.as_union at m
  rwa domain at union

theorem liftPerm.bijective (m : Nat) (u : Set)
  requires updBijective : BijectiveBetween(u, m, m)
  shows BijectiveBetween(liftPerm(m, u), succ(m), succ(m))
proof
  have union := liftPerm.union_bijective with updBijective
  have same : liftPerm(m, u) = unrotated(succ(m), u ∪ constantFunction({m}, {m}, m))
    unfold liftPerm
    rfl
  rw same
  exact unrotated.bijective with union

theorem liftPerm.zero (m : Nat) (u : Set)
  requires updBijective : BijectiveBetween(u, m, m)
  shows liftPerm(m, u)(0) = m
proof
  have union := liftPerm.union_bijective with updBijective
  have uFunction := Injective.function with (Bijective.injective with updBijective)
  have mIn : m ∈ {m}
    apply singleton.membership.backward
    rfl
  have constant := constantFunction.function at {m}, {m}, m with mIn
  have disjoint := cubeExtend.disjoint at m
  have mInSucc := bootstrap.succ.self at m
  have top : succ(m) = succ(m) by rfl
  have same : liftPerm(m, u) = unrotated(succ(m), u ∪ constantFunction({m}, {m}, m))
    unfold liftPerm
    rfl
  rw same
  have wrap := unrotated.wrap with union, mInSucc, top
  rw wrap
  have value := FunctionBetween.union_apply_right with uFunction, constant, disjoint, mIn
  rw value
  exact constantFunction.apply with mIn, mIn

theorem liftPerm.step (m : Nat) (u l : Set)
  requires updBijective : BijectiveBetween(u, m, m)
  requires hl : l ∈ m
  shows liftPerm(m, u)(succ(l)) = u(l)
proof
  have union := liftPerm.union_bijective with updBijective
  have uFunction := Injective.function with (Bijective.injective with updBijective)
  have mIn : m ∈ {m}
    apply singleton.membership.backward
    rfl
  have constant := constantFunction.function at {m}, {m}, m with mIn
  have disjoint := cubeExtend.disjoint at m
  have lIn := membership.into_successor with hl
  have lNat := natural.member at m with hl
  have lLess := Less.of_membership with hl
  have sLess : succ(l) < succ(m) by arith
  have sIn := membership.of_Less with sLess
  have same : liftPerm(m, u) = unrotated(succ(m), u ∪ constantFunction({m}, {m}, m))
    unfold liftPerm
    rfl
  rw same
  have step := unrotated.step with union, lIn, sIn
  rw step
  exact FunctionBetween.union_apply_left with uFunction, constant, disjoint, hl

theorem wallPerm.lift (m : Nat) (u : Set)
  requires updBijective : BijectiveBetween(u, m, m)
  shows wallPerm(m, liftPerm(m, u)) = u
proof
  have lifted := liftPerm.bijective with updBijective
  have zero := liftPerm.zero with updBijective
  have wallBijective := wallPerm.bijective with lifted, zero
  have wallFunction := Injective.function with (Bijective.injective with wallBijective)
  have uFunction := Injective.function with (Bijective.injective with updBijective)
  apply FunctionBetween.ext with wallFunction, uFunction
  intros l hl
  have value := wallPerm.value with lifted, hl
  rw value
  exact liftPerm.step with updBijective, hl

theorem kuhnSimplex.lift_in_family (m q : Nat) (b u : Set)
  requires baseFunction : b : m → ℕ
  requires bounds : ∀ j ∈ m, b(j) < succ(q)
  requires updBijective : BijectiveBetween(u, m, m)
  shows kuhnSimplex(succ(m), liftBase(m, q, b), liftPerm(m, u)) ∈ kuhnSimplices(succ(q), succ(m))
proof
  have liftFunction := liftBase.function at m, q, b with baseFunction
  have lifted := liftPerm.bijective with updBijective
  have liftBounds : ∀ j ∈ succ(m), liftBase(m, q, b)(j) < succ(q)
    intros j hj
    have split := bootstrap.succ.membership.forward with hj
    cases split with below | same
    . have value := liftBase.apply_below at m, q, b, j with baseFunction, below
      rw value
      exact bounds with below
    . rw same
      have value := liftBase.apply_top at m, q, b with baseFunction
      rw value
      arith
  exact kuhnSimplex.in_family with liftFunction, lifted, liftBounds

-- Every vertex but the bottom of the lifted simplex lies on the far wall.
theorem kuhnSimplex.lift_top_wall (m q : Nat) (b u x : Set)
  requires baseFunction : b : m → ℕ
  requires updBijective : BijectiveBetween(u, m, m)
  requires hx : x ∈ kuhnSimplex(succ(m), liftBase(m, q, b), liftPerm(m, u)) ∖
    {kuhnVertex(succ(m), liftBase(m, q, b), liftPerm(m, u))(0)}
  shows x(m) = succ(q)
proof
  have liftFunction := liftBase.function at m, q, b with baseFunction
  have mIn := bootstrap.succ.self at m
  have parts := difference.singleton_membership.forward with hx
  obtain ⟨inSpace, i, hi, shape⟩ := kuhnSimplex.membership.forward with parts.left
  have iNat := natural.member at succ(succ(m)) with hi
  have zeroInI : 0 ∈ i
    have split := natural.zero_or_successor with iNat
    cases split with zero | successor
    . exfalso
      apply parts.right
      rw [shape, zero]
      rfl
    . obtain ⟨k, kNat, iEq⟩ := successor
      rw iEq
      have less : 0 < succ(k) by arith
      exact membership.of_Less with less
  have zero := liftPerm.zero with updBijective
  have raisedM : m ∈ raised(succ(m), liftPerm(m, u), i)
    apply raised.membership.backward
    constructor
    . exact mIn
    . witness 0
      constructor
      . exact zeroInI
      . symm
        exact zero
  have coordinate := kuhnVertex.coordinate at succ(m), liftBase(m, q, b), liftPerm(m, u), i, m with liftFunction, hi, mIn
  have one := indicator.member with raisedM
  have top := liftBase.apply_top at m, q, b with baseFunction
  rw [shape, coordinate, one, top]
  exact add.right_one