Twitter.lean

import Mathlib.Data.Nat.Basic
import Mathlib.Data.Nat.Parity
import MIL.Common
open Nat List


@[simp]
def fact : Nat -> Nat
| 0 => 1
| n + 1 => (n + 1) * fact n

@[simp]
def sigma : Nat -> List Nat
| 0 => []
| n + 1 => sigma n ++ [n + 1]

@[simp]
def add : List Nat -> Nat
| [] => 0
| h :: t => h + add t

theorem add_dist : ∀ xs ys : List Nat, add (xs ++ ys) = add xs + add ys := by
  intro xs
  induction xs with
  | nil => simp
  | cons x xs ihx =>
    intro ys; simp
    rw [ihx]; ring


theorem twitter_puzzle :
  ∀ n : Nat, add (map (fun w => w * fact w) (sigma n)) + 1 = fact (n + 1) := by
  intro n
  induction n with
  | zero =>
    simp
  | succ n ihn =>
    simp; rw [add_dist]
    have ha : add (map (fun w ↦ w * fact w) (sigma n)) +
      add [(n + 1) * ((n + 1) * fact n)] + 1 =
      add (map (fun w ↦ w * fact w) (sigma n)) + 1 +
      add [(n + 1) * ((n + 1) * fact n)] := by ring
    rw [ha, ihn]; simp; clear ha
    have ha : (n + 2) = succ n + 1 := by rw [Nat.add_succ]
    rw [<-ha]; clear ha
    have ha : (n + 1) * fact n + (n + 1) * ((n + 1) * fact n) =
      ((n + 1) + (n + 1) * (n + 1)) * fact n := by ring
    rw [ha]; clear ha; ring