import Mathlib
import SymbolicModalLogic.FixedPoints

class Modalities (Λ : Type) where
  signature : Λ → Type
  [signature_finenum : ∀ l, FinEnum (signature l)]
  dual : Λ → Λ
  dualSignature {α : Type} : (l : Λ) → (signature l → α) → (signature (dual l) → α)
  dual_involution : ∀ l, dual (dual l) = l

def Modalities.args {Λ : Type} [Modalities Λ] (l : Λ) : List (Modalities.signature l) :=
  have : FinEnum (signature l) := Modalities.signature_finenum l
  FinEnum.toList (signature l)

def Modalities.args_mem {Λ : Type} [Modalities Λ] (l : Λ) (a : Modalities.signature l) : a ∈ args l := by
  rw [args]
  apply @FinEnum.mem_toList _ (Modalities.signature_finenum l)

inductive Formula (Λ : Type) [Modalities Λ] where
  | top : Formula Λ
  | bot : Formula Λ
  | or : Formula Λ → Formula Λ → Formula Λ
  | and : Formula Λ → Formula Λ → Formula Λ
  | modality : (l : Λ) → (args : Modalities.signature l → Formula Λ) → Formula Λ
  | var : ℕ → Formula Λ
  | lfp : Formula Λ → Formula Λ
  | gfp : Formula Λ → Formula Λ

namespace Formula

scoped notation "⊤" => Formula.top
scoped notation "⊥" => Formula.bot
scoped infixl:30 " ∨ " => Formula.or
scoped infixl:35 " ∧ " => Formula.and
scoped notation "^" n:max => Formula.var n
scoped notation "μ, " φ => Formula.lfp φ
scoped notation "ν, " φ => Formula.gfp φ

variable {Λ : Type} [Modalities Λ]

instance [Repr Λ] : ToString (Formula Λ) where
  toString :=
    let rec impl : Formula Λ → String
      | .top => "⊤"
      | .bot => "⊥"
      | .or φ ψ => s!"({impl φ}) ∨ ({impl ψ})"
      | .and φ ψ => s!"({impl φ}) ∧ ({impl ψ})"
      | .var x => s!"^{x}"
      | .modality l args =>
          have : FinEnum (Modalities.signature l) := Modalities.signature_finenum l
          let argsStr := FinEnum.toList (Modalities.signature l)
            |>.map (fun a ↦ impl (args a))
            |> String.intercalate ", "
          s!"{reprStr l}({argsStr})"
      | .lfp φ => s!"μ, {impl φ}"
      | .gfp φ => s!"ν, {impl φ}"
    impl

def not : Formula Λ → Formula Λ
  | .top             => ⊥
  | .bot             => ⊤
  | .or φ ψ          => (not φ) ∧ (not ψ)
  | .and φ ψ         => (not φ) ∨ (not ψ)
  | .modality l args => .modality (Modalities.dual l) (Modalities.dualSignature l fun a ↦ not (args a))
  | .var x           => ^x
  | .lfp φ           => ν, not φ
  | .gfp φ           => μ, not φ

scoped notation "¬" φ:60 => Formula.not φ
scoped infixl:25 " ⇒ " => fun φ ψ ↦ ¬φ ∨ ψ

def up : Formula Λ → (start : ℕ := 0) → Formula Λ
  | .top,             _     => ⊤
  | .bot,             _     => ⊥
  | .or φ ψ,          start => up φ start ∨ up ψ start
  | .and φ ψ,         start => up φ start ∧ up ψ start
  | .modality l args, start => .modality l fun a ↦ up (args a) start
  | .var x,           start => if x < start then ^x else ^(x + 1)
  | .lfp φ,           start => μ, up φ (start + 1)
  | .gfp φ,           start => ν, up φ (start + 1)

scoped notation "⇑ " φ => Formula.up φ

