import Init

namespace GoldbachMinimal

inductive Letter
  | a
  | b
  | c
  | d
deriving DecidableEq, Repr

abbrev Word := List Letter

-- ===================================================================
-- LES 16 TRANSITIONS (seule "vraie" règle de réécriture du système)
-- ===================================================================

def applyRuleToPair (x y : Letter) : Letter :=
  match (x, y) with
  | (Letter.a, Letter.a) => Letter.a
  | (Letter.a, Letter.b) => Letter.b
  | (Letter.a, Letter.c) => Letter.a
  | (Letter.a, Letter.d) => Letter.b
  | (Letter.b, Letter.a) => Letter.a
  | (Letter.b, Letter.b) => Letter.b
  | (Letter.b, Letter.c) => Letter.a
  | (Letter.b, Letter.d) => Letter.b
  | (Letter.c, Letter.a) => Letter.c
  | (Letter.c, Letter.b) => Letter.d
  | (Letter.c, Letter.c) => Letter.c
  | (Letter.c, Letter.d) => Letter.d
  | (Letter.d, Letter.a) => Letter.c
  | (Letter.d, Letter.b) => Letter.d
  | (Letter.d, Letter.c) => Letter.c
  | (Letter.d, Letter.d) => Letter.d

-- TRANSITIONS SUR TOUTES LES PAIRES ADJACENTES
def nextGeneration (w : Word) : Word :=
  match w with
  | [] => []
  | [_] => []
  | x :: y :: rest =>
    applyRuleToPair x y :: nextGeneration (y :: rest)

-- ===================================================================
-- PRIMALITÉ (test par division, auto-contenu)
-- ===================================================================

-- Version à "carburant" : récursivité structurelle sur fuel,
-- pas besoin de preuve de terminaison.
-- Fuel = n suffit largement : dès que d dépasse n, d*d > n.
def isPrimeAux (n : Nat) (fuel : Nat) (d : Nat) : Bool :=
  match fuel with
  | 0 => true
  | f + 1 =>
    if d * d > n then true
    else if n % d == 0 then false
    else isPrimeAux n f (d + 1)

def isPrime (n : Nat) : Bool :=
  if n < 2 then false
  else if n == 2 then true
  else if n % 2 == 0 then false
  else isPrimeAux n n 3

-- ===================================================================
-- PREMIÈRE LETTRE : a si (n - 3) premier, c sinon.
-- (SEULE dépendance externe du système : la primalité de n - 3.)
-- Pour n < 3 on prend la convention n - 3 = 0, non premier.
-- ===================================================================

def headLetterFor (n : Nat) : Letter :=
  if isPrime (n - 3) then Letter.a else Letter.c

-- ===================================================================
-- DERNIÈRE LETTRE EN CAS DE PALIER (la longueur augmente de 1) :
--   a ou b en fin du mot de n  ->  a en fin du mot de n + 2
--   c ou d en fin du mot de n  ->  d en fin du mot de n + 2
-- (Quand la longueur ne change pas, la dernière lettre est simplement
--  la dernière transition conservée : rien à injecter.)
-- ===================================================================

def lastLetterStep (l : Letter) : Letter :=
  match l with
  | Letter.a => Letter.a
  | Letter.b => Letter.a
  | Letter.c => Letter.d
  | Letter.d => Letter.d

-- ===================================================================
-- LONGUEUR CIBLE : floor((n - 2) / 4)
-- ===================================================================

def targetLength (n : Nat) : Nat :=
  (n - 2) / 4

-- ===================================================================
-- CONSTRUCTION DU MOT SUIVANT, 100% AUTO-DÉRIVÉE
--
-- Soit L la longueur du mot de n (et T = L - 1 le nombre de
-- transitions du mot de n).
--
--   Si la longueur augmente (mot de n+2 de longueur L + 1) :
--     mot de n+2 = nouvelle tête ++ toutes les transitions
--                  ++ lastLetterStep (dernière lettre du mot de n)
--
--   Sinon (même longueur L) :
--     mot de n+2 = nouvelle tête ++ toutes les transitions
--   (la dernière lettre n'est rien d'autre que la transition
--    conservée en position finale)
-- ===================================================================

def transitionWithContext (n : Nat) (w : Word) : Word :=
  let nextN := n + 2
  let newHead := headLetterFor nextN
  let transitions := nextGeneration w
  let wantedLength := targetLength nextN
  let currentLength := targetLength n
  if wantedLength == currentLength + 1 then
    -- palier : la longueur augmente de 1
    match w.getLast? with
    | none => newHead :: transitions ++ [lastLetterStep Letter.d]
    | some l => newHead :: transitions ++ [lastLetterStep l]
  else
    -- longueur stable : tête + transitions, c'est tout
    newHead :: transitions

