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 β]

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

end FiniteHeight

end Tuple

end Spa