import Init

def OA_statement : Prop := ((fun (run : List (Bool × Bool) → Nat → Option Nat → Sum Nat (Nat × Option Nat)) =>
(fun (active : List (Bool × Bool) → List (Bool × Bool)) =>
(fun (lastClear : List (Bool × Bool) → Option Nat) =>
(fun (firstStart : List (Bool × Bool) → Option Nat → Option Nat) =>
(fun (summarize : List (Bool × Bool) → Nat → Option Nat → Sum Nat (Nat × Option Nat)) =>
∀ actions position saved, run actions position saved = summarize actions position saved
) (fun actions position saved =>
(fun chunk =>
(fun clear =>
(fun start =>
(fun finalSaved =>
if chunk.length < actions.length then .inl (finalSaved.getD (position + chunk.length))
else .inr (position + chunk.length, finalSaved))
((if clear.isSome then none else saved).or (start.map (position + ·))))
(firstStart chunk clear)) (lastClear chunk)) (active actions))
) (fun chunk clear =>
(((chunk.zipIdx).drop (clear.map Nat.succ |>.getD 0)).find? (fun pair => !pair.1.1 && pair.1.2)).map Prod.snd)
) (fun chunk =>
((chunk.zipIdx.reverse).find? (fun pair => pair.1.1 && !pair.1.2)).map Prod.snd)
) (fun actions => actions.takeWhile (fun action => !action.1 || !action.2))
) (fun actions =>
List.rec (fun position saved => .inr (position, saved))
(fun action _ rest position saved =>
if action.1 then
  if action.2 then .inl (saved.getD position) else rest (position + 1) none
else
  if action.2 then rest (position + 1) (some (saved.getD position))
  else rest (position + 1) saved) actions))

