import Init

def OA_statement : Prop := ((fun (born : List Nat → Nat → List Nat) =>
(fun (position : List Nat → Nat → Nat) =>
(fun (insertAt : List Nat → Nat → Nat → List Nat) =>
(fun (replay : List Nat → Nat → List Nat) =>
∀ (order : List Nat) (n : Nat), order.Nodup →
(∀ time < n, time ∈ order) → (∀ time ∈ order, time < n) →
(∀ cut ≤ n, replay order cut = born order cut) ∧
(∀ time < n, position order time ≤ (replay order time).length) ∧
replay order n = order
) (fun order cut => Nat.rec []
  (fun time state => insertAt state (position order time) time) cut)
) (fun state pos value => state.take pos ++ value :: state.drop pos)
) (fun order time => (born (order.takeWhile (fun x => x != time)) time).length)
) (fun order cut => order.filter (fun time => time < cut)))

theorem OA_target : OA_statement :=
  by
    let born := fun (order : List Nat) (cut : Nat) =>
      order.filter (fun time => time < cut)
    let position := fun (order : List Nat) (time : Nat) =>
      (born (order.takeWhile (fun x => x != time)) time).length
    let insertAt := fun (state : List Nat) (pos value : Nat) =>
      state.take pos ++ value :: state.drop pos
    let replay : List Nat → Nat → List Nat := fun (order : List Nat) (cut : Nat) =>
      Nat.rec [] (fun time state => insertAt state (position order time) time) cut
    change ∀ (order : List Nat) (n : Nat), order.Nodup →
      (∀ time < n, time ∈ order) → (∀ time ∈ order, time < n) →
      (∀ cut ≤ n, replay order cut = born order cut) ∧
      (∀ time < n, position order time ≤ (replay order time).length) ∧
      replay order n = order
    have one_birth (order : List Nat) (time : Nat)
        (hu : order.Nodup) (hm : time ∈ order) :
        insertAt (born order time) (position order time) time = born order (time + 1) := by
      obtain ⟨pre, post, he⟩ := List.append_of_mem hm
      subst order
      have hu' := List.nodup_append.mp hu
      have hpre : time ∉ pre := by
        intro h
        exact hu'.2.2 time h time List.mem_cons_self rfl
      have hpost : time ∉ post := (List.nodup_cons.mp hu'.2.1).1
      have hf (xs : List Nat) (hx : time ∉ xs) : born xs (time + 1) = born xs time := by
        unfold born
        apply List.filter_congr
        intro x hmem
        have hn : x ≠ time := by intro heq; subst x; exact hx hmem
        simp only [decide_eq_decide]
        omega
      have htake : (pre ++ time :: post).takeWhile (fun x => x != time) = pre := by
        rw [List.takeWhile_append_of_pos (by
          intro x hx
          simp only [bne_iff_ne]
          intro heq
          subst x
          exact hpre hx)]
        simp
      have hstate : born (pre ++ time :: post) time = born pre time ++ born post time := by
        simp [born]
      have hnext : born (pre ++ time :: post) (time + 1) =
          born pre time ++ time :: born post time := by
        simp only [born, List.filter_append, List.filter_cons]
        simp only [show time < time + 1 by omega, decide_true, ↓reduceIte]
        change born pre (time + 1) ++ time :: born post (time + 1) = _
        rw [hf pre hpre, hf post hpost]
      rw [hstate, hnext]
      unfold insertAt position
      rw [htake]
      simp
  
    have heq (order : List Nat) (hu : order.Nodup) (cut : Nat)
        (hc : ∀ time < cut, time ∈ order) : replay order cut = born order cut := by
      induction cut with
      | zero =>
        change [] = born order 0
        simp [born]
      | succ cut ih =>
        change insertAt (replay order cut) (position order cut) cut = _
        rw [ih (by intro time ht; exact hc time (by omega))]
        exact one_birth order cut hu (hc cut (by omega))
    have hpos (order : List Nat) (time : Nat) :
        position order time ≤ (born order time).length := by
      have hs := congrArg (fun xs : List Nat => (born xs time).length)
        (List.takeWhile_append_dropWhile (p := fun x : Nat => x != time) (l := order))
      simp only [born, List.filter_append, List.length_append] at hs
      dsimp [position, born]
      omega
    intro order n hu hc hb
    have hall (cut : Nat) (hcut : cut ≤ n) : replay order cut = born order cut :=
      heq order hu cut (by intro time ht; exact hc time (by omega))
    refine ⟨hall, ?_, ?_⟩
    · intro time ht
      rw [hall time (by omega)]
      exact hpos order time
    · rw [hall n (Nat.le_refl n)]
      dsimp [born]
      apply List.filter_eq_self.mpr
      intro time ht
      simpa using hb time ht

#print axioms OA_target
#check OA_target