import Erdos1135.KrasikovLagarias.AdaptiveAuxiliaryTermination
import Mathlib.Order.WellFounded

/-!
# Path extraction for the adaptive KL policy

An ordinary infinite adaptive path has a first depth at every auxiliary count.
Those depths form the compressed sequence ruled out by the selected-lift
potential.
-/

namespace Erdos1135
namespace KrasikovLagarias
namespace AdaptiveEliminationPolicy

open EliminationPolicy
open EliminationResidue

/-- An ordinary infinite path through the adaptive elimination relation. -/
structure InfinitePath {k : Nat} {hk : 2 ≤ k}
    (potential : ForcedPotential k hk) where
  state : Nat → State k
  rootIndex : PrincipalIndex k
  state_zero : state 0 = State.root rootIndex
  legal : ∀ n, LegalChild potential (state (n + 1)) (state n)

namespace LegalChild

theorem prior_eq {k : Nat} {hk : 2 ≤ k} {potential : ForcedPotential k hk}
    {child parent : State k} (edge : LegalChild potential child parent) :
    child.prior = parent.current :: parent.prior := by
  cases edge <;> rfl

theorem parent_nonnegative {k : Nat} {hk : 2 ≤ k}
    {potential : ForcedPotential k hk} {child parent : State k}
    (edge : LegalChild potential child parent) :
    0 ≤ parent.current.value := by
  cases edge with
  | principal _ hadvanced => exact hadvanced
  | auxiliary _ hadvanced _ => exact hadvanced

theorem alphaCount_bounds {k : Nat} {hk : 2 ≤ k}
    {potential : ForcedPotential k hk} {child parent : State k}
    (edge : LegalChild potential child parent) :
    parent.current.shift.alphaCount ≤ child.current.shift.alphaCount ∧
      child.current.shift.alphaCount ≤ parent.current.shift.alphaCount + 1 := by
  cases edge with
  | principal parent hadvanced =>
      simp [State.descend, Label.principalChild]
  | auxiliary parent hadvanced hallowed =>
      cases hallowed <;>
        simp [State.descend, Label.d1Child, Label.d3Child]

theorem twoCount_add_one_le {k : Nat} {hk : 2 ≤ k}
    {potential : ForcedPotential k hk} {child parent : State k}
    (edge : LegalChild potential child parent) :
    parent.current.shift.twoCount + 1 ≤ child.current.shift.twoCount := by
  cases edge with
  | principal parent hadvanced =>
      simp [State.descend, Label.principalChild]
  | auxiliary parent hadvanced hallowed =>
      cases hallowed <;>
        simp [State.descend, Label.d1Child, Label.d3Child]

theorem current_mem_prior {k : Nat} {hk : 2 ≤ k}
    {potential : ForcedPotential k hk} {child parent : State k}
    (edge : LegalChild potential child parent) :
    parent.current ∈ child.prior := by
  rw [edge.prior_eq]
  simp

theorem eq_principal_of_alphaCount_eq {k : Nat} {hk : 2 ≤ k}
    {potential : ForcedPotential k hk} {child parent : State k}
    (edge : LegalChild potential child parent)
    (hcount : child.current.shift.alphaCount =
      parent.current.shift.alphaCount) :
    ∃ hadvanced : 0 ≤ parent.current.value,
      child = parent.descend
        (parent.current.principalChild (Nat.le_trans (by omega) hk)) := by
  cases edge with
  | principal parent hadvanced => exact ⟨hadvanced, rfl⟩
  | auxiliary parent hadvanced hallowed =>
      cases hallowed <;>
        simp [State.descend, Label.d1Child, Label.d3Child] at hcount

theorem exists_allowed_of_alphaCount_eq_add_one {k : Nat} {hk : 2 ≤ k}
    {potential : ForcedPotential k hk} {child parent : State k}
    (edge : LegalChild potential child parent)
    (hcount : child.current.shift.alphaCount =
      parent.current.shift.alphaCount + 1) :
    ∃ (hadvanced : 0 ≤ parent.current.value) (label : Label k),
      AllowedAuxiliary potential parent label ∧
        child = parent.descend label := by
  cases edge with
  | principal parent hadvanced =>
      simp [State.descend, Label.principalChild] at hcount
  | auxiliary parent hadvanced hallowed =>
      exact ⟨hadvanced, _, hallowed, rfl⟩

