module

public import Papers.OaiCapacityCriterionState.Helpers.FiniteExperiments
public import Papers.OaiCapacityCriterionState.Helpers.InformationFlow
public import Mathlib.Data.Fin.Tuple.Basic
public import Mathlib.Algebra.BigOperators.Fin

/-! # Products of experiments with different observation spaces

Splitting off the first coordinate reduces a dependent finite product to the
binary posterior product. This transfers companion-class closure and its cubic
information saving to arbitrary finite families of observations.
-/

public section

open MeasureTheory
open scoped BigOperators

namespace OaiCapacityCriterionState

-- @node: dependentLikelihood_sum
/-- Given [a finite family's likelihoods for each experiment, state, and
observation](hyp:lik), if [each experiment's likelihoods sum to one over its observations for every
state](hyp:hnorm), then [the product likelihood sums to one over all observation tuples
for the chosen state](goal). -/
lemma dependentLikelihood_sum {ι : Type} [Fintype ι] [DecidableEq ι]
    (O : ι → Type) [∀ a, Fintype (O a)] (lik : ∀ a, Fin 4 → O a → ℝ)
    (hnorm : ∀ a i, ∑ o, lik a i o = 1) (i : Fin 4) :
    ∑ obs : ∀ a, O a, ∏ a, lik a i (obs a) = 1 := by
  rw [← Fintype.prod_sum]
  simp only [hnorm, Finset.prod_const_one]

-- @node: experimentLaw_fin_zero
/-- For [a likelihood family indexed by zero experiments](hyp:lik), [the posterior
law generated by their product is the uninformative law](goal). -/
lemma experimentLaw_fin_zero (O : Fin 0 → Type) [∀ a, Fintype (O a)]
    (lik : ∀ a, Fin 4 → O a → ℝ) :
    experimentLaw (fun i (obs : ∀ a, O a) => ∏ a, lik a i (obs a)) =
      uninformativeLaw := by
  classical
  let e : Unit ≃ (∀ a, O a) := Equiv.ofUnique _ _
  rw [← experimentLaw_reindex e]
  simpa using experimentLaw_unit

