import RoleFlow.Sig

/-! # Live states (`docs/roleflow.tex`, `sec:live`)

The density proof runs on *live states*: the values of the positions in a finite set `A`,
`X_A = ∏_{v ∈ A} X_{σ v}`. Here the sorts `σ` are read from one fixed context `Γ` (the
context of all variables of a program), so a live state is a function on the positions of `Γ`
that lie in `A` (`LState Γ A`), measured by the product `lvol Γ A`. Positions of `A` beyond
`Γ` are ignored.

The reorders between live states and values of contexts are measure-preserving measurable
equivalences, built from `reorderEquiv`:
* reading a live state at a duplicate-free list of positions that covers `A` (`LState.pick`)
  gives `X_A ≅ X_as` (`LState.pickEquiv`, `measurePreserving_pick`);
* if `A` is the disjoint union of `B` and the positions of such a list, keeping `B` and reading
  the list gives `X_A ≅ X_B × X_as` (`LState.extract`, `measurePreserving_extract`), with
  Mathlib's `measurePreserving_piEquivPiSubtypeProd`. -/

open MeasureTheory

namespace RoleFlow

variable {S : Type} [Sorts S]

variable (Γ : List S)

/-- The positions of `Γ` that lie in `A`. -/
abbrev Pos (A : Finset ℕ) : Type := {i : Fin Γ.length // i.1 ∈ A}

/-- A live state: a value for each position of `Γ` in `A`. -/
abbrev LState (A : Finset ℕ) : Type := ∀ i : Pos Γ A, X (Γ.get i.1)

/-- The product measure on live states. -/
noncomputable def lvol (A : Finset ℕ) : Measure (LState Γ A) :=
  Measure.pi fun i : Pos Γ A => μ (Γ.get i.1)

instance (A : Finset ℕ) : SigmaFinite (lvol Γ A) := by unfold lvol; infer_instance

variable {Γ}

namespace LState
variable {A B C : Finset ℕ} {a : S}

/-- The live state of a value of `Γ`. -/
def restrict (A : Finset ℕ) (v : Val Γ) : LState Γ A := fun i => v i.1

/-- Forget the positions outside `A`. -/
def sub (h : A ⊆ B) (s : LState Γ B) : LState Γ A := fun i => s ⟨i.1, h i.2⟩

@[simp] theorem sub_restrict (h : A ⊆ B) (v : Val Γ) : sub h (restrict B v) = restrict A v := rfl

@[fun_prop] theorem measurable_sub (h : A ⊆ B) : Measurable (sub (Γ := Γ) h) :=
  measurable_pi_lambda _ fun _ => measurable_pi_apply _

/-- The value at the position `i ∈ A`. -/
def get (s : LState Γ A) (i : ℕ) (hi : i ∈ A) (h : Γ[i]? = some a) : X a :=
  tr X (Val.get_eq_of_getElem? h) (s ⟨⟨i, Val.lt_of_getElem? h⟩, hi⟩)

@[fun_prop] theorem measurable_get {i : ℕ} (hi : i ∈ A) (h : Γ[i]? = some a) :
    Measurable fun s : LState Γ A => s.get i hi h :=
  (measurable_tr _).comp (measurable_pi_apply _)

@[simp] theorem get_restrict (v : Val Γ) {i : ℕ} (hi : i ∈ A) (h : Γ[i]? = some a) :
    (restrict A v).get i hi h = v.get i h := rfl

/-! ## Reading a live state at a list of positions -/

section Pick
variable {xs : List ℕ} {as : List S}

/-- The position of the `j`-th entry of `xs`. -/
def pickPos (hA : ∀ x ∈ xs, x ∈ A) (h : Args Γ xs as) (j : Fin as.length) : Pos Γ A :=
  ⟨⟨xs[j.1]'(h.lt_length j.2), Val.lt_of_getElem? (h.getElem? j.1 (h.lt_length j.2))⟩,
    hA _ (List.getElem_mem _)⟩

omit [Sorts S] in
theorem pickPos_sort (hA : ∀ x ∈ xs, x ∈ A) (h : Args Γ xs as) (j : Fin as.length) :
    Γ.get (pickPos hA h j).1 = as.get j :=
  Val.get_eq_of_getElem? (h.getElem? j.1 (h.lt_length j.2))

/-- Read a live state at the positions `xs`. -/
def pick (s : LState Γ A) (xs : List ℕ) (hA : ∀ x ∈ xs, x ∈ A) (h : Args Γ xs as) : Val as :=
  fun j => tr X (pickPos_sort hA h j) (s (pickPos hA h j))

@[fun_prop] theorem measurable_pick (hA : ∀ x ∈ xs, x ∈ A) (h : Args Γ xs as) :
    Measurable fun s : LState Γ A => s.pick xs hA h :=
  measurable_pi_lambda _ fun _ => (measurable_tr _).comp (measurable_pi_apply _)

@[simp] theorem pick_restrict (v : Val Γ) (hA : ∀ x ∈ xs, x ∈ A) (h : Args Γ xs as) :
    (restrict A v).pick xs hA h = v.pick xs h := by
  funext j
  rw [Val.apply_eq_get (v.pick xs h), Val.get_pick]
  rfl

omit [Sorts S] in
/-- If `xs` has no duplicates and covers `A`, the positions of `xs` are the positions of `A`. -/
theorem pickPos_bijective (hA : ∀ x ∈ xs, x ∈ A) (h : Args Γ xs as) (hnd : xs.Nodup)
    (hcov : ∀ x ∈ A, x ∈ xs) : Function.Bijective (pickPos hA h) := by
  refine ⟨fun j j' e => ?_, fun ⟨⟨i, hi⟩, hiA⟩ => ?_⟩
  · have := congrArg (fun p : Pos Γ A => p.1.1) e
    exact Fin.ext ((hnd.getElem_inj_iff).1 this)
  · obtain ⟨j, hj, rfl⟩ := List.mem_iff_getElem.1 (hcov _ hiA)
    exact ⟨⟨j, by rw [← h.length_eq]; exact hj⟩, rfl⟩

variable (hA : ∀ x ∈ xs, x ∈ A) (h : Args Γ xs as) (hnd : xs.Nodup) (hcov : ∀ x ∈ A, x ∈ xs)

/-- **Reorder**: reading a live state at a duplicate-free list that covers `A`, as a measurable
equivalence `X_A ≅ X_as`. -/
noncomputable def pickEquiv : LState Γ A ≃ᵐ Val as :=
  reorderEquiv (fun j : Fin as.length => as.get j) (fun p : Pos Γ A => Γ.get p.1)
    (.ofBijective _ (pickPos_bijective hA h hnd hcov)) (pickPos_sort hA h)

theorem pickEquiv_apply (s : LState Γ A) : pickEquiv hA h hnd hcov s = s.pick xs hA h := rfl

theorem measurePreserving_pick : MeasurePreserving (pickEquiv hA h hnd hcov) (lvol Γ A) (vol as) :=
  measurePreserving_reorder μ _ _ _ _

end Pick

/-! ## Extracting positions from a live state -/

section Extract
variable {xs : List ℕ} {as : List S} (hB : B ⊆ A) (hA : ∀ x ∈ xs, x ∈ A) (h : Args Γ xs as)
  (hnd : xs.Nodup) (hBx : ∀ i ∈ A, i ∈ B ↔ i ∉ xs)

/-- The positions of `B ⊆ A` among those of `A`. -/
def posIn : Pos Γ B ≃ {i : Pos Γ A // i.1.1 ∈ B} where
  toFun i := ⟨⟨i.1, hB i.2⟩, i.2⟩
  invFun i := ⟨i.1.1, i.2⟩
  left_inv _ := rfl
  right_inv _ := rfl

/-- The positions of `xs` among those of `A`: those not in `B`. -/
def posOut : Pos Γ xs.toFinset ≃ {i : Pos Γ A // i.1.1 ∉ B} where
  toFun i := ⟨⟨i.1, hA _ (List.mem_toFinset.1 i.2)⟩,
    fun hi => (hBx _ (hA _ (List.mem_toFinset.1 i.2))).1 hi (List.mem_toFinset.1 i.2)⟩
  invFun i := ⟨i.1.1, List.mem_toFinset.2 (not_not.1 fun h => i.2 ((hBx _ i.1.2).2 h))⟩
  left_inv _ := rfl
  right_inv _ := rfl

/-- **Splitting**: `X_A ≅ X_B × X_as` when `A` is the disjoint union of `B` and the positions
`xs`; it keeps the positions of `B` and reads those of `xs`. -/
noncomputable def extract : LState Γ A ≃ᵐ LState Γ B × Val as :=
  (MeasurableEquiv.piEquivPiSubtypeProd _ _).trans <| .prodCongr
    (reorderEquiv (fun i : Pos Γ B => Γ.get i.1) _ (posIn hB) fun _ => rfl)
    ((reorderEquiv (fun i : Pos Γ xs.toFinset => Γ.get i.1) _ (posOut hA hBx) fun _ => rfl).trans
      (pickEquiv (fun _ => List.mem_toFinset.2) h hnd fun _ => List.mem_toFinset.1))

theorem extract_apply (s : LState Γ A) : extract hB hA h hnd hBx s = (sub hB s, s.pick xs hA h) :=
  rfl

theorem extract_restrict (v : Val Γ) :
    extract hB hA h hnd hBx (restrict A v) = (restrict B v, v.pick xs h) := by
  rw [extract_apply, pick_restrict]; rfl

theorem measurePreserving_extract :
    MeasurePreserving (extract hB hA h hnd hBx) (lvol Γ A) ((lvol Γ B).prod (vol as)) :=
  ((measurePreserving_reorder μ _ _ (posIn hB) _).prod
    ((measurePreserving_pick _ h hnd fun _ => List.mem_toFinset.1).comp
      (measurePreserving_reorder μ _ _ (posOut hA hBx) _))).comp
    (measurePreserving_piEquivPiSubtypeProd (fun i : Pos Γ A => μ (Γ.get i.1)) _)

end Extract

end LState

end RoleFlow