import RoleFlow.InvTyping

/-! # Inference of roles for invertibility

Unlike role typing, invertibility does not determine the roles: a line whose results nobody
needs may take either role. And the roles of the backward pass (`backward`) may fail where
others succeed: a line whose results are needed is marked `spine`, but if its data is returned,
only its conditioner recipe helps (see `Examples.lean`).

Without roles, every line contributes both of its recipes (`xrecsAll`). By monotonicity of
`Known`, these reach at least what the recipes of any roles reach. Conversely, in one closure
only one recipe of a line can add anything: after the conditioner recipe, its arguments, among
them the data, are known; after the spine recipe, its sources, among them the results, are known
(`known_pair_spine`, `known_pair_cond`). `inferAll` picks that recipe line by line, deciding
`Known` by the schedule (`knownSet`), so its roles reach everything that `xrecsAll` reaches
(`inferAll_known`). Hence `P` is invertible under some roles iff the inputs are `Known` under
`xrecsAll` (`Prog.exists_invertible_iff`), and then under `Prog.rolesAll`
(`Prog.invertible_rolesAll`).

`Prog.inferInv` tries the roles of the backward pass first, so it agrees with `Prog.infer`
where that succeeds (`inferInv_of_infer`), and falls back to `Prog.rolesAll`; it is sound and
complete (`inferInv_sound`, `inferInv_complete`, `inferInv_isSome_iff`). -/

namespace RoleFlow

variable {S : Type} {Sg : Sig S}

/-! ## Deciding `Known` -/

/-- The schedule's known set from `B`, ignoring the emitted lines. -/
def knownSet (R : List (XRecipe Sg)) (B : List ℕ) : Std.HashSet ℕ :=
  (fire R.length R 0 id (insertAll ∅ B)).2.2

theorem mem_knownSet {R : List (XRecipe Sg)} {B : List ℕ} {v : ℕ} :
    v ∈ knownSet R B ↔ Known R {w | w ∈ B} v := by
  constructor
  · refine fun h => fire_known _ _ _ _ _ (fun _ hx => hx) (fun w hw => .base ?_) v h
    simpa [mem_insertAll] using hw
  · exact known_fire le_rfl fun w hw => mem_insertAll.2 (.inl hw)

/-! ## All recipes -/

/-- The recipes of the lines under every role: each line's conditioner recipe and, if it has a
partial inverse, its spine recipe. -/
def xrecsAll : ℕ → List (Line Sg) → List (XRecipe Sg)
  | _, [] => []
  | n, l :: ls => ⟨l, l.ys n⟩ ::
      ((match Sg.inv l.op with
        | some si => [⟨⟨si.op, l.ys n ++ l.args.drop si.k⟩, l.args.take si.k⟩]
        | none => []) ++ xrecsAll (n + l.bs.length) ls)

theorem mem_xrecsAll_of_mem {n : ℕ} {ls : List (Line Sg)} {ρ : List Role} {x : XRecipe Sg}
    (h : x ∈ xrecs n ls ρ) : x ∈ xrecsAll n ls := by
  induction ls generalizing n ρ with
  | nil => cases ρ <;> simp [xrecs] at h
  | cons l ls ih =>
    rcases mem_xrecs h with h | rfl | ⟨si, hsi, rfl⟩
    · simp only [xrecsAll]; split <;> simp [ih h]
    · simp [xrecsAll]
    · simp [xrecsAll, hsi]

/-- **Roles never beat all recipes.** -/
theorem Known.xrecsAll {n : ℕ} {ls : List (Line Sg)} {ρ : List Role} {B : Set ℕ} {v : ℕ}
    (h : Known (xrecs n ls ρ) B v) : Known (xrecsAll n ls) B v :=
  h.of_subset fun _ => mem_xrecsAll_of_mem

/-! ## One line's two recipes -/

section Pair
variable {rc rs : XRecipe Sg} {R : List (XRecipe Sg)} {A : Set ℕ}

