import Spa.Lattice

namespace Spa

section Unzip

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

theorem 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
    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 := by
      have h := c.step ⟨0, by omega⟩
      have h1 : (⟨0, by omega⟩ : Fin c.length).succ = 1 := by
        ext; simp [Fin.val_one, Nat.mod_eq_of_lt (by omega : 1 < c.length + 1)]
      rwa [h1] at h
    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*} [Lattice α] [Lattice β]

instance prod [A : FiniteHeightLattice α] [B : FiniteHeightLattice β] :
    FiniteHeightLattice (α × β) where
  bot := ((⊥ : α), (⊥ : β))
  top := ((⊤ : α), (⊤ : β))
  height := A.height + B.height
  longestChain :=
    { series :=
        RelSeries.smash
          (A.longestChain.series.map (fun a => (a, (⊥ : β)))
            (fun _ _ h => Prod.mk_lt_mk_iff_left.mpr h))
          (B.longestChain.series.map (fun b => ((⊤ : α), b))
            (fun _ _ h => Prod.mk_lt_mk_iff_right.mpr h))
          (by simp [A.longestChain.last_series, B.longestChain.head_series])
      head_series :=
        (RelSeries.head_smash _).trans
          ((LTSeries.head_map _ _ _).trans
            (congrArg (·, (⊥ : β)) A.longestChain.head_series))
      last_series :=
        (RelSeries.last_smash _).trans
          ((LTSeries.last_map _ _ _).trans
            (congrArg ((⊤ : α), ·) B.longestChain.last_series))
      length_series := by
        show A.longestChain.series.length + B.longestChain.series.length = _
        rw [A.longestChain.length_series, B.longestChain.length_series] }
  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₂
    omega

end FixedHeight

end Spa