-- Restricting far-wall faces to the lower cube is a bijection onto the
-- completely labelled simplices there, so the two counts agree
-- 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.KuhnWallCount
import KuhnLift

defgraph wallRestrict (p m lab : Set) (f ∈ topWallFaces(p, m, lab)) : 𝒫(functionSpace(m, succ(p))) :=
  functionImage(cubeRestrict(m, p), f, functionSpace(m, succ(p)))
proof
  exact powerset.membership_iff.backward with (functionImage.subset at cubeRestrict(m, p), f, functionSpace(m, succ(p)))

theorem topWallFaces.cube (m p : Nat) (lab f : Set)
  requires hf : f ∈ topWallFaces(p, m, lab)
  shows f ⊆ functionSpace(succ(m), succ(p)) ∧ ∀ x ∈ f, x(m) = p
proof
  obtain ⟨fIn, top⟩ := topWallFaces.membership.forward with hf
  have parts := labelledFaces.membership.forward with fIn
  constructor
  . exact powerset.membership_iff.forward with parts.left
  . exact top

theorem wallRestrict.subset_of_image (m p : Nat) (f g : Set)
  requires fCube : f ⊆ functionSpace(succ(m), succ(p))
  requires fTop : ∀ x ∈ f, x(m) = p
  requires gCube : g ⊆ functionSpace(succ(m), succ(p))
  requires gTop : ∀ x ∈ g, x(m) = p
  requires inclusion : functionImage(cubeRestrict(m, p), f, functionSpace(m, succ(p))) ⊆
    functionImage(cubeRestrict(m, p), g, functionSpace(m, succ(p)))
  shows f ⊆ g
proof
  have function := cubeRestrict.function at m, p with *
  intros x hx
  have xIn := fCube with x, hx
  have value := cubeRestrict.value with xIn
  have inImage : cubeRestrict(m, p)(x) ∈ functionImage(cubeRestrict(m, p), f, functionSpace(m, succ(p)))
    apply functionImage.membership.backward
    constructor
    . exact FunctionBetween.apply_closed with function, xIn
    . witness x
      constructor
      . exact hx
      . rfl
  have inG := inclusion with inImage
  obtain ⟨yIn, x2, hx2, shape⟩ := functionImage.membership.forward with inG
  have x2In := gCube with x2, hx2
  have value2 := cubeRestrict.value with x2In
  rw [value, value2] at shape
  have back := cubeRestrict.extend with xIn, (fTop with hx)
  have back2 := cubeRestrict.extend with x2In, (gTop with hx2)
  rwa [← back, shape, back2]

theorem wallRestrict.injective (m p : Nat) (lab : Set)
  shows InjectiveBetween(wallRestrict(p, m, lab), topWallFaces(p, m, lab), 𝒫(functionSpace(m, succ(p))))
proof
  have function := wallRestrict.function at p, m, lab
  apply Injective.of_cancel with function
  intros f hf g hg same
  have fValue := wallRestrict.value with hf
  have gValue := wallRestrict.value with hg
  rw [fValue, gValue] at same
  obtain ⟨fCube, fTop⟩ := topWallFaces.cube with hf
  obtain ⟨gCube, gTop⟩ := topWallFaces.cube with hg
  have forward : functionImage(cubeRestrict(m, p), f, functionSpace(m, succ(p))) ⊆ functionImage(cubeRestrict(m, p), g, functionSpace(m, succ(p)))
    rw same
    intros y hy
    exact hy
  have backward : functionImage(cubeRestrict(m, p), g, functionSpace(m, succ(p))) ⊆ functionImage(cubeRestrict(m, p), f, functionSpace(m, succ(p)))
    rw same
    intros y hy
    exact hy
  apply Subset.antisymmetric
  . exact wallRestrict.subset_of_image with fCube, fTop, gCube, gTop, forward
  . exact wallRestrict.subset_of_image with gCube, gTop, fCube, fTop, backward

-- The restriction of a far-wall face is a completely labelled simplex of the
-- lower cube.
theorem wallRestrict.complete (m p : Nat) (lab L f : Set)
  requires labFunction : lab : functionSpace(succ(m), succ(p)) → L
  requires fullWall : ∀ x ∈ functionSpace(succ(m), succ(p)), ∀ j ∈ succ(m), x(j) = p → lab(x)(j) = 1
  requires hf : f ∈ topWallFaces(p, m, lab)
  shows wallRestrict(p, m, lab)(f) ∈ completeSimplices(p, m, wallLabels(m, p, lab, L))
