Brackets.lean

module

import Std.Tactic.Do
public section

inductive Paren where
  | open : Paren
  | close : Paren

def Paren.toInt : Paren → Int
  | .open => 1
  | .close => -1

@[simp, grind =] theorem Paren.toInt_open : Paren.open.toInt = 1 := (rfl)
@[simp, grind =] theorem Paren.toInt_close : Paren.close.toInt = -1 := (rfl)

@[simp] theorem Paren.toInt_nonneg_iff {p : Paren} : 0 ≤ p.toInt ↔ p = .open := by cases p <;> simp
@[simp] theorem Paren.toInt_pos_iff {p : Paren} : 0 < p.toInt ↔ p = .open := by cases p <;> simp
@[simp] theorem Paren.toInt_nonpos_iff {p : Paren} : p.toInt ≤ 0 ↔ p = .close := by cases p <;> simp
@[simp] theorem Paren.toInt_neg_iff {p : Paren} : p.toInt < 0 ↔ p = .close := by cases p <;> simp
@[simp] theorem Paren.toInt_eq_one_iff {p : Paren} : p.toInt = 1 ↔ p = .open := by cases p <;> simp
@[simp] theorem Paren.toInt_eq_neg_one_iff {p : Paren} : p.toInt = -1 ↔ p = .close := by cases p <;> simp

inductive IsBalanced : List Paren → Prop where
  | empty : IsBalanced []
  | append (l₁ l₂) : IsBalanced l₁ → IsBalanced l₂ → IsBalanced (l₁ ++ l₂)
  | enclose (l) : IsBalanced l → IsBalanced (.open :: l ++ [.close])

attribute [simp] IsBalanced.empty

def balance (l : List Paren) : Int :=
  l.map (·.toInt) |>.sum

@[simp, grind =]
theorem balance_nil : balance [] = 0 := by
  simp [balance]

@[simp, grind =]
theorem balance_append : balance (l₁ ++ l₂) = balance l₁ + balance l₂ := by
  simp [balance]

@[simp, grind =]
theorem balance_cons : balance (p :: l) = p.toInt + balance l := by
  simp [balance]

@[simp]
theorem balance_flatMap {l : List α} {f : α → List Paren} : balance (l.flatMap f) = (l.map (balance ∘ f)).sum := by
  induction l with simp_all

theorem List.sum_eq_zero {l : List α} [Add α] [Zero α] [Std.LawfulIdentity (α := α) (· + ·) 0] (hl : ∀ a ∈ l, a = 0) : l.sum = 0 := by
  induction l with
  | nil => simp
  | cons hd tl ih =>
    simp only [mem_cons, forall_eq_or_imp] at hl
    simp [hl.1, ih hl.2, Std.LawfulLeftIdentity.left_id]

theorem List.take_cons_eq_if {l : List α} {a : α} {n : Nat} :
    (a::l).take n = if 0 < n then a :: l.take (n - 1) else [] := by
  split
  · exact take_cons ‹_›
  · obtain rfl : n = 0 := by omega
    simp

theorem List.take_singleton {a : α} {n : Nat} : [a].take n = if 0 < n then [a] else [] := by
  rw [List.take_cons_eq_if, List.take_nil]

def maxBalance (l : List Paren) : Int :=
  (0...=l.length).toList.map (fun k => balance (l.take k)) |>.max (by simp)

attribute [-simp] Nat.toList_rcc_eq_append
attribute [simp] Std.Rcc.mem_toList_iff_mem Std.Rcc.mem_iff

@[grind! .]
theorem maxBalance_nonneg (l : List Paren) : 0 ≤ maxBalance l := by
  rw [maxBalance]
  apply List.le_max_of_mem
  simp only [List.mem_map, Std.Rcc.mem_toList_iff_mem, Std.Rcc.mem_iff, Nat.zero_le, true_and]
  exact ⟨0, by simp⟩

@[simp]
theorem maxBalance_nil : maxBalance [] = 0 := by
  simp [maxBalance]

attribute [simp] Nat.le_max_left Nat.le_max_right

@[simp]
theorem maxBalance_cons {l : List Paren} {p : Paren} :
    maxBalance (p :: l) = max 0 (p.toInt + maxBalance l) := by
  rw [maxBalance]
  simp [Nat.toList_rcc_succ_right_eq_cons_map]
  rw [List.max_cons, List.max?_eq_some_max (by simp)]
  simp only [Option.elim_some, maxBalance, ← (List.max_map_eq_max (f := (p.toInt + ·)) (by simp)),
    List.map_map]
  congr 1

