import Spa.Lattice

namespace Spa

inductive AboveBelow (α : Type*) where
  | bot
  | top
  | mk (x : α)
  deriving DecidableEq

namespace AboveBelow

instance {α : Type*} [ToString α] : ToString (AboveBelow α) where
  toString
    | bot => "⊥"
    | top => "⊤"
    | mk x => toString x

variable {α : Type*} [DecidableEq α]

instance : Max (AboveBelow α) where
  max
    | bot, x => x
    | top, _ => top
    | mk x, mk y => if x = y then mk x else top
    | mk x, bot => mk x
    | mk _, top => top

instance : Min (AboveBelow α) where
  min
    | bot, _ => bot
    | top, x => x
    | mk x, mk y => if x = y then mk x else bot
    | mk _, bot => bot
    | mk x, top => mk x

@[simp] lemma bot_sup (x : AboveBelow α) : bot ⊔ x = x := rfl
@[simp] lemma top_sup (x : AboveBelow α) : top ⊔ x = top := rfl
@[simp] lemma sup_bot (x : AboveBelow α) : x ⊔ bot = x := by cases x <;> rfl
@[simp] lemma sup_top (x : AboveBelow α) : x ⊔ top = top := by cases x <;> rfl
@[simp] lemma mk_sup_mk (x y : α) :
    (mk x ⊔ mk y : AboveBelow α) = if x = y then mk x else top := rfl

@[simp] lemma bot_inf (x : AboveBelow α) : bot ⊓ x = bot := rfl
@[simp] lemma top_inf (x : AboveBelow α) : top ⊓ x = x := rfl
@[simp] lemma inf_bot (x : AboveBelow α) : x ⊓ bot = bot := by cases x <;> rfl
@[simp] lemma inf_top (x : AboveBelow α) : x ⊓ top = x := by cases x <;> rfl
@[simp] lemma mk_inf_mk (x y : α) :
    (mk x ⊓ mk y : AboveBelow α) = if x = y then mk x else bot := rfl

protected lemma sup_comm (a b : AboveBelow α) : a ⊔ b = b ⊔ a := by
  rcases a with _ | _ | x <;> rcases b with _ | _ | y <;> simp only
    [bot_sup, sup_bot, top_sup, sup_top, mk_sup_mk]
  split_ifs with h₁ h₂ h₂ <;> simp_all

protected lemma sup_assoc (a b c : AboveBelow α) : a ⊔ b ⊔ c = a ⊔ (b ⊔ c) := by
  rcases a with _ | _ | x <;> rcases b with _ | _ | y <;> rcases c with _ | _ | z <;>
    simp only [bot_sup, sup_bot, top_sup, sup_top, mk_sup_mk]
  split_ifs <;> simp_all

protected lemma inf_comm (a b : AboveBelow α) : a ⊓ b = b ⊓ a := by
  rcases a with _ | _ | x <;> rcases b with _ | _ | y <;> simp only
    [bot_inf, inf_bot, top_inf, inf_top, mk_inf_mk]
  split_ifs with h₁ h₂ h₂ <;> simp_all

protected lemma inf_assoc (a b c : AboveBelow α) : a ⊓ b ⊓ c = a ⊓ (b ⊓ c) := by
  rcases a with _ | _ | x <;> rcases b with _ | _ | y <;> rcases c with _ | _ | z <;>
    simp only [bot_inf, inf_bot, top_inf, inf_top, mk_inf_mk]
  split_ifs <;> simp_all

protected lemma sup_inf_self (a b : AboveBelow α) : a ⊔ a ⊓ b = a := by
  rcases a with _ | _ | x <;> rcases b with _ | _ | y <;>
    simp only [bot_sup, sup_bot, top_sup, sup_top, mk_sup_mk,
               bot_inf, inf_bot, top_inf, inf_top, mk_inf_mk] <;>
    try (split_ifs <;> simp_all)