-- ===================================================================
-- GÉNÉRATION D'UNE SUITE
-- ===================================================================

def generateWordSequence
  (n : Nat) (w : Word) (steps : Nat) :
  List (Nat × Word) :=
  match steps with
  | 0 => [(n, w)]
  | s + 1 =>
    let nextW := transitionWithContext n w
    (n, w) :: generateWordSequence (n + 2) nextW s

-- =====================================================================
-- PARTIE 2 : la colonne définie directement par la primalité
-- lettre de (p, q) : a = (premier, premier), b = (composé, premier),
-- c = (premier, composé), d = (composé, composé)
-- =====================================================================
def letterB : Bool → Bool → Letter
| true, true => Letter.a
| false, true => Letter.b
| true, false => Letter.c
| false, false => Letter.d

def letterOf (p q : Nat) : Letter := letterB (isPrime p) (isPrime q)

/-- lettres des positions i, i+1, ..., i+m-1 (le petit sommant vaut 3+2i) -/
def mkCol (n : Nat) : Nat → Nat → Word
| _, 0 => []
| i, m + 1 => letterOf (3 + 2 * i) (n - (3 + 2 * i)) :: mkCol n (i + 1) m

def colOf (n : Nat) : Word := mkCol n 0 (targetLength n)

-- =====================================================================
-- PARTIE 3 : lemmes élémentaires
-- =====================================================================
theorem mkCol_succ (n i m : Nat) :
  mkCol n i (m + 1) = letterOf (3 + 2 * i) (n - (3 + 2 * i)) :: mkCol n (i + 1) m := by
  first | rfl | simp [mkCol]

theorem mkCol_nil (n i : Nat) : mkCol n i 0 = [] := by
  first | rfl | simp [mkCol]

theorem ng_single (x : Letter) : nextGeneration [x] = [] := by
  first | rfl | simp [nextGeneration]

theorem ng_cons2 (x y : Letter) (r : Word) :
  nextGeneration (x :: y :: r) = applyRuleToPair x y :: nextGeneration (y :: r) := by
  first | rfl | simp [nextGeneration]

/-- La loi du treillis : la lettre résultante garde le statut du q du premier
opérande et le statut du p du second (bande rectangulaire). -/
theorem rule_letterB (bp bq br bs : Bool) :
  applyRuleToPair (letterB bp bq) (letterB br bs) = letterB br bq := by
  cases bp <;> cases bq <;> cases br <;> cases bs <;> rfl

theorem lastStep_letterB (bp bq : Bool) :
  lastLetterStep (letterB bp bq) = letterB bq bq := by
  cases bp <;> cases bq <;> rfl

theorem prime3 : isPrime 3 = true := by
  first | rfl | decide

theorem head_eq (k : Nat) : headLetterFor k = letterOf 3 (k - 3) := by
  unfold headLetterFor letterOf
  rw [prime3]
  cases h : isPrime (k - 3) <;> simp [letterB]

-- =====================================================================
-- PARTIE 4 : décomposition de mkCol
-- =====================================================================
theorem mkCol_snoc (n : Nat) : ∀ (m i : Nat),
  mkCol n i (m + 1)
  = mkCol n i m ++ [letterOf (3 + 2 * (i + m)) (n - (3 + 2 * (i + m)))] := by
  intro m
  induction m with
  | zero =>
    intro i
    first | rfl | simp [mkCol]
  | succ m ih =>
    intro i
    rw [mkCol_succ n i (m + 1)]
    rw [ih (i + 1)]
    rw [mkCol_succ n i m]
    have e : i + 1 + m = i + (m + 1) := by omega
    rw [e]
    all_goals (first | rfl | simp)

theorem getLast_mkCol (n : Nat) : ∀ (m i : Nat),
  (mkCol n i (m + 1)).getLast?
  = some (letterOf (3 + 2 * (i + m)) (n - (3 + 2 * (i + m)))) := by
  intro m
  induction m with
  | zero =>
    intro i
    first | rfl | simp [mkCol]
  | succ m ih =>
    intro i
    have h := ih (i + 1)
    have e : i + 1 + m = i + (m + 1) := by omega
    rw [e] at h
    rw [mkCol_succ n i (m + 1)]
    rw [← h]
    rw [mkCol_succ n (i + 1) m]
    all_goals (first | rfl | simp)