end LegalChild

namespace InfinitePath

theorem nonnegative {k : Nat} {hk : 2 ≤ k}
    {potential : ForcedPotential k hk} (path : InfinitePath potential) (n : Nat) :
    0 ≤ (path.state n).current.value :=
  (path.legal n).parent_nonnegative

theorem alphaCount_mono {k : Nat} {hk : 2 ≤ k}
    {potential : ForcedPotential k hk} (path : InfinitePath potential) :
    Monotone (fun n => (path.state n).current.shift.alphaCount) :=
  monotone_nat_of_le_succ fun n => (path.legal n).alphaCount_bounds.1

theorem alphaCount_succ_le {k : Nat} {hk : 2 ≤ k}
    {potential : ForcedPotential k hk} (path : InfinitePath potential) (n : Nat) :
    (path.state (n + 1)).current.shift.alphaCount ≤
      (path.state n).current.shift.alphaCount + 1 :=
  (path.legal n).alphaCount_bounds.2

theorem depth_le_twoCount {k : Nat} {hk : 2 ≤ k}
    {potential : ForcedPotential k hk} (path : InfinitePath potential) (n : Nat) :
    n ≤ (path.state n).current.shift.twoCount := by
  induction n with
  | zero => simp
  | succ n ih =>
      have hstep := (path.legal n).twoCount_add_one_le
      omega

theorem current_mem_prior_of_lt {k : Nat} {hk : 2 ≤ k}
    {potential : ForcedPotential k hk} (path : InfinitePath potential)
    {i j : Nat} (hij : i < j) :
    (path.state i).current ∈ (path.state j).prior := by
  induction j with
  | zero => omega
  | succ j ih =>
      rw [(path.legal j).prior_eq]
      by_cases h : i = j
      · subst i
        simp
      · exact List.mem_cons_of_mem _ (ih (by omega))

theorem depth_le_two_mul_alphaCount {k : Nat} {hk : 2 ≤ k}
    {potential : ForcedPotential k hk} (path : InfinitePath potential) (n : Nat) :
    n ≤ 2 * (path.state n).current.shift.alphaCount := by
  have hnonnegative := path.nonnegative n
  rw [Label.value, EliminationShift.value] at hnonnegative
  have htwoReal : (n : Real) ≤
      (path.state n).current.shift.twoCount := by
    exact_mod_cast path.depth_le_twoCount n
  have halpha : alpha ≤ 2 := by
    linarith [three_mul_alpha_lt_five]
  have hmul :
      ((path.state n).current.shift.alphaCount : Real) * alpha ≤
        ((path.state n).current.shift.alphaCount : Real) * 2 :=
    mul_le_mul_of_nonneg_left halpha (Nat.cast_nonneg _)
  have hreal : (n : Real) ≤
      2 * (path.state n).current.shift.alphaCount := by
    push_cast
    linarith
  exact_mod_cast hreal

theorem exists_depth_alphaCount_ge {k : Nat} {hk : 2 ≤ k}
    {potential : ForcedPotential k hk} (path : InfinitePath potential)
    (count : Nat) :
    ∃ depth, count ≤ (path.state depth).current.shift.alphaCount := by
  refine ⟨2 * count, ?_⟩
  have := path.depth_le_two_mul_alphaCount (2 * count)
  omega

noncomputable def firstDepth {k : Nat} {hk : 2 ≤ k}
    {potential : ForcedPotential k hk} (path : InfinitePath potential)
    (count : Nat) : Nat :=
  Nat.find (path.exists_depth_alphaCount_ge count)

