import CoalgebraicMuCalculus.MuCalculus
import Mathlib

inductive FixedPointOp where
  | μ
  | ν
  deriving DecidableEq

variable {Λ : Type} [Modalities Λ]
open Formula

inductive IsAlternationDepthBound : FixedPointOp → ℕ → Formula Λ → Prop where
  | top {η : FixedPointOp} : IsAlternationDepthBound η 0 .top
  | bot {η : FixedPointOp} : IsAlternationDepthBound η 0 .bot
  | var {η : FixedPointOp} (x : ℕ) : IsAlternationDepthBound η 0 (.var x)
  | or {η : FixedPointOp} {n : ℕ} {φ ψ : Formula Λ} :
      IsAlternationDepthBound η n φ →
      IsAlternationDepthBound η n ψ →
      IsAlternationDepthBound η n (.or φ ψ)
  | and {η : FixedPointOp} {n : ℕ} {φ ψ : Formula Λ} :
      IsAlternationDepthBound η n φ →
      IsAlternationDepthBound η n ψ →
      IsAlternationDepthBound η n (.and φ ψ)
  | modality {η : FixedPointOp} {n : ℕ} {l : Λ} {args : Modalities.signature l → Formula Λ} :
      (∀ l, IsAlternationDepthBound η n (args l)) →
      IsAlternationDepthBound η n (.modality l args)
  | lfp {n : ℕ} {φ : Formula Λ} :
      IsAlternationDepthBound .ν n φ →
      IsAlternationDepthBound .ν n (μ, φ)
  | gfp {n : ℕ} {φ : Formula Λ} :
      IsAlternationDepthBound .μ n φ →
      IsAlternationDepthBound .μ n (ν, φ)
  | subst {η : FixedPointOp} {n : ℕ} {φ ψ : Formula Λ} {x : ℕ} :
      IsAlternationDepthBound η n φ →
      IsAlternationDepthBound η n ψ →
      IsAlternationDepthBound η n (φ.subst (x, ψ))
  | up {η₁ η₂ : FixedPointOp} {n : ℕ} {φ : Formula Λ} :
      IsAlternationDepthBound η₁ n φ →
      IsAlternationDepthBound η₂ (n + 1) φ

noncomputable def alternationDepth (φ : Formula Λ) : ℕ :=
  sInf { n : ℕ | IsAlternationDepthBound .μ n φ ∧ IsAlternationDepthBound .ν n φ }

def Formula.closure : Formula Λ → List (ℕ × Formula Λ)
  | .top => [(0, ⊤)]
  | .bot => [(0, ⊥)]
  | .or φ ψ => [(0, φ ∨ ψ)] ++ closure φ ++ closure ψ
  | .and φ ψ => [(0, φ ∧ ψ)] ++ closure φ ++ closure ψ
  | .modality l args =>
      [(0, .modality l args)] ++ (Modalities.args l |>.flatMap fun a ↦ closure (args a))
  | .var x => [(0, ^x)]
  | .lfp φ => [(0, μ, φ)] ++ (closure φ).map (fun ⟨n, ψ⟩ ↦ (n + 1, ψ))
  | .gfp φ => [(0, ν, φ)] ++ (closure φ).map (fun ⟨n, ψ⟩ ↦ (n + 1, ψ))

def mentionsVar (n : ℕ) : Formula Λ → Bool
  | .top => false
  | .bot => false
  | .or φ ψ => mentionsVar n φ || mentionsVar n ψ
  | .and φ ψ => mentionsVar n φ || mentionsVar n ψ
  | .modality l args =>
      Modalities.args l |>.any fun a ↦ mentionsVar n (args a)
  | .var x => x == n
  | .lfp φ => mentionsVar (n + 1) φ
  | .gfp φ => mentionsVar (n + 1) φ

def size : Formula Λ → ℕ
  | .top => 1
  | .bot => 1
  | .or φ ψ => 1 + size φ + size ψ
  | .and φ ψ => 1 + size φ + size ψ
  | .modality l args =>
      1 + (Modalities.args l |>.map (fun a ↦ size (args a)) |>.sum)
  | .var _ => 1
  | .lfp φ => 1 + size φ
  | .gfp φ => 1 + size φ

