import Spa.Lattice

namespace Spa

section Unzip

variable {α β : Type*} [PartialOrder α] [PartialOrder β]

lemma LTSeries.exists_unzip (c : LTSeries (α × β)) :
    ∃ (c₁ : LTSeries α) (c₂ : LTSeries β),
      c₁.head = c.head.1 ∧ c₁.last = c.last.1 ∧
      c₂.head = c.head.2 ∧ c₂.last = c.last.2 ∧
      c.length ≤ c₁.length + c₂.length := by
  suffices H : ∀ (n : ℕ) (c : LTSeries (α × β)), c.length = n →
      ∃ (c₁ : LTSeries α) (c₂ : LTSeries β),
        c₁.head = c.head.1 ∧ c₁.last = c.last.1 ∧
        c₂.head = c.head.2 ∧ c₂.last = c.last.2 ∧
        c.length ≤ c₁.length + c₂.length from H c.length c rfl
  intro n
  induction n with
  | zero =>
    intro c hn
    refine ⟨RelSeries.singleton _ c.head.1, RelSeries.singleton _ c.head.2,
      rfl, ?_, rfl, ?_, by simp [hn]⟩ <;>
    · have hlast : Fin.last c.length = 0 := by ext; simp [hn]
      simp [RelSeries.last, RelSeries.head, hlast]
  | succ n ih =>
    intro c hn
    have h0 : c.length ≠ 0 := by omega
    haveI : NeZero c.length := ⟨h0⟩
    obtain ⟨c₁, c₂, hh₁, hl₁, hh₂, hl₂, hlen⟩ :=
      ih (c.tail h0) (by simp [RelSeries.tail_length, hn])
    rw [RelSeries.last_tail] at hl₁ hl₂
    rw [RelSeries.head_tail] at hh₁ hh₂
    rw [RelSeries.tail_length] at hlen
    have hstep : c.head < c 1 := c.strictMono Fin.one_pos'
    obtain ⟨hle1, hle2⟩ := Prod.le_def.mp hstep.le
    rcases eq_or_lt_of_le hle1 with heq1 | hlt1 <;>
      rcases eq_or_lt_of_le hle2 with heq2 | hlt2
    · exact absurd (Prod.ext heq1 heq2) hstep.ne
    · refine ⟨c₁, c₂.cons c.head.2 (hh₂ ▸ hlt2),
        hh₁.trans heq1.symm, hl₁, RelSeries.head_cons .., by
          rw [RelSeries.last_cons]; exact hl₂, by
          simp only [RelSeries.cons_length]; omega⟩
    · refine ⟨c₁.cons c.head.1 (hh₁ ▸ hlt1), c₂,
        RelSeries.head_cons .., by
          rw [RelSeries.last_cons]; exact hl₁,
        hh₂.trans heq2.symm, hl₂, by
          simp only [RelSeries.cons_length]; omega⟩
    · refine ⟨c₁.cons c.head.1 (hh₁ ▸ hlt1), c₂.cons c.head.2 (hh₂ ▸ hlt2),
        RelSeries.head_cons .., by
          rw [RelSeries.last_cons]; exact hl₁,
        RelSeries.head_cons .., by
          rw [RelSeries.last_cons]; exact hl₂, by
          simp only [RelSeries.cons_length]; omega⟩

end Unzip

section FixedHeight

variable {α β : Type*}

instance prod [A : FiniteHeightLattice α] [B : FiniteHeightLattice β] :
    FiniteHeightLattice (α × β) where
  toLattice := inferInstance
  longestChain :=
    RelSeries.smash
      (A.longestChain.map (fun a => (a, (⊥ : β)))
        (fun _ _ h => Prod.mk_lt_mk_iff_left.mpr h))
      (B.longestChain.map (fun b => ((⊤ : α), b))
        (fun _ _ h => Prod.mk_lt_mk_iff_right.mpr h))
      rfl
  chains_bounded := fun c => by
    obtain ⟨c₁, c₂, -, -, -, -, hlen⟩ := LTSeries.exists_unzip c
    have h₁ := A.chains_bounded c₁
    have h₂ := B.chains_bounded c₂
    show c.length ≤ A.longestChain.length + B.longestChain.length
    omega

end FixedHeight

end Spa