def subst : Formula Λ → (ℕ × Formula Λ) → Formula Λ
  | .top,             _      => ⊤
  | .bot,             _      => ⊥
  | .or φ ψ,          σ      => subst φ σ ∨ subst ψ σ
  | .and φ ψ,         σ      => subst φ σ ∧ subst ψ σ
  | .modality l args, σ      => .modality l fun a ↦ subst (args a) σ
  | .var x,           (y, ψ) => if x < y then ^x else
                                if x > y then ^(x - 1) else
                                ψ
  | .lfp φ,           (y, ψ) => μ, subst φ (y + 1, ⇑ψ)
  | .gfp φ,           (y, ψ) => ν, subst φ (y + 1, ⇑ψ)

macro_rules
  | `($φ[$ψ/^$n]) => `(Formula.subst $φ ($n, $ψ))

end Formula

section Semantics

class Model (Λ : Type) [Modalities Λ] (F : Type → Type) [Functor F] where
  lifting {X : Type} (l : Λ) : (args : Modalities.signature l → Set X) → Set (F X)

abbrev Model.lift {Λ : Type} [Modalities Λ] (l : Λ) (F : Type → Type) [Functor F] [Model Λ F] {X : Type} : (args : Modalities.signature l → Set X) → Set (F X) :=
  Model.lifting l