@[simp]
theorem maxBalance_append_singleton {l : List Paren} {p : Paren} :
    maxBalance (l ++ [p]) = max (maxBalance l) (balance l + p.toInt) := by
  induction l with
  | nil => simp
  | cons hd tl ih =>
    simp only [List.cons_append, maxBalance_cons, ih, balance_cons, Int.max_assoc]
    grind

theorem maxBalance_append {l m : List Paren} :
    maxBalance (l ++ m) = max (maxBalance l) (balance l + maxBalance m) := by
  induction l with
  | nil => simpa using (Int.max_eq_right (maxBalance_nonneg _)).symm
  | cons p l ih => grind [maxBalance_cons]

def minBalance (l : List Paren) : Int :=
  (0...=l.length).toList.map (fun k => balance (l.take k)) |>.min (by simp)

@[grind! .]
theorem minBalance_nonpos (l : List Paren) : minBalance l ≤ 0 := by
  rw [minBalance]
  apply List.min_le_of_mem
  simp only [List.mem_map, Std.Rcc.mem_toList_iff_mem, Std.Rcc.mem_iff, Nat.zero_le, true_and]
  exact ⟨0, by simp⟩

theorem minBalance_le_balance (l : List Paren) : minBalance l ≤ balance l := by
  rw [minBalance]
  apply List.min_le_of_mem
  simp only [List.mem_map, Std.Rcc.mem_toList_iff_mem, Std.Rcc.mem_iff, Nat.zero_le, true_and]
  exact ⟨l.length, by simp⟩

@[simp]
theorem minBalance_nil : minBalance [] = 0 := by
  simp [minBalance]

attribute [simp] Nat.min_le_left Nat.min_le_right

theorem minBalance_eq_zero_iff {l : List Paren} : minBalance l = 0 ↔ ∀ k, 0 ≤ balance (l.take k) := by
  simp only [← Std.le_antisymm_iff, minBalance_nonpos, true_and]
  simp only [minBalance, List.le_min_iff, List.mem_map, Std.Rcc.mem_toList_iff_mem, Std.Rcc.mem_iff,
    Nat.zero_le, true_and, forall_exists_index, and_imp, forall_apply_eq_imp_iff₂]
  exact ⟨fun h n => List.take_eq_take_min ▸ h (min n l.length) (by simp), fun h n _ => h n⟩

@[simp]
theorem minBalance_cons {l : List Paren} {p : Paren} :
    minBalance (p :: l) = min 0 (p.toInt + minBalance l) := by
  rw [minBalance]
  simp [Nat.toList_rcc_succ_right_eq_cons_map]
  rw [List.min_cons, List.min?_eq_some_min (by simp)]
  simp only [Option.elim_some, minBalance, ← (List.min_map_eq_min (f := (p.toInt + ·)) (by simp)),
    List.map_map]
  congr 1

@[simp]
theorem minBalance_append_singleton {l : List Paren} {p : Paren} :
    minBalance (l ++ [p]) = min (minBalance l) (balance l + p.toInt) := by
  induction l with
  | nil => simp
  | cons hd tl ih =>
    simp only [List.cons_append, minBalance_cons, ih, balance_cons, Int.min_assoc]
    grind

theorem minBalance_append {l m : List Paren} :
    minBalance (l ++ m) = min (minBalance l) (balance l + minBalance m) := by
  induction l with
  | nil => simpa using (Int.min_eq_right (minBalance_nonpos _)).symm
  | cons p l ih => grind [minBalance_cons]

