import Spa.Lattice
import Mathlib.Data.Fin.Tuple.Basic
import Mathlib.Algebra.Order.BigOperators.Group.Finset

/-!

# Finite Tuple Lattices

This file provides a proof that, in addition to being a lattice, the function
space `Fin n → β` is itself a `Spa.FiniteHeightLattice` if the element type
`β` is a lattice.

Finite tuple lattices are the workhorse behind `FiniteMap`, whose carrier is
`Fin ks.length → β`.

The proof proceeds by "unzipping" a chain (`LTSeries`):

$$
(a_1, b_1, c_1) < \ldots < (a_1, b_1, c_o) < \ldots < (a_1, b_m, c_o) <
  \ldots < (a_n, b_m, c_o)
$$

In which, at each step, at least one of the components must have increased
(otherwise, the chain is not striclty increasing), into `n` chains
(`LTSeries`).

$$
\begin{aligned}
a_1 < \ldots < a_n \\
b_1 < \ldots < b_m \
c_1 < \ldots < c_o \
\end{aligned}
$$

Because at least one of the two "unzipped" chains grows with each element of
the product chain, the full chain length can't exceed the sum of the
components. By the definition of finite height, these two chains are bounded,
and therefore, the product chain is bounded too. -/

namespace Spa

namespace Tuple

variable {β : Type*}

section Unzip

variable [PartialOrder β]

open Classical in -- chain bounds are in Prop, so classical helps here.
/-- The generalized unzip: any chain in `Fin n → β` decomposes into a family of
    per-tuple-coordinate chains in `β`, agreeing with the original at each end, whose
    lengths sum to an upper bound on the original chain's length. -/