class LawfulModel (Λ : Type) [Modalities Λ] (F : Type → Type) [Functor F] [LawfulFunctor F] [Model Λ F] where
  lifting_natural {X Y : Type} (f : X → Y) (l : Λ) (args : Modalities.signature l → Set Y) :
      (Model.lift l F fun a ↦ f ⁻¹' (args a)) = ((Functor.map f) ⁻¹' (Model.lift l F args))
  lifting_monotone {X : Type} (l : Λ) (args args' : Modalities.signature l → Set X) :
      (∀ a, args a ⊆ args' a) → Model.lift l F args ⊆ Model.lift l F args'
  lifting_dual {X : Type} (l : Λ) (args : Modalities.signature l → Set X) :
      Model.lift (Modalities.dual l) F (Modalities.dualSignature l args) = (Model.lift l F fun a ↦ (args a)ᶜ)ᶜ

variable {Λ : Type} [Modalities Λ]
variable {F : Type → Type} [Functor F] [LawfulFunctor F] [Model Λ F]

def Environment (X : Type) : Type := ℕ → Set X

namespace Environment

variable {X : Type}

def empty : Environment X := λ _ ↦ ∅

def up (xs : Environment X) : Environment X
  | 0     => ∅
  | n + 1 => xs n

def subst (xs : Environment X) (n : ℕ) (x : Set X) : Environment X :=
  λ m ↦ if m = n then x else xs m

scoped notation "⇑ " φ => Environment.up φ

def add (x : Set X) (xs : Environment X) : Environment X := subst (⇑ xs) 0 x

infixr:67 " :: " => Environment.add

variable {xs : Environment X} {x y : Set X} {n m : ℕ}

@[simp]
lemma subst_eq : subst xs n x n = x := by simp [subst]

@[simp]
lemma subst_neq : m ≠ n → subst xs n x m = xs m := by intro h; simp [subst, h]

@[simp]
lemma subst_subst_eq : subst (subst xs n x) n y = subst xs n y := by
  apply funext
  intro m
  by_cases h : m = n <;> simp [h]

lemma subst_subst_neq : m ≠ n → subst (subst xs n x) m y = subst (subst xs m y) n x := by
  intro hneq
  apply funext
  intro k
  by_cases h₁ : k = m <;> by_cases h₂ : k = n <;> simp [h₁, h₂, hneq, Ne.symm hneq]

lemma subst_up : (⇑(subst xs n x)) = subst (⇑xs) (n + 1) x := by
  apply funext
  intro m
  cases m with
  | zero   => simp [up]
  | succ m => by_cases h : m = n <;> simp [h, up]

lemma subst_add : x :: subst xs n y = subst (x :: xs) (n + 1) y := by
  rw [add, add, subst_subst_neq, subst_up]
  simp

end Environment

def interp {X : Type} (γ : X → F X) : Formula Λ → Environment X → Set X
  | .top,             _  => Set.univ
  | .bot,             _  => ∅
  | .or φ ψ,          xs => (interp γ φ xs) ∪ (interp γ ψ xs)
  | .and φ ψ,         xs => (interp γ φ xs) ∩ (interp γ ψ xs)
  | .modality l args, xs => γ ⁻¹' Model.lift l F fun a ↦ interp γ (args a) xs
  | .var x,           xs => xs x
  | .lfp φ,           xs => FixedPoints.lfp fun S ↦ interp γ φ (S :: xs)
  | .gfp φ,           xs => FixedPoints.gfp fun S ↦ interp γ φ (S :: xs)

def interp_subst_monotone [LawfulModel Λ F] {X : Type} {γ : X → F X} (φ : Formula Λ) (xs : Environment X) (n : ℕ) :
    Monotone fun x ↦ interp γ φ (xs.subst n x) := by
  induction φ generalizing xs n with
  | top =>
      intro S S' hle
      rfl
  | bot =>
      intro S S' hle
      rfl
  | or φ ψ ihφ ihψ =>
      intro S S' hle
      simp only [Set.le_eq_subset, Set.union_subset_iff, interp]
      constructor
      · apply Set.subset_union_of_subset_left
        exact ihφ xs n hle
      · apply Set.subset_union_of_subset_right
        exact ihψ xs n hle
  | and φ ψ ihφ ihψ =>
      intro S S' hle
      simp only [Set.le_eq_subset, Set.subset_inter_iff, interp]
      constructor
      · intro x ⟨hx, _⟩
        exact ihφ xs n hle hx
      · intro x ⟨_, hx⟩
        exact ihψ xs n hle hx
  | modality l args ih =>
      intro S S' hle
      simp only [Set.le_eq_subset]
      apply Set.preimage_mono
      apply LawfulModel.lifting_monotone
      intro a
      exact ih a xs n hle
  | var x =>
      intro S S' hle
      by_cases h : x = n
      · simp only [interp, h, Environment.subst_eq, Set.le_eq_subset]
        exact hle
      · simp [h, interp]
  | lfp φ ihφ =>
      intro S S' hle
      simp only [interp, Set.le_eq_subset, FixedPoints.lfp, Set.sInf_eq_sInter, Set.subset_sInter_iff, Set.mem_setOf_eq]
      intro U hU x hx
      apply hx
      calc
        interp γ φ (U :: xs.subst n S) = interp γ φ ((U :: xs).subst (n + 1) S)  := by rw [Environment.subst_add]
        _                              ⊆ interp γ φ ((U :: xs).subst (n + 1) S') := ihφ _ _ hle
        _                              = interp γ φ (U :: xs.subst n S')         := by rw [Environment.subst_add]
        _                              ⊆ U                                       := hU
  | gfp φ ihφ =>
      intro S S' hle
      simp only [interp, Set.le_eq_subset, FixedPoints.gfp, Set.sSup_eq_sUnion, Set.sUnion_subset_iff, Set.mem_setOf_eq]
      intro U hU x hx
      use U; constructor
      · calc
          U ⊆ interp γ φ (U :: xs.subst n S)          := hU
          _ = interp γ φ ((U :: xs).subst (n + 1) S)  := by rw [Environment.subst_add]
          _ ⊆ interp γ φ ((U :: xs).subst (n + 1) S') := ihφ _ _ hle
          _ = interp γ φ (U :: xs.subst n S')         := by rw [Environment.subst_add]
      · exact hx

lemma interp_lfp_isLeastFixedPt [LawfulModel Λ F] {X : Type} (γ : X → F X) (φ : Formula Λ) (xs : Environment X) :
    Function.IsLeastFixedPt (fun S ↦ interp γ φ (S :: xs)) (interp γ (.lfp φ) xs) := by
  apply FixedPoints.lfp_isLeastFixedPt
  simp [Environment.add, interp_subst_monotone]

lemma interp_gfp_isGreatestFixedPt [LawfulModel Λ F] {X : Type} (γ : X → F X) (φ : Formula Λ) (xs : Environment X) :
    Function.IsGreatestFixedPt (fun S ↦ interp γ φ (S :: xs)) (interp γ (.gfp φ) xs) := by
  apply FixedPoints.gfp_isGreatestFixedPt
  simp [Environment.add, interp_subst_monotone]

end Semantics