import MaxLeafFormal.RSCCheck
import MaxLeafFormal.RSCStages

/-!
# Checkers for the exact `c_Z` and the stage sets

`checkCZ` and `checkStages` are computable Boolean checkers over the exported data of `RSCData`; the
generated shard theorems of `MaxLeafFormal.RSCData.Extras.Shards` decide them by `native_decide`.
Their soundness theorems, and the bridge `isStages_of_facts` from the stage facts of `RSCData`
to `RSCCert.IsStages`, are kernel-checked and use no data.
-/

namespace MaxLeafFormal.RSCCheck

/-! ### Computable checkers -/

/-- `fin` attains its minimum `a / b` over `Z` at index `k` (all cross-multiplied):
`b · finN(N_k) = a · 2000 D_k` and `a · 2000 D ≤ b · finN(N)` for every point `(N, D)`. -/
def checkCZ (Z : List IPt) (k a b : ℕ) : Bool :=
  okPts Z && decide (k < Z.length) && decide (0 < b) &&
    decide (b * finN (Z.getD k zeroI).1 = a * (2000 * (Z.getD k zeroI).2)) &&
    Z.all (fun P => decide (a * (2000 * P.2) ≤ b * finN P.1))

/-- Each stage list has the given length, well-formed points, and strictly increasing values. -/
def checkStages : List (List IPt) → List ℕ → Bool
  | [], [] => true
  | Y :: Ys, n :: ns => (Y.length == n) && okPts Y && sortedZ Y && checkStages Ys ns
  | _, _ => false

open MaxLeafFormal.RSCCert MaxLeafFormal.RSCParams
open MaxLeafFormal.Antinorm8 (Kpos PZ)
open Pointwise

noncomputable section

/-! ### Soundness -/

theorem fin_val {P : IPt} (h : okPt P = true) :
    Antinorm8.fin p₂ (val P) = (finN P.1 : ℝ) / (2000 * (P.2 : ℝ)) := by
  have hDpos : 0 < P.2 := okPt_pos h
  have hL : P.1.length = 8 := by
    simp only [okPt, Bool.and_eq_true, beq_iff_eq] at h
    exact h.1
  obtain ⟨L, D⟩ := P
  obtain ⟨N0, N1, N2, N3, N4, N5, N6, N7, rfl⟩ := list8_of_length hL
  have hD0 : (D : ℝ) ≠ 0 := by exact_mod_cast hDpos.ne'
  simp only [finN]
  push_cast
  rw [p2v]
  simp [Antinorm8.fin, val]
  field_simp
  ring

theorem nodup_val_of_sorted {Y : List IPt} (hok : okPts Y = true) (hs : sortedZ Y = true) :
    (Y.map val).Nodup := by
  have hokL : ∀ P ∈ Y, okPt P = true := by
    simp only [okPts, List.all_eq_true] at hok
    exact hok
  have hpos : ∀ P ∈ Y, 0 < P.2 := fun P hP => okPt_pos (hokL P hP)
  have hnV := nodup_vlist hpos hs
  have hY : Y.Nodup := List.Nodup.of_map vlist hnV
  refine List.Nodup.map_on ?_ hY
  intro x hx y hy hxy
  exact List.inj_on_of_nodup_map hnV hx hy (vlist_eq_of_val (hokL x hx) (hokL y hy) hxy)

theorem card_finsetOf_of_sorted {Y : List IPt} (hok : okPts Y = true) (hs : sortedZ Y = true) :
    (finsetOf Y).card = Y.length := by
  unfold finsetOf
  rw [List.toFinset_card_of_nodup (nodup_val_of_sorted hok hs), List.length_map]

