import Mathlib.Combinatorics.SimpleGraph.Basic
import Mathlib.Combinatorics.SimpleGraph.Finite
import Mathlib.Combinatorics.SimpleGraph.Coloring
import Mathlib.Combinatorics.SimpleGraph.Subgraph
import Mathlib.Combinatorics.SimpleGraph.Path
import Mathlib.Combinatorics.SimpleGraph.Walk
import Mathlib.Combinatorics.SimpleGraph.Maps
import Mathlib.Order.Disjoint
import Mathlib.Data.List.Basic
import Mathlib.Data.List.Lemmas
import Mathlib.Data.Fintype.Basic
import Mathlib.Data.Fintype.Card
import Mathlib.Data.Fintype.Sets
import Mathlib.Data.Set.Basic
import Mathlib.Tactic.Linarith

open SimpleGraph List Walk Set.Notation
open Classical

variable
  {α : Type}
  {V : Type} [Fintype V]
  (G : SimpleGraph V) [DecidableRel G.Adj]

noncomputable def extend_coloring
  (S : Set V) (S_coloring : Coloring (G.induce S) α)
  (v : V) (c : α)
  (h : c ∉ S_coloring '' (S ↓∩ G.neighborSet v))
  : Coloring (G.induce (S ∪ {v})) α
  := by
  simp at h
  apply Coloring.mk
  case color =>
    intro ⟨x, mem⟩
    by_cases hx : x ∈ S
    · exact S_coloring ⟨x, hx⟩
    · exact c
  case valid =>
    rintro ⟨x₁, mem_x₁⟩ ⟨x₂, mem_x₂⟩
    intro adj
    by_cases hx₁ : x₁ ∈ S <;> by_cases hx₂ : x₂ ∈ S
    -- 4 cases according to whether x₁ and/or x₂ are in S or not
    · simp [hx₁, hx₂]
      exact S_coloring.valid adj
    · simp [hx₁, hx₂]
      replace hx₂ : x₂ = v := mem_x₂.resolve_left hx₂
      simp [hx₂] at adj
      exact h x₁ (G.adj_symm adj) hx₁
    · simp [hx₁, hx₂]
      replace hx₁ : x₁ = v := mem_x₁.resolve_left hx₁
      simp [hx₁] at adj
      -- TODO the following 2 lines are kind of cumbersome
      intro wrong
      exact (h x₂ adj hx₂) (Eq.symm wrong)
    · intro a
      simp_all only [comap_adj, Function.Embedding.coe_subtype]
      simp_all only [Set.union_singleton, Set.mem_insert_iff, or_false, SimpleGraph.irrefl]

lemma preimage_val_card (A B : Set V) : (A ↓∩ B).toFinset.card = (A ∩ B).toFinset.card := by
  apply Finset.card_bij'
  case i => exact fun a _ => ↑a
  case j =>
    intro a h
    simp only [Set.toFinset_inter, Finset.mem_inter, Set.mem_toFinset] at h
    exact ⟨a, h.left⟩
  case hi =>
    intro a ha
    simp_all only [not_false_eq_true, true_and, Set.mem_toFinset, Set.mem_preimage, Set.toFinset_inter,
      Finset.mem_inter, Subtype.coe_prop, and_self]
  case hj =>
    intro a ha
    simp_all only [not_false_eq_true, true_and, Set.mem_toFinset, Set.mem_preimage]
    simp_all only [Set.toFinset_inter, Finset.mem_inter, Set.mem_toFinset]
  case left_inv =>
    intro a ha
    simp_all only [not_false_eq_true, true_and, Subtype.coe_eta]
  case right_inv =>
    intro a ha
    simp_all only [not_false_eq_true, true_and]

lemma extend_coloring'
  {k : ℕ}
  (S : Set V) (S_colorable : Colorable (G.induce S) k)
  (v : V) (h : (G.neighborSet v ∩ S).toFinset.card < k)
  : Colorable (G.induce (S ∪ {v})) k := by
  have S_coloring := S_colorable.some
  let N := G.neighborSet v

  have card_fact : (S ↓∩ N).toFinset.card = (S ∩ N).toFinset.card := by
    -- without "convert" this gives a timeout error
    have foo := preimage_val_card S N
    convert foo

  have unused_color : ∃ c : Fin k, c ∉ S_coloring '' (S ↓∩ N) := by
    suffices (S_coloring '' (S ↓∩ N)).toFinset.card < (Set.univ : Set (Fin k)).toFinset.card from by
      have ⟨c, hc⟩ := Finset.exists_mem_not_mem_of_card_lt_card this
      simp only [Set.mem_toFinset] at hc
      exact ⟨c, hc.right⟩
    simp only [Set.toFinset_image, Set.toFinset_univ, Finset.card_univ, Fintype.card_fin]
    calc
      (Finset.image (⇑S_coloring) (S ↓∩ N).toFinset).card ≤ (S ↓∩ N).toFinset.card := Finset.card_image_le
      _ = (S ∩ N).toFinset.card := card_fact
      _ = (N ∩ S).toFinset.card := by simp only [Set.inter_comm]
      _ < k := h
  have ⟨c, hc⟩ := unused_color
  exact Nonempty.intro (extend_coloring G S S_coloring v c hc)