-- =====================================================================
-- PARTIE 5 : nextGeneration décale la colonne de n vers celle de n+2
-- nextGeneration (lettres i..i+m) = lettres (i+1)..(i+m) pour n+2
-- =====================================================================
theorem ng_mkCol (n : Nat) : ∀ (m i : Nat),
  nextGeneration (mkCol n i (m + 1)) = mkCol (n + 2) (i + 1) m := by
  intro m
  induction m with
  | zero =>
    intro i
    first
    | (rw [mkCol_succ n i 0, mkCol_nil n (i + 1), ng_single, mkCol_nil (n + 2) (i + 1)])
    | simp [mkCol, nextGeneration]
  | succ m ih =>
    intro i
    have h1 :
      nextGeneration (letterOf (3 + 2 * i) (n - (3 + 2 * i)) :: mkCol n (i + 1) (m + 1))
      = applyRuleToPair (letterOf (3 + 2 * i) (n - (3 + 2 * i)))
        (letterOf (3 + 2 * (i + 1)) (n - (3 + 2 * (i + 1))))
        :: nextGeneration (mkCol n (i + 1) (m + 1)) := by
      rw [mkCol_succ n (i + 1) m]
      all_goals exact ng_cons2 _ _ _
    rw [mkCol_succ n i (m + 1), h1, ih (i + 1), mkCol_succ (n + 2) (i + 1) m]
    have e : (n + 2) - (3 + 2 * (i + 1)) = n - (3 + 2 * i) := by omega
    rw [e]
    unfold letterOf
    rw [rule_letterB]

-- =====================================================================
-- PARTIE 6 : les deux cas de transitionWithContext
-- =====================================================================
theorem trans_palier (n : Nat) (w : Word) (l : Letter)
  (h : targetLength (n + 2) = targetLength n + 1)
  (hl : w.getLast? = some l) :
  transitionWithContext n w
  = headLetterFor (n + 2) :: nextGeneration w ++ [lastLetterStep l] := by
  first
  | simp [transitionWithContext, h, hl]
  | (unfold transitionWithContext; simp [h, hl])

theorem trans_stable (n : Nat) (w : Word)
  (h : ¬ targetLength (n + 2) = targetLength n + 1) :
  transitionWithContext n w = headLetterFor (n + 2) :: nextGeneration w := by
  first
  | simp [transitionWithContext, h]
  | (unfold transitionWithContext; simp [h])

-- =====================================================================
-- PARTIE 7 : un pas de la récurrence = colonne suivante
-- =====================================================================
theorem step_mkCol (n u : Nat) (hn : n % 2 = 0) (hu : (n - 2) / 4 = u + 1) :
  transitionWithContext n (mkCol n 0 (u + 1)) = mkCol (n + 2) 0 (targetLength (n + 2)) := by
  have hg := ng_mkCol n u 0
  have h1 : headLetterFor (n + 2) = letterOf (3 + 2 * 0) ((n + 2) - (3 + 2 * 0)) := by
    first | exact head_eq (n + 2) | simpa using head_eq (n + 2)
  by_cases hp : targetLength (n + 2) = targetLength n + 1
  · -- cas du palier : la colonne gagne une lettre (n = 4u+8)
    have hn4 : n = 4 * u + 8 := by
      unfold targetLength at hp
      omega
    have hl := getLast_mkCol n u 0
    have ht : targetLength (n + 2) = u + 1 + 1 := by
      unfold targetLength
      omega
    have h2 :
      lastLetterStep (letterOf (3 + 2 * (0 + u)) (n - (3 + 2 * (0 + u))))
      = letterOf (3 + 2 * (0 + 1 + u)) ((n + 2) - (3 + 2 * (0 + 1 + u))) := by
      unfold letterOf
      rw [lastStep_letterB]
      have e1 : n - (3 + 2 * (0 + u)) = 3 + 2 * (0 + 1 + u) := by omega
      have e2 : (n + 2) - (3 + 2 * (0 + 1 + u)) = 3 + 2 * (0 + 1 + u) := by
        omega
      rw [e1, e2]
    rw [trans_palier n (mkCol n 0 (u + 1)) _ hp hl, hg]
    rw [ht, mkCol_succ (n + 2) 0 (u + 1), mkCol_snoc (n + 2) u (0 + 1), h1, h2]
    -- `a :: b ++ c` se lit `(a :: b) ++ c` ; égal à `a :: (b ++ c)` par
    -- définition de `++`
    all_goals (first | rfl | simp)
  · -- cas stable : même longueur
    have ht : targetLength (n + 2) = u + 1 := by
      unfold targetLength at hp ⊢
      omega
    rw [trans_stable n (mkCol n 0 (u + 1)) hp, hg]
    rw [ht, mkCol_succ (n + 2) 0 u, h1]