-- @node: experimentLaw_fin_succ
/-- Given [likelihoods for every experiment, state, and observation in a family of
n plus one experiments](hyp:lik), if [all likelihoods in the family are nonnegative](hyp:hnonneg)
and [each experiment's likelihoods sum to one over observations for every state](hyp:hnorm),
then [the posterior law of their product equals the law obtained by combining the first
experiment's posterior law with the posterior law of the remaining n experiments](goal). -/
lemma experimentLaw_fin_succ {n : ℕ} (O : Fin (n + 1) → Type)
    [∀ a, Fintype (O a)] (lik : ∀ a, Fin 4 → O a → ℝ)
    (hnonneg : ∀ a i o, 0 ≤ lik a i o)
    (hnorm : ∀ a i, ∑ o, lik a i o = 1) :
    experimentLaw (fun i (obs : ∀ a, O a) => ∏ a, lik a i (obs a)) =
      starLaw (experimentLaw (lik 0))
        (experimentLaw (fun i (obs : ∀ a : Fin n, O a.succ) =>
          ∏ a : Fin n, lik a.succ i (obs a))) := by
  classical
  rw [← experimentLaw_reindex (Fin.consEquiv O)]
  have heq : (fun i (obs : O 0 × (∀ a : Fin n, O a.succ)) =>
      ∏ a, lik a i (Fin.consEquiv O obs a)) =
      productLikelihood (lik 0) (fun i (obs : ∀ a : Fin n, O a.succ) =>
        ∏ a : Fin n, lik a.succ i (obs a)) := by
    funext i obs
    simp [Fin.prod_univ_succ, productLikelihood]
  rw [heq]
  exact experimentLaw_prod _ _ (hnonneg 0)
    (fun i obs => Finset.prod_nonneg (fun a _ => hnonneg a.succ i (obs a)))
    (hnorm 0) (fun i => dependentLikelihood_sum (fun a : Fin n => O a.succ)
      (fun a : Fin n => lik a.succ) (fun a : Fin n => hnorm a.succ) i)

-- @node: companion_fin_product_bound
/-- Suppose [a class contains the uninformative law](hyp:hzero), [is closed under
combining conditionally independent experiments](hyp:hstar), and [quadratic information
augmented by its cube divided by 3000 is subadditive for any two class members](hyp:hsave).
For n experiments with [likelihoods specified for every state and observation](hyp:lik),
[nonnegative likelihoods](hyp:hnonneg), [likelihoods summing to
one over observations for every state](hyp:hnorm), and [individual posterior laws in the
class](hyp:hmem), [the product posterior law also belongs to the class and its augmented
information is at most the sum of the individual augmented informations](goal). -/
lemma companion_fin_product_bound (C : Set (Measure (Fin 4 → ℝ)))
    (hzero : uninformativeLaw ∈ C)
    (hstar : ∀ mu ∈ C, ∀ nu ∈ C, starLaw mu nu ∈ C)
    (hsave : ∀ mu ∈ C, ∀ nu ∈ C,
      cubicInfoTransform (quadInfo (starLaw mu nu)) ≤
        cubicInfoTransform (quadInfo mu) + cubicInfoTransform (quadInfo nu))
    (n : ℕ) (O : Fin n → Type) [∀ a, Fintype (O a)]
    (lik : ∀ a, Fin 4 → O a → ℝ)
    (hnonneg : ∀ a i o, 0 ≤ lik a i o)
    (hnorm : ∀ a i, ∑ o, lik a i o = 1)
    (hmem : ∀ a, experimentLaw (lik a) ∈ C) :
    experimentLaw (fun i (obs : ∀ a, O a) => ∏ a, lik a i (obs a)) ∈ C ∧
      cubicInfoTransform (quadInfo (experimentLaw
        (fun i (obs : ∀ a, O a) => ∏ a, lik a i (obs a)))) ≤
        ∑ a, cubicInfoTransform (quadInfo (experimentLaw (lik a))) := by
  induction n with
  | zero =>
      rw [experimentLaw_fin_zero]
      refine ⟨hzero, ?_⟩
      rw [extreme_moments.2.1]
      norm_num [cubicInfoTransform]
  | succ n ih =>
      obtain ⟨htail, hbound⟩ := ih (fun a => O a.succ) (fun a => lik a.succ)
        (fun a => hnonneg a.succ) (fun a => hnorm a.succ) (fun a => hmem a.succ)
      rw [experimentLaw_fin_succ O lik hnonneg hnorm, Fin.sum_univ_succ]
      exact ⟨hstar _ (hmem 0) _ htail,
        (hsave _ (hmem 0) _ htail).trans (add_le_add_right hbound _)⟩

-- @node: experimentLaw_pi_reindex
/-- Given [a finite family's likelihoods for each experiment, state, and
observation](hyp:lik), [reindexing the experiments by a bijection leaves the posterior law
of their product unchanged, even when their observation spaces differ](goal). -/
lemma experimentLaw_pi_reindex {ι κ : Type} [Fintype ι] [Fintype κ]
    [DecidableEq ι] [DecidableEq κ] (e : κ ≃ ι)
    (O : ι → Type) [∀ a, Fintype (O a)] (lik : ∀ a, Fin 4 → O a → ℝ) :
    experimentLaw (fun i (obs : ∀ a, O a) => ∏ a, lik a i (obs a)) =
      experimentLaw (fun i (obs : ∀ a, O (e a)) => ∏ a, lik (e a) i (obs a)) := by
  classical
  have h := experimentLaw_reindex (Equiv.piCongrLeft O e).symm
    (fun i (obs : ∀ a, O (e a)) => ∏ a, lik (e a) i (obs a))
  calc
    _ = experimentLaw (fun i (obs : ∀ a, O a) =>
        ∏ a : κ, lik (e a) i (obs (e a))) := by
      congr 1
      funext i obs
      exact (e.prod_comp (fun a => lik a i (obs a))).symm
    _ = _ := by simpa only [Equiv.piCongrLeft_symm_apply] using h

-- @node: companion_dependent_product_bound
/-- Suppose [a class contains the uninformative law](hyp:hzero), [is closed under
combining conditionally independent experiments](hyp:hstar), and [quadratic information
augmented by its cube divided by 3000 is subadditive for any two class members](hyp:hsave).
For any finite family of experiments with [likelihoods specified for every state and
observation](hyp:lik), [nonnegative likelihoods](hyp:hnonneg),
[likelihoods summing to one over observations for every state](hyp:hnorm), and [individual
posterior laws in the class](hyp:hmem), [the product posterior law belongs to the class
and its augmented information is at most the sum of the individual augmented informations](goal). -/
lemma companion_dependent_product_bound (C : Set (Measure (Fin 4 → ℝ)))
    (hzero : uninformativeLaw ∈ C)
    (hstar : ∀ mu ∈ C, ∀ nu ∈ C, starLaw mu nu ∈ C)
    (hsave : ∀ mu ∈ C, ∀ nu ∈ C,
      cubicInfoTransform (quadInfo (starLaw mu nu)) ≤
        cubicInfoTransform (quadInfo mu) + cubicInfoTransform (quadInfo nu))
    {ι : Type} [Fintype ι] [DecidableEq ι]
    (O : ι → Type) [∀ a, Fintype (O a)] (lik : ∀ a, Fin 4 → O a → ℝ)
    (hnonneg : ∀ a i o, 0 ≤ lik a i o)
    (hnorm : ∀ a i, ∑ o, lik a i o = 1)
    (hmem : ∀ a, experimentLaw (lik a) ∈ C) :
    experimentLaw (fun i (obs : ∀ a, O a) => ∏ a, lik a i (obs a)) ∈ C ∧
      cubicInfoTransform (quadInfo (experimentLaw
        (fun i (obs : ∀ a, O a) => ∏ a, lik a i (obs a)))) ≤
        ∑ a, cubicInfoTransform (quadInfo (experimentLaw (lik a))) := by
  classical
  let e := (Fintype.equivFin ι).symm
  rw [experimentLaw_pi_reindex e O lik, ← e.sum_comp]
  exact companion_fin_product_bound C hzero hstar hsave _ (fun a => O (e a))
    (fun a => lik (e a)) (fun a => hnonneg (e a)) (fun a => hnorm (e a))
    (fun a => hmem (e a))

end OaiCapacityCriterionState