import Mathlib.Tactic.Ring

inductive WFF : Type
  | p_atm : String → WFF
  | p_not : WFF → WFF
  | p_and : WFF → WFF → WFF
  | p_or  : WFF → WFF → WFF
  | p_imp : WFF → WFF → WFF
  | p_iff : WFF → WFF → WFF

open WFF

/-1.1 递归定义 num_binary(φ)：公式中二元联结词 (∧, ∨, →, ↔) 出现的次数-/
def num_binary : WFF → Nat
  | p_atm _   => 0
  | p_not φ   => num_binary φ
  | p_and φ ψ => num_binary φ + num_binary ψ + 1
  | p_or  φ ψ => num_binary φ + num_binary ψ + 1
  | p_imp φ ψ => num_binary φ + num_binary ψ + 1
  | p_iff φ ψ => num_binary φ + num_binary ψ + 1

/-
1.2 递归定义 num_sentence(φ)：公式中命题符号（原子命题）出现的次数
    注意：同一个命题符号多次出现应计算多次
    例如：若 α = (A → (¬ A))，则 num_binary α = 1, num_sentence α = 2
-/
def num_sentence : WFF → Nat
  | p_atm _   => 1
  | p_not φ   => num_sentence φ
  | p_and φ ψ => num_sentence φ + num_sentence ψ
  | p_or  φ ψ => num_sentence φ + num_sentence ψ
  | p_imp φ ψ => num_sentence φ + num_sentence ψ
  | p_iff φ ψ => num_sentence φ + num_sentence ψ

/-
1.3 证明：对任意公式 α，num_sentence α = num_binary α + 1
    即：命题符号出现次数 = 二元联结词出现次数 + 1
-/
theorem sentence_eq_binary_plus_one :
  ∀ α : WFF, num_sentence α = num_binary α + 1 := by
  intro α
  induction α with
  | p_atm =>
      simp [num_sentence, num_binary]
  | p_not φ ih =>
      simp [num_sentence, num_binary]
      exact ih
  | p_and φ ψ ihφ ihψ =>
      simp [num_sentence, num_binary]
      rw [ihφ, ihψ]
      ring
  | p_or φ ψ ihφ ihψ =>
      simp [num_sentence, num_binary]
      rw [ihφ, ihψ]
      ring
  | p_imp φ ψ ihφ ihψ =>
      simp [num_sentence, num_binary]
      rw [ihφ, ihψ]
      ring
  | p_iff φ ψ ihφ ihψ =>
      simp [num_sentence, num_binary]
      rw [ihφ, ihψ]
      ring


/-
2.1 递归定义 symbolLength(φ)：公式作为符号串的长度（符号总数）
    约定：每个原子命题算 1 个符号，¬ 算 1 个符号，二元联结词算 1 个符号，每个括号算 1 个符号。
    例如：p_atm "A" 对应串 "A"，长度为 1
    p_not φ  对应串 "(¬ φ)"，长度为 symbolLength φ + 3
    p_and φ ψ 对应串 "(φ ∧ ψ)"，长度为 symbolLength φ + symbolLength ψ + 3
-/
def symbolLength : WFF → Nat
  | p_atm _   => 1
  | p_not φ   => symbolLength φ + 3
  | p_and φ ψ => symbolLength φ + symbolLength ψ + 3
  | p_or  φ ψ => symbolLength φ + symbolLength ψ + 3
  | p_imp φ ψ => symbolLength φ + symbolLength ψ + 3
  | p_iff φ ψ => symbolLength φ + symbolLength ψ + 3

/-
2.2 证明：若 α 不含 ¬，则 symbolLength α 为奇数。
    提示：为简化证明和方便后面的证明，可证明更强的结论：对任意不含 ¬ 的公式 α，存在 k 使得 symbolLength α = 4 * k + 1
-/

def has_no_not : WFF → Prop
  | p_atm _   => True
  | p_not _   => False
  | p_and φ ψ => has_no_not φ ∧ has_no_not ψ
  | p_or  φ ψ => has_no_not φ ∧ has_no_not ψ
  | p_imp φ ψ => has_no_not φ ∧ has_no_not ψ
  | p_iff φ ψ => has_no_not φ ∧ has_no_not ψ

theorem symbolLength_4k1 :
  ∀ α : WFF, has_no_not α → ∃ k, symbolLength α = 4 * k + 1 := by
    intro α h
    induction α with
    | p_atm _ =>
        simp [symbolLength]
    | p_not φ =>
        contradiction
    | p_and φ ψ ihφ ihψ | p_or φ ψ ihφ ihψ | p_imp φ ψ ihφ ihψ | p_iff φ ψ ihφ ihψ =>
        simp [symbolLength]
        rcases h with ⟨hφ, hψ⟩
        specialize ihφ hφ
        specialize ihψ hψ
        rcases ihφ with ⟨kφ, hkφ⟩
        rcases ihψ with ⟨kψ, hkψ⟩
        use kφ + kψ + 1
        rw [hkφ, hkψ]
        ring

theorem symbolLength_no_not_odd :
  ∀ α : WFF, has_no_not α → symbolLength α % 2 = 1 := by
    intro α h
    rcases symbolLength_4k1 α h with ⟨k, hk⟩
    rw [hk]
    omega

/-
2.3 证明：若 α 不含 ¬，则命题符号出现次数 > symbolLength α / 4
    即 4 * num_sentence α > symbolLength α
    提示：证明若 symbolLength α = 4 * k + 1，则 num_sentence α = k + 1