theorem firstDepth_spec {k : Nat} {hk : 2 ≤ k}
    {potential : ForcedPotential k hk} (path : InfinitePath potential)
    (count : Nat) :
    count ≤ (path.state (path.firstDepth count)).current.shift.alphaCount :=
  Nat.find_spec (path.exists_depth_alphaCount_ge count)

theorem alphaCount_lt_of_lt_firstDepth {k : Nat} {hk : 2 ≤ k}
    {potential : ForcedPotential k hk} (path : InfinitePath potential)
    {depth count : Nat} (hdepth : depth < path.firstDepth count) :
    (path.state depth).current.shift.alphaCount < count := by
  have hnot := Nat.find_min (path.exists_depth_alphaCount_ge count) hdepth
  omega

theorem firstDepth_alphaCount {k : Nat} {hk : 2 ≤ k}
    {potential : ForcedPotential k hk} (path : InfinitePath potential)
    (count : Nat) :
    (path.state (path.firstDepth count)).current.shift.alphaCount = count := by
  have hlower := path.firstDepth_spec count
  cases hdepth : path.firstDepth count with
  | zero =>
      rw [hdepth] at hlower
      rw [path.state_zero] at hlower ⊢
      simp [State.root, Label.root] at hlower ⊢
      omega
  | succ depth =>
      have hprevious :
          (path.state depth).current.shift.alphaCount < count :=
        path.alphaCount_lt_of_lt_firstDepth (by omega)
      have hstep := path.alphaCount_succ_le depth
      rw [hdepth] at hlower
      omega

