Algorithm.lean

/-
Copyright (c) 2025 Rémy Degenne. All rights reserved.
Released under Apache 2.0 license as described in the file LICENSE.
Authors: Rémy Degenne, Paulo Rauber
-/
module

public import LeanMachineLearning.ForMathlib.MeasureTheory.Measurable
public import LeanMachineLearning.ForMathlib.Probability.Kernel.IonescuTulcea.Traj
public import LeanMachineLearning.ForMathlib.Probability.Kernel.MeasurableSpace

/-!
# Algorithms and environments

We define structures for stochastic, sequential algorithms and environments, and the notion of an
algorithm-environment sequence, which is a sequence of observations, actions and feedbacks generated
by an algorithm interacting with an environment.

A round of interaction consists of an observation in `𝓞`, an action in `𝓐` and a feedback in `𝓨`,
in that order: it is an element of `Round 𝓞 𝓐 𝓨 = 𝓞 × 𝓐 × 𝓨`.
At time `n`, the algorithm has seen the `n` complete rounds at times `0, ..., n - 1`.
That history is an element of `Hist 𝓞 𝓐 𝓨 n = Fin n → Round 𝓞 𝓐 𝓨`. The environment then draws
an observation `O n` given the history, the algorithm chooses an action `A n` according to its
policy applied to the history and `O n`, and the environment returns a feedback `Y n` given the
history, `O n` and `A n`.
In particular, the first observation is drawn according to the observation kernel at time `0`,
applied to the empty history (the unique element of `Hist 𝓞 𝓐 𝓨 0`).

## Main definitions

* `Round 𝓞 𝓐 𝓨`: one round of interaction, an observation-action-feedback triple.
* `Hist 𝓞 𝓐 𝓨 n`: a history of `n` complete rounds.
* `Algorithm 𝓞 𝓐 𝓨`: a stochastic, sequential algorithm.
* `Environment 𝓞 𝓐 𝓨`: a stochastic environment.
* `history O A Y n`: the history before time `n`, a random variable with values in `Hist 𝓞 𝓐 𝓨 n`.
* `IsAlgEnvSeq O A Y alg env P`: an algorithm-environment sequence. That is, a sequence of
  observations `O`, actions `A` and feedbacks `Y` that have the correct conditional distributions
  to be generated by an algorithm `alg` interacting with an environment `env`, defined on a
  probability space `(Ω, P)`.
* `IsAlgEnvSeqUntil O A Y alg env P N`: `O`, `A` and `Y` form an algorithm-environment sequence for
  the times `n < N`.

-/

@[expose] public section

open MeasureTheory ProbabilityTheory Filter Real Finset

open scoped ENNReal NNReal

-- TODO: this belongs to Mathlib.
instance Prod.instUnique {α β : Type*} [Unique α] [Unique β] : Unique (α × β) where
  default := (default, default)
  uniq p := by simp [Prod.ext_iff, Unique.eq_default]

namespace Learning

variable {𝓞 𝓐 𝓨 Ω : Type*} {m𝓞 : MeasurableSpace 𝓞} {m𝓐 : MeasurableSpace 𝓐}
  {m𝓨 : MeasurableSpace 𝓨} {mΩ : MeasurableSpace Ω}

/-- One round of interaction: an observation, then an action, then a feedback. -/
abbrev Round (𝓞 𝓐 𝓨 : Type*) := 𝓞 × 𝓐 × 𝓨

/-- The observation of a round. -/
def Round.obs (r : Round 𝓞 𝓐 𝓨) : 𝓞 := r.1
/-- The action of a round. -/
def Round.action (r : Round 𝓞 𝓐 𝓨) : 𝓐 := r.2.1
/-- The feedback of a round. -/
def Round.feedback (r : Round 𝓞 𝓐 𝓨) : 𝓨 := r.2.2

namespace Round

@[simp] lemma obs_mk (o : 𝓞) (a : 𝓐) (y : 𝓨) : Round.obs (o, a, y) = o := rfl
@[simp] lemma action_mk (o : 𝓞) (a : 𝓐) (y : 𝓨) : Round.action (o, a, y) = a := rfl
@[simp] lemma feedback_mk (o : 𝓞) (a : 𝓐) (y : 𝓨) : Round.feedback (o, a, y) = y := rfl

@[simp]
lemma mk_obs_action_feedback (r : Round 𝓞 𝓐 𝓨) : (r.obs, r.action, r.feedback) = r := rfl

lemma obs_eq_fst (r : Round 𝓞 𝓐 𝓨) : r.obs = r.1 := rfl
lemma action_eq_snd_fst (r : Round 𝓞 𝓐 𝓨) : r.action = r.2.1 := rfl
lemma feedback_eq_snd_snd (r : Round 𝓞 𝓐 𝓨) : r.feedback = r.2.2 := rfl

@[fun_prop]
lemma measurable_obs : Measurable (Round.obs (𝓞 := 𝓞) (𝓐 := 𝓐) (𝓨 := 𝓨)) := measurable_fst

@[fun_prop]
lemma measurable_action : Measurable (Round.action (𝓞 := 𝓞) (𝓐 := 𝓐) (𝓨 := 𝓨)) :=
  measurable_snd.fst

@[fun_prop]
lemma measurable_feedback : Measurable (Round.feedback (𝓞 := 𝓞) (𝓐 := 𝓐) (𝓨 := 𝓨)) :=
  measurable_snd.snd

end Round

/-- History of `n` complete rounds; `n = 0` is the empty history. -/
abbrev Hist (𝓞 𝓐 𝓨 : Type*) (n : ℕ) := Fin n → Round 𝓞 𝓐 𝓨

/-- A stochastic, sequential algorithm.
At each round, it sees an observation in `𝓞`, then takes an action in `𝓐`, and finally receives
feedback in `𝓨`. The action is a random function of the past rounds and the current observation. -/
@[ext]
structure Algorithm (𝓞 𝓐 𝓨 : Type*) [MeasurableSpace 𝓞] [MeasurableSpace 𝓐] [MeasurableSpace 𝓨]
    where
  /-- Law of the action of round `n` given the past rounds and the current observation. -/
  policy : (n : ℕ) → Kernel (Hist 𝓞 𝓐 𝓨 n × 𝓞) 𝓐
  /-- The policy is a Markov kernel. -/
  [isMarkovKernel_policy : ∀ n, IsMarkovKernel (policy n)]

