import Mathlib.Data.Fintype.Sigma
import HardnessReductions.Graph.Colorability
import HardnessReductions.Graph.EdgeList
import HardnessReductions.Logic.Sat3

/-!
# SAT → 3-Colorability

The classical palette-and-gadget reduction. Three special vertices `T`, `F`,
`B` form a triangle, so in any proper 3-colouring they carry the three colours;
naming a colour then means naming the palette vertex that wears it. Each
variable slot gets a pair of vertices, one per polarity, joined to each other
and to `B`: the pair therefore wears `T`'s and `F`'s colours in some order,
which *is* a truth assignment. Each clause gets a chain of OR components
verifying that at least one of its literals wears `T`'s colour.

## The OR component

Three fresh vertices `u`, `v`, `w` per literal occurrence, forming a triangle,
with `u` joined to the chain so far, `v` joined to the occurrence's literal
vertex, and `w` joined to `B`. Writing `x` for the chain input and `y` for the
literal vertex — both of which avoid `B`'s colour, hence wear `T`'s or `F`'s —
the triangle forces exactly the OR:

* if `x` and `y` both wear `F`'s colour then `u` and `v` must take the other
  two, leaving `w` wearing `F`'s colour (`or_forcing`);
* conversely every combination of input colours extends, with `w` wearing
  `T`'s colour iff some input does (`colorOf`, and the case check inside
  `colorable_of_sat`).

The chain of a clause starts at a fourth special vertex `S`, joined to `T` and
`B` and so forced to wear `F`'s colour, and its final output is joined to `F`
and `B`, forcing it to wear `T`'s. So a proper colouring exists exactly when
some literal of every clause can wear `T`'s colour.

## Why the seed vertex, and why this works for unrestricted SAT

Textbook presentations reduce from 3SAT and start each clause's chain at its
first literal, which needs the clause to be non-empty and makes the gadget's
shape depend on the clause width. Seeding the chain at `S` instead buys
uniformity: a clause of *any* width, including zero, gets the same treatment.
An empty clause's chain never leaves `S`, so `S` ends up joined to `T`, `B`
*and* `F` — a `K₄` on the palette, and the graph is not 3-colourable. That is
the right answer, an empty clause being unsatisfiable, and it arrives with no
case split anywhere in the construction.

The result is a reduction from unrestricted `SAT`, with `SAT3` following by
composition through `sat3ToSat`. It also keeps the construction *flat*: the
chain is a fact about the edges, described positionally by `chainAt`, not a
recursion building them. Only the correctness proof inducts along it, which is
where the induction belongs — and which is what lets the certificate be a
straight-line enumeration rather than needing a closed form (contrast
`splitClause_eq` in `Logic/Sat3/FromSat.lean`).
-/

namespace Hardness

open SimpleGraph

variable {V : Type*}

/-! ## Vertex names

Fresh names are tagged tuples of input objects, as everywhere in this library.
The variable component is indexed by *slots* (`varSlot`) rather than by `V`,
which is what keeps the vertex type finite without a `Fintype V` hypothesis;
the gadget component is indexed by literal occurrences, reusing `CNF.Occ`. -/

/-- Vertices of the colouring gadget: four special vertices (`0` true, `1`
false, `2` base, `3` seed), a polarity pair per variable slot, and three per
literal occurrence. -/
abbrev ColVtx (φ : CNF V) (vs : List V) : Type :=
  Fin 4 ⊕ (Fin (vs.length + 1) × Bool) ⊕ (φ.Occ × Fin 3)

/-! The vertex type is built from `Fin`, `Bool`, sums, products and `CNF.Occ`,
so its `DecidableEq` and `Fintype` instances exist already. They are *named*
here because instance search will not rediscover them at depth: synthesising
`DecidableEq (ColVtx φ vs × ColVtx φ vs)` from scratch exceeds the search
budget, and adjacency is a membership test on a list of such pairs. Same
precaution, same reason, as `instFintypeExt` and friends in
`Sets/Matching3D/FromSat.lean`. -/

instance instDecidableEqColVtx (φ : CNF V) (vs : List V) : DecidableEq (ColVtx φ vs) :=
  inferInstance

instance instFintypeColVtx (φ : CNF V) (vs : List V) : Fintype (ColVtx φ vs) :=
  inferInstance

instance instDecidableEqColVtxPair (φ : CNF V) (vs : List V) :
    DecidableEq (ColVtx φ vs × ColVtx φ vs) :=
  instDecidableEqProd

/-- Special vertex `i`: `0` true, `1` false, `2` base, `3` seed. -/
def palVtx (φ : CNF V) (vs : List V) (i : Fin 4) : ColVtx φ vs := Sum.inl i

/-- The vertex of variable slot `k` at polarity `b`. -/
def varVtx (φ : CNF V) (vs : List V) (k : Fin (vs.length + 1)) (b : Bool) : ColVtx φ vs :=
  Sum.inr (Sum.inl (k, b))

/-- Vertex `r` of the OR component at occurrence `o`. -/
def gadVtx (φ : CNF V) (vs : List V) (o : φ.Occ) (r : Fin 3) : ColVtx φ vs :=
  Sum.inr (Sum.inr (o, r))

/-- The vertex carrying literal `l`: its variable's slot vertex at its
polarity. -/
def litVtx [DecidableEq V] (φ : CNF V) (vs : List V) (l : Literal V) : ColVtx φ vs :=
  varVtx φ vs (varSlot vs l.var) l.polarity

variable {φ : CNF V} {vs : List V}