theorem checkCZ_sound {Z : List IPt} {k a b : ℕ} (h : checkCZ Z k a b = true)
    (hne : (finsetOf Z).Nonempty) : Antinorm8.cZ p₂ (finsetOf Z) hne = (a : ℝ) / (b : ℝ) := by
  simp only [checkCZ, Bool.and_eq_true, decide_eq_true_eq, List.all_eq_true] at h
  obtain ⟨⟨⟨⟨hok, hk⟩, hb⟩, heq⟩, hall⟩ := h
  have hokL : ∀ P ∈ Z, okPt P = true := by
    simp only [okPts, List.all_eq_true] at hok
    exact hok
  have hbr : (0 : ℝ) < b := by exact_mod_cast hb
  have hmemk : Z.getD k zeroI ∈ Z := by
    rw [List.getD_eq_getElem _ _ hk]
    exact List.getElem_mem hk
  unfold Antinorm8.cZ
  apply le_antisymm
  · have h1 : val (Z.getD k zeroI) ∈ finsetOf Z := fam_mem_finsetOf.2 ⟨_, hmemk, rfl⟩
    refine (Finset.inf'_le _ h1).trans (le_of_eq ?_)
    have hok' := hokL _ hmemk
    have hD : (0 : ℝ) < ((Z.getD k zeroI).2 : ℝ) := by exact_mod_cast okPt_pos hok'
    rw [fin_val hok', div_eq_div_iff (mul_pos (by norm_num) hD).ne' hbr.ne']
    have := congrArg (fun n : ℕ => (n : ℝ)) heq
    push_cast at this
    linarith
  · apply Finset.le_inf'
    intro z hz
    obtain ⟨P, hP, rfl⟩ := fam_mem_finsetOf.1 hz
    have hok' := hokL P hP
    have hD : (0 : ℝ) < (P.2 : ℝ) := by exact_mod_cast okPt_pos hok'
    rw [fin_val hok', div_le_div_iff₀ hbr (mul_pos (by norm_num) hD)]
    have := hall P hP
    have h' : ((a * (2000 * P.2) : ℕ) : ℝ) ≤ ((b * finN P.1 : ℕ) : ℝ) := by exact_mod_cast this
    push_cast at h'
    linarith

theorem checkStages_sound {stages : List (List IPt)} {ns : List ℕ}
    (h : checkStages stages ns = true) :
    stages.all okPts = true ∧ (stages.map finsetOf).map Finset.card = ns := by
  induction stages generalizing ns with
  | nil =>
    cases ns with
    | nil => simp
    | cons n ns => simp [checkStages] at h
  | cons Y Ys ih =>
    cases ns with
    | nil => simp [checkStages] at h
    | cons n ns =>
      simp only [checkStages, Bool.and_eq_true, beq_iff_eq] at h
      obtain ⟨⟨⟨hl, hok⟩, hs⟩, hrest⟩ := h
      obtain ⟨ih1, ih2⟩ := ih hrest
      refine ⟨by simp [hok, ih1], ?_⟩
      simp only [List.map_cons, ih2, card_finsetOf_of_sorted hok hs, hl]

theorem finsetOf_nil : finsetOf [] = ∅ := by
  simp [finsetOf]

theorem finsetOf_initI (ν c₂ c₃ : ℕ) :
    finsetOf [initI ν c₂ c₃] = {(ê ^ (1 + ν))⁻¹ • init p₂ p₃ y ϖ₂ ϖ₃ ν c₂ c₃} := by
  simp [finsetOf, val_initI]

theorem getD_map_finsetOf : ∀ (l : List (List IPt)) (n : ℕ),
    (l.map finsetOf).getD n ∅ = finsetOf (l.getD n [])
  | [], n => by simp [finsetOf_nil]
  | a :: l, 0 => by simp
  | a :: l, n + 1 => by simpa using getD_map_finsetOf l n

theorem isStages_of_facts {Z : List IPt} (hZ : checkZ Z = true) (key : ℕ × ℕ × ℕ)
    {stages : List (List IPt)} (hst : stages.all okPts = true)
    (hfacts : List.Forall (fun i => StageFact Z (([initI key.1 key.2.1 key.2.2] :: stages).getD i [])
      ((stages ++ [Z]).getD i [])) (List.range (stages.length + 1))) :
    RSCCert.IsStages p₂ p₃ (finsetOf Z)
      ((ê ^ (1 + key.1))⁻¹ • init p₂ p₃ y ϖ₂ ϖ₃ key.1 key.2.1 key.2.2) (stages.map finsetOf) := by
  have _hZ := hZ
  have _hst := hst
  refine ⟨?_, ?_⟩
  · intro Y hY x hx
    obtain ⟨S, -, rfl⟩ := List.mem_map.1 hY
    obtain ⟨P, -, rfl⟩ := fam_mem_finsetOf.1 hx
    exact ⟨val_nonneg P, fun s => ⟨((P.1.getD s 0 : ℕ) : ℚ) / (P.2 : ℚ), by simp [val]⟩⟩
  · intro i hi1 hi2 x hx z hz
    rw [List.length_map] at hi2
    have hF := forall_lt_of_forall_range hfacts (i - 1) (by omega)
    have e1 : (({(ê ^ (1 + key.1))⁻¹ • init p₂ p₃ y ϖ₂ ϖ₃ key.1 key.2.1 key.2.2} :
        Finset (Fin 8 → ℝ)) :: stages.map finsetOf) =
        ([initI key.1 key.2.1 key.2.2] :: stages).map finsetOf := by
      rw [List.map_cons, finsetOf_initI]
    have e2 : stages.map finsetOf ++ [finsetOf Z] = (stages ++ [Z]).map finsetOf := by
      simp
    rw [e1, getD_map_finsetOf] at hx
    rw [e2, getD_map_finsetOf]
    obtain ⟨X, hX, rfl⟩ := fam_mem_finsetOf.1 hx
    obtain ⟨W, hW, rfl⟩ := fam_mem_finsetOf.1 hz
    unfold StageFact at hF
    exact hF X hX W hW

/-- A checked chain with one stage set of `n` points gives stages `[Y₁]` with `|Y₁| = n`. -/
theorem stages_one {Z : List IPt} (hZ : checkZ Z = true) (key : ℕ × ℕ × ℕ)
    {stages : List (List IPt)} {n : ℕ} (hsz : checkStages stages [n] = true)
    (hfacts : List.Forall (fun i => StageFact Z (([initI key.1 key.2.1 key.2.2] :: stages).getD i [])
      ((stages ++ [Z]).getD i [])) (List.range (stages.length + 1))) :
    ∃ Y₁ : Finset (Fin 8 → ℝ), Y₁.card = n ∧
      RSCCert.IsStages p₂ p₃ (finsetOf Z)
        ((ê ^ (1 + key.1))⁻¹ • init p₂ p₃ y ϖ₂ ϖ₃ key.1 key.2.1 key.2.2) [Y₁] := by
  obtain ⟨hst, hc⟩ := checkStages_sound hsz
  have hS := isStages_of_facts hZ key hst hfacts
  rcases hm : stages.map finsetOf with _ | ⟨Y, _ | ⟨Y', l⟩⟩
  · rw [hm] at hc; simp at hc
  · rw [hm] at hc hS
    simp only [List.map_cons, List.map_nil, List.cons.injEq, and_true] at hc
    exact ⟨Y, hc, hS⟩
  · rw [hm] at hc; simp at hc

/-- A checked chain with two stage sets of `n₁`, `n₂` points gives stages `[Y₁, Y₂]`. -/
theorem stages_two {Z : List IPt} (hZ : checkZ Z = true) (key : ℕ × ℕ × ℕ)
    {stages : List (List IPt)} {n₁ n₂ : ℕ} (hsz : checkStages stages [n₁, n₂] = true)
    (hfacts : List.Forall (fun i => StageFact Z (([initI key.1 key.2.1 key.2.2] :: stages).getD i [])
      ((stages ++ [Z]).getD i [])) (List.range (stages.length + 1))) :
    ∃ Y₁ Y₂ : Finset (Fin 8 → ℝ), Y₁.card = n₁ ∧ Y₂.card = n₂ ∧
      RSCCert.IsStages p₂ p₃ (finsetOf Z)
        ((ê ^ (1 + key.1))⁻¹ • init p₂ p₃ y ϖ₂ ϖ₃ key.1 key.2.1 key.2.2) [Y₁, Y₂] := by
  obtain ⟨hst, hc⟩ := checkStages_sound hsz
  have hS := isStages_of_facts hZ key hst hfacts
  rcases hm : stages.map finsetOf with _ | ⟨Y, _ | ⟨Y', _ | ⟨Y'', l⟩⟩⟩
  · rw [hm] at hc; simp at hc
  · rw [hm] at hc; simp at hc
  · rw [hm] at hc hS
    simp only [List.map_cons, List.map_nil, List.cons.injEq, and_true] at hc
    exact ⟨Y, Y', hc.1, hc.2, hS⟩
  · rw [hm] at hc; simp at hc

end

end MaxLeafFormal.RSCCheck