instance (alg : Algorithm 𝓞 𝓐 𝓨) (n : ℕ) : IsMarkovKernel (alg.policy n) :=
  alg.isMarkovKernel_policy n

instance : MeasurableSpace (Algorithm 𝓞 𝓐 𝓨) :=
  MeasurableSpace.comap (fun alg ↦ alg.policy) inferInstance

lemma measurable_algorithm_iff (f : Ω → Algorithm 𝓞 𝓐 𝓨) :
    Measurable f ↔ ∀ n, Measurable fun x ↦ (f x).policy n := by
  unfold instMeasurableSpaceAlgorithm
  rw [measurable_comap_iff, measurable_pi_iff]
  simp

/-- A stochastic environment.
At each round, an observation is drawn prior to the algorithm taking an action. Then the environment
provides feedback based on the observation and the action. -/
@[ext]
structure Environment (𝓞 𝓐 𝓨 : Type*) [MeasurableSpace 𝓞] [MeasurableSpace 𝓐] [MeasurableSpace 𝓨]
    where
  /-- Law of the observation of round `n` given the past rounds. -/
  obs : (n : ℕ) → Kernel (Hist 𝓞 𝓐 𝓨 n) 𝓞
  /-- Law of the feedback of round `n` given the past rounds, the observation and the action. -/
  feedback : (n : ℕ) → Kernel ((Hist 𝓞 𝓐 𝓨 n × 𝓞) × 𝓐) 𝓨
  /-- The observation kernel is a Markov kernel. -/
  [isMarkovKernel_obs : ∀ n, IsMarkovKernel (obs n)]
  /-- The feedback kernel is a Markov kernel. -/
  [isMarkovKernel_feedback : ∀ n, IsMarkovKernel (feedback n)]

instance (env : Environment 𝓞 𝓐 𝓨) (n : ℕ) : IsMarkovKernel (env.obs n) := env.isMarkovKernel_obs n
instance (env : Environment 𝓞 𝓐 𝓨) (n : ℕ) : IsMarkovKernel (env.feedback n) :=
  env.isMarkovKernel_feedback n

instance : MeasurableSpace (Environment 𝓞 𝓐 𝓨) :=
  MeasurableSpace.comap (fun env ↦ (env.obs, env.feedback)) inferInstance

lemma measurable_environment_iff (f : Ω → Environment 𝓞 𝓐 𝓨) :
    Measurable f ↔
      ∀ n, Measurable (fun x ↦ (f x).obs n) ∧ Measurable (fun x ↦ (f x).feedback n) := by
  simp_rw [measurable_comap_iff, measurable_fun_prod, forall_and, measurable_pi_iff,
    measurable_kernel_iff]
  rfl

/-- Distribution of the first observation: the observation kernel at time `0` applied to the empty
history. -/
def Environment.obs0 (env : Environment 𝓞 𝓐 𝓨) : Measure 𝓞 :=
  env.obs 0 default
deriving IsProbabilityMeasure

lemma Environment.obs0_def (env : Environment 𝓞 𝓐 𝓨) : env.obs0 = env.obs 0 default := rfl

lemma Environment.obs_zero (env : Environment 𝓞 𝓐 𝓨) (h : Hist 𝓞 𝓐 𝓨 0) :
    env.obs 0 h = env.obs0 := by
  rw [Unique.eq_default h]
  rfl

/-- Distribution of the first action given the first observation: the policy at time `0` applied to
the empty history. -/
noncomputable def Algorithm.p0 (alg : Algorithm 𝓞 𝓐 𝓨) : Kernel 𝓞 𝓐 :=
  (alg.policy 0).sectR default
deriving IsMarkovKernel

lemma Algorithm.p0_def (alg : Algorithm 𝓞 𝓐 𝓨) : alg.p0 = (alg.policy 0).sectR default := rfl

lemma Algorithm.p0_apply (alg : Algorithm 𝓞 𝓐 𝓨) (o : 𝓞) :
    alg.p0 o = alg.policy 0 (default, o) := rfl

lemma Algorithm.policy_zero (alg : Algorithm 𝓞 𝓐 𝓨) (h : Hist 𝓞 𝓐 𝓨 0) (o : 𝓞) :
    alg.policy 0 (h, o) = alg.p0 o := by
  rw [Unique.eq_default h]
  rfl

/-- Distribution of the first feedback given the first observation and action: the feedback kernel
at time `0` applied to the empty history. -/
noncomputable def Environment.ν0 (env : Environment 𝓞 𝓐 𝓨) : Kernel (𝓞 × 𝓐) 𝓨 :=
  (env.feedback 0).comap (fun p ↦ ((default, p.1), p.2)) (by fun_prop)
deriving IsMarkovKernel

lemma Environment.ν0_def (env : Environment 𝓞 𝓐 𝓨) :
    env.ν0 = (env.feedback 0).comap (fun p ↦ ((default, p.1), p.2)) (by fun_prop) := rfl

lemma Environment.ν0_apply (env : Environment 𝓞 𝓐 𝓨) (o : 𝓞) (a : 𝓐) :
    env.ν0 (o, a) = env.feedback 0 ((default, o), a) := rfl

lemma Environment.feedback_zero (env : Environment 𝓞 𝓐 𝓨) (h : Hist 𝓞 𝓐 𝓨 0) (o : 𝓞) (a : 𝓐) :
    env.feedback 0 ((h, o), a) = env.ν0 (o, a) := by
  rw [Unique.eq_default h]
  rfl

/-- Kernel describing the distribution of the round at time `n` given the history before `n`. -/
noncomputable
def stepKernel (alg : Algorithm 𝓞 𝓐 𝓨) (env : Environment 𝓞 𝓐 𝓨) (n : ℕ) :
    Kernel (Hist 𝓞 𝓐 𝓨 n) (Round 𝓞 𝓐 𝓨) :=
  env.obs n ⊗ₖ (alg.policy n ⊗ₖ env.feedback n)
deriving IsMarkovKernel