proof
  have value := wallRestrict.value with hf
  rw value
  obtain ⟨fIn, top⟩ := topWallFaces.membership.forward with hf
  have fParts := labelledFaces.membership.forward with fIn
  have fCube := powerset.membership_iff.forward with fParts.left
  obtain ⟨s, hs, a, ha, shape⟩ := fParts.right.left
  have labels := fParts.right.right
  obtain ⟨b, baseFunction, bounds, u, updBijective, sShape⟩ := kuhnSimplices.shape with hs
  have updInjective := Bijective.injective with updBijective
  rw sShape at ha
  rw sShape at shape
  have mIn := bootstrap.succ.self at m
  have zeroLess : 0 < succ(m) by arith
  have positive := membership.of_Less with zeroLess
  have zeroLess2 : 0 < succ(succ(m)) by arith
  have zeroIn := membership.of_Less with zeroLess2
  obtain ⟨inSpace, i, hi, aShape⟩ := kuhnSimplex.membership.forward with ha
  have flat : ∀ x ∈ kuhnSimplex(succ(m), b, u) ∖ {kuhnVertex(succ(m), b, u)(i)}, x(m) = p
    rw under ← aShape
    rwa under ← shape
  have below := bounds with mIn
  have bottom := kuhnSimplex.full_face_bottom with baseFunction, updBijective, hi, mIn, below, flat
  rw bottom at aShape
  rw aShape at shape
  have oneLess : succ(0) < succ(succ(m)) by arith
  have oneIn := membership.of_Less with oneLess
  have oneMember := kuhnVertex.member at succ(m), b, u, succ(0) with baseFunction, oneIn
  have oneInFace : kuhnVertex(succ(m), b, u)(succ(0)) ∈ kuhnSimplex(succ(m), b, u) ∖ {kuhnVertex(succ(m), b, u)(0)}
    apply difference.singleton_membership.backward
    constructor
    . exact oneMember
    . intro equal
      have same := kuhnVertex.injective with baseFunction, updInjective, oneIn, zeroIn, equal
      arith
  rw ← shape at oneInFace
  have oneTop := top with oneInFace
  have first : u(0) = m
    have coordinate := kuhnVertex.coordinate_succ with baseFunction, updInjective, positive, mIn
    have zeroValue := kuhnVertex.zero_coordinate at succ(m), b, u, m with baseFunction, mIn
    rw [coordinate, zeroValue] at oneTop
    have baseNat := nat.omega.backward with (FunctionBetween.apply_closed with baseFunction, mIn)
    by_contra different
    have zero := indicator.singleton_other with different
    rw zero at oneTop
    arith
  have face := kuhnSimplex.wall_face with baseFunction, bounds, updBijective, first
  have imageEq : functionImage(cubeRestrict(m, p), f, functionSpace(m, succ(p))) =
      kuhnSimplex(m, restrict(b, m), wallPerm(m, u))
    rwa shape
  rw imageEq
  apply completeSimplices.membership.backward
  constructor
  . have mSubset := successor.subset at m
    have restrictedBase := restrict.domain_function with baseFunction, mSubset
    have restrictedBounds : ∀ j ∈ m, restrict(b, m)(j) < p
      intros j hj
      have baseValue := restrict.domain_apply with baseFunction, mSubset, hj
      rw baseValue
      exact bounds with (membership.into_successor with hj)
    have wallBijective := wallPerm.bijective with updBijective, first
    exact kuhnSimplex.in_family with restrictedBase, wallBijective, restrictedBounds
  . have transfer := labelImage.restrict_wall with labFunction, fullWall, fCube, top
    rwa [imageEq, labels] at transfer

-- Every completely labelled simplex of the lower cube is the restriction of
-- a far-wall face.
theorem wallRestrict.onto (m q : Nat) (lab L t : Set)
  requires labFunction : lab : functionSpace(succ(m), succ(succ(q))) → L
  requires fullWall : ∀ x ∈ functionSpace(succ(m), succ(succ(q))), ∀ j ∈ succ(m), x(j) = succ(q) → lab(x)(j) = 1
  requires ht : t ∈ completeSimplices(succ(q), m, wallLabels(m, succ(q), lab, L))
  shows ∃ f ∈ topWallFaces(succ(q), m, lab), wallRestrict(succ(q), m, lab)(f) = t