theorem step_col (n : Nat) (hn : n % 2 = 0) (h6 : 6 ≤ n) :
  transitionWithContext n (colOf n) = colOf (n + 2) := by
  obtain ⟨u, hu⟩ : ∃ u, (n - 2) / 4 = u + 1 := ⟨(n - 2) / 4 - 1, by omega⟩
  have ht : targetLength n = u + 1 := hu
  unfold colOf
  rw [ht]
  exact step_mkCol n u hn hu

-- =====================================================================
-- PARTIE 8 : THÉORÈME PRINCIPAL
-- =====================================================================
theorem seq_correct : ∀ (steps n : Nat) (w : Word),
  n % 2 = 0 → 6 ≤ n → w = colOf n →
  ∀ e ∈ generateWordSequence n w steps, e.2 = colOf e.1 := by
  intro steps
  induction steps with
  | zero =>
    intro n w hn h6 hw e he
    have he' : e = (n, w) := by
      simpa [generateWordSequence] using he
    rw [he']
    exact hw
  | succ s ih =>
    intro n w hn h6 hw e he
    have he' : e = (n, w) ∨ e ∈ generateWordSequence (n + 2)
        (transitionWithContext n w) s := by
      simpa [generateWordSequence] using he
    rcases he' with he' | he'
    · rw [he']
      exact hw
    · exact ih (n + 2) (transitionWithContext n w) (by omega) (by omega)
        (by rw [hw]; exact step_col n hn h6) e he'

/-- Le mot de [a] (n = 6) engendré par la récurrence est, à chaque étape,
exactement la colonne définie par la primalité. -/
theorem recurrence_eq_colonne (steps : Nat) :
  ∀ e ∈ generateWordSequence 6 [Letter.a] steps, e.2 = colOf e.1 :=
  seq_correct steps 6 [Letter.a] (by decide) (by decide) (by first | rfl | decide)

-- Contrôles de cohérence (affichage seulement)
#eval colOf 98
#eval (generateWordSequence 6 [Letter.a] 46).getLast?.map (fun e => e.2)
#print axioms GoldbachMinimal.recurrence_eq_colonne

-- ===================================================================
-- AFFICHAGE ET DENSITÉS
-- ===================================================================

def letterToString : Letter → String
  | Letter.a => "a"
  | Letter.b => "b"
  | Letter.c => "c"
  | Letter.d => "d"

def wordToString : Word → String :=
  fun w =>
    w.foldl
    (fun acc l => acc ++ letterToString l)
    ""

-- COMPTAGE DES a
def countA (w : Word) : Nat :=
    w.foldl
    (fun acc l =>
    if l == Letter.a then acc + 1 else acc)
    0

def densityA (w : Word) : Float :=
    if w.length == 0 then
    0.0
    else
    Float.ofNat (countA w) / Float.ofNat w.length

def printComputedWordsAndDensities
  (generated : List (Nat × Word)) : IO Unit := do
  IO.println "======================================================"
  IO.println "MOTS CALCULÉS PAR LE PROGRAMME (et densité de a)"
  IO.println "======================================================"
  for (n, w) in generated do
    let s := wordToString w
    let cnt := countA w
    let total := w.length
    let dens := densityA w
    IO.println s!"n = {n} : {s}   | longueur = {total}, nombre de a = {cnt}, densité de a = {dens}"

end GoldbachMinimal

-- PROGRAMME PRINCIPAL

def main : IO Unit := do
  IO.println "=============================="
  IO.println "GÉNÉRATION AUTONOME DE 6 À 118"
  IO.println "(aucune table : uniquement la primalité de n-3,"
  IO.println " les 16 transitions et la règle de palier)"
  IO.println "=============================="
  IO.println ""

  -- mot initial : n = 6, "a"
  let initialWord : GoldbachMinimal.Word :=
    [GoldbachMinimal.Letter.a]

  -- 57 mots : 6, 8, ..., 118
  let generated :=
    GoldbachMinimal.generateWordSequence 6 initialWord 56

  GoldbachMinimal.printComputedWordsAndDensities generated