FiniteEnum.lean

import Std
import FPRDTransfer.LMR

namespace FPRD

universe u v

/-- An executable exhaustive list for a finite type.  Duplicate entries are
allowed at the construction boundary and removed before cardinality or
transport to `Fin n`. -/
class FiniteEnum (α : Type u) where
  elems : List α
  complete : ∀ a, a ∈ elems

/-- The minimal bundled equivalence needed by this dependency-light package. -/
structure FiniteEquiv (α : Type u) (β : Type v) where
  toFun : α → β
  invFun : β → α
  left_inv : ∀ a, invFun (toFun a) = a
  right_inv : ∀ b, toFun (invFun b) = b

namespace FiniteEquiv

instance : CoeFun (FiniteEquiv α β) (fun _ => α → β) := ⟨FiniteEquiv.toFun⟩

def symm (e : FiniteEquiv α β) : FiniteEquiv β α where
  toFun := e.invFun
  invFun := e.toFun
  left_inv := e.right_inv
  right_inv := e.left_inv

@[simp] theorem symm_apply_apply (e : FiniteEquiv α β) (a : α) :
    e.symm (e a) = a := e.left_inv a

@[simp] theorem apply_symm_apply (e : FiniteEquiv α β) (b : β) :
    e (e.symm b) = b := e.right_inv b

end FiniteEquiv

namespace FiniteEnum

variable {α : Type u} {β : Type v}

/-- Executable membership driven directly by `DecidableEq`, independent of a
potentially unrelated `BEq` implementation. -/
def containsEq [DecidableEq α] (a : α) : List α → Bool
  | [] => false
  | b :: bs => if b = a then true else containsEq a bs

@[simp] theorem containsEq_eq_true [DecidableEq α] {a : α} {as : List α} :
    containsEq a as = true ↔ a ∈ as := by
  induction as with
  | nil => simp [containsEq]
  | cons b bs ih =>
      by_cases h : b = a
      · subst b
        simp [containsEq]
      · constructor
        · intro hc
          have ht : containsEq a bs = true := by simpa [containsEq, h] using hc
          exact List.mem_cons.mpr (Or.inr (ih.mp ht))
        · intro hm
          have ht : a ∈ bs :=
            (List.mem_cons.mp hm).resolve_left (Ne.symm h)
          simpa [containsEq, h] using ih.mpr ht

/-- Propositional-equality-based deduplication, kept local so the transport
does not depend on a potentially unrelated `BEq` implementation. -/
def dedup [DecidableEq α] : List α → List α
  | [] => []
  | a :: as =>
      let rest := dedup as
      if containsEq a rest then rest else a :: rest

@[simp] theorem mem_dedup [DecidableEq α] {a : α} {as : List α} :
    a ∈ dedup as ↔ a ∈ as := by
  induction as generalizing a with
  | nil => simp [dedup]
  | cons b bs ih =>
      simp only [dedup]
      split <;> simp_all

theorem nodup_dedup [DecidableEq α] (as : List α) : (dedup as).Nodup := by
  induction as with
  | nil => simp [dedup]
  | cons a as ih =>
      simp only [dedup]
      split <;> simp_all

/-- The canonical duplicate-free list associated with an enumeration. -/
def unique (α : Type u) [FiniteEnum α] [DecidableEq α] : List α :=
  dedup (FiniteEnum.elems (α := α))

theorem complete_unique (α : Type u) [FiniteEnum α] [DecidableEq α] (a : α) :
    a ∈ unique α := by
  simp [unique, FiniteEnum.complete]

/-- Transport an exhaustive enumeration across a proved type equivalence. -/
@[reducible] def ofEquiv (source : FiniteEnum α) (e : FiniteEquiv α β) :
    FiniteEnum β where
  elems := source.elems.map e
  complete b := by
    apply List.mem_map.mpr
    exact ⟨e.symm b, source.complete (e.symm b), e.right_inv b⟩

/-- Executable cardinality of the duplicate-free enumeration. -/
def card (α : Type u) [FiniteEnum α] [DecidableEq α] : Nat :=
  (unique α).length

/-- First position of an element, or the list length when it is absent. -/
def indexOf [DecidableEq α] (a : α) : List α → Nat
  | [] => 0
  | b :: bs => if b = a then 0 else indexOf a bs + 1

theorem indexOf_lt_length_of_mem [DecidableEq α] {a : α} {as : List α}
    (h : a ∈ as) : indexOf a as < as.length := by
  induction as with
  | nil => simp at h
  | cons b bs ih =>
      simp only [indexOf, List.length_cons]
      by_cases hba : b = a
      · simp [hba]
      · have : a ∈ bs := (List.mem_cons.mp h).resolve_left (Ne.symm hba)
        simp [hba, ih this]