lemma extend_coloring''
  {k : ℕ}
  (S : Set V) (S_colorable : Colorable (G.induce S) k)
  (v : V) (h : G.degree v < k)
  : Colorable (G.induce (S ∪ {v})) k := by
  apply extend_coloring'
  · assumption
  · calc
    (G.neighborSet v ∩ S).toFinset.card ≤ (G.neighborSet v).toFinset.card := by
      simp only [Set.toFinset_inter]
      apply Finset.card_le_card
      simp only [Set.subset_toFinset, Finset.coe_inter, Set.coe_toFinset, Set.inter_subset_left]
    _ = G.degree v := by simp only [Set.toFinset_card, G.card_neighborSet_eq_degree]
    _ < k := h

theorem color_path
    (k : ℕ) (maxDeg_le_k : G.maxDegree ≤ k)
    {u v : V} (P : G.Path u v)
    (S : Set V)
    (P_outside_S : ∀ x ∈ P.val.support, x ∉ S)
    (S_colorable : Colorable (G.induce S) k)
    : Colorable (G.induce (S ∪ P.val.support.tail.toFinset)) k := by
  classical
  let ⟨W, W_is_path⟩ := P
  cases W with
  | nil =>
    simp
    rw [Finset.coe_empty, Set.union_empty]
    exact S_colorable
  | @cons u w v adj W' =>
    simp at *
    let P' : G.Path w v := ⟨W', IsPath.of_cons W_is_path⟩
    have P'_colorable := color_path k maxDeg_le_k P' S (P_outside_S.right) S_colorable

    let N := G.neighborSet w
    -- set of vertices already colored by P'_coloring
    let C := S ∪ W'.support.tail.toFinset

    have u_not_in_W' : u ∉ W'.support := ((Walk.cons_isPath_iff adj W').mp W_is_path).right
    have u_not_in_W'_tail : u ∉ W'.support.tail := by
      intro wrong
      have u_in_W' := List.Mem.tail w wrong
      rw [←(Walk.support_eq_cons W')] at u_in_W'
      contradiction
    have hu : u ∈ (N \ C) := by
      simp [N, C]
      exact ⟨G.adj_symm adj, P_outside_S.left, u_not_in_W'_tail⟩

    have h : (N ∩ C).toFinset.card < k := by
      have ineq1 : (N ∩ C).toFinset.card + (N \ C).toFinset.card ≤ k := calc
        _ = N.toFinset.card := by simp only [Set.toFinset_inter, Set.toFinset_diff, Finset.card_inter_add_card_sdiff]
        _ ≤ G.maxDegree := G.degree_le_maxDegree w
        _ ≤ k := maxDeg_le_k
      have ineq2 : 1 ≤ (N \ C).toFinset.card := by
        simp only [Finset.one_le_card, Set.toFinset_nonempty]
        exact ⟨u, hu⟩
      linarith

    have extended_coloring := by
      apply extend_coloring' G C P'_colorable w
      -- TODO without "convert" this gives a timeout error
      convert h

    have fact : (W'.support.tail.toFinset ∪ {w} : Set V) = W'.support.toFinset := by
      simp only [coe_toFinset, Set.union_def, Set.mem_setOf_eq, Set.mem_singleton_iff, Or.comm, mem_support_iff]

    rw [Set.union_assoc, fact] at extended_coloring
    exact extended_coloring

lemma induce_maxDegree_mono (S : Set V) : (G.induce S).maxDegree ≤ G.maxDegree := by
  simp only [G.induce_eq_coe_induce_top S]
  apply maxDegree_le_of_forall_degree_le
  intro v
  rw [Subgraph.coe_degree]
  calc
    ((⊤ : G.Subgraph).induce S).degree v ≤ G.degree v := Subgraph.degree_le ((⊤ : G.Subgraph).induce S) v
    _ ≤ G.maxDegree := G.degree_le_maxDegree v

noncomputable def extend_path_maximally {u v : V} (P : G.Path u v)
  : ∃ w : V, ∃ Q : G.Walk w u, IsPath (Walk.append Q P.1) ∧ G.neighborFinset w ⊆ (Walk.append Q P.1).support.toFinset := by
  replace ⟨P, hP⟩ := P
  induction hn : Fintype.card V - P.support.length generalizing u P

  case zero =>
    have fact : P.support.toFinset = Finset.univ := by
      apply Finset.eq_univ_of_card
      rw [isPath_def] at hP
      rw [←(List.toFinset_card_of_nodup hP)] at hn
      apply Nat.le_antisymm
      · exact Finset.card_le_univ P.support.toFinset
      · omega
    use u, Walk.nil
    simp only [Walk.nil_append, fact, Finset.coe_univ, Finset.subset_univ, and_true]
    assumption

  case succ n ih =>
    by_cases hu : G.neighborFinset u ⊆ P.support.toFinset
    case pos =>
      use u, Walk.nil
      simp only [Walk.nil_append]
      exact ⟨hP, hu⟩
    case neg =>
      simp [Finset.subset_iff] at hu
      have ⟨x, e, hx⟩ := hu
      replace e : G.Adj x u := G.adj_symm e
      let P' := Walk.cons e P
      have hP' : IsPath P' := IsPath.cons hP hx
      have hn' : Fintype.card V - P'.support.length = n := by
        simp only [P', Walk.support_cons, List.length_cons]
        omega
      have ⟨w, Q', hQ'⟩ := ih P' hP' hn'
      let Q := Walk.concat Q' e
      use w, Q
      simp only [Q, Walk.concat_append]
      simp only [P'] at hQ'
      assumption


theorem brooks'
  (k : ℕ) (k_ge_3 : k ≥ 3)
  (maxDeg_le_k : G.maxDegree ≤ k)
  (no_max_clique : CliqueFree G (k + 1))
  : Colorable G k := by
  induction hn : Fintype.card V using Nat.strong_induction_on generalizing V G
  case h n ih _ _ =>
    by_cases h : n ≤ k
    · have n_colorable := G.colorable_of_fintype
      rw [hn] at n_colorable
      exact Colorable.mono h n_colorable
    · by_cases k_reg : G.IsRegularOfDegree k
      -- Assume G is not k-regular.
      case neg =>
        simp [IsRegularOfDegree] at k_reg
        have ⟨x, hx⟩ := k_reg
        let S : Set V := Set.univ \ {x}
        let G' := G.induce S

        -- Apply induction
        have S_maxDeg_le_k : G'.maxDegree ≤ k := sorry -- induce_maxDegree_mono G S
        have S_no_max_clique : G'.CliqueFree (k + 1) := CliqueFree.comap (Embedding.induce S) no_max_clique
        have S_card : Fintype.card (↑S) = n - 1 := by
          simp only [S, ←Set.toFinset_card, Set.toFinset_diff, Set.toFinset_univ, Set.toFinset_singleton]
          have fact := Finset.card_inter_add_card_sdiff Finset.univ {x}
          simp [hn] at fact
          omega
        have fact : n - 1 < n := by omega
        have S_colorable : (G.induce S).Colorable k := ih (n - 1) fact (G.induce S) S_maxDeg_le_k S_no_max_clique S_card

        replace hx : G.degree x < k := by
          have hx' : G.degree x ≤ k := Nat.le_trans (G.degree_le_maxDegree x) maxDeg_le_k
          omega

        have G_colorable := extend_coloring'' G S S_colorable x hx
        have fact : S ∪ {x} = Set.univ := by simp only [Set.union_singleton,
          Set.insert_diff_singleton, Set.mem_univ, Set.insert_eq_of_mem, S]
        rw [fact] at G_colorable
        exact Colorable.of_embedding (G.induceUnivIso.symm.toEmbedding) G_colorable

      -- Assume G is k-regular.
      case pos =>
        -- Pick any vertex v.
        have V_nonempty : Nonempty V := by
          apply Fintype.card_pos_iff.mp
          omega
        let v := V_nonempty.some

        -- Since G does not contain a clique on k+1 vertices, there exist
        -- two neighbours x, y of v that are not adjacent in G.
        have claim : ∃ x y : V, x ≠ y ∧ G.Adj v x ∧ G.Adj v y ∧ ¬G.Adj x y := by
          by_contra wrong
          apply no_max_clique
          case t => exact (G.neighborSet v ∪ {v}).toFinset
          case a =>
            apply IsNClique.mk
            case card_eq =>
              simp only [Set.union_singleton, Set.toFinset_insert, Set.mem_toFinset,
                mem_neighborSet, SimpleGraph.irrefl, not_false_eq_true,
                Finset.card_insert_of_not_mem, Set.toFinset_card, Fintype.card_ofFinset,
                Nat.add_left_inj]
              have fact : ({x | G.Adj v x} : Finset V) = G.neighborFinset v := by
                ext a
                simp only [Finset.mem_filter, Finset.mem_univ, true_and, mem_neighborFinset]
              simp only [fact, card_neighborFinset_eq_degree]
              exact k_reg v
            case isClique =>
              intro x hx y hy x_ne_y
              simp at hx hy
              cases' hx with x_eq_v hx <;> cases' hy with y_eq_v hy
              · rw [x_eq_v, y_eq_v] at x_ne_y
                contradiction
              · rw [←x_eq_v] at hy
                assumption
              · rw [←y_eq_v] at hx
                exact G.adj_symm hx
              · simp at wrong
                exact wrong x y x_ne_y hx hy

        have ⟨x, y, x_ne_y, hx, hy, x_nadj_y⟩ := claim
        sorry