protected lemma inf_sup_self (a b : AboveBelow α) : a ⊓ (a ⊔ b) = a := by
  rcases a with _ | _ | x <;> rcases b with _ | _ | y <;>
    simp only [bot_sup, sup_bot, top_sup, sup_top, mk_sup_mk,
               bot_inf, inf_bot, top_inf, inf_top, mk_inf_mk] <;>
    try (split_ifs <;> simp_all)

instance : Lattice (AboveBelow α) :=
  Lattice.mk' AboveBelow.sup_comm AboveBelow.sup_assoc
    AboveBelow.inf_comm AboveBelow.inf_assoc
    AboveBelow.sup_inf_self AboveBelow.inf_sup_self

lemma le_iff {a b : AboveBelow α} : a ≤ b ↔ a ⊔ b = b := sup_eq_right.symm

lemma bot_le' (a : AboveBelow α) : (bot : AboveBelow α) ≤ a :=
  le_iff.mpr (bot_sup a)

lemma le_top' (a : AboveBelow α) : a ≤ (top : AboveBelow α) :=
  le_iff.mpr (sup_top a)

instance : OrderBot (AboveBelow α) where
  bot := bot
  bot_le := bot_le'

instance : OrderTop (AboveBelow α) where
  top := top
  le_top := le_top'