/-- If the conditioner recipe `rc` only fires once its targets are known anyway, it can be
dropped. -/
theorem known_pair_spine
    (hc : (∀ a ∈ rc.line.args, Known (rs :: R) A a) → ∀ y ∈ rc.tgt, Known (rs :: R) A y)
    {v : ℕ} (h : Known (rc :: rs :: R) A v) : Known (rs :: R) A v := by
  induction h with
  | base hv => exact .base hv
  | @fire x _ hx _ hv ih =>
    rcases List.mem_cons.1 hx with rfl | hx
    · exact hc ih _ hv
    · exact .fire hx ih hv

/-- If the sources of the conditioner recipe `rc` are known without it and its targets are not,
the spine recipe `rs`, whose sources contain those targets, never fires, and can be dropped. -/
theorem known_pair_cond (hcs : ∀ t ∈ rs.tgt, t ∈ rc.line.args)
    (hsc : ∀ t ∈ rc.tgt, t ∈ rs.line.args)
    (ha : ∀ a ∈ rc.line.args, Known (rs :: R) A a) (hy : ¬ ∀ y ∈ rc.tgt, Known (rs :: R) A y)
    {v : ℕ} (h : Known (rc :: rs :: R) A v) : Known (rc :: R) A v := by
  obtain ⟨y, hyt, hyn⟩ := not_forall₂.1 hy
  -- without its targets, `rs` never fires
  have hdrop : ∀ {u}, Known (rs :: R) A u → Known R A u := fun hu => by
    induction hu with
    | base hv => exact .base hv
    | @fire x _ hx _ hv ih =>
      rcases List.mem_cons.1 hx with rfl | hx
      · exact absurd ((ih y (hsc y hyt)).of_subset fun _ => List.mem_cons_of_mem _) hyn
      · exact .fire hx ih hv
  induction h with
  | base hv => exact .base hv
  | @fire x _ hx _ hv ih =>
    rcases List.mem_cons.1 hx with rfl | hx
    · exact .fire (List.mem_cons_self ..) ih hv
    rcases List.mem_cons.1 hx with rfl | hx
    · exact (hdrop (ha _ (hcs _ hv))).of_subset fun _ => List.mem_cons_of_mem _
    · exact .fire (List.mem_cons_of_mem _ hx) ih hv

end Pair

/-! ## Choosing the roles -/

/-- The roles, line by line, from the base `B`: `F` holds the recipes chosen so far. A line
becomes a conditioner if, with all later recipes, its arguments are known but its results are
not; otherwise a spine step. A line without a partial inverse becomes a conditioner. -/
def inferAll (B : List ℕ) : List (XRecipe Sg) → ℕ → List (Line Sg) → List Role
  | _, _, [] => []
  | F, n, l :: ls =>
    match Sg.inv l.op with
    | none => .cond :: inferAll B (⟨l, l.ys n⟩ :: F) (n + l.bs.length) ls
    | some si =>
      let rs : XRecipe Sg := ⟨⟨si.op, l.ys n ++ l.args.drop si.k⟩, l.args.take si.k⟩
      let K := knownSet (rs :: (F ++ xrecsAll (n + l.bs.length) ls)) B
      if l.args.all K.contains && !(l.ys n).all K.contains then
        .cond :: inferAll B (⟨l, l.ys n⟩ :: F) (n + l.bs.length) ls
      else .spine :: inferAll B (rs :: F) (n + l.bs.length) ls

theorem length_inferAll (B : List ℕ) :
    ∀ (F : List (XRecipe Sg)) (n : ℕ) (ls : List (Line Sg)), (inferAll B F n ls).length = ls.length
  | _, _, [] => rfl
  | F, n, l :: ls => by
    simp only [inferAll]
    split
    · simp [length_inferAll]
    · split <;> simp [length_inferAll]