lemma closure_mem_size (φ : Formula Λ) : ∀ {n ψ}, (n, ψ) ∈ φ.closure → size ψ ≤ size φ := by
  intro n χ
  induction φ generalizing n χ with
  | top => grind [Formula.closure]
  | bot => grind [Formula.closure]
  | or φ ψ hφ hψ =>
      intro h
      simp only [Formula.closure, List.cons_append, List.nil_append, List.mem_cons, Prod.mk.injEq, List.mem_append] at h
      cases h with
      | inl eq => rw [eq.right]
      | inr h => cases h <;> grind [size]
  | and φ ψ hφ hψ =>
      intro h
      simp only [Formula.closure, List.cons_append, List.nil_append, List.mem_cons, Prod.mk.injEq, List.mem_append] at h
      cases h with
      | inl eq => rw [eq.right]
      | inr h => cases h <;> grind [size]
  | modality l args hargs =>
      intro h
      simp only [Formula.closure, List.cons_append, List.nil_append, List.mem_cons, Prod.mk.injEq, List.mem_flatMap] at h
      cases h with
      | inl eq => rw [eq.right]
      | inr h =>
          obtain ⟨a, _, ha⟩ := h
          simp only [size]
          grw [hargs _ ha]
          apply Nat.le_add_left_of_le
          apply List.le_sum_of_mem
          apply List.mem_map_of_mem
          apply Modalities.args_mem
  | var x => grind [Formula.closure]
  | lfp φ hφ =>
      intro h
      simp only [Formula.closure, List.cons_append, List.nil_append, List.mem_cons, Prod.mk.injEq, List.mem_map, Prod.exists, exists_eq_right_right] at h
      cases h with
      | inl eq => rw [eq.right]
      | inr h => obtain ⟨_, h, _⟩ := h; grind [size]
  | gfp φ hφ =>
      intro h
      simp only [Formula.closure, List.cons_append, List.nil_append, List.mem_cons, Prod.mk.injEq, List.mem_map, Prod.exists, exists_eq_right_right] at h
      cases h with
      | inl eq => rw [eq.right]
      | inr h => obtain ⟨_, h, _⟩ := h; grind [size]

def ad : Formula Λ → ℕ
  | .top => 0
  | .bot => 0
  | .or φ ψ => Nat.max (ad φ) (ad ψ)
  | .and φ ψ => Nat.max (ad φ) (ad ψ)
  | .modality l args =>
      Modalities.args l
        |>.map (fun a ↦ ad (args a))
        |> List.foldr Nat.max 0
  | .var _ => 0
  | .lfp φ => Id.run do
      let mut n := Nat.max 1 (ad φ)
      for h : (offset, ψ) in closure φ do
        if h' : mentionsVar offset ψ then
          if let Formula.gfp χ := ψ then
            n := Nat.max n (ad (ν, χ) + 1)
      return n
  | .gfp φ => Id.run do
      let mut n := Nat.max 1 (ad φ)
      for h : (offset, ψ) in closure φ do
        if h' : mentionsVar offset ψ then
          if let Formula.lfp χ := ψ then
            n := Nat.max n (ad (μ, χ) + 1)
      return n
  termination_by φ => size φ
  decreasing_by
    · grind [size]
    · grind [size]
    · grind [size]
    · grind [size]
    · rw [size, Nat.lt_one_add_iff]
      apply List.le_sum_of_mem
      apply List.mem_map_of_mem
      apply Modalities.args_mem
    · grind [size]
    · rw [show size (μ, φ) = 1 + size φ from rfl, Nat.lt_one_add_iff]
      apply closure_mem_size
      exact h
    · grind [size]
    · rw [show size (ν, φ) = 1 + size φ from rfl, Nat.lt_one_add_iff]
      apply closure_mem_size
      exact h

lemma ad_eq_alternationDepth (φ : Formula Λ) : ad φ = alternationDepth φ := by
  sorry