import CubicPathwidth.Alg.Plan2
import CubicPathwidth.Alg.Proofs.Targets

/-!
# U4: `UndoWidth` (width bridge for `undo`, `lem:undo`)

Proof of `undoWidth_of`, from the hypothesis `PdTw2Width` (task U1). Every bag of
`undo` is `X ++ D` for a bag `X` of the list `outBags ++ tailBags` built in `undoF`;
each such `X` has at most `maxBagCard bagsR + 3 ⌊log₂ |V(F)|⌋ + 1` vertices, and `D`
adds at most `|D|`. The local definitions `plansOf`, `outBags`, `tailBags`, `undoOf`
restate the body of `undoF` with the plan list as a parameter (equal by `rfl`).
Helpers live in the namespace `CubicPathwidth.Alg.U4`.
-/

noncomputable section

namespace CubicPathwidth.Alg

open Graph

namespace U4

/-! ## 1. Restated body of `undoF` -/

/-- The plan list of `undoF`: one entry `(S, P, fP, idx)` per component `C` of `rest`. -/
def plansOf (F : Graph) (ops : List Op) (R : Graph) (bagsR : List Bag) :
    List (List Nat × List Bag × Bool × Option Nat) :=
  (F.induced fun v => !(deletedSet ops).contains v && !R.mem v).components.map fun C =>
    (undoAttach F (deletedSet ops) C, (pdTw2F (F.induced C.contains)).1,
      (pdTw2F (F.induced C.contains)).2,
      if (undoAttach F (deletedSet ops) C).isEmpty = true then none
      else bagsR.findIdx? (fun x => subsetB (undoAttach F (deletedSet ops) C) x))

/-- The inserted bags of `undoF`, before the union with `D`. -/
def outBags (bagsR : List Bag)
    (plans : List (List Nat × List Bag × Bool × Option Nat)) : List Bag :=
  bagsR.zipIdx.flatMap fun ⟨B, i⟩ =>
    B :: plans.flatMap fun ⟨_, P, _, idx⟩ => if idx == some i then P.map (bagUnion · B) else []

/-- The tail bags of `undoF` (plans with empty attachment), before the union with `D`. -/
def tailBags (plans : List (List Nat × List Bag × Bool × Option Nat)) : List Bag :=
  plans.flatMap fun ⟨S, P, _, _⟩ => if S.isEmpty then P else []

/-- `undoF` with the plan list as a parameter. -/
def undoOf (F : Graph) (D : List Nat) (R : Graph) (bagsR : List Bag)
    (plans : List (List Nat × List Bag × Bool × Option Nat)) : List Bag × Bool :=
  let validS : List Nat → Bool := fun S => S.all R.mem && match S with
    | [] => true
    | [_] => true
    | [a, b] => R.hasEdge a b
    | _ => false
  let ok := plans.all fun ⟨S, _, _, idx⟩ => validS S && (S.isEmpty || idx.isSome)
  if !ok then ([F.nodes], true) else
  let fl := plans.any fun ⟨_, _, f, _⟩ => f
  let res := ((outBags bagsR plans) ++ tailBags plans).map (bagUnion · D)
  (if res.isEmpty then [D] else res, fl)

theorem undoF_eq (F : Graph) (ops : List Op) (R : Graph) (bagsR : List Bag) :
    undoF F ops R bagsR = undoOf F (deletedSet ops) R bagsR (plansOf F ops R bagsR) := rfl

theorem undo_eq (F : Graph) (ops : List Op) (R : Graph) (bagsR : List Bag) :
    undo F ops R bagsR = (undoOf F (deletedSet ops) R bagsR (plansOf F ops R bagsR)).1 := rfl

/-! ## 2. Cardinality helpers -/

theorem card_bagUnion_le (X Y : Bag) :
    (bagUnion X Y).toFinset.card ≤ X.toFinset.card + Y.toFinset.card := by
  unfold bagUnion
  rw [List.toFinset_append]
  exact Finset.card_union_le _ _