/-- `litVtx` unfolded. A `rw` cannot use `litVtx`'s equation theorems directly
(it is a plain definition, not an equation-compiler match), so proofs that need
to see through it use this. -/
theorem litVtx_eq [DecidableEq V] (l : Literal V) :
    litVtx φ vs l = varVtx φ vs (varSlot vs l.var) l.polarity := rfl

@[simp] theorem palVtx_inj {i j : Fin 4} :
    palVtx φ vs i = palVtx φ vs j ↔ i = j := by simp [palVtx]

@[simp] theorem varVtx_inj {k k' : Fin (vs.length + 1)} {b b' : Bool} :
    varVtx φ vs k b = varVtx φ vs k' b' ↔ k = k' ∧ b = b' := by
  simp [varVtx, Prod.ext_iff]

@[simp] theorem gadVtx_inj {o o' : φ.Occ} {r r' : Fin 3} :
    gadVtx φ vs o r = gadVtx φ vs o' r' ↔ o = o' ∧ r = r' := by
  simp [gadVtx, Prod.ext_iff]

@[simp] theorem palVtx_ne_varVtx {i : Fin 4} {k : Fin (vs.length + 1)} {b : Bool} :
    palVtx φ vs i ≠ varVtx φ vs k b := by simp [palVtx, varVtx]

@[simp] theorem palVtx_ne_gadVtx {i : Fin 4} {o : φ.Occ} {r : Fin 3} :
    palVtx φ vs i ≠ gadVtx φ vs o r := by simp [palVtx, gadVtx]

@[simp] theorem varVtx_ne_gadVtx {k : Fin (vs.length + 1)} {b : Bool}
    {o : φ.Occ} {r : Fin 3} : varVtx φ vs k b ≠ gadVtx φ vs o r := by
  simp [varVtx, gadVtx]

@[simp] theorem varVtx_ne_palVtx {i : Fin 4} {k : Fin (vs.length + 1)} {b : Bool} :
    varVtx φ vs k b ≠ palVtx φ vs i := Ne.symm palVtx_ne_varVtx

@[simp] theorem gadVtx_ne_palVtx {i : Fin 4} {o : φ.Occ} {r : Fin 3} :
    gadVtx φ vs o r ≠ palVtx φ vs i := Ne.symm palVtx_ne_gadVtx

@[simp] theorem gadVtx_ne_varVtx {k : Fin (vs.length + 1)} {b : Bool}
    {o : φ.Occ} {r : Fin 3} : gadVtx φ vs o r ≠ varVtx φ vs k b := Ne.symm varVtx_ne_gadVtx

/-- The literal vertex is a variable vertex, never special or gadget. -/
@[simp] theorem palVtx_ne_litVtx [DecidableEq V] {i : Fin 4} {l : Literal V} :
    palVtx φ vs i ≠ litVtx φ vs l := palVtx_ne_varVtx

@[simp] theorem litVtx_ne_gadVtx [DecidableEq V] {l : Literal V} {o : φ.Occ} {r : Fin 3} :
    litVtx φ vs l ≠ gadVtx φ vs o r := varVtx_ne_gadVtx

/-! ## The chain

`chainAt i j` is the vertex carrying "some literal among the first `j` of
clause `i` holds": the seed for `j = 0`, and the output vertex of the `j`-th OR
component afterwards. The `else` branch is unreachable junk, entered only for
`j` past the clause's width; it is there so the definition is total in `j`,
which keeps the chain a plain `ℕ`-indexed family and so keeps `omega` and `rw`
away from dependent indices. -/