theorem get_indexOf [DecidableEq α] {a : α} {as : List α} (h : a ∈ as) :
    as.get ⟨indexOf a as, indexOf_lt_length_of_mem h⟩ = a := by
  induction as with
  | nil => simp at h
  | cons b bs ih =>
      by_cases hba : b = a
      · subst b
        simp [indexOf]
      · have ht : a ∈ bs := (List.mem_cons.mp h).resolve_left (Ne.symm hba)
        simpa [indexOf, hba] using ih ht

theorem indexOf_get [DecidableEq α] {as : List α} (hn : as.Nodup)
    (i : Fin as.length) : indexOf (as.get i) as = i.val := by
  induction as with
  | nil => exact Fin.elim0 i
  | cons a as ih =>
      have hnot : a ∉ as := (List.nodup_cons.mp hn).1
      have htail : as.Nodup := (List.nodup_cons.mp hn).2
      refine Fin.cases ?_ (fun j => ?_) i
      · simp [indexOf]
      · have hne : a ≠ as.get j := by
          intro h
          apply hnot
          rw [h]
          exact as.get_mem j
        rw [List.get_cons_succ', indexOf, if_neg hne, ih htail j]
        rfl

/-- Every finite enumeration yields an equivalence with a concrete `Fin n`.
The target size is computed after deduplication. -/
def equivFin (α : Type u) [FiniteEnum α] [DecidableEq α] :
    FiniteEquiv α (Fin (card α)) where
  toFun a := ⟨indexOf a (unique α), indexOf_lt_length_of_mem (complete_unique α a)⟩
  invFun i := (unique α).get i
  left_inv a := get_indexOf (complete_unique α a)
  right_inv i := by
    apply Fin.ext
    exact indexOf_get (nodup_dedup (FiniteEnum.elems (α := α))) i

instance : FiniteEnum Unit where
  elems := [()]
  complete a := by cases a; simp

instance : FiniteEnum Bool where
  elems := [false, true]
  complete a := by cases a <;> simp

instance {n : Nat} : FiniteEnum (Fin n) where
  elems := List.finRange n
  complete a := by simp

instance [FiniteEnum α] : FiniteEnum (Option α) where
  elems := none :: (FiniteEnum.elems (α := α)).map some
  complete a := by
    cases a with
    | none => simp
    | some a => simp [FiniteEnum.complete]

instance [FiniteEnum α] [FiniteEnum β] : FiniteEnum (Sum α β) where
  elems :=
    (FiniteEnum.elems (α := α)).map Sum.inl ++
      (FiniteEnum.elems (α := β)).map Sum.inr
  complete a := by
    cases a with
    | inl a => simp [FiniteEnum.complete]
    | inr b => simp [FiniteEnum.complete]

instance [FiniteEnum α] [FiniteEnum β] : FiniteEnum (α × β) where
  elems := (FiniteEnum.elems (α := α)).flatMap fun a =>
    (FiniteEnum.elems (α := β)).map fun b => (a, b)
  complete p := by
    rcases p with ⟨a, b⟩
    simp [FiniteEnum.complete]

/-- Enumerate functions on `Fin n` by choosing a head and then a tail. -/
def finFunctions (values : List α) : (n : Nat) → List (Fin n → α)
  | 0 => [fun i => Fin.elim0 i]
  | n + 1 => values.flatMap fun a =>
      (finFunctions values n).map fun f => Fin.cases a f

theorem mem_finFunctions (values : List α) (hvalues : ∀ a, a ∈ values) :
    ∀ {n : Nat} (f : Fin n → α), f ∈ finFunctions values n := by
  intro n
  induction n with
  | zero =>
      intro f
      simp only [finFunctions, List.mem_singleton]
      funext i
      exact Fin.elim0 i
  | succ n ih =>
      intro f
      let tail : Fin n → α := fun i => f i.succ
      have htail : tail ∈ finFunctions values n := ih tail
      simp only [finFunctions, List.mem_flatMap, List.mem_map]
      refine ⟨f 0, hvalues (f 0), tail, htail, ?_⟩
      funext i
      exact Fin.cases rfl (fun j => rfl) i

@[reducible] def finFunctionEnum (n : Nat) (e : FiniteEnum α) :
    FiniteEnum (Fin n → α) where
  elems := finFunctions e.elems n
  complete := mem_finFunctions e.elems e.complete

instance {n : Nat} [FiniteEnum α] : FiniteEnum (Fin n → α) :=
  finFunctionEnum n (inferInstance : FiniteEnum α)

/-- Decidable equality for finite functions, proved by splitting the domain at
zero.  Core Lean deliberately supplies no general function equality instance. -/
protected def finFunctionDecEq [DecidableEq α] :
    (n : Nat) → DecidableEq (Fin n → α)
  | 0 => fun _ _ => isTrue (funext fun i => Fin.elim0 i)
  | n + 1 => fun f g =>
      match decEq (f 0) (g 0),
          FiniteEnum.finFunctionDecEq n (fun i => f i.succ) (fun i => g i.succ) with
      | isTrue hzero, isTrue htail =>
          isTrue (funext fun i => Fin.cases hzero (fun j => congrFun htail j) i)
      | isFalse hzero, _ =>
          isFalse fun h => hzero (congrFun h 0)
      | _, isFalse htail =>
          isFalse fun h => htail (funext fun i => congrFun h i.succ)

instance {n : Nat} [DecidableEq α] : DecidableEq (Fin n → α) :=
  FiniteEnum.finFunctionDecEq n

end FiniteEnum

/-- Port words whose length is bounded by the automaton distance. -/
structure BoundedPath (d k : Nat) where
  elems : List (Fin d)
  length_le : elems.length ≤ k
  deriving DecidableEq

namespace BoundedPath

@[reducible] def enumeration (d : Nat) : (k : Nat) → FiniteEnum (BoundedPath d k)
  | 0 =>
      { elems := [⟨[], by simp⟩]
        complete := fun p => by
          have hp : p.elems = [] := List.eq_nil_of_length_eq_zero (Nat.eq_zero_of_le_zero p.length_le)
          cases p
          simp_all }
  | k + 1 =>
      let previous := enumeration d k
      { elems :=
          ⟨[], by simp⟩ ::
            (FiniteEnum.elems (α := Fin d)).flatMap fun a =>
              previous.elems.map fun p =>
                ⟨a :: p.elems, Nat.succ_le_succ p.length_le⟩
        complete := fun p => by
          cases he : p.elems with
          | nil =>
              have hp : p = ⟨[], by simp⟩ := by cases p; simp_all
              simp [hp]
          | cons a as =>
              have has : as.length ≤ k := by
                simpa [he] using p.length_le
              let tail : BoundedPath d k := ⟨as, has⟩
              have htail : tail ∈ previous.elems := previous.complete tail
              have hp : p = ⟨a :: tail.elems, Nat.succ_le_succ tail.length_le⟩ := by
                cases p
                simp_all [tail]
              simp only [List.mem_cons, List.mem_flatMap, List.mem_map]
              exact Or.inr ⟨a, FiniteEnum.complete a, tail, htail, hp.symm⟩ }

instance {d k : Nat} : FiniteEnum (BoundedPath d k) := enumeration d k

end BoundedPath

namespace Unfold

/-- Decidable equality follows the fixed observation depth. -/
protected def decEq {Gamma : Type u} [DecidableEq Gamma] {d : Nat} :
    (k : Nat) → DecidableEq (Unfold d Gamma k)
  | 0, .missing, .missing => isTrue rfl
  | 0, .missing, .leaf _ => isFalse fun h => by cases h
  | 0, .leaf _, .missing => isFalse fun h => by cases h
  | 0, .leaf a, .leaf b =>
      match decEq a b with
      | isTrue h => isTrue (by cases h; rfl)
      | isFalse h => isFalse fun hab => by cases hab; exact h rfl
  | k + 1, .missing, .missing => isTrue rfl
  | k + 1, .missing, .node _ _ => isFalse fun h => by cases h
  | k + 1, .node _ _, .missing => isFalse fun h => by cases h
  | k + 1, .node a childrenA, .node b childrenB => by
      letI : DecidableEq (Unfold d Gamma k) := Unfold.decEq k
      exact match decEq a b with
        | isFalse h => isFalse fun hab => by cases hab; exact h rfl
        | isTrue hlabel =>
            match decEq childrenA childrenB with
            | isTrue hchildren => isTrue (by cases hlabel; cases hchildren; rfl)
            | isFalse hchildren =>
                isFalse fun hab => by cases hab; exact hchildren rfl

instance {Gamma : Type u} [DecidableEq Gamma] {d k : Nat} :
    DecidableEq (Unfold d Gamma k) := Unfold.decEq k

/-- Exhaustive finite unfolding trees at an exact observation depth. -/
@[reducible] def enumeration {Gamma : Type u} [FiniteEnum Gamma] (d : Nat) :
    (k : Nat) → FiniteEnum (Unfold d Gamma k)
  | 0 =>
      let labels : FiniteEnum (Option Gamma) := inferInstance
      { elems := Unfold.missing :: labels.elems.map Unfold.leaf
        complete := fun tree => by
          cases tree with
          | missing => simp
          | leaf label => simp [labels.complete label] }
  | k + 1 =>
      let labels : FiniteEnum (Option Gamma) := inferInstance
      let previous := enumeration d k
      let children := FiniteEnum.finFunctionEnum d previous
      { elems := Unfold.missing :: labels.elems.flatMap fun label =>
          children.elems.map fun child => Unfold.node label child
        complete := fun tree => by
          cases tree with
          | missing => simp
          | node label child =>
              simp only [List.mem_cons, List.mem_flatMap, List.mem_map]
              exact Or.inr
                ⟨label, labels.complete label, child, children.complete child, rfl⟩ }

instance {d k : Nat} {Gamma : Type u} [FiniteEnum Gamma] :
    FiniteEnum (Unfold d Gamma k) := enumeration d k

end Unfold

/- The executable checks below were compiled once before these definitions;
that red run failed on the deliberately absent API. -/

theorem test_boundedPath_card : FiniteEnum.card (BoundedPath 1 2) = 3 := by decide

theorem test_fin_card : FiniteEnum.card (Fin 3) = 3 := by decide

theorem test_option_card : FiniteEnum.card (Option Bool) = 3 := by decide

theorem test_sum_card : FiniteEnum.card (Sum Bool Unit) = 3 := by decide

theorem test_product_card : FiniteEnum.card (Bool × Bool) = 4 := by decide

theorem test_finFunction_card : FiniteEnum.card (Fin 2 → Bool) = 4 := by decide

/-- The exact shape of the scaffold-label carrier produced by the T89 stack
compiler, at a small executable calibration. -/
abbrev StackTupleFixture := Fin 2 → Option Bool

theorem test_stackTuple_card : FiniteEnum.card StackTupleFixture = 9 := by decide

theorem test_stackTuple_equivFin_left (value : StackTupleFixture) :
    (FiniteEnum.equivFin StackTupleFixture).symm
        (FiniteEnum.equivFin StackTupleFixture value) = value := by
  simp

theorem test_stackTuple_equivFin_right
    (i : Fin (FiniteEnum.card StackTupleFixture)) :
    FiniteEnum.equivFin StackTupleFixture
        ((FiniteEnum.equivFin StackTupleFixture).symm i) = i := by
  simp

/-- Small record with the same field shape as T101's root-sharing cell. -/
structure RootCellFixture where
  symbol : Bool
  next : Fin 2
  deriving DecidableEq

def rootCellFixtureEquiv : FiniteEquiv (Bool × Fin 2) RootCellFixture where
  toFun value := ⟨value.1, value.2⟩
  invFun cell := (cell.symbol, cell.next)
  left_inv value := by cases value; rfl
  right_inv cell := by cases cell; rfl

@[reducible] instance : FiniteEnum RootCellFixture :=
  FiniteEnum.ofEquiv (inferInstance : FiniteEnum (Bool × Fin 2)) rootCellFixtureEquiv

/-- The exact finite-function/optional-record shape of T101's root label. -/
abbrev RootLabelShapeFixture := Fin 2 → Option RootCellFixture

theorem test_rootLabelShape_card : FiniteEnum.card RootLabelShapeFixture = 25 := by decide

theorem test_rootLabelShape_equivFin_left (value : RootLabelShapeFixture) :
    (FiniteEnum.equivFin RootLabelShapeFixture).symm
        (FiniteEnum.equivFin RootLabelShapeFixture value) = value := by
  simp

theorem test_rootLabelShape_equivFin_right
    (i : Fin (FiniteEnum.card RootLabelShapeFixture)) :
    FiniteEnum.equivFin RootLabelShapeFixture
        ((FiniteEnum.equivFin RootLabelShapeFixture).symm i) = i := by
  simp

theorem test_unfold_depth0_card : FiniteEnum.card (Unfold 1 Unit 0) = 3 := by decide

theorem test_unfold_depth1_card : FiniteEnum.card (Unfold 1 Unit 1) = 7 := by decide

theorem test_unfold_depth2_card : FiniteEnum.card (Unfold 1 Unit 2) = 15 := by decide

theorem test_unfold_equivFin_left (value : Unfold 1 Unit 2) :
    (FiniteEnum.equivFin (Unfold 1 Unit 2)).symm
        (FiniteEnum.equivFin (Unfold 1 Unit 2) value) = value := by
  simp

theorem test_unfold_equivFin_right
    (i : Fin (FiniteEnum.card (Unfold 1 Unit 2))) :
    FiniteEnum.equivFin (Unfold 1 Unit 2)
        ((FiniteEnum.equivFin (Unfold 1 Unit 2)).symm i) = i := by
  simp

theorem test_equivFin_left (p : BoundedPath 1 2) :
    (FiniteEnum.equivFin (BoundedPath 1 2)).symm
        (FiniteEnum.equivFin (BoundedPath 1 2) p) = p := by
  simp

theorem test_equivFin_right (i : Fin (FiniteEnum.card (BoundedPath 1 2))) :
    FiniteEnum.equivFin (BoundedPath 1 2)
        ((FiniteEnum.equivFin (BoundedPath 1 2)).symm i) = i := by
  simp

end FPRD