-/
theorem more_than_quarter_sentence :
  ∀ α : WFF, has_no_not α → 4 * num_sentence α > symbolLength α := by
  intro α h
  have : ∀k, symbolLength α = 4 * k + 1 → num_sentence α = k + 1 := by
    induction α with
    | p_atm _ =>
        intro k hk
        simp [symbolLength] at hk
        simp [num_sentence]
        exact hk
    | p_not φ =>
        contradiction
    | p_and φ ψ ihφ ihψ | p_or φ ψ ihφ ihψ | p_imp φ ψ ihφ ihψ | p_iff φ ψ ihφ ihψ =>
        intro k hk
        simp [num_sentence]
        simp [symbolLength] at hk
        rcases h with ⟨hφ, hψ⟩
        specialize ihφ hφ
        specialize ihψ hψ
        have h1 : ∃ kφ, symbolLength φ = 4 * kφ + 1 := by
          apply symbolLength_4k1
          exact hφ
        have h2 : ∃ kψ, symbolLength ψ = 4 * kψ + 1 := by
          apply symbolLength_4k1
          exact hψ
        rcases h1 with ⟨kφ, hkφ⟩
        rcases h2 with ⟨kψ, hkψ⟩
        rw [hkφ, hkψ] at hk
        have : k = kφ + kψ + 1 := by
          omega
        specialize ihφ kφ hkφ
        specialize ihψ kψ hkψ
        rw [ihφ, ihψ]
        omega
  have : ∃ k, symbolLength α = 4 * k + 1 := symbolLength_4k1 α h
  rcases this with ⟨k, hk⟩
  specialize this k hk
  rw [hk, this]
  omega

/-
 3.0 目标：实现核心判定器 isTautology：判断一个公式是否为重言式
     思路：首先收集公式中所有原子命题的名字，生成所有可能的真值表赋值，然后对每种赋值计算公式的真值，最后检查是否所有赋值下公式都为真。
-/

-- 3.1 实现辅助函数 getVars：收集公式中所有原子命题的名字并去重
def getVars : WFF → List String
  | WFF.p_atm s   => [s]
  | WFF.p_not p   => getVars p
  | WFF.p_and p q => (getVars p ++ getVars q).eraseDups
  | WFF.p_or  p q => (getVars p ++ getVars q).eraseDups
  | WFF.p_imp p q => (getVars p ++ getVars q).eraseDups
  | WFF.p_iff p q => (getVars p ++ getVars q).eraseDups

-- 3.2 实现辅助函数 lookup：在赋值环境中查找变量的值
def lookup (env : List (String × Bool)) (s : String) : Bool :=
  match env.find? (λ x => x.1 == s) with
  | some (_, val) => val
  | none          => false -- 理论上不应发生，因为环境包含所有变量

-- 3.3 实现辅助函数 evaluate：根据赋值环境计算公式的真值
def evaluate (env : List (String × Bool)) : WFF → Bool
  | WFF.p_atm s   => lookup env s
  | WFF.p_not p   => !(evaluate env p)
  | WFF.p_and p q => (evaluate env p) && (evaluate env q)
  | WFF.p_or  p q => (evaluate env p) || (evaluate env q)
  | WFF.p_imp p q => !(evaluate env p) || (evaluate env q)
  | WFF.p_iff p q => (evaluate env p) == (evaluate env q)

-- 3.4 实现辅助函数 generateEnvironments：生成所有可能的真值表赋值
/-
给定 n 个变量，生成 2ⁿ  种赋值组合。
通过递归来实现：如果列表为空，只有一种空组合；
如果加入一个新变量 x，就把已有的每种组合分别复制两份，一份配上 (x, true)，另一份配上 (x, false)。
-/

def generateEnvironments : List String → List (List (String × Bool))
  | [] => [[]]
  | s :: ss =>
    let subEnvs := generateEnvironments ss
    let withTrue  := subEnvs.map (λ env => (s, true) :: env)
    let withFalse := subEnvs.map (λ env => (s, false) :: env)
    withTrue ++ withFalse

-- 3.5 组装核心判定器
def isTautology (f : WFF) : Bool :=
  let vars := getVars f
  let allEnvs := generateEnvironments vars
  allEnvs.all (λ env => evaluate env f == true)

/-3.6 用 #eval 验证以下重言式（结果应全部为 true）-/

def A : WFF := p_atm "A"
def B : WFF := p_atm "B"
def C : WFF := p_atm "C"

-- (a) 排中律: A ∨ (¬ A)
#eval isTautology (p_or A (p_not A))

-- (b) 矛盾律: ¬(A ∧ (¬ A))
#eval isTautology (p_not (p_and A (p_not A)))

-- (c) 蕴涵自反: A → A
#eval isTautology (p_imp A A)

-- (d) De Morgan 1: (¬(A ∧ B)) ↔ ((¬ A) ∨ (¬ B))
#eval isTautology (p_iff (p_not (p_and A B)) (p_or (p_not A) (p_not B)))

-- (e) De Morgan 2: (¬(A ∨ B)) ↔ ((¬ A) ∧ (¬ B))
#eval isTautology (p_iff (p_not (p_or A B)) (p_and (p_not A) (p_not B)))

-- (f) 逆否命题: (A → B) ↔ ((¬ B) → (¬ A))
#eval isTautology (p_iff (p_imp A B) (p_imp (p_not B) (p_not A)))

-- (g) 输出律 (Exportation): ((A ∧ B) → C) ↔ (A → (B → C))
#eval isTautology (p_iff (p_imp (p_and A B) C) (p_imp A (p_imp B C)))

-- (h) 双重否定: (¬(¬ A)) ↔ A
#eval isTautology (p_iff (p_not (p_not A)) A)

-- (i) 验证一个非重言式（结果应为 false）
#eval isTautology (p_imp A B)