import BFPP.StepFunctions
import Mathlib.Topology.ContinuousMap.ZeroAtInftyUnitization
import Mathlib.Topology.ContinuousMap.StoneWeierstrass

/-! # The algebra generated by initial-segment functions -/

namespace BFPP

set_option autoImplicit false

open Set Order Filter Topology
open scoped ZeroAtInfty

universe u
variable {ι : Type u} [LinearOrder ι] [OrderBot ι] [SuccOrder ι] [NoMaxOrder ι]
  [TopologicalSpace ι] [OrderTopology ι] [CompactIccSpace ι]

theorem initialStep_mul (i j : ι) : initialStep i * initialStep j = initialStep (min i j) := by
  ext k
  simp only [ZeroAtInftyContinuousMap.mul_apply, initialStep_apply, le_min_iff]
  split_ifs <;> simp_all

noncomputable def stepSpan : Submodule ℝ C₀(ι, ℝ) := LinearMap.range stepCombination

theorem initialStep_mem_stepSpan (i : ι) : initialStep i ∈ (stepSpan : Submodule ℝ C₀(ι, ℝ)) :=
  ⟨Finsupp.single i 1, stepCombination_single i⟩

theorem stepCombination_eq_sum (c : ι →₀ ℝ) :
    stepCombination c = ∑ i ∈ c.support, c i • initialStep i := rfl

theorem stepSpan_mul_mem (f g : C₀(ι, ℝ)) (hf : f ∈ stepSpan) (hg : g ∈ stepSpan) :
    f * g ∈ stepSpan := by
  obtain ⟨c, rfl⟩ := hf
  obtain ⟨d, rfl⟩ := hg
  rw [stepCombination_eq_sum, stepCombination_eq_sum]
  simp_rw [Finset.sum_mul, Finset.mul_sum]
  apply Submodule.sum_mem
  intro i _hi
  apply Submodule.sum_mem
  intro j _hj
  rw [smul_mul_smul_comm, initialStep_mul]
  exact Submodule.smul_mem _ _ (initialStep_mem_stepSpan _)

/-- Add constants after extending the finite step functions by zero at infinity. -/
noncomputable def stepAlgebra : Subalgebra ℝ C(OnePoint ι, ℝ) where
  carrier := {f | ∃ r : ℝ, ∃ g : C₀(ι, ℝ), g ∈ stepSpan ∧
    f = ContinuousMap.const (OnePoint ι) r + g.toOnePoint}
  algebraMap_mem' r := by
    refine ⟨r, 0, Submodule.zero_mem _, ?_⟩
    simp
    ext x
    rfl
  zero_mem' := by
    refine ⟨0, 0, Submodule.zero_mem _, ?_⟩
    simp
  add_mem' := by
    rintro f h ⟨r, g, hg, rfl⟩ ⟨s, k, hk, rfl⟩
    refine ⟨r + s, g + k, Submodule.add_mem _ hg hk, ?_⟩
    ext x
    simp
    ring
  one_mem' := by
    refine ⟨1, 0, Submodule.zero_mem _, ?_⟩
    simp
  mul_mem' := by
    rintro f h ⟨r, g, hg, rfl⟩ ⟨s, k, hk, rfl⟩
    refine ⟨r * s, r • k + s • g + g * k, ?_, ?_⟩
    · exact Submodule.add_mem _
        (Submodule.add_mem _ (Submodule.smul_mem _ _ hk) (Submodule.smul_mem _ _ hg))
        (stepSpan_mul_mem g k hg hk)
    · ext x
      simp
      ring

theorem initialStep_toOnePoint_mem_stepAlgebra (i : ι) :
    (initialStep i).toOnePoint ∈ (stepAlgebra : Subalgebra ℝ C(OnePoint ι, ℝ)) := by
  refine ⟨0, initialStep i, initialStep_mem_stepSpan i, ?_⟩
  simp

theorem stepAlgebra_separatesPoints : (stepAlgebra : Subalgebra ℝ C(OnePoint ι, ℝ)).SeparatesPoints := by
  intro x y hxy
  have hsep (i : ι) (hi : (initialStep i).toOnePoint x ≠ (initialStep i).toOnePoint y) :
      ∃ f ∈ (fun f : C(OnePoint ι, ℝ) => (f : OnePoint ι → ℝ)) ''
        (stepAlgebra (ι := ι) : Set C(OnePoint ι, ℝ)), f x ≠ f y := by
    refine ⟨((initialStep i).toOnePoint : OnePoint ι → ℝ), ?_, hi⟩
    exact Set.mem_image_of_mem _ (initialStep_toOnePoint_mem_stepAlgebra i)
  induction x using OnePoint.rec with
  | infty =>
    induction y using OnePoint.rec with
    | infty => exact (hxy rfl).elim
    | coe j => apply hsep j; simp
  | coe i =>
    induction y using OnePoint.rec with
    | infty => apply hsep i; simp
    | coe j =>
      have hij : i ≠ j := fun h => hxy (congrArg (fun z : ι => (z : OnePoint ι)) h)
      rcases lt_or_gt_of_ne hij with hij | hji
      · apply hsep i
        simp [not_le.mpr hij]
      · apply hsep j
        simp [not_le.mpr hji]

theorem stepCombination_denseRange : DenseRange (stepCombination : (ι →₀ ℝ) → C₀(ι, ℝ)) := by
  apply Metric.denseRange_iff.mpr
  intro f ε hε
  obtain ⟨h, hh⟩ := ContinuousMap.exists_mem_subalgebra_near_continuousMap_of_separatesPoints
    stepAlgebra stepAlgebra_separatesPoints f.toOnePoint (ε / 3) (by linarith)
  obtain ⟨r, g, hg, he⟩ := h.property
  obtain ⟨c, rfl⟩ := hg
  refine ⟨c, ?_⟩
  have hr : |r| < ε / 3 := by
    have hinf := (h.val - f.toOnePoint).norm_coe_le_norm (OnePoint.infty : OnePoint ι)
    have hinf' : |r| ≤ ‖h.val - f.toOnePoint‖ := by
      simpa [he] using hinf
    exact hinf'.trans_lt hh
  have hnorm : ‖stepCombination c - f‖ ≤ 2 * (ε / 3) := by
    apply c0_norm_le _ _ (by linarith)
    intro i
    have hp := (h.val - f.toOnePoint).norm_coe_le_norm (i : OnePoint ι)
    have hp' : |r + stepCombination c i - f i| < ε / 3 := by
      have h₁ : |r + stepCombination c i - f i| ≤ ‖h.val - f.toOnePoint‖ := by
        simpa [he] using hp
      exact h₁.trans_lt hh
    change |stepCombination c i - f i| ≤ 2 * (ε / 3)
    have h₁ := abs_le.mp hp'.le
    have h₂ := abs_le.mp hr.le
    exact abs_le.mpr ⟨by linarith, by linarith⟩
  rw [dist_comm, dist_eq_norm]
  exact hnorm.trans_lt (by linarith)

end BFPP