@[simp]
theorem firstDepth_zero {k : Nat} {hk : 2 ≤ k}
    {potential : ForcedPotential k hk} (path : InfinitePath potential) :
    path.firstDepth 0 = 0 := by
  exact Nat.eq_zero_of_le_zero
    (Nat.find_min' (path.exists_depth_alphaCount_ge 0) (Nat.zero_le _))

theorem firstDepth_strictMono {k : Nat} {hk : 2 ≤ k}
    {potential : ForcedPotential k hk} (path : InfinitePath potential) :
    StrictMono path.firstDepth := by
  intro i j hij
  by_contra hnot
  have hdepth : path.firstDepth j ≤ path.firstDepth i := by omega
  have hcount := path.alphaCount_mono hdepth
  change (path.state (path.firstDepth j)).current.shift.alphaCount ≤
    (path.state (path.firstDepth i)).current.shift.alphaCount at hcount
  rw [path.firstDepth_alphaCount, path.firstDepth_alphaCount] at hcount
  omega

theorem alphaCount_eq_of_firstDepth_le_lt {k : Nat} {hk : 2 ≤ k}
    {potential : ForcedPotential k hk} (path : InfinitePath potential)
    {count depth : Nat} (hlower : path.firstDepth count ≤ depth)
    (hupper : depth < path.firstDepth (count + 1)) :
    (path.state depth).current.shift.alphaCount = count := by
  have hge := path.alphaCount_mono hlower
  change (path.state (path.firstDepth count)).current.shift.alphaCount ≤
    (path.state depth).current.shift.alphaCount at hge
  rw [path.firstDepth_alphaCount] at hge
  have hlt := path.alphaCount_lt_of_lt_firstDepth hupper
  omega

theorem principalSteps_of_constant_alphaCount {k : Nat} {hk : 2 ≤ k}
    {potential : ForcedPotential k hk} (path : InfinitePath potential)
    {start finish count : Nat} (hle : start ≤ finish)
    (hcount : ∀ depth, start ≤ depth → depth ≤ finish →
      (path.state depth).current.shift.alphaCount = count) :
    PrincipalSteps potential (path.state start) (path.state finish) := by
  induction finish generalizing start with
  | zero =>
      have : start = 0 := by omega
      subst start
      exact PrincipalSteps.refl _
  | succ finish ih =>
      by_cases heq : start = finish + 1
      · subst start
        exact PrincipalSteps.refl _
      · have hle' : start ≤ finish := by omega
        have previous := ih hle' (fun depth hstart hfinish =>
          hcount depth hstart (by omega))
        have hedge := path.legal finish
        have halpha :
            (path.state (finish + 1)).current.shift.alphaCount =
              (path.state finish).current.shift.alphaCount := by
          rw [hcount (finish + 1) (by omega) (by omega),
            hcount finish hle' (by omega)]
        obtain ⟨hadvanced, hstate⟩ :=
          hedge.eq_principal_of_alphaCount_eq halpha
        rw [hstate]
        exact PrincipalSteps.step previous hadvanced

noncomputable def auxiliaryStep_firstDepth {k : Nat} {hk : 2 ≤ k}
    {potential : ForcedPotential k hk} (path : InfinitePath potential)
    (count : Nat) :
    AuxiliaryStep potential (path.state (path.firstDepth count))
      (path.state (path.firstDepth (count + 1))) := by
  let start := path.firstDepth count
  let finish := path.firstDepth (count + 1)
  let branchDepth := finish - 1
  have hstartFinish : start < finish := path.firstDepth_strictMono (by omega)
  have hbranchSucc : branchDepth + 1 = finish := by
    dsimp [branchDepth]
    omega
  have hstartBranch : start ≤ branchDepth := by
    dsimp [branchDepth]
    omega
  have hbranchFinish : branchDepth < finish := by omega
  have hprincipals :
      PrincipalSteps potential (path.state start) (path.state branchDepth) :=
    path.principalSteps_of_constant_alphaCount hstartBranch
      (fun depth hstart hbranch =>
        path.alphaCount_eq_of_firstDepth_le_lt hstart
          (hbranch.trans_lt hbranchFinish))
  have hbranchCount :
      (path.state branchDepth).current.shift.alphaCount = count :=
    path.alphaCount_eq_of_firstDepth_le_lt hstartBranch hbranchFinish
  have hfinishCount :
      (path.state finish).current.shift.alphaCount = count + 1 := by
    dsimp [finish]
    exact path.firstDepth_alphaCount (count + 1)
  have hedge : LegalChild potential (path.state finish)
      (path.state branchDepth) := by
    rw [← hbranchSucc]
    exact path.legal branchDepth
  have hexists := hedge.exists_allowed_of_alphaCount_eq_add_one (by omega)
  let hadvanced := Classical.choose hexists
  have hexistsLabel := Classical.choose_spec hexists
  let label := Classical.choose hexistsLabel
  have hlabel := Classical.choose_spec hexistsLabel
  exact
    { branch := path.state branchDepth
      principals := hprincipals
      advanced := hadvanced
      child := label
      allowed := hlabel.1
      finish_eq := hlabel.2 }

noncomputable def toAuxiliarySequence {k : Nat} {hk : 2 ≤ k}
    {potential : ForcedPotential k hk} (path : InfinitePath potential) :
    AuxiliarySequence potential where
  state count := path.state (path.firstDepth count)
  rootIndex := path.rootIndex
  state_zero := by rw [path.firstDepth_zero, path.state_zero]
  alphaCount := path.firstDepth_alphaCount
  nonnegative count := path.nonnegative (path.firstDepth count)
  earlier_mem_prior i j hij :=
    path.current_mem_prior_of_lt (path.firstDepth_strictMono hij)
  step := path.auxiliaryStep_firstDepth

theorem false_of_forcedPotential {k : Nat} {hk : 2 ≤ k}
    {potential : ForcedPotential k hk} (path : InfinitePath potential) : False :=
  path.toAuxiliarySequence.false_of_forcedPotential

end InfinitePath

/-- Every adaptive root is accessible in the presence of its potential. -/
theorem root_accessible {k : Nat} {hk : 2 ≤ k}
    (potential : ForcedPotential k hk) (rootIndex : PrincipalIndex k) :
    Acc (LegalChild potential) (State.root rootIndex) := by
  by_contra hnot
  obtain ⟨state, hzero, hlegal⟩ :=
    (not_acc_iff_exists_descending_chain.mp hnot)
  exact (InfinitePath.mk state rootIndex hzero hlegal).false_of_forcedPotential

end AdaptiveEliminationPolicy
end KrasikovLagarias
end Erdos1135