theorem isBalanced_iff {l : List Paren} :
    IsBalanced l ↔ (balance l = 0 ∧ minBalance l = 0) := by
  rw [minBalance_eq_zero_iff]
  refine ⟨fun h => ?_, fun h => ?_⟩
  · induction h with
    | empty => simp
    | append l₁ l₂ ih₁ ih₂ h₁ h₂ => exact ⟨by grind, by grind [List.take_append]⟩ -- https://github.com/leanprover/lean4/issues/12581
    | enclose l h ih =>
      refine ⟨by grind, fun k => ?_⟩
      simp only [List.cons_append, List.take_cons_eq_if]
      grind [List.take_append, List.take_singleton]
  · rcases h with ⟨h₁, h₂⟩
    generalize h : l.length = n
    induction n using Nat.strongRecOn generalizing l with | ind n ih
    subst h
    by_cases h : ∃ k, 0 < k ∧ k < l.length ∧ balance (l.take k) = 0
    · obtain ⟨k, hk₁, hk₂, hk₃⟩ := h
      rw [← List.take_append_drop k l]
      have := List.take_append_drop k l
      refine IsBalanced.append (l.take k) (l.drop k)
        (ih k (by grind) (by grind) (by grind) (by grind))
        (ih (l.length - k) (by grind) (by grind) (fun n => ?_) (by grind))
      suffices balance (List.take n (List.drop k l)) = balance (List.take ((l.take k).length + n) l) by
        simpa [this] using h₂ _
      rw (occs := [3]) [← this]
      rw [List.take_length_add_append, balance_append, hk₃, Int.zero_add]
    · by_cases h₀ : l.length = 0
      · simp_all
      · have h' : ∀ k, 0 < k → k < l.length → 0 < balance (l.take k) := by grind
        obtain ⟨l, rfl⟩ : ∃ l', l = .open :: l' ++ [.close] := by
          obtain ⟨l, rfl⟩ : ∃ l', l = .open :: l' := by
            refine ⟨l.tail, ?_⟩
            rw (occs := [1]) [← List.cons_head_tail (l := l) (by grind), List.cons.injEq]
            have := h₂ 1
            simp_all [List.take_one, List.head?_eq_some_head (l := l) (by grind)]
          refine ⟨l.take (l.length - 1), ?_⟩
          have : l ≠ [] := by grind
          have hl : l = l.take (l.length - 1) ++ [l[l.length - 1]'(by grind)] := by
            rw (occs := [1]) [← List.take_append_drop (l.length - 1) l, List.append_right_inj]
            rw [List.drop_sub_one (by grind), List.drop_length, List.append_nil,
              getElem?_pos _ _ (by grind), Option.toList_some]
          suffices l[l.length - 1]'(by grind) = .close by
            rw (occs := [1]) [hl, this, List.cons_append]
          have : balance [l[l.length - 1]'(by grind)] ≤ 0 := by grind
          simpa using this
        refine IsBalanced.enclose l (ih _ (by grind) (by grind) (fun k => ?_) rfl)
        rw [List.take_eq_take_min]
        have := h' (min k l.length + 1)
        grind

theorem not_isBalanced_append_of_balance_neg {l m : List Paren} (h : balance l < 0) :
    ¬ IsBalanced (l ++ m) := by
  simpa [isBalanced_iff, minBalance_eq_zero_iff] using fun _ => ⟨l.length, by simp [h]⟩

theorem balance_nonneg_of_isBalanced_append {l m : List Paren} (h : IsBalanced (l ++ m)) :
    0 ≤ balance l := by
  simp only [isBalanced_iff, balance_append, minBalance_eq_zero_iff] at h
  simpa using h.2 l.length

theorem IsBalanced.balance_eq_zero {l : List Paren} (h : IsBalanced l) : balance l = 0 :=
  (isBalanced_iff.1 h).1

def parens (openBracket closeBracket : Char) (s : String) : List Paren :=
  s.toList.filterMap (fun c => if c = openBracket then some .open else if c = closeBracket then some .close else none)

@[simp]
theorem parens_empty {o c : Char} : parens o c "" = [] := by
  simp [parens]

@[simp]
theorem parens_append {o c : Char} {s t : String} :
    parens o c (s ++ t) = parens o c s ++ parens o c t := by
  simp [parens]

@[simp]
theorem parens_push {o c : Char} {s : String} {t : Char} :
    parens o c (s.push t) =
      parens o c s ++ (if t = o then some Paren.open else if t = c then some Paren.close else none).toList := by
  simp only [parens, String.toList_push, List.filterMap_append, List.filterMap_cons,
    List.filterMap_nil, List.append_cancel_left_eq]
  grind

def isBalanced (openBracket closeBracket : Char) (s : String) : Bool := Id.run do
  let mut depth := 0
  for c in s do
    if c == openBracket then
      depth := depth + 1
    else if c == closeBracket then
      if depth == 0 then
        return false
      depth := depth - 1
  return depth == 0

open Std.Do

set_option mvcgen.warning false
theorem isBalanced_eq_true_iff {openBracket closeBracket : Char} {s : String} (h : openBracket ≠ closeBracket) :
    isBalanced openBracket closeBracket s = true ↔ IsBalanced (parens openBracket closeBracket s) := by
  generalize hwp : isBalanced openBracket closeBracket s = w
  apply Std.Do.Id.of_wp_run_eq hwp
  mvcgen
  case inv1 =>
    exact Std.Do.StringInvariant.withEarlyReturn
      (fun pos depth => ⌜∀ t₁ t₂, pos.Splits t₁ t₂ → minBalance (parens openBracket closeBracket t₁) = 0 ∧ depth = balance (parens openBracket closeBracket t₁)⌝)
      (fun res depth => ⌜res = false ∧ ¬ IsBalanced (parens openBracket closeBracket s)⌝)
  next pos _ hp depth h₁ ih =>
    simp only [SPred.and_nil, SPred.down_pure, SPred.exists_nil, Bool.exists_bool, true_and,
      Bool.true_eq_false, false_and, and_false, or_false, SPred.or_nil] at ih
    have := ih.resolve_right (by simp [hp])
    simp only [Int.natCast_add, Int.cast_ofNat_Int, SPred.and_nil, SPred.down_pure, true_and,
      reduceCtorEq, false_and, SPred.exists_nil, exists_const, SPred.or_nil, or_false]
    intro t₁ t₂ hsp
    obtain ⟨t₁, rfl⟩ := hsp.exists_eq_append_singleton
    simp only [beq_iff_eq] at h₁
    simp only [h₁, String.append_singleton, parens_push, ↓reduceIte, Option.toList_some,
      minBalance_append, minBalance_cons, Paren.toInt_open, minBalance_nil, Int.add_zero,
      balance_append, balance_cons, balance_nil, Int.add_left_inj]
    have := this.2 _ _ hsp.of_next
    grind
  next pos _ hp depth h₁ h₂ h₃ ih =>
    simp only [SPred.and_nil, SPred.down_pure, SPred.exists_nil, Bool.exists_bool, true_and,
      Bool.true_eq_false, false_and, and_false, or_false, SPred.or_nil, reduceCtorEq,
      String.splits_endPos_iff, and_imp, forall_eq_apply_imp_iff, forall_eq, Option.some.injEq,
      Bool.false_eq, and_self_left, exists_eq_left, false_or] at ⊢ ih
    have := ih.resolve_right (by simp [hp])
    obtain ⟨t₁, t₂, hsp⟩ : ∃ t₁ t₂, (pos.next hp).Splits t₁ t₂ := ⟨_, _, (pos.next hp).splits⟩
    obtain ⟨t₁, rfl⟩ := hsp.exists_eq_append_singleton
    rw [hsp.eq_append, parens_append]
    apply not_isBalanced_append_of_balance_neg
    simp only [beq_iff_eq] at h₂ h₃
    simp only [h₂, String.append_singleton, parens_push, Ne.symm h, ↓reduceIte, Option.toList_some,
      balance_append, balance_cons, Paren.toInt_close, Int.reduceNeg, balance_nil, Int.add_zero]
    have := this.2 _ _ hsp.of_next
    grind
  next pos _ hp depth h₁ h₂ h₃ ih =>
    simp only [SPred.and_nil, SPred.down_pure, SPred.exists_nil, Bool.exists_bool, true_and,
      Bool.true_eq_false, false_and, and_false, or_false, SPred.or_nil, reduceCtorEq,
      exists_const] at ⊢ ih
    have := ih.resolve_right (by simp [hp])
    intro t₁ t₂ hsp
    obtain ⟨t₁, rfl⟩ := hsp.exists_eq_append_singleton
    simp only [beq_iff_eq] at h₂ h₃
    simp only [h₂, String.append_singleton, parens_push, Ne.symm h, ↓reduceIte, Option.toList_some,
      minBalance_append, minBalance_cons, Paren.toInt_close, Int.reduceNeg, minBalance_nil,
      Int.add_zero, balance_append, balance_cons, balance_nil]
    have := this.2 _ _ hsp.of_next
    grind
  next pos _ hp depth h₁ h₂ ih =>
    simp only [SPred.and_nil, SPred.down_pure, SPred.exists_nil, Bool.exists_bool, true_and,
      Bool.true_eq_false, false_and, and_false, or_false, SPred.or_nil, reduceCtorEq,
      exists_const] at ⊢ ih
    have := ih.resolve_right (by simp [hp])
    intro t₁ t₂ hsp
    obtain ⟨t₁, rfl⟩ := hsp.exists_eq_append_singleton
    simp only [beq_iff_eq] at h₁ h₂
    simpa [h₁, h₂] using this.2 _ _ hsp.of_next
  next => simp
  next p h ih =>
    simp only [h, String.splits_endPos_iff, and_imp, forall_eq_apply_imp_iff, forall_eq,
      SPred.and_nil, SPred.down_pure, true_and, reduceCtorEq, false_and, SPred.exists_nil,
      exists_const, SPred.or_nil, or_false, beq_iff_eq] at ⊢ ih
    rw [isBalanced_iff]
    grind
  next p b h ih => simp_all