import Mathlib.Order.Lattice
import Mathlib.Order.RelSeries

/-!

# Lattice Definitions

This file provides some definitions for lattices. It used to be more critical
when this was an Agda project, since it defined (semi)lattices, the ordering
relation, etc. However, these have been lifted into `Mathlib.Order.Lattice`
etc.. What remains are a couple of theorems about folds, as well
as `FiniteHeightLattice`, the core concept of lattice-based static
program analyses. See the documentation on that class for more information. -/

namespace Spa

/-- Predicate for binary functions independently monotone in both their arguments. -/
def Monotone₂ {α β γ : Type*} [Preorder α] [Preorder β] [Preorder γ]
    (f : α → β → γ) : Prop :=
  (∀ b, Monotone (f · b)) ∧ (∀ a, Monotone (f a ·))

section Folds

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

/-- (right) folds are monotonic in both their arguments if the underlying accumulator function is. -/
lemma foldr_mono {l₁ l₂ : List α} (f : α → β → β) {b₁ b₂ : β}
    (hl : List.Forall₂ (· ≤ ·) l₁ l₂) (hb : b₁ ≤ b₂)
    (hf₁ : ∀ b, Monotone (f · b)) (hf₂ : ∀ a, Monotone (f a ·)) :
    l₁.foldr f b₁ ≤ l₂.foldr f b₂ := by
  induction hl with
  | nil => exact hb
  | cons hxy _ ih =>
    exact le_trans (hf₁ _ hxy) (hf₂ _ ih)

/-- (left) folds are monotinic in both their arguments if the underlying accumulator function is. -/
lemma foldl_mono {l₁ l₂ : List α} (f : β → α → β) {b₁ b₂ : β}
    (hl : List.Forall₂ (· ≤ ·) l₁ l₂) (hb : b₁ ≤ b₂)
    (hf₁ : ∀ a, Monotone (f · a)) (hf₂ : ∀ b, Monotone (f b ·)) :
    l₁.foldl f b₁ ≤ l₂.foldl f b₂ := by
  induction hl generalizing b₁ b₂ with
  | nil => exact hb
  | cons hxy _ ih =>
    exact ih (le_trans (hf₁ _ hb) (hf₂ _ hxy))

omit [Preorder α] in
/-- (right) folds on a particular list are monotonic if the underlying accumulator is monotonic in its accumulator argument. -/
lemma foldr_mono' (l : List α) (f : α → β → β)
    (hf : ∀ a, Monotone (f a ·)) : Monotone (l.foldr f ·) := by
  intro b₁ b₂ hb
  induction l with
  | nil => exact hb
  | cons x xs ih => exact hf x ih

omit [Preorder α] in
/-- (left) folds on a particular list are monotonic if the underlying accumulator is monotonic in its accumulator argument. -/
lemma foldl_mono' (l : List α) (f : β → α → β)
    (hf : ∀ a, Monotone (f · a)) : Monotone fun b => l.foldl f b := by
  intro b₁ b₂ hb
  induction l generalizing b₁ b₂ with
  | nil => exact hb
  | cons x xs ih => exact ih (hf x hb)

end Folds

/-- Predicate on types with `Preorder` that claims all $<$ chains in the type have at most `n` comparisons. -/
def BoundedChains (α : Type*) [Preorder α] (n : ℕ) : Prop :=
  ∀ c : LTSeries α, c.length ≤ n

/-- Since a singleton type's preorder has no nonempty `<` chains,
    they are vacuously bounded by any minimum height. -/
lemma boundedChains_of_subsingleton (α : Type*) [Preorder α] [Subsingleton α]
    (n : ℕ) : BoundedChains α n := fun c => by
  by_contra hc
  push_neg at hc
  exact (c.step ⟨0, by omega⟩).ne (Subsingleton.elim _ _)

/-- A finite height lattice is a lattice in which all chains $a < \ldots < z$ have a maximum height `height`. -/
class FiniteHeightLattice (α : Type*) extends Lattice α, OrderBot α, OrderTop α where
  longestChain : LTSeries α
  chains_bounded : BoundedChains α longestChain.length

-- a < ... < z
-- ----------- length <= height

namespace FiniteHeightLattice

def height (α : Type*) [FiniteHeightLattice α] : ℕ :=
  (longestChain (α := α)).length

variable (α : Type*) [FiniteHeightLattice α]

/-- Any maximum-length chain in a bounded finite-height lattice starts at `⊥`. -/
lemma longestChain_head : (longestChain (α := α)).head = ⊥ := by
  by_contra hne
  have hbound := chains_bounded ((longestChain (α := α)).cons ⊥ (bot_lt_iff_ne_bot.mpr hne))
  rw [RelSeries.cons_length] at hbound
  omega

/-- Any maximum-length chain in a bounded finite-height lattice ends at `⊤`. -/
lemma longestChain_last : (longestChain (α := α)).last = ⊤ := by
  by_contra hne
  have hbound := chains_bounded ((longestChain (α := α)).snoc ⊤ (lt_top_iff_ne_top.mpr hne))
  rw [RelSeries.snoc_length] at hbound
  omega

/-- This is something like a lemma about isomorphic types having the same height.
    Given a finite-height lattice `α`, lattice `β`, and a `Monotone` bijection
    between the two, we can show that lattice `β` also has a finite height.

    The proof is fairly trivial: the longest chain in `α` can be transported
    to be a longest chain in `β` (by monotonicity), establishing a height witness.
    At the same time, any chain in `β` can be transported to a chain in `α`,
    and must be bounded by the same height by `FiniteHeightLattice.chains_bounded`. -/
def transport {α β : Type*} [Lattice β]
    [I : FiniteHeightLattice α] (f : α → β) (g : β → α)
    (hf : Monotone f) (hg : Monotone g)
    (hgf : Function.LeftInverse g f) (hfg : Function.LeftInverse f g) :
    FiniteHeightLattice β where
  toLattice := inferInstance
  toOrderBot := {
    bot := f (⊥ : α)
    bot_le := fun b => by
      rw [← hfg b]
      exact hf (_root_.bot_le : (⊥ : α) ≤ g b) }
  toOrderTop := {
    top := f (⊤ : α)
    le_top := fun b => by
      rw [← hfg b]
      exact hf (_root_.le_top : g b ≤ (⊤ : α)) }
  longestChain :=
    I.longestChain.map f (hf.strictMono_of_injective hgf.injective)
  chains_bounded := fun c =>
    I.chains_bounded (c.map g (hg.strictMono_of_injective hfg.injective))

/-- A `Unique` lattice trivially has finite height: its only chain is the singleton
    `[default]`, and there are no nontrivial `<` chains in a subsingleton. -/
def ofUnique (α : Type*) [Lattice α] [Unique α] :
    FiniteHeightLattice α where
  toLattice := inferInstance
  toOrderBot := {
    bot := default
    bot_le := fun _ => le_of_eq (Subsingleton.elim _ _) }
  toOrderTop := {
    top := default
    le_top := fun _ => le_of_eq (Subsingleton.elim _ _) }
  longestChain := RelSeries.singleton _ default
  chains_bounded := boundedChains_of_subsingleton α 0

end FiniteHeightLattice

end Spa