import Spa.Lattice

/-!

# The Above-Below Lattice

This file defines the `AboveBelow` lattice, which takes a flat domain
$a_1, \ldots, a_n \in \alpha$ and lifts it into a lattice bounded
above by a synthetic $\top$ element, and below by a synthetic $\bot$
element.

$$
\begin{array}{ccccc}
  && \top && \\
  & \swarrow & \downarrow & \searrow & \\
  a_1 & & … & & a_n \\
  & \searrow & \downarrow & \swarrow & \\
  && \bot &&
\end{array}
$$

This lattice is also a `Spa.FiniteHeightLattice`, because no chain can
exceed the bottom-to-top chain $\bot < a_i < \top$.

The above-below lattice is helpful for for analyses such as
`Spa/Analysis/Sign.lean` and `Spa/Analysis/Constant.lean`, whose
classifications of values (by sign or by exact value) do not have
any inherent structure beyond "matching exactly".

-/

namespace Spa

/-- The above-below lattice, with bottom element `bot` and top element `top`. -/
@[aesop safe cases]
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 aesop
protected lemma sup_assoc (a b c : AboveBelow α) : a ⊔ b ⊔ c = a ⊔ (b ⊔ c) := by aesop
protected lemma inf_comm (a b : AboveBelow α) : a ⊓ b = b ⊓ a := by aesop
protected lemma inf_assoc (a b c : AboveBelow α) : a ⊓ b ⊓ c = a ⊓ (b ⊓ c) := by aesop
protected lemma sup_inf_self (a b : AboveBelow α) : a ⊔ a ⊓ b = a := by aesop
protected lemma inf_sup_self (a b : AboveBelow α) : a ⊓ (a ⊔ b) = a := by aesop

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
  rw [le_iff] at h
  rcases a with _ | _ | x <;> rcases b with _ | _ | y <;> simp_all

/-- If `f` sends `⊥` to `⊥` (in both arguments) and `⊤` to `⊤`
    (against any non-`⊥` argument), it is monotone in both arguments.
    The values of the the elements in `α` are irrelevant since they
    are always incomparable. This makes it easy to prove monotonicity
    for operations that "just" combine their flat elements, or give up. -/
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 c a b hab <;>
    rcases eq_or_ne c bot with rfl | hc <;>
    rcases le_cases hab with rfl | rfl | rfl <;>
    simp [hbotl, hbotr, htopl, htopr, bot_le', le_top', *]

section Interp

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

/-- As long as the interpretation of a the above-below lattice respects the
    fact that `bot` means "impossible", interpreting the above-below
    lattice agrees with its `⊔`. -/
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 aesop

/-- As long as two distinct values in the flat domain don't overlap,
    interpreting the above-below lattice agrees with its `⊔` -/
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 <;> rcases s₂ with _ | _ | y <;> simp_all
  split
  · exact h.2
  · next hne => exact (hdisj hne v h.1 h.2).elim

end Interp

/-- synthetic rank of an element, used to prove chain bounds. -/
private def rank : AboveBelow α → ℕ
  | bot => 0
  | mk _ => 1
  | top => 2

/-- It's not possible for any two lifted flat-domain elements to be less
    than one another. -/
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
  rcases le_cases hle with h | h | h <;> simp_all

/-- The rank of elements is strictly monotonic. -/
lemma rank_strictMono : StrictMono (rank : AboveBelow α → ℕ) := by
  intro a b hab
  rcases a with _ | _ | x <;> rcases b with _ | _ | y <;>
    simp_all [rank, not_mk_lt_mk, (bot_le' _).not_lt, (le_top' _).not_lt]

/-- All chains in the above-below lattice have at most 2 comparisons. -/
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
  height := 2
  chains_bounded := boundedChains

end AboveBelow

end Spa