/-- **The chosen roles reach what all recipes reach.** -/
theorem inferAll_known (B : List ℕ) :
    ∀ (F : List (XRecipe Sg)) (n : ℕ) (ls : List (Line Sg)) {v : ℕ},
      Known (F ++ xrecsAll n ls) {w | w ∈ B} v →
        Known (F ++ xrecs n ls (inferAll B F n ls)) {w | w ∈ B} v
  | F, n, [], _, h => by simpa [inferAll, xrecs, xrecsAll] using h
  | F, n, l :: ls, v, h => by
    set n' := n + l.bs.length
    set rc : XRecipe Sg := ⟨l, l.ys n⟩
    cases hsi : Sg.inv l.op with
    | none =>
      have h' : Known ((rc :: F) ++ xrecsAll n' ls) {w | w ∈ B} v :=
        h.of_subset fun x hx => by simp [xrecsAll, hsi] at hx ⊢; tauto
      have ih := inferAll_known B (rc :: F) n' ls h'
      simp only [inferAll, hsi]
      exact ih.of_subset fun x hx => by simp [xrecs_cond] at hx ⊢; tauto
    | some si =>
      set rs : XRecipe Sg := ⟨⟨si.op, l.ys n ++ l.args.drop si.k⟩, l.args.take si.k⟩
      set R := F ++ xrecsAll n' ls
      have h' : Known (rc :: rs :: R) {w | w ∈ B} v :=
        h.of_subset fun x hx => by simp [R, rc, rs, xrecsAll, hsi] at hx ⊢; tauto
      have hcs : ∀ t ∈ rs.tgt, t ∈ rc.line.args := fun _ ht => List.mem_of_mem_take ht
      have hsc : ∀ t ∈ rc.tgt, t ∈ rs.line.args := fun _ ht => List.mem_append_left _ ht
      have hK : ∀ u, (knownSet (rs :: R) B).contains u = true ↔ Known (rs :: R) {w | w ∈ B} u :=
        fun u => by rw [Std.HashSet.contains_iff_mem, mem_knownSet]
      have hc_iff : (l.args.all (knownSet (rs :: R) B).contains &&
            !(l.ys n).all (knownSet (rs :: R) B).contains) = true ↔
          (∀ a ∈ rc.line.args, Known (rs :: R) {w | w ∈ B} a) ∧
            ¬ ∀ y ∈ rc.tgt, Known (rs :: R) {w | w ∈ B} y := by
        simp only [Bool.and_eq_true, List.all_eq_true, Bool.not_eq_true', ← Bool.not_eq_true,
          hK]
        rfl
      simp only [inferAll, hsi]
      split
      · next hc =>
        have hc := hc_iff.1 hc
        have ih := inferAll_known B (rc :: F) n' ls
          ((known_pair_cond hcs hsc hc.1 hc.2 h').of_subset fun x hx => by
            simp [R] at hx ⊢; tauto)
        exact ih.of_subset fun x hx => by simp [xrecs_cond] at hx ⊢; tauto
      · next hc =>
        have hc' : (∀ a ∈ rc.line.args, Known (rs :: R) {w | w ∈ B} a) →
            ∀ y ∈ rc.tgt, Known (rs :: R) {w | w ∈ B} y := fun ha => by
          by_contra hn
          exact hc (hc_iff.2 ⟨ha, hn⟩)
        have ih := inferAll_known B (rs :: F) n' ls
          ((known_pair_spine hc' h').of_subset fun x hx => by simp [R] at hx ⊢; tauto)
        exact ih.of_subset fun x hx => by simp [xrecs_spine _ _ _ _ hsi] at hx ⊢; tauto

/-! ## Programs -/

namespace Prog

variable {P : Prog Sg} {ρ : List Role}

/-- The roles chosen from the weights and the returned variables. -/
def rolesAll (P : Prog Sg) : List Role :=
  inferAll (List.range P.nw ++ P.ret) [] P.ctx₀.length P.lines

theorem length_rolesAll (P : Prog Sg) : P.rolesAll.length = P.lines.length :=
  length_inferAll _ _ _ _

theorem base_eq (P : Prog Sg) :
    {w | w ∈ List.range P.nw ++ P.ret} = ({w | w < P.nw} ∪ ↑P.ret.toFinset : Set ℕ) := by
  ext; simp

theorem invertible_rolesAll_of_known (hwf : P.WF) (hk : ∀ x ∈ List.range' P.nw P.ins.length,
      Known (xrecsAll P.ctx₀.length P.lines) ({w | w < P.nw} ∪ ↑P.ret.toFinset) x) :
    P.Invertible P.rolesAll := by
  refine invertible_iff_known.2 ⟨hwf, fun x hx => ?_⟩
  have := inferAll_known (List.range P.nw ++ P.ret) [] P.ctx₀.length P.lines
    (v := x) (by rw [List.nil_append, base_eq]; exact hk x hx)
  rw [List.nil_append, base_eq] at this
  exact this

/-- **Invertibility under some roles** is invertibility under all recipes at once. -/
theorem exists_invertible_iff : (∃ ρ, P.Invertible ρ) ↔ P.WF ∧
    ∀ x ∈ List.range' P.nw P.ins.length,
      Known (xrecsAll P.ctx₀.length P.lines) ({w | w < P.nw} ∪ ↑P.ret.toFinset) x := by
  refine ⟨fun ⟨ρ, hI⟩ => ?_, fun h => ⟨_, invertible_rolesAll_of_known h.1 h.2⟩⟩
  obtain ⟨hwf, hk⟩ := invertible_iff_known.1 hI
  exact ⟨hwf, fun x hx => (hk x hx).xrecsAll⟩

/-- **`rolesAll` is complete**: if any roles invert `P`, these do. -/
theorem invertible_rolesAll (h : ∃ ρ, P.Invertible ρ) : P.Invertible P.rolesAll :=
  let h := exists_invertible_iff.1 h
  invertible_rolesAll_of_known h.1 h.2

theorem length_backward : ∀ (n : ℕ) (ls : List (Line Sg)) (N : Finset ℕ),
    (backward n ls N).1.length = ls.length
  | _, [], _ => rfl
  | n, l :: ls, N => by
    simp only [backward]
    split <;> simp [length_backward]

/-- **Inverse inference without density**: the backward pass's roles if they invert `P`, and
otherwise `rolesAll` if they do. -/
def inferInv [DecidableEq S] (P : Prog Sg) : Option (List Role) :=
  let ρ := (backward P.ctx₀.length P.lines P.ret.toFinset).1
  if P.Invertible ρ then some ρ
  else if P.Invertible P.rolesAll then some P.rolesAll else none

variable [DecidableEq S]

theorem inferInv_sound (h : P.inferInv = some ρ) : P.Invertible ρ := by
  unfold inferInv at h; simp only at h
  split_ifs at h with h₁ h₂ <;> cases h
  · exact h₁
  · exact h₂

/-- The roles found have one role per line, so they also give an invertibility typing. -/
theorem inferInv_invTyped (h : P.inferInv = some ρ) : P.InvTyped ρ := by
  refine invTyped_iff.2 ⟨inferInv_sound h, ?_⟩
  unfold inferInv at h; simp only at h
  split_ifs at h <;> cases h
  · exact length_backward ..
  · exact P.length_rolesAll

theorem inferInv_of_infer (h : P.infer = some ρ) : P.inferInv = some ρ := by
  have hT := (P.infer_iff ρ).1 h
  unfold infer at h; simp only at h; split_ifs at h; cases h
  simp only [inferInv, if_pos hT.invertible]

/-- **Completeness**: if any roles invert `P`, inference succeeds. -/
theorem inferInv_complete (h : P.Invertible ρ) : P.inferInv.isSome := by
  unfold inferInv; simp only
  split_ifs with h₁ h₂
  · rfl
  · rfl
  · exact absurd (invertible_rolesAll ⟨ρ, h⟩) h₂

theorem inferInv_isSome_iff : P.inferInv.isSome ↔ ∃ ρ, P.InvTyped ρ := by
  constructor
  · intro h
    obtain ⟨ρ, hρ⟩ := Option.isSome_iff_exists.1 h
    exact ⟨ρ, inferInv_invTyped hρ⟩
  · rintro ⟨ρ, hρ⟩
    exact inferInv_complete (invTyped_iff.1 hρ).1

end Prog

end RoleFlow