lemma exists_unzip {n : ℕ} (c : LTSeries (Fin n → β)) :
    ∃ cs : Fin n → LTSeries β,
      (∀ i, (cs i).head = c.head i) ∧ (∀ i, (cs i).last = c.last i) ∧
      c.length ≤ ∑ i, (cs i).length := by
  suffices H : ∀ (m : ℕ) (c : LTSeries (Fin n → β)), c.length = m →
      ∃ cs : Fin n → LTSeries β,
        (∀ i, (cs i).head = c.head i) ∧ (∀ i, (cs i).last = c.last i) ∧
        c.length ≤ ∑ i, (cs i).length from H c.length c rfl
  intro m
  induction m with
  | zero =>
    intro c hn
    have hlast : (Fin.last c.length) = 0 := by ext; simp [hn]
    have hhl : c.last = c.head := by rw [RelSeries.last, RelSeries.head, hlast]
    refine ⟨fun i => RelSeries.singleton _ (c.head i), fun i => ?_, fun i => ?_, ?_⟩
    · exact RelSeries.head_singleton _
    · rw [RelSeries.last_singleton, hhl]
    · simp [hn, RelSeries.singleton]
  | succ m ih =>
    intro c hn
    have h0 : c.length ≠ 0 := by omega
    haveI : NeZero c.length := ⟨h0⟩
    obtain ⟨cs', hh', hl', hlen'⟩ := ih (c.tail h0) (by rw [RelSeries.tail_length]; omega)
    have hstep : c.head < c 1 := c.strictMono Fin.one_pos'
    obtain ⟨hle, j, hjlt⟩ := Pi.lt_def.mp hstep
    have hh'1 : ∀ i, (cs' i).head = c 1 i := fun i => by rw [hh' i, RelSeries.head_tail]
    refine ⟨fun i =>
        if hlt : c.head i < c 1 i then
          (cs' i).cons (c.head i) (by rw [hh'1 i]; exact hlt)
        else cs' i,
      fun i => ?_, fun i => ?_, ?_⟩
    · by_cases hlt : c.head i < c 1 i
      · simp only [dif_pos hlt, RelSeries.head_cons]
      · simp only [dif_neg hlt]
        rw [hh'1 i]
        exact ((lt_or_eq_of_le (hle i)).resolve_left hlt).symm
    · by_cases hlt : c.head i < c 1 i
      · simp only [dif_pos hlt, RelSeries.last_cons, hl' i, RelSeries.last_tail]
      · simp only [dif_neg hlt, hl' i, RelSeries.last_tail]
    · calc c.length
          = (c.tail h0).length + 1 := by rw [RelSeries.tail_length]; omega
        _ ≤ (∑ i, (cs' i).length) + 1 := Nat.add_le_add_right hlen' 1
        _ ≤ ∑ i, (if hlt : c.head i < c 1 i then
              (cs' i).cons (c.head i) (by rw [hh'1 i]; exact hlt) else cs' i).length :=
          Nat.succ_le_of_lt (Finset.sum_lt_sum (fun i _ => by
              split
              · rw [RelSeries.cons_length]; omega
              · exact le_rfl)
            ⟨j, Finset.mem_univ j, by rw [dif_pos hjlt, RelSeries.cons_length]; omega⟩)

end Unzip

section FiniteHeight

variable [FiniteHeightLattice β]

private lemma consBot_strictMono {n : ℕ} :
    StrictMono (fun b : β => (Fin.cons b (⊥ : Fin n → β) : Fin (n + 1) → β)) := by
  intro a b hab
  refine lt_iff_le_and_ne.mpr ⟨?_, ?_⟩
  · refine Pi.le_def.mpr (fun i => Fin.cases ?_ (fun j => ?_) i)
    · simpa using hab.le
    · simp
  · exact fun h => hab.ne (by simpa using congrFun h 0)

private lemma consTop_strictMono {n : ℕ} :
    StrictMono (fun f : Fin n → β => (Fin.cons (⊤ : β) f : Fin (n + 1) → β)) := by
  intro f g hfg
  refine lt_iff_le_and_ne.mpr ⟨?_, ?_⟩
  · refine Pi.le_def.mpr (fun i => Fin.cases ?_ (fun j => ?_) i)
    · simp
    · simpa using Pi.le_def.mp hfg.le j
  · intro h
    apply hfg.ne
    funext j
    simpa using congrFun h j.succ

/-- The maximal chain in `Fin n → β`: walk the first tuple element from `⊥` to `⊤`
    through `β`'s longest chain, then do that with the second element, and so on. -/
private def stdChain : (n : ℕ) →
    { s : LTSeries (Fin n → β) //
        s.head = (⊥ : Fin n → β) ∧
        s.length = n * (FiniteHeightLattice.longestChain (α := β)).length }
  | 0 => ⟨RelSeries.singleton _ ⊥, by rw [RelSeries.head_singleton], by simp⟩
  | n + 1 =>
      let prev := stdChain n
      ⟨RelSeries.smash
          ((FiniteHeightLattice.longestChain (α := β)).map
            (fun b => (Fin.cons b (⊥ : Fin n → β) : Fin (n + 1) → β)) consBot_strictMono)
          (prev.1.map (fun f => (Fin.cons (⊤ : β) f : Fin (n + 1) → β)) consTop_strictMono)
          (by rw [LTSeries.last_map, LTSeries.head_map, prev.2.1]; rfl),
        by
          simp only [RelSeries.head_smash, LTSeries.head_map]
          rw [show (FiniteHeightLattice.longestChain (α := β)).head = (⊥ : β) from rfl]
          funext i
          refine Fin.cases ?_ (fun j => ?_) i <;> simp [Pi.bot_apply],
        by
          show (FiniteHeightLattice.longestChain (α := β)).length + prev.1.length
              = (n + 1) * (FiniteHeightLattice.longestChain (α := β)).length
          rw [prev.2.2, Nat.succ_mul]; exact Nat.add_comm _ _⟩

instance instFiniteHeight {n : ℕ} : FiniteHeightLattice (Fin n → β) where
  toLattice := inferInstance
  longestChain := (stdChain n).1
  chains_bounded := fun c => by
    obtain ⟨cs, _, _, hbound⟩ := exists_unzip c
    refine hbound.trans ?_
    rw [(stdChain n).2.2]
    calc ∑ i, (cs i).length
        ≤ ∑ _i : Fin n, (FiniteHeightLattice.longestChain (α := β)).length :=
          Finset.sum_le_sum (fun i _ => FiniteHeightLattice.chains_bounded (cs i))
      _ = n * (FiniteHeightLattice.longestChain (α := β)).length := by
          simp [Finset.sum_const, Finset.card_univ, Fintype.card_fin]

end FiniteHeight

end Tuple

end Spa