proof
  obtain ⟨tIn, complete⟩ := completeSimplices.membership.forward with ht
  obtain ⟨b, baseFunction, bounds, u, updBijective, tShape⟩ := kuhnSimplices.shape with tIn
  have liftFunction := liftBase.function at m, q, b with baseFunction
  have lifted := liftPerm.bijective with updBijective
  have zero := liftPerm.zero 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
  have sIn := kuhnSimplex.lift_in_family with baseFunction, bounds, updBijective
  have face := kuhnSimplex.wall_face with liftFunction, liftBounds, lifted, zero
  have restrictBase := liftBase.restrict at m, q, b with baseFunction
  have restrictPerm := wallPerm.lift with updBijective
  rw [restrictBase, restrictPerm, ← tShape] at face
  have zeroLess : 0 < succ(succ(m)) by arith
  have zeroIn := membership.of_Less with zeroLess
  have bottomMember := kuhnVertex.member at succ(m), liftBase(m, q, b), liftPerm(m, u), 0 with liftFunction, zeroIn
  have faceCube : kuhnSimplex(succ(m), liftBase(m, q, b), liftPerm(m, u)) ∖ {kuhnVertex(succ(m), liftBase(m, q, b), liftPerm(m, u))(0)} ⊆ functionSpace(succ(m), succ(succ(q)))
    intros x hx
    have parts := difference.singleton_membership.forward with hx
    obtain ⟨inSpace, i, hi, xShape⟩ := kuhnSimplex.membership.forward with parts.left
    rw xShape
    exact kuhnVertex.into_cube with liftFunction, liftBounds, hi
  have faceTop : ∀ x ∈ kuhnSimplex(succ(m), liftBase(m, q, b), liftPerm(m, u)) ∖ {kuhnVertex(succ(m), liftBase(m, q, b), liftPerm(m, u))(0)}, x(m) = succ(q)
    intros x hx
    exact kuhnSimplex.lift_top_wall with baseFunction, updBijective, hx
  have faceIn : kuhnSimplex(succ(m), liftBase(m, q, b), liftPerm(m, u)) ∖ {kuhnVertex(succ(m), liftBase(m, q, b), liftPerm(m, u))(0)} ∈ topWallFaces(succ(q), m, lab)
    apply topWallFaces.membership.backward
    constructor
    . apply labelledFaces.membership.backward
      constructor
      . exact powerset.membership_iff.backward with faceCube
      . constructor
        . witness kuhnSimplex(succ(m), liftBase(m, q, b), liftPerm(m, u))
          constructor
          . exact sIn
          . witness kuhnVertex(succ(m), liftBase(m, q, b), liftPerm(m, u))(0)
            constructor
            . exact bottomMember
            . rfl
        . have transfer := labelImage.restrict_wall with labFunction, fullWall, faceCube, faceTop
          rw [face, complete] at transfer
          symm
          exact transfer
    . exact faceTop
  witness kuhnSimplex(succ(m), liftBase(m, q, b), liftPerm(m, u)) ∖ {kuhnVertex(succ(m), liftBase(m, q, b), liftPerm(m, u))(0)}
  constructor
  . exact faceIn
  . have value := wallRestrict.value at succ(q), m, lab, kuhnSimplex(succ(m), liftBase(m, q, b), liftPerm(m, u)) ∖ {kuhnVertex(succ(m), liftBase(m, q, b), liftPerm(m, u))(0)} with faceIn
    rwa value

theorem topWallFaces.count (m q : Nat) (lab L : Set)
  requires labFunction : lab : functionSpace(succ(m), succ(succ(q))) → L
  requires fullWall : ∀ x ∈ functionSpace(succ(m), succ(succ(q))), ∀ j ∈ succ(m), x(j) = succ(q) → lab(x)(j) = 1
  shows cardinality(topWallFaces(succ(q), m, lab)) =
    cardinality(completeSimplices(succ(q), m, wallLabels(m, succ(q), lab, L)))
proof
  have injective := wallRestrict.injective at m, succ(q), lab
  have function := wallRestrict.function at succ(q), m, lab
  have selfSubset : topWallFaces(succ(q), m, lab) ⊆ topWallFaces(succ(q), m, lab)
    intros f hf
    exact hf
  have finite : Finite(topWallFaces(succ(q), m, lab))
    have inside : topWallFaces(succ(q), m, lab) ⊆ labelledFaces(succ(q), succ(m), lab)
      intros f hf
      have parts := topWallFaces.membership.forward with hf
      exact parts.left
    exact Finite.subset with (labelledFaces.finite at succ(m), succ(q), lab), inside
  have count := cardinality.injective_image with injective, selfSubset, finite
  have same : functionImage(wallRestrict(succ(q), m, lab), topWallFaces(succ(q), m, lab), 𝒫(functionSpace(m, succ(succ(q))))) =
      completeSimplices(succ(q), m, wallLabels(m, succ(q), lab, L))
    ext t
    constructor
    . intro ht
      obtain ⟨tIn, f, hf, shape⟩ := functionImage.membership.forward with ht
      rw shape
      exact wallRestrict.complete with labFunction, fullWall, hf
    . intro ht
      obtain ⟨f, hf, shape⟩ := wallRestrict.onto with labFunction, fullWall, ht
      apply functionImage.membership.backward
      constructor
      . rw ← shape
        exact FunctionBetween.apply_closed with function, hf
      . witness f
        constructor
        . exact hf
        . symm
          exact shape
  rw same at count
  symm
  exact count