theorem OA_target : OA_statement :=
  by
    dsimp only [OA_statement]
    let run : List (Bool × Bool) → Nat → Option Nat → Sum Nat (Nat × Option Nat) := fun actions =>
      List.rec (fun position saved => .inr (position, saved))
        (fun action _ rest position saved =>
          if action.1 then
            if action.2 then .inl (saved.getD position) else rest (position + 1) none
          else
            if action.2 then rest (position + 1) (some (saved.getD position))
            else rest (position + 1) saved) actions
    let active : List (Bool × Bool) → List (Bool × Bool) :=
      fun actions => actions.takeWhile (fun action => !action.1 || !action.2)
    let lastClear : List (Bool × Bool) → Option Nat :=
      fun chunk => ((chunk.zipIdx.reverse).find? (fun pair => pair.1.1 && !pair.1.2)).map Prod.snd
    let firstStart : List (Bool × Bool) → Option Nat → Option Nat :=
      fun chunk clear => (((chunk.zipIdx).drop (clear.map Nat.succ |>.getD 0)).find?
        (fun pair => !pair.1.1 && pair.1.2)).map Prod.snd
    let summarize : List (Bool × Bool) → Nat → Option Nat → Sum Nat (Nat × Option Nat) :=
      fun actions position saved =>
        let chunk := active actions
        let clear := lastClear chunk
        let start := firstStart chunk clear
        let finalSaved := (if clear.isSome then none else saved).or (start.map (position + ·))
        if chunk.length < actions.length then .inl (finalSaved.getD (position + chunk.length))
        else .inr (position + chunk.length, finalSaved)
    change ∀ actions position saved, run actions position saved = summarize actions position saved
    have lc (a : Bool × Bool) (xs : List (Bool × Bool)) :
        lastClear (a :: xs) = ((lastClear xs).map (· + 1)).or
          (if a.1 && !a.2 then some 0 else none) := by
      simp only [lastClear, List.zipIdx_cons', List.reverse_cons,
        ← List.map_reverse, List.find?_append]
      rcases a with ⟨a, b⟩
      cases a <;> cases b <;>
        simp [-List.map_reverse, List.find?_map, Option.map_map, Function.comp_def]
      cases xs.zipIdx.reverse.find? (fun pair => pair.1.1 && !pair.1.2) <;> simp
    have fs_none (a : Bool × Bool) (xs : List (Bool × Bool)) :
        firstStart (a :: xs) none =
          if !a.1 && a.2 then some 0 else (firstStart xs none).map (· + 1) := by
      rcases a with ⟨a, b⟩
      cases a <;> cases b <;> simp [firstStart, List.zipIdx_succ,
        List.find?_map, Option.map_map, Function.comp_def]
    have fs_zero (a : Bool × Bool) (xs : List (Bool × Bool)) :
        firstStart (a :: xs) (some 0) = (firstStart xs none).map (· + 1) := by
      simp [firstStart, List.zipIdx_succ, List.find?_map, Option.map_map, Function.comp_def]
    have fs_succ (a : Bool × Bool) (xs : List (Bool × Bool)) (c : Nat) :
        firstStart (a :: xs) (some (c + 1)) =
          (firstStart xs (some c)).map (· + 1) := by
      simp [firstStart, List.zipIdx_succ, ← List.map_drop,
        List.find?_map, Option.map_map, Function.comp_def]
    have ac_skip (xs : List (Bool × Bool)) :
        active ((false, false) :: xs) = (false, false) :: active xs := by simp [active]
    have ac_start (xs : List (Bool × Bool)) :
        active ((false, true) :: xs) = (false, true) :: active xs := by simp [active]
    have ac_clear (xs : List (Bool × Bool)) :
        active ((true, false) :: xs) = (true, false) :: active xs := by simp [active]
    have sn (p : Nat) (s : Option Nat) : summarize [] p s = .inr (p, s) := by
      cases s <;> rfl
    have st (xs : List (Bool × Bool)) (p : Nat) (s : Option Nat) :
        summarize ((true, true) :: xs) p s = .inl (s.getD p) := by
      cases s <;> simp [summarize, active, lastClear, firstStart]
    have sk (xs : List (Bool × Bool)) (p : Nat) (s : Option Nat) :
        summarize ((false, false) :: xs) p s = summarize xs (p + 1) s := by
      simp only [summarize, ac_skip, List.length_cons, Nat.succ_lt_succ_iff, lc]
      cases lastClear (active xs) with
      | none => simp [fs_none, Option.map_map, Function.comp_def, Nat.add_assoc, Nat.add_comm, Nat.add_left_comm]
      | some c => simp [fs_succ, Option.map_map, Function.comp_def, Nat.add_assoc, Nat.add_comm, Nat.add_left_comm]
    have ss (xs : List (Bool × Bool)) (p : Nat) (s : Option Nat) :
        summarize ((false, true) :: xs) p s = summarize xs (p + 1) (some (s.getD p)) := by
      simp only [summarize, ac_start, List.length_cons, Nat.succ_lt_succ_iff, lc]
      cases lastClear (active xs) with
      | none => cases s <;> simp [fs_none, Nat.add_assoc, Nat.add_comm, Nat.add_left_comm]
      | some c => simp [fs_succ, Option.map_map, Function.comp_def, Nat.add_assoc, Nat.add_comm, Nat.add_left_comm]
    have sc (xs : List (Bool × Bool)) (p : Nat) (s : Option Nat) :
        summarize ((true, false) :: xs) p s = summarize xs (p + 1) none := by
      simp only [summarize, ac_clear, List.length_cons, Nat.succ_lt_succ_iff, lc]
      cases lastClear (active xs) with
      | none => simp [fs_zero, Option.map_map, Function.comp_def, Nat.add_assoc, Nat.add_comm, Nat.add_left_comm]
      | some c => simp [fs_succ, Option.map_map, Function.comp_def, Nat.add_assoc, Nat.add_comm, Nat.add_left_comm]
    intro actions position saved
    induction actions generalizing position saved with
    | nil => simp only [run, sn]
    | cons a xs ih =>
      rcases a with ⟨a, b⟩
      cases a <;> cases b <;> simp only [run, Bool.false_eq_true, ↓reduceIte, sk, ss, sc, st, ih]

#print axioms OA_target
#check OA_target