theorem card_le_maxBagCard {bags : List Bag} {B : Bag} (h : B ∈ bags) :
    B.toFinset.card ≤ maxBagCard bags := by
  induction bags with
  | nil => simp at h
  | cons C L ih =>
    unfold maxBagCard at ih ⊢
    simp only [List.map_cons, List.foldr_cons]
    rcases List.mem_cons.mp h with rfl | h'
    · exact le_max_left _ _
    · exact le_trans (ih h') (le_max_right _ _)

theorem maxBagCard_le {bags : List Bag} {K : Nat} (h : ∀ B ∈ bags, B.toFinset.card ≤ K) :
    maxBagCard bags ≤ K := by
  induction bags with
  | nil => simp [maxBagCard]
  | cons C L ih =>
    unfold maxBagCard at ih ⊢
    simp only [List.map_cons, List.foldr_cons]
    exact max_le (h C (by simp)) (ih fun B hB => h B (by simp [hB]))

theorem induced_numNodes_le (hInd : InducedSpec) {F : Graph} (hF : F.WF) (p : Nat → Bool) :
    (F.induced p).numNodes ≤ F.numNodes := by
  obtain ⟨hW, hmem, -⟩ := hInd F p hF
  have h1 : (F.induced p).nodes.toFinset.card = (F.induced p).nodes.length :=
    List.toFinset_card_of_nodup hW.nodup
  have h2 : (F.induced p).nodes.toFinset ⊆ F.nodes.toFinset := by
    intro x hx
    rw [List.mem_toFinset] at hx ⊢
    exact ((hmem x).mp hx).1
  have h3 : F.nodes.toFinset.card ≤ F.nodes.length := List.toFinset_card_le _
  unfold Graph.numNodes
  rw [← h1]
  exact le_trans (Finset.card_le_card h2) h3

/-- Every bag of `pd_tw2` on an induced subgraph of `F` has at most
`3 ⌊log₂ |V(F)|⌋ + 1` vertices, when the flag is `false`. -/
theorem pd_card_le (h : PdTw2Width) (hInd : InducedSpec) {F : Graph} (hF : F.WF)
    (p : Nat → Bool) (hfl : (pdTw2F (F.induced p)).2 = false) :
    ∀ B ∈ (pdTw2F (F.induced p)).1, B.toFinset.card ≤ 3 * Nat.log 2 F.numNodes + 1 := by
  intro B hB
  have hH : (F.induced p).WF := (hInd F p hF).1
  have hw : maxBagCard (pdTw2F (F.induced p)).1 ≤ 3 * Nat.log 2 (F.induced p).numNodes + 1 :=
    h (F.induced p).numNodes (F.induced p) hH hfl
  have hlog : Nat.log 2 (F.induced p).numNodes ≤ Nat.log 2 F.numNodes :=
    Nat.log_mono_right (induced_numNodes_le hInd hF p)
  have hmax := card_le_maxBagCard hB
  omega

theorem plansOf_bound (h : PdTw2Width) (hInd : InducedSpec) {F : Graph} (hF : F.WF)
    (ops : List Op) (R : Graph) (bagsR : List Bag) :
    ∀ pl ∈ plansOf F ops R bagsR, pl.2.2.1 = false →
      ∀ P ∈ pl.2.1, P.toFinset.card ≤ 3 * Nat.log 2 F.numNodes + 1 := by
  intro pl hpl hf P hP
  unfold plansOf at hpl
  obtain ⟨C, -, rfl⟩ := List.mem_map.mp hpl
  exact pd_card_le h hInd hF C.contains hf P hP

/-- Membership bound for the bag list `outBags ++ tailBags`: every bag has at most
`M + K` vertices, where `M` bounds `bagsR` and `K` bounds the `pd_tw2` bags. -/
theorem card_out_le {bagsR : List Bag}
    {plans : List (List Nat × List Bag × Bool × Option Nat)} {M K : Nat}
    (hM : ∀ B ∈ bagsR, B.toFinset.card ≤ M)
    (hK : ∀ pl ∈ plans, ∀ P ∈ pl.2.1, P.toFinset.card ≤ K) :
    ∀ X ∈ outBags bagsR plans ++ tailBags plans, X.toFinset.card ≤ M + K := by
  unfold outBags tailBags
  intro X hX
  rw [List.mem_append] at hX
  rcases hX with hX | hX
  · rw [List.mem_flatMap] at hX
    obtain ⟨⟨B, i⟩, hBi, hX⟩ := hX
    have hB : B ∈ bagsR := by
      have := List.mem_map_of_mem (f := Prod.fst) hBi
      simpa [List.zipIdx_map_fst] using this
    rcases List.mem_cons.mp hX with hXB | hX
    · rw [hXB]
      exact le_trans (hM B hB) (Nat.le_add_right _ _)
    · rw [List.mem_flatMap] at hX
      obtain ⟨⟨S, P, f, idx⟩, hpl, hX⟩ := hX
      dsimp only at hX
      split at hX
      · rw [List.mem_map] at hX
        obtain ⟨P', hP', rfl⟩ := hX
        have h1 := card_bagUnion_le P' B
        have h2 := hK _ hpl P' hP'
        have h3 := hM B hB
        omega
      · simp at hX
  · rw [List.mem_flatMap] at hX
    obtain ⟨⟨S, P, f, idx⟩, hpl, hX⟩ := hX
    dsimp only at hX
    split at hX
    · have h2 := hK _ hpl X hX
      omega
    · simp at hX

/-! ## 3. Width of `undoOf` -/

theorem undoOf_width (F : Graph) (D : List Nat) (R : Graph) (bagsR : List Bag)
    (plans : List (List Nat × List Bag × Bool × Option Nat)) (K : Nat)
    (hK : ∀ pl ∈ plans, pl.2.2.1 = false → ∀ P ∈ pl.2.1, P.toFinset.card ≤ K)
    (hfl : (undoOf F D R bagsR plans).2 = false) :
    maxBagCard (undoOf F D R bagsR plans).1 ≤ D.toFinset.card + maxBagCard bagsR + K := by
  unfold undoOf at hfl ⊢
  dsimp only at hfl ⊢
  split at hfl
  · simp at hfl
  · rename_i hok
    rw [if_neg hok]
    dsimp only
    have hf : ∀ pl ∈ plans, pl.2.2.1 = false := by
      intro ⟨S, P, f, idx⟩ hpl
      have h1 := (List.any_eq_false.mp hfl) _ hpl
      simpa using h1
    have hK' : ∀ pl ∈ plans, ∀ P ∈ pl.2.1, P.toFinset.card ≤ K :=
      fun pl hpl P hP => hK pl hpl (hf pl hpl) P hP
    have hcard := card_out_le (bagsR := bagsR) (M := maxBagCard bagsR)
      (fun B hB => card_le_maxBagCard hB) hK'
    split
    · have h0 : maxBagCard [D] = D.toFinset.card := by simp [maxBagCard]
      rw [h0]
      omega
    · apply maxBagCard_le
      intro X hX
      rw [List.mem_map] at hX
      obtain ⟨Y, hY, hYX⟩ := hX
      subst hYX
      have h1 : (bagUnion Y D).toFinset.card ≤ Y.toFinset.card + D.toFinset.card :=
        card_bagUnion_le Y D
      have h2 := hcard Y hY
      show (bagUnion Y D).toFinset.card ≤ _
      omega

end U4

open U4 in
/-- `UndoWidth` (the width bound of `lem:undo` for `undo`), from the `pd_tw2` width bound
`PdTw2Width`. -/
theorem undoWidth_of (h : PdTw2Width) : UndoWidth := by
  intro F R ops bagsR hF hfl
  have hInd : InducedSpec := inducedSpec
  rw [undoF_eq] at hfl
  rw [undo_eq]
  have hK := plansOf_bound h hInd hF ops R bagsR
  have hmax := undoOf_width F (deletedSet ops) R bagsR (plansOf F ops R bagsR)
    (3 * Nat.log 2 F.numNodes + 1) hK hfl
  have hreal := (Nat.cast_le (α := ℝ)).mpr hmax
  unfold pdWidth
  push_cast at hreal ⊢
  linarith

end CubicPathwidth.Alg

#print axioms CubicPathwidth.Alg.undoWidth_of