lemma stepKernel_def (alg : Algorithm 𝓞 𝓐 𝓨) (env : Environment 𝓞 𝓐 𝓨) (n : ℕ) :
    stepKernel alg env n = env.obs n ⊗ₖ (alg.policy n ⊗ₖ env.feedback n) := rfl

@[simp]
lemma fst_stepKernel (alg : Algorithm 𝓞 𝓐 𝓨) (env : Environment 𝓞 𝓐 𝓨) (n : ℕ) :
    (stepKernel alg env n).fst = env.obs n := by
  rw [stepKernel, Kernel.fst_compProd]

lemma stepKernel_zero (alg : Algorithm 𝓞 𝓐 𝓨) (env : Environment 𝓞 𝓐 𝓨) (h : Hist 𝓞 𝓐 𝓨 0) :
    stepKernel alg env 0 h = env.obs0 ⊗ₘ (alg.p0 ⊗ₖ env.ν0) := by
  rw [Unique.eq_default h, stepKernel, Kernel.compProd_apply_eq_compProd_sectR]
  congr 1
  ext o s hs
  rw [Kernel.sectR_apply, Kernel.compProd_apply hs, Kernel.compProd_apply hs]
  rfl

section IsAlgEnvSeq

variable {O : ℕ → Ω → 𝓞} {A : ℕ → Ω → 𝓐} {Y : ℕ → Ω → 𝓨}
    {alg : Algorithm 𝓞 𝓐 𝓨} {env : Environment 𝓞 𝓐 𝓨}
    {P : Measure Ω} [IsFiniteMeasure P] {N : ℕ}

/-- Step of the algorithm-environment sequence: the round at time `n`. -/
def step (O : ℕ → Ω → 𝓞) (A : ℕ → Ω → 𝓐) (Y : ℕ → Ω → 𝓨) (n : ℕ) (ω : Ω) : Round 𝓞 𝓐 𝓨 :=
  (O n ω, A n ω, Y n ω)

@[simp] lemma obs_step (n : ℕ) (ω : Ω) : (step O A Y n ω).obs = O n ω := rfl
@[simp] lemma action_step (n : ℕ) (ω : Ω) : (step O A Y n ω).action = A n ω := rfl
@[simp] lemma feedback_step (n : ℕ) (ω : Ω) : (step O A Y n ω).feedback = Y n ω := rfl

@[fun_prop]
lemma measurable_step (n : ℕ) (hO : Measurable (O n)) (hA : Measurable (A n))
    (hY : Measurable (Y n)) :
    Measurable (step O A Y n) := by
  unfold step
  fun_prop

/-- A random variable that gives the sequence of rounds. -/
def trajectory (O : ℕ → Ω → 𝓞) (A : ℕ → Ω → 𝓐) (Y : ℕ → Ω → 𝓨) (ω : Ω) : ℕ → Round 𝓞 𝓐 𝓨 :=
  fun n ↦ (O n ω, A n ω, Y n ω)

@[fun_prop]
lemma measurable_trajectory {O : ℕ → Ω → 𝓞} {A : ℕ → Ω → 𝓐} {Y : ℕ → Ω → 𝓨}
    (hO : ∀ n, Measurable (O n)) (hA : ∀ n, Measurable (A n))
    (hR : ∀ n, Measurable (Y n)) : Measurable (trajectory O A Y) := by
  unfold trajectory
  fun_prop

/-- History of the algorithm-environment sequence before time `n`: the rounds at
times `0, ..., n - 1`. -/
def history (O : ℕ → Ω → 𝓞) (A : ℕ → Ω → 𝓐) (Y : ℕ → Ω → 𝓨) (n : ℕ) (ω : Ω) : Hist 𝓞 𝓐 𝓨 n :=
  fun i ↦ (O i ω, A i ω, Y i ω)

lemma history_apply (n : ℕ) (ω : Ω) (i : Fin n) :
    history O A Y n ω i = (O i ω, A i ω, Y i ω) := rfl

@[simp]
lemma history_zero : history O A Y 0 = fun _ ↦ (default : Hist 𝓞 𝓐 𝓨 0) := Unique.eq_default _

@[fun_prop]
lemma measurable_history (hO : ∀ n, Measurable (O n)) (hA : ∀ n, Measurable (A n))
    (hY : ∀ n, Measurable (Y n)) (n : ℕ) :
    Measurable (history O A Y n) := by
  unfold history
  fun_prop