/-- The chain vertex of clause `i` after `j` literals. -/
def chainAt (φ : CNF V) (vs : List V) (i : Fin φ.length) : ℕ → ColVtx φ vs
  | 0 => palVtx φ vs 3
  | j + 1 =>
    if h : j < (φ[i.1]'i.isLt).length then gadVtx φ vs ⟨i, ⟨j, h⟩⟩ 2 else palVtx φ vs 3

/-- The output vertex of clause `i`'s chain: where the whole clause's verdict
ends up, and the vertex that gets forced to wear `T`'s colour. For an empty
clause this is the seed itself, which is what makes an empty clause complete
the palette to a `K₄`. -/
def lastOut (φ : CNF V) (vs : List V) (i : Fin φ.length) : ColVtx φ vs :=
  chainAt φ vs i (φ[i.1]'i.isLt).length

/-- Every chain vertex is either the seed or an OR output, and both wear a
colour other than `B`'s — which is what the forcing argument needs of the
chain input. -/
theorem chainAt_eq (i : Fin φ.length) (j : ℕ) :
    chainAt φ vs i j = palVtx φ vs 3 ∨ ∃ o : φ.Occ, chainAt φ vs i j = gadVtx φ vs o 2 := by
  cases j with
  | zero => exact Or.inl rfl
  | succ m =>
    by_cases h : m < (φ[i.1]'i.isLt).length
    · exact Or.inr ⟨⟨i, ⟨m, h⟩⟩, by simp [chainAt, h]⟩
    · exact Or.inl (by simp [chainAt, h])

/-! ## The edges

The gadget is presented as an explicit edge list in one orientation and handed
to `ofEdgeList` (`Graph/EdgeList.lean`), which supplies symmetry and drops the
diagonal. `mem_colorEdges` names the families and is the only way any proof
below looks at the list. -/

/-- The edges of the colouring gadget, in one orientation. -/
def colorEdges [DecidableEq V] (φ : CNF V) (vs : List V) :
    List (ColVtx φ vs × ColVtx φ vs) :=
  -- the palette triangle T—F, T—B, F—B, and the seed's S—T, S—B
  [(palVtx φ vs 0, palVtx φ vs 1), (palVtx φ vs 0, palVtx φ vs 2),
    (palVtx φ vs 1, palVtx φ vs 2), (palVtx φ vs 3, palVtx φ vs 0),
    (palVtx φ vs 3, palVtx φ vs 2)] ++
  -- a variable slot's two polarities disagree, and neither wears B's colour
  (List.finRange (vs.length + 1)).flatMap (fun k ↦
    [(varVtx φ vs k true, varVtx φ vs k false), (varVtx φ vs k true, palVtx φ vs 2),
      (varVtx φ vs k false, palVtx φ vs 2)]) ++
  -- one OR component per occurrence: a triangle, its output off B, and its two
  -- inputs — the occurrence's literal and the chain so far
  (allOccs φ).flatMap (fun o ↦
    [(gadVtx φ vs o 0, gadVtx φ vs o 1), (gadVtx φ vs o 0, gadVtx φ vs o 2),
      (gadVtx φ vs o 1, gadVtx φ vs o 2), (gadVtx φ vs o 2, palVtx φ vs 2),
      (litVtx φ vs (φ.occLit o), gadVtx φ vs o 1),
      (chainAt φ vs o.1 o.2.1, gadVtx φ vs o 0)]) ++
  -- every clause's chain must end wearing T's colour
  (List.finRange φ.length).flatMap (fun i ↦ [(lastOut φ vs i, palVtx φ vs 1)])

/-- Membership in `allOccs` is no constraint — every occurrence is listed — so
it can be simped away when unfolding the edge list. -/
private theorem mem_allOccs_iff_true (φ : CNF V) (o : φ.Occ) : (o ∈ allOccs φ) ↔ True :=
  iff_true_intro (mem_allOccs φ o)

/-- The edge families of the gadget, named. Every proof about the graph goes
through this lemma rather than through `colorEdges` itself. -/
theorem mem_colorEdges [DecidableEq V] (u w : ColVtx φ vs) :
    (u, w) ∈ colorEdges φ vs ↔
      (u = palVtx φ vs 0 ∧ w = palVtx φ vs 1) ∨
      (u = palVtx φ vs 0 ∧ w = palVtx φ vs 2) ∨
      (u = palVtx φ vs 1 ∧ w = palVtx φ vs 2) ∨
      (u = palVtx φ vs 3 ∧ w = palVtx φ vs 0) ∨
      (u = palVtx φ vs 3 ∧ w = palVtx φ vs 2) ∨
      (∃ k, u = varVtx φ vs k true ∧ w = varVtx φ vs k false) ∨
      (∃ k, u = varVtx φ vs k true ∧ w = palVtx φ vs 2) ∨
      (∃ k, u = varVtx φ vs k false ∧ w = palVtx φ vs 2) ∨
      (∃ o : φ.Occ, u = gadVtx φ vs o 0 ∧ w = gadVtx φ vs o 1) ∨
      (∃ o : φ.Occ, u = gadVtx φ vs o 0 ∧ w = gadVtx φ vs o 2) ∨
      (∃ o : φ.Occ, u = gadVtx φ vs o 1 ∧ w = gadVtx φ vs o 2) ∨
      (∃ o : φ.Occ, u = gadVtx φ vs o 2 ∧ w = palVtx φ vs 2) ∨
      (∃ o : φ.Occ, u = litVtx φ vs (φ.occLit o) ∧ w = gadVtx φ vs o 1) ∨
      (∃ o : φ.Occ, u = chainAt φ vs o.1 o.2.1 ∧ w = gadVtx φ vs o 0) ∨
      (∃ i : Fin φ.length, u = lastOut φ vs i ∧ w = palVtx φ vs 1) := by
  simp only [colorEdges, List.mem_append, List.mem_flatMap, List.mem_cons,
    List.not_mem_nil, or_false, Prod.mk.injEq, List.mem_finRange, mem_allOccs_iff_true,
    true_and, exists_or, or_assoc]

/-- The colouring gadget graph of `φ` over the variable list `vs`. -/
def satColorGraph [DecidableEq V] (φ : CNF V) (vs : List V) : SimpleGraph (ColVtx φ vs) :=
  ofEdgeList (colorEdges φ vs)

instance [DecidableEq V] (φ : CNF V) (vs : List V) :
    DecidableRel (satColorGraph φ vs).Adj := by
  unfold satColorGraph
  infer_instance

/-- Adjacency in the gadget graph, unfolded once and for all. -/
theorem satColorGraph_adj [DecidableEq V] (u w : ColVtx φ vs) :
    (satColorGraph φ vs).Adj u w ↔
      u ≠ w ∧ ((u, w) ∈ colorEdges φ vs ∨ (w, u) ∈ colorEdges φ vs) :=
  ofEdgeList_adj _ u w

/-- Building an adjacency: exhibit the edge in its stated orientation. -/
theorem adj_of_mem_colorEdges [DecidableEq V] {u w : ColVtx φ vs}
    (hne : u ≠ w) (h : (u, w) ∈ colorEdges φ vs) : (satColorGraph φ vs).Adj u w :=
  adj_of_mem_ofEdgeList hne h

/-! ## From an assignment to a colouring

The forward direction is a construction: name the three colours `0`, `1`, `2`
after `T`, `F`, `B`, colour each variable slot by what the assignment says, and
colour each OR component by the case table `gadColor`. Properness is then a
finite check per edge family, which is what `decide` is for — the only families
needing an argument are the ones mentioning the chain. -/

/-- What the chain of clause `i` carries after `j` literals: whether one of the
first `j` literals of the clause holds. -/
def chainVal (σ : V → Bool) (φ : CNF V) (i : Fin φ.length) (j : ℕ) : Bool :=
  ((φ[i.1]'i.isLt).take j).any fun l ↦ decide (l.Holds σ)

@[simp] theorem chainVal_zero (σ : V → Bool) (φ : CNF V) (i : Fin φ.length) :
    chainVal σ φ i 0 = false := rfl

/-- The chain recurrence: one more literal is OR-ed in. -/
theorem chainVal_succ {σ : V → Bool} {φ : CNF V} {i : Fin φ.length} {j : ℕ}
    (h : j < (φ[i.1]'i.isLt).length) :
    chainVal σ φ i (j + 1) =
      (chainVal σ φ i j || decide (((φ[i.1]'i.isLt)[j]'h).Holds σ)) := by
  rw [chainVal, chainVal, List.take_succ_eq_append_getElem h, List.any_append]
  simp

/-- A clause's chain ends carrying exactly whether the clause is satisfied. -/
theorem chainVal_length {σ : V → Bool} {φ : CNF V} (i : Fin φ.length) :
    chainVal σ φ i (φ[i.1]'i.isLt).length = true ↔ Clause.Sat σ (φ[i.1]'i.isLt) := by
  rw [chainVal, List.take_length]
  simp [Clause.Sat]

/-- The colours of an OR component, as a function of its two inputs' truth
values: `a` the chain so far, `b` the occurrence's literal. Vertex `2` is the
output, and carries `a || b`. -/
def gadColor (a b : Bool) (r : Fin 3) : Fin 3 :=
  if r = 0 then (if a then 1 else if b then 2 else 0)
  else if r = 1 then (if b && !a then 1 else 2)
  else (if a || b then 0 else 1)

/-- The colours of the four special vertices: `T`, `F`, `B` get the three
colours, and the seed `S` shares `F`'s. -/
def palColor (i : Fin 4) : Fin 3 := if i = 0 then 0 else if i = 2 then 2 else 1

/-- The colouring determined by an assignment. -/
def colorOf [DecidableEq V] (σ : V → Bool) (φ : CNF V) (vs : List V) :
    ColVtx φ vs → Fin 3
  | Sum.inl i => palColor i
  | Sum.inr (Sum.inl (k, b)) => if slotValue vs σ k = b then 0 else 1
  | Sum.inr (Sum.inr (o, r)) =>
      gadColor (chainVal σ φ o.1 o.2.1) (decide ((φ.occLit o).Holds σ)) r

variable {σ : V → Bool}

@[simp] theorem colorOf_palVtx [DecidableEq V] (i : Fin 4) :
    colorOf σ φ vs (palVtx φ vs i) = palColor i := rfl

@[simp] theorem colorOf_varVtx [DecidableEq V] (k : Fin (vs.length + 1)) (b : Bool) :
    colorOf σ φ vs (varVtx φ vs k b) = if slotValue vs σ k = b then 0 else 1 := rfl

@[simp] theorem colorOf_gadVtx [DecidableEq V] (o : φ.Occ) (r : Fin 3) :
    colorOf σ φ vs (gadVtx φ vs o r) =
      gadColor (chainVal σ φ o.1 o.2.1) (decide ((φ.occLit o).Holds σ)) r := rfl

/-- A literal's vertex wears `T`'s colour exactly when the literal holds —
provided the literal's variable is listed, so that its slot is its own. -/
theorem colorOf_litVtx [DecidableEq V] {l : Literal V} (hl : l.var ∈ vs) :
    colorOf σ φ vs (litVtx φ vs l) = if l.Holds σ then 0 else 1 := by
  rw [litVtx, colorOf_varVtx, slotValue_varSlot hl]
  rfl

/-- The chain vertex after `j` literals wears `T`'s colour exactly when one of
those literals holds. The junk branch of `chainAt` is never reached, `j` being
bounded by the clause's width. -/
theorem colorOf_chainAt [DecidableEq V] (i : Fin φ.length) :
    ∀ j : ℕ, j ≤ (φ[i.1]'i.isLt).length →
      colorOf σ φ vs (chainAt φ vs i j) = if chainVal σ φ i j then 0 else 1
  | 0, _ => by simp [chainAt, palColor]
  | j + 1, hj => by
    have h : j < (φ[i.1]'i.isLt).length := Nat.lt_of_lt_of_le (Nat.lt_succ_self j) hj
    rw [show chainAt φ vs i (j + 1) = gadVtx φ vs ⟨i, ⟨j, h⟩⟩ 2 by simp [chainAt, h],
      colorOf_gadVtx, chainVal_succ h]
    rcases Bool.eq_false_or_eq_true (chainVal σ φ i j) with ha | ha <;>
      by_cases hb : ((φ[i.1]'i.isLt)[j]'h).Holds σ <;>
      simp [gadColor, CNF.occLit, ha, hb]

/-- The clause output vertex wears `T`'s colour exactly when the clause is
satisfied. -/
theorem colorOf_lastOut [DecidableEq V] (i : Fin φ.length) :
    colorOf σ φ vs (lastOut φ vs i) =
      if chainVal σ φ i (φ[i.1]'i.isLt).length then 0 else 1 :=
  colorOf_chainAt i _ le_rfl

/-- **Properness**, family by family. The only hypotheses used are that the
assignment satisfies `φ` (for the clause-output family) and that every variable
occurring is listed (for the literal-input family). -/
theorem colorOf_ne_of_mem_colorEdges [DecidableEq V] (hsat : φ.Sat σ)
    (hvs : ∀ c ∈ φ, ∀ l ∈ c, l.var ∈ vs) {u w : ColVtx φ vs}
    (h : (u, w) ∈ colorEdges φ vs) : colorOf σ φ vs u ≠ colorOf σ φ vs w := by
  rw [mem_colorEdges] at h
  rcases h with ⟨rfl, rfl⟩ | ⟨rfl, rfl⟩ | ⟨rfl, rfl⟩ | ⟨rfl, rfl⟩ | ⟨rfl, rfl⟩ |
    ⟨k, rfl, rfl⟩ | ⟨k, rfl, rfl⟩ | ⟨k, rfl, rfl⟩ |
    ⟨o, rfl, rfl⟩ | ⟨o, rfl, rfl⟩ | ⟨o, rfl, rfl⟩ | ⟨o, rfl, rfl⟩ |
    ⟨o, rfl, rfl⟩ | ⟨o, rfl, rfl⟩ | ⟨i, rfl, rfl⟩
  -- the palette triangle and the seed's two edges
  · simp [palColor]
  · simp [palColor]
  · simp [palColor]
  · simp [palColor]
  · simp [palColor]
  -- a variable slot's two polarities, and each against B
  · rcases Bool.eq_false_or_eq_true (slotValue vs σ k) with hk | hk <;> simp [hk]
  · rcases Bool.eq_false_or_eq_true (slotValue vs σ k) with hk | hk <;> simp [hk, palColor]
  · rcases Bool.eq_false_or_eq_true (slotValue vs σ k) with hk | hk <;> simp [hk, palColor]
  -- the OR triangle, and its output against B
  · simp only [colorOf_gadVtx]
    rcases Bool.eq_false_or_eq_true (chainVal σ φ o.1 o.2.1) with ha | ha <;>
      rcases Bool.eq_false_or_eq_true (decide ((φ.occLit o).Holds σ)) with hb | hb <;>
      simp [gadColor, ha, hb]
  · simp only [colorOf_gadVtx]
    rcases Bool.eq_false_or_eq_true (chainVal σ φ o.1 o.2.1) with ha | ha <;>
      rcases Bool.eq_false_or_eq_true (decide ((φ.occLit o).Holds σ)) with hb | hb <;>
      simp [gadColor, ha, hb]
  · simp only [colorOf_gadVtx]
    rcases Bool.eq_false_or_eq_true (chainVal σ φ o.1 o.2.1) with ha | ha <;>
      rcases Bool.eq_false_or_eq_true (decide ((φ.occLit o).Holds σ)) with hb | hb <;>
      simp [gadColor, ha, hb]
  · simp only [colorOf_gadVtx, colorOf_palVtx]
    rcases Bool.eq_false_or_eq_true (chainVal σ φ o.1 o.2.1) with ha | ha <;>
      rcases Bool.eq_false_or_eq_true (decide ((φ.occLit o).Holds σ)) with hb | hb <;>
      simp [gadColor, palColor, ha, hb]
  -- the literal input: this is where `hvs` is needed
  · have hmem : (φ.occLit o).var ∈ vs :=
      hvs _ (List.getElem_mem _) _ (CNF.occLit_mem φ o)
    rw [colorOf_litVtx hmem, colorOf_gadVtx]
    rcases Bool.eq_false_or_eq_true (chainVal σ φ o.1 o.2.1) with ha | ha <;>
      by_cases hb : (φ.occLit o).Holds σ <;> simp [gadColor, ha, hb]
  -- the chain input
  · rw [colorOf_chainAt o.1 o.2.1 (Nat.le_of_lt o.2.isLt), colorOf_gadVtx]
    rcases Bool.eq_false_or_eq_true (chainVal σ φ o.1 o.2.1) with ha | ha <;>
      rcases Bool.eq_false_or_eq_true (decide ((φ.occLit o).Holds σ)) with hb | hb <;>
      simp [gadColor, ha, hb]
  -- the clause output, forced to T's colour: this is where `hsat` is needed
  · rw [colorOf_lastOut, if_pos ((chainVal_length i).mpr (hsat _ (List.getElem_mem _)))]
    simp [palColor]

/-- A satisfying assignment colours the gadget. -/
theorem colorable_of_sat [DecidableEq V] (hsat : φ.Sat σ)
    (hvs : ∀ c ∈ φ, ∀ l ∈ c, l.var ∈ vs) : (satColorGraph φ vs).Colorable 3 :=
  (colorable_iff_exists_fun _ _).mpr
    ⟨colorOf σ φ vs, fun u w hadj ↦ by
      rcases (satColorGraph_adj u w).mp hadj with ⟨-, h | h⟩
      · exact colorOf_ne_of_mem_colorEdges hsat hvs h
      · exact (colorOf_ne_of_mem_colorEdges hsat hvs h).symm⟩

/-! ## From a colouring to an assignment

The backward direction is the one that has to reason about an *arbitrary*
witness: we are handed some proper 3-colouring we did not build, and must
extract a satisfying assignment from it. Nothing is known about the colouring
except that adjacent vertices differ, so every step is a pigeonhole in three
colours — which is what `fin3_pin` packages.

The shape of the argument is: the palette triangle fixes three distinct
colours, hence *all* of them; the seed is pinned to `F`'s; each variable pair
splits between `T`'s and `F`'s, which is the assignment; each OR component
propagates a disjunction *backwards* (`or_forcing`); and the chain, whose end
is pinned to `T`'s colour, is walked back down to the seed, where it must have
picked up a holding literal on the way (`exists_holds_of_chainAt`). -/

/-- In three colours, avoiding two of three distinct colours pins the third.
The whole backward direction is assembled from this one pigeonhole. -/
theorem fin3_pin (a b c d : Fin 3) (hab : a ≠ b) (hac : a ≠ c) (hbc : b ≠ c)
    (hda : d ≠ a) (hdc : d ≠ c) : d = b := by
  have ha := a.isLt
  have hb := b.isLt
  have hc := c.isLt
  have hd := d.isLt
  have hab' : a.val ≠ b.val := fun h ↦ hab (Fin.ext h)
  have hac' : a.val ≠ c.val := fun h ↦ hac (Fin.ext h)
  have hbc' : b.val ≠ c.val := fun h ↦ hbc (Fin.ext h)
  have hda' : d.val ≠ a.val := fun h ↦ hda (Fin.ext h)
  have hdc' : d.val ≠ c.val := fun h ↦ hdc (Fin.ext h)
  exact Fin.ext (by omega)

section Backward

variable [DecidableEq V]

/-! ### The edges, as adjacencies

`mem_colorEdges` gives membership; these turn the families that the backward
argument consumes into adjacency facts, ready to feed a colouring's validity. -/

theorem adj_pal_TF : (satColorGraph φ vs).Adj (palVtx φ vs 0) (palVtx φ vs 1) :=
  adj_of_mem_colorEdges (by simp) (by simp [mem_colorEdges])

theorem adj_pal_TB : (satColorGraph φ vs).Adj (palVtx φ vs 0) (palVtx φ vs 2) :=
  adj_of_mem_colorEdges (by simp) (by simp [mem_colorEdges])

theorem adj_pal_FB : (satColorGraph φ vs).Adj (palVtx φ vs 1) (palVtx φ vs 2) :=
  adj_of_mem_colorEdges (by simp)
    (by simp [mem_colorEdges])

theorem adj_pal_ST : (satColorGraph φ vs).Adj (palVtx φ vs 3) (palVtx φ vs 0) :=
  adj_of_mem_colorEdges (by simp)
    (by simp [mem_colorEdges])

theorem adj_pal_SB : (satColorGraph φ vs).Adj (palVtx φ vs 3) (palVtx φ vs 2) :=
  adj_of_mem_colorEdges (by simp)
    (by simp [mem_colorEdges])

theorem adj_varPair (k : Fin (vs.length + 1)) :
    (satColorGraph φ vs).Adj (varVtx φ vs k true) (varVtx φ vs k false) :=
  adj_of_mem_colorEdges (by simp)
    (by simp [mem_colorEdges])

theorem adj_varVtx_base (k : Fin (vs.length + 1)) (b : Bool) :
    (satColorGraph φ vs).Adj (varVtx φ vs k b) (palVtx φ vs 2) := by
  refine adj_of_mem_colorEdges (Ne.symm palVtx_ne_varVtx) ?_
  cases b <;> simp [mem_colorEdges]

theorem adj_gad_01 (o : φ.Occ) :
    (satColorGraph φ vs).Adj (gadVtx φ vs o 0) (gadVtx φ vs o 1) :=
  adj_of_mem_colorEdges (by simp) (by simp [mem_colorEdges])

theorem adj_gad_02 (o : φ.Occ) :
    (satColorGraph φ vs).Adj (gadVtx φ vs o 0) (gadVtx φ vs o 2) :=
  adj_of_mem_colorEdges (by simp) (by simp [mem_colorEdges])

theorem adj_gad_12 (o : φ.Occ) :
    (satColorGraph φ vs).Adj (gadVtx φ vs o 1) (gadVtx φ vs o 2) :=
  adj_of_mem_colorEdges (by simp) (by simp [mem_colorEdges])

theorem adj_gad_base (o : φ.Occ) :
    (satColorGraph φ vs).Adj (gadVtx φ vs o 2) (palVtx φ vs 2) :=
  adj_of_mem_colorEdges (Ne.symm palVtx_ne_gadVtx) (by simp [mem_colorEdges])

theorem adj_lit_gad (o : φ.Occ) :
    (satColorGraph φ vs).Adj (litVtx φ vs (φ.occLit o)) (gadVtx φ vs o 1) :=
  adj_of_mem_colorEdges litVtx_ne_gadVtx (by simp [mem_colorEdges])

theorem adj_chain_gad (o : φ.Occ) :
    (satColorGraph φ vs).Adj (chainAt φ vs o.1 o.2.1) (gadVtx φ vs o 0) := by
  refine adj_of_mem_colorEdges ?_ (by simp [mem_colorEdges])
  rcases chainAt_eq o.1 o.2.1 with h | ⟨o', h⟩ <;> rw [h] <;> simp

theorem adj_last_pal (i : Fin φ.length) :
    (satColorGraph φ vs).Adj (lastOut φ vs i) (palVtx φ vs 1) := by
  refine adj_of_mem_colorEdges ?_ (by simp [mem_colorEdges])
  rw [lastOut]
  rcases chainAt_eq i (φ[i.1]'i.isLt).length with h | ⟨o', h⟩ <;> rw [h] <;> simp

/-! ### Reading off the assignment -/

variable {f : ColVtx φ vs → Fin 3}
  (hf : ∀ u w : ColVtx φ vs, (satColorGraph φ vs).Adj u w → f u ≠ f w)

include hf

/-- The seed wears `F`'s colour: it is adjacent to `T` and `B`, and there are
only three colours. -/
theorem f_seed : f (palVtx φ vs 3) = f (palVtx φ vs 1) :=
  fin3_pin _ _ _ _ (hf _ _ adj_pal_TF) (hf _ _ adj_pal_TB) (hf _ _ adj_pal_FB)
    (hf _ _ adj_pal_ST) (hf _ _ adj_pal_SB)

/-- Every chain vertex avoids `B`'s colour: it is the seed or an OR output, and
both are adjacent to `B`. -/
theorem f_chainAt_ne_base (i : Fin φ.length) (j : ℕ) :
    f (chainAt φ vs i j) ≠ f (palVtx φ vs 2) := by
  rcases chainAt_eq i j with h | ⟨o, h⟩ <;> rw [h]
  · exact hf _ _ adj_pal_SB
  · exact hf _ _ (adj_gad_base o)

/-- The assignment read off a colouring: a variable is true exactly when its
positive vertex wears `T`'s colour. Variables that do not occur share the dummy
slot and so share a value, which no clause can see. -/
def assignmentOfColoring (f : ColVtx φ vs → Fin 3) (v : V) : Bool :=
  decide (f (varVtx φ vs (varSlot vs v) true) = f (palVtx φ vs 0))

/-- A variable slot's two vertices split `T`'s and `F`'s colours between them:
whichever one does not wear `T`'s wears `F`'s. -/
theorem f_varVtx_false (k : Fin (vs.length + 1)) :
    f (varVtx φ vs k false) = f (palVtx φ vs 0) ↔
      f (varVtx φ vs k true) ≠ f (palVtx φ vs 0) := by
  constructor
  · intro hfalse htrue
    exact hf _ _ (adj_varPair k) (htrue.trans hfalse.symm)
  · intro htrue
    have h1 : f (varVtx φ vs k true) = f (palVtx φ vs 1) :=
      fin3_pin _ _ _ _ (hf _ _ adj_pal_TF) (hf _ _ adj_pal_TB) (hf _ _ adj_pal_FB)
        htrue (hf _ _ (adj_varVtx_base k true))
    refine fin3_pin _ _ _ _ (Ne.symm (hf _ _ adj_pal_TF)) (hf _ _ adj_pal_FB)
      (hf _ _ adj_pal_TB) ?_ (hf _ _ (adj_varVtx_base k false))
    rw [← h1]
    exact fun hc ↦ hf _ _ (adj_varPair k) hc.symm

/-- A literal's vertex wears `T`'s colour exactly when the literal holds under
the extracted assignment. -/
theorem f_litVtx_iff (l : Literal V) :
    f (litVtx φ vs l) = f (palVtx φ vs 0) ↔ l.Holds (assignmentOfColoring f) := by
  have hpol : l.Holds (assignmentOfColoring f) ↔
      (decide (f (varVtx φ vs (varSlot vs l.var) true) = f (palVtx φ vs 0)) = l.polarity) :=
    Iff.rfl
  rw [litVtx_eq, hpol]
  cases hp : l.polarity
  · rw [f_varVtx_false hf, decide_eq_false_iff_not]
  · rw [decide_eq_true_iff]

/-- **The OR component forces the disjunction.** If a component's output wears
`T`'s colour then one of its two inputs does. This is the local step the chain
induction repeats, and the only place the three-colour pigeonhole does real
work: both inputs avoid `B`'s colour, so if neither wore `T`'s they would both
wear `F`'s, forcing the component's other two vertices onto the same colour. -/
theorem or_forcing (o : φ.Occ) (hout : f (gadVtx φ vs o 2) = f (palVtx φ vs 0)) :
    f (chainAt φ vs o.1 o.2.1) = f (palVtx φ vs 0) ∨
      f (litVtx φ vs (φ.occLit o)) = f (palVtx φ vs 0) := by
  by_contra hcon
  have hx : f (chainAt φ vs o.1 o.2.1) ≠ f (palVtx φ vs 0) := fun h ↦ hcon (Or.inl h)
  have hy : f (litVtx φ vs (φ.occLit o)) ≠ f (palVtx φ vs 0) := fun h ↦ hcon (Or.inr h)
  -- both inputs avoid B's colour, so both wear F's
  have hxF : f (chainAt φ vs o.1 o.2.1) = f (palVtx φ vs 1) :=
    fin3_pin _ _ _ _ (hf _ _ adj_pal_TF) (hf _ _ adj_pal_TB) (hf _ _ adj_pal_FB)
      hx (f_chainAt_ne_base hf o.1 o.2.1)
  have hyF : f (litVtx φ vs (φ.occLit o)) = f (palVtx φ vs 1) := by
    rw [litVtx_eq]
    exact fin3_pin _ _ _ _ (hf _ _ adj_pal_TF) (hf _ _ adj_pal_TB) (hf _ _ adj_pal_FB)
      (by rwa [litVtx_eq] at hy) (hf _ _ (adj_varVtx_base _ _))
  -- so the component's other two vertices are both pushed onto B's colour
  have hu : f (gadVtx φ vs o 0) = f (palVtx φ vs 2) := by
    refine fin3_pin _ _ _ _ (hf _ _ adj_pal_TB) (hf _ _ adj_pal_TF)
      (Ne.symm (hf _ _ adj_pal_FB)) ?_ ?_
    · rw [← hout]; exact hf _ _ (adj_gad_02 o)
    · rw [← hxF]; exact Ne.symm (hf _ _ (adj_chain_gad o))
  have hv : f (gadVtx φ vs o 1) = f (palVtx φ vs 2) := by
    refine fin3_pin _ _ _ _ (hf _ _ adj_pal_TB) (hf _ _ adj_pal_TF)
      (Ne.symm (hf _ _ adj_pal_FB)) ?_ ?_
    · rw [← hout]; exact hf _ _ (adj_gad_12 o)
    · rw [← hyF]; exact Ne.symm (hf _ _ (adj_lit_gad o))
  exact hf _ _ (adj_gad_01 o) (hu.trans hv.symm)

/-- **Walking the chain back.** If clause `i`'s chain wears `T`'s colour after
`j` literals, one of those `j` literals holds. The induction bottoms out at the
seed, which wears `F`'s colour and so can never have started the chain off
true — which is exactly why an empty clause has no colouring. -/
theorem exists_holds_of_chainAt (i : Fin φ.length) :
    ∀ j : ℕ, j ≤ (φ[i.1]'i.isLt).length →
      f (chainAt φ vs i j) = f (palVtx φ vs 0) →
      ∃ l ∈ (φ[i.1]'i.isLt).take j, l.Holds (assignmentOfColoring f)
  | 0, _, h => by
    rw [show chainAt φ vs i 0 = palVtx φ vs 3 from rfl, f_seed hf] at h
    exact absurd h.symm (hf _ _ adj_pal_TF)
  | j + 1, hj, h => by
    have hlt : j < (φ[i.1]'i.isLt).length := Nat.lt_of_lt_of_le (Nat.lt_succ_self j) hj
    rw [show chainAt φ vs i (j + 1) = gadVtx φ vs ⟨i, ⟨j, hlt⟩⟩ 2 by simp [chainAt, hlt]] at h
    have ht : (φ[i.1]'i.isLt).take (j + 1) =
        (φ[i.1]'i.isLt).take j ++ [(φ[i.1]'i.isLt)[j]'hlt] :=
      List.take_succ_eq_append_getElem hlt
    rcases or_forcing hf ⟨i, ⟨j, hlt⟩⟩ h with hchain | hlit
    · obtain ⟨l, hl, hlh⟩ := exists_holds_of_chainAt i j (Nat.le_of_lt hlt) hchain
      exact ⟨l, by rw [ht]; exact List.mem_append_left _ hl, hlh⟩
    · refine ⟨(φ[i.1]'i.isLt)[j]'hlt,
        by rw [ht]; exact List.mem_append_right _ (List.mem_singleton_self _), ?_⟩
      exact (f_litVtx_iff hf _).mp hlit

/-- A proper 3-colouring satisfies the formula. -/
theorem sat_assignmentOfColoring : φ.Sat (assignmentOfColoring f) := by
  intro c hc
  obtain ⟨i, hi, rfl⟩ := List.getElem_of_mem hc
  -- the clause's output is pinned to T's colour
  have hlast : f (lastOut φ vs ⟨i, hi⟩) = f (palVtx φ vs 0) :=
    fin3_pin _ _ _ _ (Ne.symm (hf _ _ adj_pal_TF)) (hf _ _ adj_pal_FB) (hf _ _ adj_pal_TB)
      (hf _ _ (adj_last_pal ⟨i, hi⟩)) (f_chainAt_ne_base hf ⟨i, hi⟩ _)
  obtain ⟨l, hl, hlh⟩ :=
    exists_holds_of_chainAt hf ⟨i, hi⟩ _ le_rfl hlast
  rw [List.take_length] at hl
  exact ⟨l, hl, hlh⟩

end Backward

/-! ## Correctness -/

/-- **The mathematical content of the reduction**: a formula is satisfiable
exactly when its colouring gadget is 3-colourable. The variable list must cover
every variable that occurs — the forward direction needs distinct occurring
variables to get distinct slots. -/
theorem satisfiable_iff_colorable [DecidableEq V] (φ : CNF V) (vs : List V)
    (hvs : ∀ c ∈ φ, ∀ l ∈ c, l.var ∈ vs) :
    φ.Satisfiable ↔ (satColorGraph φ vs).Colorable 3 := by
  constructor
  · rintro ⟨σ, hσ⟩
    exact colorable_of_sat hσ hvs
  · intro hcol
    obtain ⟨f, hf⟩ := (colorable_iff_exists_fun _ _).mp hcol
    exact ⟨assignmentOfColoring f, sat_assignmentOfColoring hf⟩

/-- The instance produced by the reduction: the colouring gadget over the
formula's own variable list, asking for three colours. -/
def satToColorInstance (x : SatInstance) : GraphThresholdInstance where
  V := ColVtx x.φ (cnfVars x.φ)
  G := satColorGraph x.φ (cnfVars x.φ)
  k := 3

/-- SAT reduces to K-COLORABILITY, landing on three colours. -/
def satToColorability : SAT ⟶ᵣ GRAPH_COLORABILITY where
  map := satToColorInstance
  correct x :=
    satisfiable_iff_colorable x.φ (cnfVars x.φ) fun _ hc _ hl ↦ mem_cnfVars hc hl

/-- **SAT reduces to 3-COLORABILITY.** The gadget always asks for three
colours, so the restriction is free. -/
def satTo3Colorability : SAT ⟶ᵣ GRAPH_3_COLORABILITY :=
  satToColorability.codRestrict _ fun _ ↦ rfl

/-- 3SAT reduces to 3-COLORABILITY, through SAT. The gadget handles clauses of
any width, so nothing here needs the width bound; this composition exists only
because 3SAT is the more usual starting point in the literature. -/
def sat3To3Colorability : SAT3 ⟶ᵣ GRAPH_3_COLORABILITY :=
  sat3ToSat.comp satTo3Colorability

end Hardness