/-
Copyright (c) 2026 Scott Armstrong. All rights reserved.
Released under Apache 2.0 license as described in the file LICENSE.
Authors: Scott Armstrong
-/
import Definitions.Def_MarkovProcessModule_MarkovProcess_DenseTime_CemeteryExtension
import Definitions.Def_MarkovProcessModule_MarkovProcess_Path_Basic
import Mathlib.MeasureTheory.Constructions.Polish.Basic

/-!
# Paths with a lifetime

This file defines paths that are continuous in the live state space strictly before a possibly
infinite lifetime and are at the cemetery state from the lifetime onward.

The canonical measurable structure is generated by the lifetime and all fixed-time coordinates.
No countable-generation or standard-Borel claim for the lifetime-path carrier is made here.
-/

open MeasureTheory
open scoped ENNReal

namespace MarkovProcess

noncomputable section

/-- A path with a possibly infinite lifetime, continuous on its open live-time interval. -/
structure LifetimePath (α : Type*) [TopologicalSpace α] where
  /-- The lifetime, with `∞` denoting a nonexplosive path. -/
  lifetime : ENNReal
  /-- The live-state path, defined precisely at times strictly before the lifetime. -/
  livePath : {t : NNReal // (t : ENNReal) < lifetime} → α
  /-- The live-state path is continuous on its natural subtype domain. -/
  continuous_livePath : Continuous livePath

namespace LifetimePath

variable {α : Type*} [TopologicalSpace α]

/-- The coordinate of a lifetime path, equal to the live path before its lifetime and to the
cemetery state from its lifetime onward. -/
def coordinate (t : NNReal) (ω : LifetimePath α) : Cemetery α :=
  if h : (t : ENNReal) < ω.lifetime then Cemetery.alive (ω.livePath ⟨t, h⟩)
  else Cemetery.delta

@[simp]
theorem coordinate_of_lt (ω : LifetimePath α) (t : NNReal)
    (ht : (t : ENNReal) < ω.lifetime) :
    coordinate t ω = Cemetery.alive (ω.livePath ⟨t, ht⟩) := by
  rw [coordinate, dif_pos ht]

@[simp]
theorem coordinate_of_le (ω : LifetimePath α) (t : NNReal)
    (ht : ω.lifetime ≤ (t : ENNReal)) : coordinate t ω = Cemetery.delta := by
  rw [coordinate, dif_neg (not_lt_of_ge ht)]

/-- Coordinates strictly before the lifetime are live. -/
theorem coordinate_ne_delta_of_lt (ω : LifetimePath α) (t : NNReal)
    (ht : (t : ENNReal) < ω.lifetime) : coordinate t ω ≠ Cemetery.delta := by
  rw [coordinate_of_lt ω t ht]
  exact Sum.inl_ne_inr

/-- A coordinate is live exactly when its time is strictly before the lifetime. -/
theorem coordinate_ne_delta_iff (ω : LifetimePath α) (t : NNReal) :
    coordinate t ω ≠ Cemetery.delta ↔ (t : ENNReal) < ω.lifetime := by
  constructor
  · intro h
    by_contra hlt
    exact h (coordinate_of_le ω t (not_lt.mp hlt))
  · exact coordinate_ne_delta_of_lt ω t

/-- Two lifetime paths with the same lifetime and coordinates are equal. -/
theorem ext_coordinate {ω η : LifetimePath α} (hlifetime : ω.lifetime = η.lifetime)
    (hcoordinate : ∀ t, coordinate t ω = coordinate t η) : ω = η := by
  cases ω with
  | mk ωlifetime ωpath hωpath =>
    cases η with
    | mk ηlifetime ηpath hηpath =>
      dsimp only at hlifetime
      subst ηlifetime
      congr
      funext t
      have h := hcoordinate t.1
      rw [coordinate_of_lt _ _ t.property, coordinate_of_lt _ _ t.property] at h
      exact Sum.inl.inj h

/-- The cemetery-valued coordinate path is continuous on its live-time subtype. -/
theorem continuous_coordinate_before (ω : LifetimePath α) :
    Continuous (fun t : {t : NNReal // (t : ENNReal) < ω.lifetime} ↦
      coordinate t ω) := by
  have hcoordinate :
      (fun t : {t : NNReal // (t : ENNReal) < ω.lifetime} ↦ coordinate t ω) =
        Cemetery.alive ∘ ω.livePath := by
    funext t
    exact coordinate_of_lt ω t t.property
  rw [hcoordinate]
  exact continuous_inl.comp ω.continuous_livePath

/-- Restrict a lifetime path to a compact time interval lying strictly before its lifetime. -/
def restrictBefore (ω : LifetimePath α) (T : NNReal)
    (hT : (T : ENNReal) < ω.lifetime) : C(Set.Iic T, α) where
  toFun t := ω.livePath ⟨t, (ENNReal.coe_le_coe.mpr t.property).trans_lt hT⟩
  continuous_toFun := by
    have hinclusion :
        Continuous (fun t : Set.Iic T ↦
          (⟨t, (ENNReal.coe_le_coe.mpr t.property).trans_lt hT⟩ :
            {t : NNReal // (t : ENNReal) < ω.lifetime})) :=
      continuous_subtype_val.subtype_mk _
    exact ω.continuous_livePath.comp hinclusion

@[simp]
theorem restrictBefore_apply (ω : LifetimePath α) (T : NNReal)
    (hT : (T : ENNReal) < ω.lifetime) (t : Set.Iic T) :
    restrictBefore ω T hT t =
      ω.livePath ⟨t, (ENNReal.coe_le_coe.mpr t.property).trans_lt hT⟩ := by
  apply congrArg ω.livePath
  apply Subtype.ext
  rfl

/-- The measurable structure generated by the lifetime and all fixed-time coordinates. -/
instance instMeasurableSpace [MeasurableSpace α] : MeasurableSpace (LifetimePath α) :=
  MeasurableSpace.comap lifetime inferInstance ⊔
    ⨆ t : NNReal, MeasurableSpace.comap (coordinate t) inferInstance

/-- The lifetime is measurable for the canonical measurable structure. -/
theorem measurable_lifetime [MeasurableSpace α] :
    Measurable (lifetime : LifetimePath α → ENNReal) := by
  apply Measurable.of_comap_le
  exact le_sup_left

/-- Every fixed-time coordinate is measurable for the canonical measurable structure. -/
theorem measurable_coordinate [MeasurableSpace α] (t : NNReal) :
    Measurable (coordinate (α := α) t) := by
  apply Measurable.of_comap_le
  exact le_sup_of_le_right (le_iSup (fun s : NNReal ↦
    MeasurableSpace.comap (coordinate s) inferInstance) t)

/-- An ordinary continuous path regarded as a lifetime path with infinite lifetime. -/
def ofContinuousPath (ω : C(NNReal, α)) : LifetimePath α where
  lifetime := ∞
  livePath := fun t ↦ ω t
  continuous_livePath := ω.continuous.comp continuous_subtype_val

@[simp]
theorem lifetime_ofContinuousPath (ω : C(NNReal, α)) :
    (ofContinuousPath ω).lifetime = ∞ := rfl

@[simp]
theorem coordinate_ofContinuousPath (ω : C(NNReal, α)) (t : NNReal) :
    coordinate t (ofContinuousPath ω) = Cemetery.alive (ω t) := by
  rw [coordinate_of_lt]
  simp only [ofContinuousPath]
  exact ENNReal.coe_lt_top

/-- The infinite-lifetime embedding of ordinary continuous paths is injective. -/
theorem ofContinuousPath_injective :
    Function.Injective (ofContinuousPath (α := α)) := by
  intro ω η h
  apply ContinuousMap.ext
  intro t
  have hcoordinate := congrArg (coordinate (α := α) t) h
  rw [coordinate_ofContinuousPath, coordinate_ofContinuousPath] at hcoordinate
  exact Sum.inl.inj hcoordinate

/-- An infinite-lifetime path regarded as an ordinary continuous path. -/
def toContinuousPath (ω : LifetimePath α) (hω : ω.lifetime = ∞) : ContinuousPath α where
  toFun t := ω.livePath ⟨t, by rw [hω]; exact ENNReal.coe_lt_top⟩
  continuous_toFun := by
    have hinclusion :
        Continuous (fun t : NNReal ↦
          (⟨t, by rw [hω]; exact ENNReal.coe_lt_top⟩ :
            {t : NNReal // (t : ENNReal) < ω.lifetime})) :=
      continuous_id.subtype_mk _
    exact ω.continuous_livePath.comp hinclusion

@[simp]
theorem toContinuousPath_ofContinuousPath (ω : ContinuousPath α) :
    toContinuousPath (ofContinuousPath ω) (lifetime_ofContinuousPath ω) = ω := by
  apply ContinuousMap.ext
  intro t
  rfl

@[simp]
theorem ofContinuousPath_toContinuousPath (ω : LifetimePath α) (hω : ω.lifetime = ∞) :
    ofContinuousPath (toContinuousPath ω hω) = ω := by
  apply ext_coordinate
  · rw [lifetime_ofContinuousPath, hω]
  · intro t
    have ht : (t : ENNReal) < ω.lifetime := by
      rw [hω]
      exact ENNReal.coe_lt_top
    rw [coordinate_ofContinuousPath, coordinate_of_lt _ _ ht]
    congr 2

section MeasurableContinuousPath

variable [MeasurableSpace α] [BorelSpace α]

/-- The infinite-lifetime embedding is measurable for the coordinate-generated structure. -/
theorem measurable_ofContinuousPath : Measurable (ofContinuousPath (α := α)) := by
  apply Measurable.of_comap_le
  rw [instMeasurableSpace, MeasurableSpace.comap_sup, MeasurableSpace.comap_iSup]
  apply sup_le
  · rw [MeasurableSpace.comap_comp]
    change MeasurableSpace.comap (fun _ : ContinuousPath α ↦ (∞ : ENNReal))
      inferInstance ≤ borel (ContinuousPath α)
    exact measurable_const.comap_le
  · rw [iSup_le_iff]
    intro t
    have hcoordinate :
        coordinate (α := α) t ∘ ofContinuousPath =
          Cemetery.alive ∘ fun ω : ContinuousPath α ↦ ω t := by
      funext ω
      exact coordinate_ofContinuousPath ω t
    rw [MeasurableSpace.comap_comp, hcoordinate]
    have heval : Measurable (fun ω : ContinuousPath α ↦ ω t) :=
      (continuous_eval_const t).borel_measurable.mono le_rfl
        (le_of_eq BorelSpace.measurable_eq)
    exact (measurable_inl.comp heval).comap_le

variable [StandardBorelSpace (ContinuousPath α)]
  [MeasurableSpace.CountablySeparated (LifetimePath α)]

/-- The infinite-lifetime embedding is a measurable embedding whenever the lifetime-path
measurable space is countably separated. -/
theorem measurableEmbedding_ofContinuousPath :
    MeasurableEmbedding (ofContinuousPath (α := α)) :=
  measurable_ofContinuousPath.measurableEmbedding ofContinuousPath_injective

end MeasurableContinuousPath

end LifetimePath

end
end MarkovProcess