lemma history_congr {O' : ℕ → Ω → 𝓞} {A' : ℕ → Ω → 𝓐} {Y' : ℕ → Ω → 𝓨} {P : Measure Ω}
    (hO : ∀ n, O' n =ᵐ[P] O n) (hA : ∀ n, A' n =ᵐ[P] A n) (hY : ∀ n, Y' n =ᵐ[P] Y n) (n : ℕ) :
    history O' A' Y' n =ᵐ[P] history O A Y n := by
  have h : ∀ᵐ ω ∂P, ∀ i : Fin n, (O' i ω = O i ω ∧ A' i ω = A i ω) ∧ Y' i ω = Y i ω := by
    rw [ae_all_iff]
    exact fun i ↦ ((hO i).and (hA i)).and (hY i)
  filter_upwards [h] with ω hω
  funext i
  rw [history_apply, history_apply, (hω i).1.1, (hω i).1.2, (hω i).2]

lemma eval_comp_history (n : ℕ) :
    (fun x ↦ x (Fin.last n)) ∘ (history O A Y (n + 1)) = step O A Y n := rfl

lemma obs_eval_comp_history (n : ℕ) :
    (fun x ↦ (x (Fin.last n)).obs) ∘ (history O A Y (n + 1)) = O n := rfl

lemma action_eval_comp_history (n : ℕ) :
    (fun x ↦ (x (Fin.last n)).action) ∘ (history O A Y (n + 1)) = A n := rfl

lemma feedback_eval_comp_history (n : ℕ) :
    (fun x ↦ (x (Fin.last n)).feedback) ∘ (history O A Y (n + 1)) = Y n := rfl

/-- The history before time `m` is a restriction of the history before time `n ≥ m`. -/
lemma history_eq_comp_history {m n : ℕ} (hmn : m ≤ n) :
    history O A Y m = (fun h (i : Fin m) ↦ h (Fin.castLE hmn i)) ∘ history O A Y n := rfl

lemma history_succ (n : ℕ) :
    history O A Y (n + 1) =
      (MeasurableEquiv.finSuccProd (Round 𝓞 𝓐 𝓨) n).symm ∘
        (fun ω ↦ (history O A Y n ω, step O A Y n ω)) := by
  funext ω
  simp only [Function.comp_apply, MeasurableEquiv.finSuccProd_symm_apply]
  funext i
  refine Fin.lastCases ?_ (fun i ↦ ?_) i
  · simp [history, step]
  · simp [history]

/-- An algorithm-environment sequence: a sequence of observations, actions and feedbacks generated
by an algorithm interacting with an environment. -/
structure IsAlgEnvSeq
    (O : ℕ → Ω → 𝓞) (A : ℕ → Ω → 𝓐) (Y : ℕ → Ω → 𝓨)
    (alg : Algorithm 𝓞 𝓐 𝓨) (env : Environment 𝓞 𝓐 𝓨)
    (P : Measure Ω) [IsFiniteMeasure P] : Prop where
  /-- The observation sequence is measurable. -/
  measurable_obs n : Measurable (O n) := by fun_prop
  /-- The action sequence is measurable. -/
  measurable_action n : Measurable (A n) := by fun_prop
  /-- The feedback sequence is measurable. -/
  measurable_feedback n : Measurable (Y n) := by fun_prop
  /-- The observation at time `n` has the correct conditional distribution given the history. -/
  hasCondDistrib_obs n :
    HasCondDistrib (O n) (history O A Y n) (env.obs n) P
  /-- The action at time `n` has the correct conditional distribution given the history and the
  observation at time `n`. -/
  hasCondDistrib_action n :
    HasCondDistrib (A n) (fun ω ↦ (history O A Y n ω, O n ω)) (alg.policy n) P
  /-- The feedback at time `n` has the correct conditional distribution given the history, the
  observation and the action at time `n`. -/
  hasCondDistrib_feedback n :
    HasCondDistrib (Y n) (fun ω ↦ ((history O A Y n ω, O n ω), A n ω)) (env.feedback n) P

/-- An algorithm-environment sequence until time `N`: a sequence of observations, actions and
feedbacks such that the rounds at times `n < N` are generated by an algorithm interacting with
an environment. In particular, the law of `history O A Y N` is determined. -/
structure IsAlgEnvSeqUntil
    (O : ℕ → Ω → 𝓞) (A : ℕ → Ω → 𝓐) (Y : ℕ → Ω → 𝓨)
    (alg : Algorithm 𝓞 𝓐 𝓨) (env : Environment 𝓞 𝓐 𝓨)
    (P : Measure Ω) [IsFiniteMeasure P] (N : ℕ) : Prop where
  /-- The observation sequence is measurable. -/
  measurable_obs n : Measurable (O n) := by fun_prop
  /-- The action sequence is measurable. -/
  measurable_action n : Measurable (A n) := by fun_prop
  /-- The feedback sequence is measurable. -/
  measurable_feedback n : Measurable (Y n) := by fun_prop
  /-- The observation at time `n < N` has the correct conditional distribution given the history. -/
  hasCondDistrib_obs n (hn : n < N) :
    HasCondDistrib (O n) (history O A Y n) (env.obs n) P
  /-- The action at time `n < N` has the correct conditional distribution given the history and the
  observation at time `n`. -/
  hasCondDistrib_action n (hn : n < N) :
    HasCondDistrib (A n) (fun ω ↦ (history O A Y n ω, O n ω)) (alg.policy n) P
  /-- The feedback at time `n < N` has the correct conditional distribution given the history, the
  observation and the action at time `n`. -/
  hasCondDistrib_feedback n (hn : n < N) :
    HasCondDistrib (Y n) (fun ω ↦ ((history O A Y n ω, O n ω), A n ω)) (env.feedback n) P

lemma IsAlgEnvSeqUntil.mono (h : IsAlgEnvSeqUntil O A Y alg env P N) {N' : ℕ} (hN : N' ≤ N) :
    IsAlgEnvSeqUntil O A Y alg env P N' where
  measurable_obs := h.measurable_obs
  measurable_action := h.measurable_action
  measurable_feedback := h.measurable_feedback
  hasCondDistrib_obs n hn := h.hasCondDistrib_obs n (hn.trans_le hN)
  hasCondDistrib_action n hn := h.hasCondDistrib_action n (hn.trans_le hN)
  hasCondDistrib_feedback n hn := h.hasCondDistrib_feedback n (hn.trans_le hN)

lemma IsAlgEnvSeq.isAlgEnvSeqUntil (h : IsAlgEnvSeq O A Y alg env P) (N : ℕ) :
    IsAlgEnvSeqUntil O A Y alg env P N where
  measurable_obs := h.measurable_obs
  measurable_action := h.measurable_action
  measurable_feedback := h.measurable_feedback
  hasCondDistrib_obs n _ := h.hasCondDistrib_obs n
  hasCondDistrib_action n _ := h.hasCondDistrib_action n
  hasCondDistrib_feedback n _ := h.hasCondDistrib_feedback n

/-- `IsAlgEnvSeq` only depends on the almost everywhere equivalence classes of the processes. -/
lemma IsAlgEnvSeq.congr {O' : ℕ → Ω → 𝓞} {A' : ℕ → Ω → 𝓐} {Y' : ℕ → Ω → 𝓨}
    (h : IsAlgEnvSeq O A Y alg env P)
    (hO' : ∀ n, Measurable (O' n)) (hA' : ∀ n, Measurable (A' n)) (hY' : ∀ n, Measurable (Y' n))
    (hO : ∀ n, O' n =ᵐ[P] O n) (hA : ∀ n, A' n =ᵐ[P] A n) (hY : ∀ n, Y' n =ᵐ[P] Y n) :
    IsAlgEnvSeq O' A' Y' alg env P where
  measurable_obs := hO'
  measurable_action := hA'
  measurable_feedback := hY'
  hasCondDistrib_obs n :=
    HasCondDistrib.congr (h.hasCondDistrib_obs n) (history_congr hO hA hY n) (hO n)
  hasCondDistrib_action n := by
    refine HasCondDistrib.congr (h.hasCondDistrib_action n) ?_ (hA n)
    filter_upwards [history_congr hO hA hY n, hO n] with ω h1 h2
    rw [h1, h2]
  hasCondDistrib_feedback n := by
    refine HasCondDistrib.congr (h.hasCondDistrib_feedback n) ?_ (hY n)
    filter_upwards [history_congr hO hA hY n, hO n, hA n] with ω h1 h2 h3
    rw [h1, h2, h3]

lemma isAlgEnvSeq_iff_forall_isAlgEnvSeqUntil :
    IsAlgEnvSeq O A Y alg env P ↔ ∀ N, IsAlgEnvSeqUntil O A Y alg env P N where
  mp h N := h.isAlgEnvSeqUntil N
  mpr h := {
    measurable_obs := (h 0).measurable_obs
    measurable_action := (h 0).measurable_action
    measurable_feedback := (h 0).measurable_feedback
    hasCondDistrib_obs n := (h (n + 1)).hasCondDistrib_obs n n.lt_succ_self
    hasCondDistrib_action n := (h (n + 1)).hasCondDistrib_action n n.lt_succ_self
    hasCondDistrib_feedback n := (h (n + 1)).hasCondDistrib_feedback n n.lt_succ_self }

lemma IsAlgEnvSeq.measurable_step (h : IsAlgEnvSeq O A Y alg env P) (n : ℕ) :
    Measurable (step O A Y n) := by
  have hO := h.measurable_obs
  have hA := h.measurable_action
  have hY := h.measurable_feedback
  fun_prop

lemma IsAlgEnvSeq.measurable_history (h : IsAlgEnvSeq O A Y alg env P) (n : ℕ) :
    Measurable (history O A Y n) := by
  have hO := h.measurable_obs
  have hA := h.measurable_action
  have hY := h.measurable_feedback
  fun_prop

lemma IsAlgEnvSeq.measurable_trajectory (h : IsAlgEnvSeq O A Y alg env P) :
    Measurable (trajectory O A Y) :=
  Learning.measurable_trajectory h.measurable_obs h.measurable_action h.measurable_feedback

lemma IsAlgEnvSeqUntil.measurable_step (h : IsAlgEnvSeqUntil O A Y alg env P N) (n : ℕ) :
    Measurable (step O A Y n) := by
  have hO := h.measurable_obs
  have hA := h.measurable_action
  have hY := h.measurable_feedback
  fun_prop

lemma IsAlgEnvSeqUntil.measurable_history (h : IsAlgEnvSeqUntil O A Y alg env P N) (n : ℕ) :
    Measurable (history O A Y n) := by
  have hO := h.measurable_obs
  have hA := h.measurable_action
  have hY := h.measurable_feedback
  fun_prop

lemma IsAlgEnvSeq.hasCondDistrib_step (h : IsAlgEnvSeq O A Y alg env P) (n : ℕ) :
    HasCondDistrib (step O A Y n) (history O A Y n) (stepKernel alg env n) P :=
  (h.hasCondDistrib_obs n).prod
    ((h.hasCondDistrib_action n).prod (h.hasCondDistrib_feedback n))

lemma IsAlgEnvSeqUntil.hasCondDistrib_step (h : IsAlgEnvSeqUntil O A Y alg env P N)
    (n : ℕ) (hn : n < N) :
    HasCondDistrib (step O A Y n) (history O A Y n) (stepKernel alg env n) P :=
  (h.hasCondDistrib_obs n hn).prod
    ((h.hasCondDistrib_action n hn).prod (h.hasCondDistrib_feedback n hn))

section Zero

/-! ### Laws at time `0`

At time `0` the history is the unique element of `Hist 𝓞 𝓐 𝓨 0`: conditioning on it is the same
as not conditioning. Note that those results need `P` to be a probability measure. -/

variable [IsProbabilityMeasure P]

omit [IsFiniteMeasure P] in
/-- The history before time `0` is a constant. -/
lemma hasLaw_history_zero (O : ℕ → Ω → 𝓞) (A : ℕ → Ω → 𝓐) (Y : ℕ → Ω → 𝓨) :
    HasLaw (history O A Y 0) (Measure.dirac default) P where
  aemeasurable := by rw [history_zero]; exact aemeasurable_const
  map_eq := by rw [history_zero, Measure.map_const, measure_univ, one_smul]

lemma IsAlgEnvSeqUntil.hasLaw_obs_zero (h : IsAlgEnvSeqUntil O A Y alg env P N) (hN : 0 < N) :
    HasLaw (O 0) env.obs0 P := by
  have h0 := h.hasCondDistrib_obs 0 hN
  rw [history_zero] at h0
  exact h0.hasLaw_of_const'

lemma IsAlgEnvSeq.hasLaw_obs_zero (h : IsAlgEnvSeq O A Y alg env P) :
    HasLaw (O 0) env.obs0 P :=
  (h.isAlgEnvSeqUntil 1).hasLaw_obs_zero zero_lt_one

omit [IsProbabilityMeasure P] in
lemma IsAlgEnvSeqUntil.hasCondDistrib_action_zero (h : IsAlgEnvSeqUntil O A Y alg env P N)
    (hN : 0 < N) :
    HasCondDistrib (A 0) (O 0) alg.p0 P :=
  hasCondDistrib_prodMk_left_unique_iff.mp (h.hasCondDistrib_action 0 hN)

omit [IsProbabilityMeasure P] in
lemma IsAlgEnvSeq.hasCondDistrib_action_zero (h : IsAlgEnvSeq O A Y alg env P) :
    HasCondDistrib (A 0) (O 0) alg.p0 P :=
  (h.isAlgEnvSeqUntil 1).hasCondDistrib_action_zero zero_lt_one

omit [IsProbabilityMeasure P] in
lemma IsAlgEnvSeqUntil.hasCondDistrib_feedback_zero (h : IsAlgEnvSeqUntil O A Y alg env P N)
    (hN : 0 < N) :
    HasCondDistrib (Y 0) (fun ω ↦ (O 0 ω, A 0 ω)) env.ν0 P := by
  have h0 := h.hasCondDistrib_feedback 0 hN
  rw [history_zero] at h0
  exact h0.of_measurableEmbedding_comp_right
    ((measurableEmbedding_prodMk_left (default : Hist 𝓞 𝓐 𝓨 0)).prodMap .id)

omit [IsProbabilityMeasure P] in
lemma IsAlgEnvSeq.hasCondDistrib_feedback_zero (h : IsAlgEnvSeq O A Y alg env P) :
    HasCondDistrib (Y 0) (fun ω ↦ (O 0 ω, A 0 ω)) env.ν0 P :=
  (h.isAlgEnvSeqUntil 1).hasCondDistrib_feedback_zero zero_lt_one

lemma IsAlgEnvSeqUntil.hasLaw_step_zero (h : IsAlgEnvSeqUntil O A Y alg env P N) (hN : 0 < N) :
    HasLaw (step O A Y 0) (env.obs0 ⊗ₘ (alg.p0 ⊗ₖ env.ν0)) P := by
  have h0 := h.hasCondDistrib_step 0 hN
  rw [history_zero] at h0
  rw [← stepKernel_zero alg env default]
  exact h0.hasLaw_of_const'

lemma IsAlgEnvSeq.hasLaw_step_zero (h : IsAlgEnvSeq O A Y alg env P) :
    HasLaw (step O A Y 0) (env.obs0 ⊗ₘ (alg.p0 ⊗ₖ env.ν0)) P :=
  (h.isAlgEnvSeqUntil 1).hasLaw_step_zero zero_lt_one

end Zero

lemma IsAlgEnvSeq.hasLaw_obs_comp (h : IsAlgEnvSeq O A Y alg env P) (n : ℕ) :
    HasLaw (O n) (env.obs n ∘ₘ (P.map (history O A Y n))) P :=
  HasCondDistrib.hasLaw_comp (h.hasCondDistrib_obs n)

lemma IsAlgEnvSeq.hasLaw_action_comp (h : IsAlgEnvSeq O A Y alg env P) (n : ℕ) :
    HasLaw (A n) (alg.policy n ∘ₘ (P.map fun ω ↦ (history O A Y n ω, O n ω))) P :=
  HasCondDistrib.hasLaw_comp (h.hasCondDistrib_action n)

lemma IsAlgEnvSeq.hasLaw_feedback_comp (h : IsAlgEnvSeq O A Y alg env P) (n : ℕ) :
    HasLaw (Y n)
      ((env.feedback n) ∘ₘ (P.map fun ω ↦ ((history O A Y n ω, O n ω), A n ω))) P :=
  HasCondDistrib.hasLaw_comp (h.hasCondDistrib_feedback n)

lemma IsAlgEnvSeq.hasLaw_step_comp (h : IsAlgEnvSeq O A Y alg env P) (n : ℕ) :
    HasLaw (step O A Y n) (stepKernel alg env n ∘ₘ (P.map (history O A Y n))) P :=
  HasCondDistrib.hasLaw_comp (h.hasCondDistrib_step n)

/-- Conditionally on the event `(O 0, A 0) = p`, the first feedback has law `env.ν0 p`. -/
lemma IsAlgEnvSeq.hasLaw_feedback_zero_cond [MeasurableSingletonClass 𝓞]
    [MeasurableSingletonClass 𝓐] (h : IsAlgEnvSeq O A Y alg env P) {p : 𝓞 × 𝓐}
    (hP : P ((fun ω ↦ (O 0 ω, A 0 ω)) ⁻¹' {p}) ≠ 0) :
    HasLaw (Y 0) (env.ν0 p) P[|(fun ω ↦ (O 0 ω, A 0 ω)) ⁻¹' {p}] :=
  h.hasCondDistrib_feedback_zero.hasLaw_cond (h.measurable_feedback 0)
    (measurableSet_singleton p) (fun a ha ↦ by rw [Set.mem_singleton_iff.1 ha]) hP

section Filtration

namespace IsAlgEnvSeq

/-- Filtration generated by the history: `h.filtration n` is the σ-algebra generated by
`history O A Y n`, that is by the `n` rounds at times `0, ..., n - 1`. In particular
`h.filtration 0` is the trivial σ-algebra, and the round at time `n` is measurable with respect
to `h.filtration (n + 1)`. -/
def filtration (h : IsAlgEnvSeq O A Y alg env P) :
    Filtration ℕ mΩ where
  seq n := MeasurableSpace.comap (history O A Y n) inferInstance
  mono' i j hij := by
    simp only
    rw [← measurable_iff_comap_le, history_eq_comp_history hij]
    exact measurable_comp_comap _ (by fun_prop)
  le' i := by
    rw [← measurable_iff_comap_le]
    exact Learning.measurable_history h.measurable_obs h.measurable_action h.measurable_feedback _

lemma filtration_eq_comap (h : IsAlgEnvSeq O A Y alg env P) (n : ℕ) :
    h.filtration n = MeasurableSpace.comap (history O A Y n) inferInstance := rfl

lemma filtration_zero (h : IsAlgEnvSeq O A Y alg env P) : h.filtration 0 = ⊥ := by
  rw [filtration_eq_comap, history_zero, MeasurableSpace.comap_const]

lemma measurable_history_filtration (h : IsAlgEnvSeq O A Y alg env P) (n : ℕ) :
    Measurable[h.filtration n] (history O A Y n) :=
  measurable_iff_comap_le.mpr le_rfl

lemma adapted_history (h : IsAlgEnvSeq O A Y alg env P) :
    Adapted h.filtration (history O A Y) :=
  h.measurable_history_filtration

lemma measurable_step_filtration_succ (h : IsAlgEnvSeq O A Y alg env P) (n : ℕ) :
    Measurable[h.filtration (n + 1)] (step O A Y n) := by
  rw [← eval_comp_history]
  exact measurable_comp_comap _ (by fun_prop)

lemma measurable_obs_filtration_succ (h : IsAlgEnvSeq O A Y alg env P) (n : ℕ) :
    Measurable[h.filtration (n + 1)] (O n) := by
  rw [← obs_eval_comp_history (O := O) (A := A) (Y := Y) n]
  exact measurable_comp_comap _ (by fun_prop)

lemma measurable_action_filtration_succ (h : IsAlgEnvSeq O A Y alg env P) (n : ℕ) :
    Measurable[h.filtration (n + 1)] (A n) := by
  rw [← action_eval_comp_history (O := O) (A := A) (Y := Y) n]
  exact measurable_comp_comap _ (by fun_prop)

lemma measurable_feedback_filtration_succ (h : IsAlgEnvSeq O A Y alg env P) (n : ℕ) :
    Measurable[h.filtration (n + 1)] (Y n) := by
  rw [← feedback_eval_comp_history (O := O) (A := A) (Y := Y) n]
  exact measurable_comp_comap _ (by fun_prop)

/-- The step at time `m < n` is measurable with respect to the σ-algebra generated by the first
`n` rounds. -/
lemma measurable_step_filtration_of_lt (h : IsAlgEnvSeq O A Y alg env P) {m n : ℕ} (hmn : m < n) :
    Measurable[h.filtration n] (step O A Y m) :=
  (h.measurable_step_filtration_succ m).mono (h.filtration.mono hmn) le_rfl

lemma measurable_obs_filtration_of_lt (h : IsAlgEnvSeq O A Y alg env P) {m n : ℕ} (hmn : m < n) :
    Measurable[h.filtration n] (O m) :=
  (h.measurable_obs_filtration_succ m).mono (h.filtration.mono hmn) le_rfl

lemma measurable_action_filtration_of_lt (h : IsAlgEnvSeq O A Y alg env P) {m n : ℕ}
    (hmn : m < n) :
    Measurable[h.filtration n] (A m) :=
  (h.measurable_action_filtration_succ m).mono (h.filtration.mono hmn) le_rfl

lemma measurable_feedback_filtration_of_lt (h : IsAlgEnvSeq O A Y alg env P) {m n : ℕ}
    (hmn : m < n) :
    Measurable[h.filtration n] (Y m) :=
  (h.measurable_feedback_filtration_succ m).mono (h.filtration.mono hmn) le_rfl

/-- Filtration generated by the history before time `n` together with the observation at
time `n`. -/
def filtrationObs (h : IsAlgEnvSeq O A Y alg env P) :
    Filtration ℕ mΩ where
  seq n := MeasurableSpace.comap (fun ω ↦ (history O A Y n ω, O n ω)) inferInstance
  mono' n m hnm := by
    simp only
    rw [← measurable_iff_comap_le]
    rcases eq_or_lt_of_le hnm with rfl | hlt
    · exact measurable_iff_comap_le.mpr le_rfl
    have : (fun ω ↦ (history O A Y n ω, O n ω)) =
        (fun p : Hist 𝓞 𝓐 𝓨 m × 𝓞 ↦
          (fun i : Fin n ↦ p.1 (Fin.castLE hnm i), (p.1 ⟨n, hlt⟩).obs)) ∘
        (fun ω ↦ (history O A Y m ω, O m ω)) := rfl
    rw [this]
    exact measurable_comp_comap _ (by fun_prop)
  le' n := by
    rw [← measurable_iff_comap_le]
    exact (Learning.measurable_history h.measurable_obs h.measurable_action
      h.measurable_feedback n).prodMk (h.measurable_obs n)

lemma filtrationObs_eq_comap (h : IsAlgEnvSeq O A Y alg env P) (n : ℕ) :
    h.filtrationObs n =
      MeasurableSpace.comap (fun ω ↦ (history O A Y n ω, O n ω)) inferInstance := rfl

/-- Filtration generated by the history before time `n` together with the observation and the
action at time `n`. -/
def filtrationAction (h : IsAlgEnvSeq O A Y alg env P) :
    Filtration ℕ mΩ where
  seq n := MeasurableSpace.comap (fun ω ↦ ((history O A Y n ω, O n ω), A n ω)) inferInstance
  mono' n m hnm := by
    simp only
    rw [← measurable_iff_comap_le]
    rcases eq_or_lt_of_le hnm with rfl | hlt
    · exact measurable_iff_comap_le.mpr le_rfl
    have : (fun ω ↦ ((history O A Y n ω, O n ω), A n ω)) =
        (fun p : (Hist 𝓞 𝓐 𝓨 m × 𝓞) × 𝓐 ↦
          ((fun i : Fin n ↦ p.1.1 (Fin.castLE hnm i), (p.1.1 ⟨n, hlt⟩).obs),
            (p.1.1 ⟨n, hlt⟩).action)) ∘
        (fun ω ↦ ((history O A Y m ω, O m ω), A m ω)) := rfl
    rw [this]
    exact measurable_comp_comap _ (by fun_prop)
  le' n := by
    rw [← measurable_iff_comap_le]
    exact ((Learning.measurable_history h.measurable_obs h.measurable_action
      h.measurable_feedback n).prodMk (h.measurable_obs n)).prodMk (h.measurable_action n)

lemma filtrationAction_eq_comap (h : IsAlgEnvSeq O A Y alg env P) (n : ℕ) :
    h.filtrationAction n =
      MeasurableSpace.comap (fun ω ↦ ((history O A Y n ω, O n ω), A n ω)) inferInstance := rfl

lemma filtrationObs_zero_eq_comap (h : IsAlgEnvSeq O A Y alg env P) :
    h.filtrationObs 0 = MeasurableSpace.comap (O 0) inferInstance := by
  rw [filtrationObs_eq_comap, history_zero]
  refine le_antisymm ?_ ?_
  · rw [← measurable_iff_comap_le]
    exact measurable_const.prodMk (measurable_iff_comap_le.mpr le_rfl)
  · rw [← measurable_iff_comap_le]
    exact measurable_snd.comp (measurable_iff_comap_le.mpr le_rfl)

lemma filtrationAction_zero_eq_comap (h : IsAlgEnvSeq O A Y alg env P) :
    h.filtrationAction 0 =
      MeasurableSpace.comap (fun ω ↦ (O 0 ω, A 0 ω)) inferInstance := by
  rw [filtrationAction_eq_comap, history_zero]
  refine le_antisymm ?_ ?_
  · rw [← measurable_iff_comap_le]
    exact ((measurable_const.prodMk
      (measurable_fst.comp (measurable_iff_comap_le.mpr le_rfl)))).prodMk
      (measurable_snd.comp (measurable_iff_comap_le.mpr le_rfl))
  · rw [← measurable_iff_comap_le]
    exact ((measurable_snd.comp measurable_fst).comp
      (measurable_iff_comap_le.mpr le_rfl)).prodMk
      (measurable_snd.comp (measurable_iff_comap_le.mpr le_rfl))

lemma measurable_history_filtrationObs (h : IsAlgEnvSeq O A Y alg env P) (n : ℕ) :
    Measurable[h.filtrationObs n] (history O A Y n) :=
  measurable_fst.comp (measurable_iff_comap_le.mpr le_rfl)

@[fun_prop]
lemma measurable_history_filtrationAction (h : IsAlgEnvSeq O A Y alg env P) (n : ℕ) :
    Measurable[h.filtrationAction n] (history O A Y n) :=
  (measurable_fst.comp measurable_fst).comp (measurable_iff_comap_le.mpr le_rfl)

lemma filtrationObs_le_filtrationAction (h : IsAlgEnvSeq O A Y alg env P) (n : ℕ) :
    h.filtrationObs n ≤ h.filtrationAction n := by
  rw [filtrationObs_eq_comap, filtrationAction_eq_comap, ← measurable_iff_comap_le]
  exact measurable_fst.comp (measurable_iff_comap_le.mpr le_rfl)

lemma filtration_le_filtrationObs (h : IsAlgEnvSeq O A Y alg env P) (n : ℕ) :
    h.filtration n ≤ h.filtrationObs n :=
  measurable_iff_comap_le.mp (h.measurable_history_filtrationObs n)

lemma filtration_le_filtrationAction (h : IsAlgEnvSeq O A Y alg env P) (n : ℕ) :
    h.filtration n ≤ h.filtrationAction n :=
  measurable_iff_comap_le.mp (h.measurable_history_filtrationAction n)

lemma filtrationAction_le_filtration_succ (h : IsAlgEnvSeq O A Y alg env P) (n : ℕ) :
    h.filtrationAction n ≤ h.filtration (n + 1) := by
  rw [filtrationAction_eq_comap, ← measurable_iff_comap_le]
  exact ((h.adapted_history.measurable_le n.le_succ).prodMk
    (h.measurable_obs_filtration_succ n)).prodMk (h.measurable_action_filtration_succ n)

lemma filtrationObs_le_filtration_succ (h : IsAlgEnvSeq O A Y alg env P) (n : ℕ) :
    h.filtrationObs n ≤ h.filtration (n + 1) :=
  (h.filtrationObs_le_filtrationAction n).trans (h.filtrationAction_le_filtration_succ n)

lemma adapted_obs_filtrationObs (h : IsAlgEnvSeq O A Y alg env P) :
    Adapted h.filtrationObs O := fun _ ↦
  measurable_snd.comp (measurable_iff_comap_le.mpr le_rfl)

lemma adapted_obs_filtrationAction (h : IsAlgEnvSeq O A Y alg env P) :
    Adapted h.filtrationAction O := fun _ ↦
  (measurable_snd.comp measurable_fst).comp (measurable_iff_comap_le.mpr le_rfl)

lemma adapted_action_filtrationAction (h : IsAlgEnvSeq O A Y alg env P) :
    Adapted h.filtrationAction A := fun _ ↦
  measurable_snd.comp (measurable_iff_comap_le.mpr le_rfl)

lemma measurable_feedback_filtrationObs_of_lt (h : IsAlgEnvSeq O A Y alg env P)
    {m n : ℕ} (hmn : m < n) :
    Measurable[h.filtrationObs n] (Y m) :=
  (h.measurable_feedback_filtration_of_lt hmn).mono (filtration_le_filtrationObs h n) le_rfl

lemma measurable_feedback_filtrationAction_of_lt (h : IsAlgEnvSeq O A Y alg env P)
    {m n : ℕ} (hmn : m < n) :
    Measurable[h.filtrationAction n] (Y m) :=
  (h.measurable_feedback_filtration_of_lt hmn).mono (filtration_le_filtrationAction h n) le_rfl

end IsAlgEnvSeq

end Filtration

end IsAlgEnvSeq

section NoObservation

/-! ### Environments without observations

An environment with `𝓞 = Unit` provides no information to the algorithm before it takes its action:
the algorithm only sees the past rounds. Since `Unit` carries a unique probability measure, the
observation kernels of such an environment are all equal to `Kernel.const _ (Measure.dirac ())`,
and the observation process of an algorithm-environment sequence is `noObs`. -/

/-- Every Markov kernel with codomain `Unit` is the constant kernel at `Measure.dirac ()`. -/
lemma Kernel.eq_const_dirac_unit {α : Type*} {mα : MeasurableSpace α} (κ : Kernel α Unit)
    [IsMarkovKernel κ] :
    κ = Kernel.const α (Measure.dirac ()) := by
  ext a s hs
  rw [Kernel.const_apply]
  rcases Set.eq_empty_or_nonempty s with rfl | ⟨u, hu⟩
  · simp
  · have hs_univ : s = Set.univ := Set.eq_univ_of_forall fun x ↦ by rwa [Subsingleton.elim x u]
    simp [hs_univ]

/-- A random variable with values in `Unit` admits any Markov kernel as conditional
distribution. -/
lemma hasCondDistrib_unit {α : Type*} {mα : MeasurableSpace α} {P : Measure Ω}
    [IsProbabilityMeasure P] {X : Ω → α} (hX : AEMeasurable X P) (U : Ω → Unit)
    (κ : Kernel α Unit) [IsMarkovKernel κ] :
    HasCondDistrib U X κ P := by
  have hU : U = fun _ ↦ () := funext fun _ ↦ rfl
  subst hU
  refine HasLaw.mk (hX.prodMk aemeasurable_const) ?_
  rw [Kernel.eq_const_dirac_unit κ, Measure.compProd_const, Measure.prod_dirac,
    AEMeasurable.map_map_of_aemeasurable (by fun_prop) hX]
  rfl

/-- The observation process of an algorithm-environment sequence without observations. -/
def noObs (Ω : Type*) : ℕ → Ω → Unit := fun _ _ ↦ ()

@[simp] lemma noObs_apply (n : ℕ) (ω : Ω) : noObs Ω n ω = () := rfl

@[fun_prop]
lemma measurable_noObs (n : ℕ) : Measurable (noObs Ω n) := measurable_const

end NoObservation

end Learning