lemma bot_lt_mk (x : α) : (bot : AboveBelow α) < mk x :=
  lt_of_le_of_ne (bot_le' _) (by simp)

lemma mk_lt_top (x : α) : (mk x : AboveBelow α) < top :=
  lt_of_le_of_ne (le_top' _) (by simp)

lemma bot_lt_top : (bot : AboveBelow α) < top :=
  lt_of_le_of_ne (bot_le' _) (by simp)

lemma le_cases {a b : AboveBelow α} (h : a ≤ b) :
    a = bot ∨ b = top ∨ a = b := by
  have hsup := le_iff.mp h
  rcases a with _ | _ | x <;> rcases b with _ | _ | y
  · exact Or.inl rfl
  · exact Or.inr (Or.inl rfl)
  · exact Or.inl rfl
  · exact absurd hsup (by simp)
  · exact Or.inr (Or.inl rfl)
  · exact absurd hsup (by simp)
  · exact absurd hsup (by simp)
  · exact Or.inr (Or.inl rfl)
  · rw [mk_sup_mk] at hsup
    by_cases hxy : x = y
    · exact Or.inr (Or.inr (by rw [hxy]))
    · rw [if_neg hxy] at hsup
      exact absurd hsup (by simp)

/-- Monotonicity for *strict* operations on flat lattices: if `f` sends `⊥` to
`⊥` (in either argument) and `⊤` to `⊤` (against any non-`⊥` argument), it is
monotone in both arguments — regardless of its values on plain elements.
`Analysis/Sign.agda` and `Analysis/Constant.agda` postulated exactly these
monotonicity facts for their `plus`/`minus`, all of which have this shape. -/
lemma monotone₂_of_strict {β γ : Type*} [DecidableEq β] [DecidableEq γ]
    (f : AboveBelow α → AboveBelow β → AboveBelow γ)
    (hbotl : ∀ y, f bot y = bot) (hbotr : ∀ x, f x bot = bot)
    (htopl : ∀ y, y ≠ bot → f top y = top)
    (htopr : ∀ x, x ≠ bot → f x top = top) : Monotone₂ f := by
  constructor
  · intro y a b hab
    show f a y ≤ f b y
    rcases le_cases hab with rfl | rfl | rfl
    · rw [hbotl]; exact bot_le' _
    · rcases eq_or_ne y bot with rfl | hy
      · rw [hbotr, hbotr]
      · rw [htopl y hy]; exact le_top' _
    · exact le_rfl
  · intro x a b hab
    show f x a ≤ f x b
    rcases le_cases hab with rfl | rfl | rfl
    · rw [hbotr]; exact bot_le' _
    · rcases eq_or_ne x bot with rfl | hx
      · rw [hbotl, hbotl]
      · rw [htopr x hx]; exact le_top' _
    · exact le_rfl

/-! ### Interpretations of flat lattices -/

section Interp

variable {V : Type*} {P : AboveBelow α → V → Prop}

lemma interp_sup_of (hbot : ∀ v, ¬P bot v) (htop : ∀ v, P top v)
    {s₁ s₂ : AboveBelow α} (v : V) (h : P s₁ v ∨ P s₂ v) : P (s₁ ⊔ s₂) v := by
  rcases s₁ with _ | _ | x
  · rw [bot_sup]; exact h.resolve_left (hbot v)
  · rw [top_sup]; exact htop v
  · rcases s₂ with _ | _ | y
    · rw [sup_bot]; exact h.resolve_right (hbot v)
    · rw [sup_top]; exact htop v
    · rw [mk_sup_mk]
      split
      · next heq => subst heq; exact h.elim id id
      · exact htop v

lemma interp_inf_of
    (hdisj : ∀ {x y : α}, x ≠ y → ∀ v, ¬(P (mk x) v ∧ P (mk y) v))
    {s₁ s₂ : AboveBelow α} (v : V) (h : P s₁ v ∧ P s₂ v) : P (s₁ ⊓ s₂) v := by
  rcases s₁ with _ | _ | x
  · rw [bot_inf]; exact h.1
  · rw [top_inf]; exact h.2
  · rcases s₂ with _ | _ | y
    · rw [inf_bot]; exact h.2
    · rw [inf_top]; exact h.1
    · rw [mk_inf_mk]
      split
      · next heq => subst heq; exact h.1
      · next hne => exact absurd h (hdisj hne v)

end Interp

/-- Rank of an element: `⊥ ↦ 0`, `[x] ↦ 1`, `⊤ ↦ 2`. Used to bound chains
(Agda's `isLongest` / `x≺[y]⇒x≡⊥` / `[x]≺y⇒y≡⊤` case analysis lives here). -/
def rank : AboveBelow α → ℕ
  | bot => 0
  | mk _ => 1
  | top => 2

/-- Agda: the impossibility of `[x] ≺ [y]` (combines `x≺[y]⇒x≡⊥` and
`[x]≺y⇒y≡⊤`: the flat middle layer is an antichain). -/
lemma not_mk_lt_mk (x y : α) : ¬(mk x : AboveBelow α) < mk y := by
  intro h
  obtain ⟨hle, hne⟩ := lt_iff_le_and_ne.mp h
  have hsup := le_iff.mp hle
  rw [mk_sup_mk] at hsup
  by_cases hxy : x = y
  · rw [if_pos hxy] at hsup
    exact hne hsup
  · rw [if_neg hxy] at hsup
    exact absurd hsup (by simp)

lemma rank_strictMono : StrictMono (rank : AboveBelow α → ℕ) := by
  intro a b hab
  rcases a with _ | _ | x <;> rcases b with _ | _ | y
  · exact absurd hab (lt_irrefl _)
  · simp [rank]
  · simp [rank]
  · exact absurd hab (bot_le' _).not_lt
  · exact absurd hab (lt_irrefl _)
  · exact absurd hab (le_top' _).not_lt
  · exact absurd hab (bot_le' _).not_lt
  · simp [rank]
  · exact absurd hab (not_mk_lt_mk x y)

lemma boundedChains : BoundedChains (AboveBelow α) 2 := fun c => by
  have h := LTSeries.head_add_length_le_nat (c.map rank rank_strictMono)
  rw [LTSeries.head_map, LTSeries.last_map, LTSeries.map_length] at h
  have h2 : rank c.last ≤ 2 := by cases c.last <;> simp [rank]
  omega

instance [Inhabited α] : FiniteHeightLattice (AboveBelow α) where
  toLattice := inferInstance
  toOrderBot := inferInstance
  toOrderTop := inferInstance
  longestChain :=
    ((RelSeries.singleton _ bot).snoc (mk default)
        (by rw [RelSeries.last_singleton]; exact bot_lt_mk default)).snoc top
      (by rw [RelSeries.last_snoc]; exact mk_lt_top default)
  chains_bounded := boundedChains

end AboveBelow

end Spa
