import Mathlib.Algebra.Field.Defs
import Mathlib.GroupTheory.GroupAction.Defs
import Mathlib.Tactic.FunProp

import SciLean.Util.SorryProof

set_option linter.unusedVariables false

-- Some missing theorems -------------------------------------------------------
--------------------------------------------------------------------------------

theorem Function.invFun_comp' [Nonempty α] {f : α → β} (hf : f.Injective) {x : α} :
    f.invFun (f x) = x

:= by
  suffices (f.invFun ∘ f) x = x by assumption
  rw[Function.invFun_comp hf]
  rfl


import SciLean.Core.Rand.Rand
import SciLean.Core.Distribution.ParametricDistribFwdDeriv

namespace SciLean

open MeasureTheory

variable
  {R} [RealScalar R]
  {W} [Vec R W]
  {X} [MeasurableSpace X] [Vec R X]
  {Y} [Vec R Y] [Module ℝ Y]

set_default_scalar R

@[fun_trans]
theorem Rand.𝔼.arg_r.cderiv_rule (r : W → Rand X) (f : X → Y) :
  cderiv R (fun w => (r w).𝔼 f)
  =
  fun w dw =>
    let d := parDistribDeriv (fun w => (r w).ℙ.toDistribution (R:=R)) w dw
    d.extAction f (fun r ⊸ fun y ⊸ ((r • y) : Y)) := sorry_proof

@[fun_trans]
theorem Rand.𝔼.arg_rf.cderiv_rule' (r : W → Rand X) (f : W → X → Y)
  (hf : ∀ x, CDifferentiable R (f · x)) :
  cderiv R (fun w => (r w).𝔼 (f w))
  =
  fun w dw =>
    let dr := parDistribFwdDeriv (fun w => (r w).ℙ.toDistribution (R:=R)) w dw
    let df := fun x => fwdDeriv R (f · x) w dw
    dr.extAction df (fun rdr ⊸ fun ydy ⊸ rdr.1•ydy.2 + rdr.2•ydy.1) := sorry_proof

theorem Rand.𝔼_deriv_as_distribDeriv {X} [Vec R X] [MeasureSpace X]
  (r : W → Rand X) (f : W → X → Y) :
  cderiv R (fun w => (r w).𝔼 (f w))
  =
  fun w dw =>
    parDistribDeriv (fun w => (fun x => ((r w).pdf R volume x) • f w x).toDistribution (R:=R)) w dw |>.integrate


import Mathlib.Data.Erased
import Mathlib.Control.Random
import Mathlib.MeasureTheory.Integral.Bochner
import Mathlib.MeasureTheory.Decomposition.Lebesgue

import SciLean.Core.FunctionPropositions.Bijective
import SciLean.Core.Objects.Scalar
import SciLean.Core.Integral.CIntegral
import SciLean.Core.Rand.SimpAttr

import Mathlib.MeasureTheory.Measure.GiryMonad

open MeasureTheory ENNReal BigOperators Finset

namespace SciLean

abbrev erase (a : α) : Erased α := .mk a

@[simp,ftrans_simp]
theorem erase_out {α} (a : α) : (erase a).out = a := by simp[erase]

/-- `x : Rand X` is a random variable of type `X`

You can:
  - generate sample with `x.get : IO X`
  - get probability measure with `x.ℙ : Measure X`

The internal fields `spec` and `rand` are just an internal implementation of `Rand` and should not
be accessed by normal users.

TODO: Hide implementation using quotients or something like that
-/
structure Rand (X : Type _)  where
  /-- `spec` defines a probability measure by computing an expectation. This means if `x : Rand X`
  corresponds to a probability measure `μ` then for `φ : X → ℝ`
  ```
  x.spec.out φ = ∫ x, φ x ∂μ
  ```

  Using `(X→ℝ)→ℝ` instead of `Measure X` for the specification of random variables has the
  advantage that we can reuse Lean's `do` notation.
  -/
  spec : Erased ((X→ℝ)→ℝ)
  /-- `rand` is a pseudo randon number generator implemented using the "Standard" number generator
  -/
  rand : StateM StdGen X

namespace Rand

def _root_.Function.IsMeasure {X} [MeasurableSpace X] (F : (X → ℝ) → ℝ) : Prop :=
  ∃ μ : Measure X, ∀ (f : X → ℝ), F f = ∫ x, f x ∂μ

open Classical in
/-- Probability measure of a random variable -/
@[pp_dot]
noncomputable
def ℙ {X} [MeasurableSpace X] (r : Rand X) : Measure X :=
  if h : r.spec.out.IsMeasure then
    choose h
  else
    0

/-- Specification of `x : Rand X` is really saying that it is a probability measure. -/
class LawfulRand (x : Rand X) [MeasurableSpace X] where
  is_measure : x.spec.out.IsMeasure
  is_prob : IsProbabilityMeasure x.ℙ

variable {X Y Z : Type _}
  [MeasurableSpace X]
  [MeasurableSpace Y]

instance instIsProbabilityMeasureℙ (x : Rand X) [inst : LawfulRand x] : IsProbabilityMeasure (x.ℙ) := inst.is_prob

/-- Extensionality of random variable.

WARNING: This theorem is inconsistent!!! The random generators `x.rand` and `y.rand` might differ.
         We are not trying to model pseudo-random numbers. We assume that every random number
         generator is a true random number generator. Thus the result of any probabilistic program
         should be independent on the exact generator up to some randomness.

TODO: We might quotient all the random number generators corresponding to the measure `x.ℙ`  under
      the assumption that they are all true random generators. I believe that such type would be
      a singleton i.e. all the random number generators are all the same.
-/
@[ext]
axiom ext (x y : Rand X) : x.spec.out = y.spec.out → x = y

/-- Generate rundom number using IO randomness -/
def get (x : Rand X) : IO X := do
  let stdGen ← IO.stdGenRef.get
  let (res, new) := x.rand stdGen
  let _ ← IO.stdGenRef.set new
  pure res

----------------------------------------------------------------------------------------------------
-- Monadic structure -------------------------------------------------------------------------------
----------------------------------------------------------------------------------------------------

instance : Monad Rand where
  pure x := {
    spec := erase (fun φ => φ x),
    rand := pure x
  }
  bind x f := {
    spec := erase (fun φ => x.spec.out (fun x => (f x).spec.out φ)),
    rand := bind x.rand (fun x => (f x).rand)
  }

instance : LawfulMonad Rand where
  bind_pure_comp := by intros; rfl
  bind_map       := by intros; rfl
  pure_bind      := by intros; ext; simp[Bind.bind,Pure.pure]
  bind_assoc     := by intros; ext; simp[Bind.bind,Pure.pure]
  map_const      := by intros; ext; rfl
  id_map         := by intros; ext; simp[Bind.bind,Pure.pure,id,Functor.map]
  seqLeft_eq     := by intros; ext; simp[Bind.bind,Pure.pure,Seq.seq,Function.const,Functor.map,SeqLeft.seqLeft]
  seqRight_eq    := by intros; ext; simp[Bind.bind,Pure.pure,Seq.seq,Function.const,Functor.map,SeqRight.seqRight]
  pure_seq       := by intros; ext; simp[Bind.bind,Pure.pure,Seq.seq,Functor.map]

-- this needs some integrability and lawfulness of Rand
theorem swap_bind (f : X → Y → Z) (x : Rand X) (y : Rand Y) :
    (do let x' ← x; let y' ← y; pure (f x' y'))
    =
    (do let y' ← y; let x' ← x; pure (f x' y')) := by
  sorry_proof

@[simp, ftrans_simp]
theorem pure_ℙ (x : X) : (pure x : Rand X).ℙ = Measure.dirac x := sorry_proof

----------------------------------------------------------------------------------------------------
-- Arithmetics -------------------------------------------------------------------------------------
----------------------------------------------------------------------------------------------------

instance [Add X] : HAdd X (Rand X) (Rand X) := ⟨fun x' x => do
  let x'' ← x
  pure (x' + x'')⟩

instance [Add X] : HAdd (Rand X) X (Rand X) := ⟨fun x x' => do
  let x'' ← x
  pure (x'' + x')⟩

-- instance [Add X] : HAdd (Rand X) (Rand X) (Rand X) := ⟨fun x y =>
--   let x' ~ x
--   let y' ~ y
--   pure (x' + y')⟩

-- todo: add simp theorems that inline these operations

----------------------------------------------------------------------------------------------------
-- Simple Random Variable functions ----------------------------------------------------------------
----------------------------------------------------------------------------------------------------

@[pp_dot]
abbrev map (r : Rand X) (f : X → Y) : Rand Y := do
  let x' ← r
  return f x'

/-- Marginal distribution for the first component of a pair. -/
@[pp_dot]
abbrev fst (r : Rand (X×Y)) : Rand X := do
  let (x,_) ← r
  return x

/-- Marginal distribution for the second component of a pair. -/
@[pp_dot]
abbrev snd (r : Rand (X×Y)) : Rand Y := do
  let (_,y) ← r
  return y

@[simp, ftrans_simp]
theorem map_ℙ  (r : Rand X) (f : X → Y) :
  (r.map f).ℙ = r.ℙ.map f := sorry_proof

----------------------------------------------------------------------------------------------------
-- Expected Value ----------------------------------------------------------------------------------
----------------------------------------------------------------------------------------------------

section ExpectedValue

variable
  {R} [RealScalar R]
  [AddCommGroup Y] [Module ℝ Y]
  [AddCommGroup Z] [Module ℝ Z]
  -- [AddCommGroup U] [TopologicalSpace U] [TopologicalAddGroup U] [Module ℝ U] [LocallyConveUSpace ℝ U]
  {U} [AddCommGroup U] [Module ℝ U]

@[pp_dot]
noncomputable
def 𝔼 (r : Rand X) (φ : X → Y) : Y := ∫' x, φ x ∂r.ℙ

@[simp, ftrans_simp, rand_push_E]
theorem pure_𝔼 (x : X) (φ : X → Y) :
    (pure (f:=Rand) x).𝔼 φ = φ x := by simp [𝔼]

-- What are the right assumptions here? Lambda lawfulness of `x` and `f x'` and integrability of `φ`
@[rand_push_E]
theorem bind_E (r : Rand X) (f : X → Rand Y) (φ : Y → Z) :
    (r >>= f).𝔼 φ = r.𝔼 (fun x' => (f x').𝔼 φ) := by simp[𝔼]; sorry_proof

-- todo: We might want this to hold without lawfulness
-- consider adding as a property inside of `Distribution` or `Rand`
@[simp, ftrans_simp, rand_push_E]
theorem E_zero (r : Rand X) :
    r.𝔼 (fun _ => (0 : U)) = 0 := by simp[𝔼]

@[simp, ftrans_simp, add_pull, rand_push_E]
theorem E_add (r : Rand X) (φ ψ : X → U)
    (hφ : CIntegrable φ r.ℙ) (hψ : CIntegrable ψ r.ℙ) :
    r.𝔼 (fun x => φ x + ψ x) = r.𝔼 φ + r.𝔼 ψ := by
  simp[𝔼]; rw[cintegral_add] <;> assumption

@[simp, ftrans_simp, smul_pull, rand_push_E]
theorem E_smul (r : Rand X) (φ : X → ℝ) (y : Y) :
    r.𝔼 (fun x' => φ x' • y) = r.𝔼 φ • y := by sorry_proof

theorem reparameterize [Nonempty X] (f : X → Y) (hf : f.Injective) {r : Rand X} {φ : X → Z} :
    r.𝔼 φ
    =
    let invf := f.invFun
    (r.map f).𝔼 (fun y => φ (invf y))

:= by
  simp [𝔼,Function.invFun_comp' hf]


import Mathlib.Data.Set.Defs
import Mathlib.Data.Set.Image

import Mathlib.Tactic.FunTrans.Attr
import Mathlib.Tactic.FunTrans.Elab

import SciLean.Core.Objects.Scalar
import SciLean.Util.SorryProof

variable {α β γ : Type _}

attribute [fun_trans] Set.preimage
attribute [fun_trans] Set.image

attribute [fun_trans] Set.preimage_id Set.preimage_id'

namespace Set

open Classical in
@[fun_trans]
theorem preimage_const' (b : β) (s : Set β) :
    (fun _ : α => b) ⁻¹' s = if b ∈ s then univ else ∅ := by apply preimage_const

@[fun_trans]
theorem preimage_comp' (f : β → γ) (g : α → β) :
    preimage (fun x => f (g x))
    =
    fun s => g ⁻¹' (f ⁻¹' s) := rfl

----------------------------------------------------------------------------------------------------

@[fun_trans]
theorem Prod.mk.arg_fstsnd.preimage_rule_prod (f : α → β) (g : α → γ) (B : Set β) (C : Set γ) :
    preimage (fun x => (f x, g x)) (B.prod C)
    =
    f ⁻¹' B ∩ g ⁻¹' C


import Mathlib.Data.Set.Defs
import Mathlib.Data.Set.Image

import Mathlib.Tactic.FunTrans.Attr
import Mathlib.Tactic.FunTrans.Elab

import SciLean.Core.Objects.Scalar
import SciLean.Util.SorryProof

variable {α β γ : Type _}

attribute [fun_trans] Set.preimage
attribute [fun_trans] Set.image

attribute [fun_trans] Set.preimage_id Set.preimage_id'

namespace Set

open Classical in
@[fun_trans]
theorem preimage_const' (b : β) (s : Set β) :
    (fun _ : α => b) ⁻¹' s = if b ∈ s then univ else ∅ := by apply preimage_const

@[fun_trans]
theorem preimage_comp' (f : β → γ) (g : α → β) :
    preimage (fun x => f (g x))
    =
    fun s => g ⁻¹' (f ⁻¹' s) := rfl

----------------------------------------------------------------------------------------------------

@[fun_trans]
theorem Prod.mk.arg_fstsnd.preimage_rule_prod (f : α → β) (g : α → γ) (B : Set β) (C : Set γ) :
    preimage (fun x => (f x, g x)) (B.prod C)
    =
    f ⁻¹' B ∩ g ⁻¹' C := sorry_proof

@[fun_trans]
theorem Prod.mk.arg_fst.preimage_rule_prod (f : α → β) (c : γ) :
    preimage (fun x => (f x, c))
    =
    fun s => f ⁻¹' {y | (y,c) ∈ s} := sorry_proof

@[fun_trans]
theorem Prod.mk.arg_snd.preimage_rule_prod (b : β) (g : α → γ) :
    preimage (fun x => (b, g x))
    =
    fun s => g ⁻¹' {z | (b,z) ∈ s}


import Mathlib.Data.Set.Defs
import Mathlib.Data.Set.Image

import Mathlib.Tactic.FunTrans.Attr
import Mathlib.Tactic.FunTrans.Elab

import SciLean.Core.Objects.Scalar
import SciLean.Util.SorryProof

variable {α β γ : Type _}

attribute [fun_trans] Set.preimage
attribute [fun_trans] Set.image

attribute [fun_trans] Set.preimage_id Set.preimage_id'

namespace Set

open Classical in
@[fun_trans]
theorem preimage_const' (b : β) (s : Set β) :
    (fun _ : α => b) ⁻¹' s = if b ∈ s then univ else ∅ := by apply preimage_const

@[fun_trans]
theorem preimage_comp' (f : β → γ) (g : α → β) :
    preimage (fun x => f (g x))
    =
    fun s => g ⁻¹' (f ⁻¹' s) := rfl

----------------------------------------------------------------------------------------------------

@[fun_trans]
theorem Prod.mk.arg_fstsnd.preimage_rule_prod (f : α → β) (g : α → γ) (B : Set β) (C : Set γ) :
    preimage (fun x => (f x, g x)) (B.prod C)
    =
    f ⁻¹' B ∩ g ⁻¹' C := sorry_proof

@[fun_trans]
theorem Prod.mk.arg_fst.preimage_rule_prod (f : α → β) (c : γ) :
    preimage (fun x => (f x, c))
    =
    fun s => f ⁻¹' {y | (y,c) ∈ s} := sorry_proof

@[fun_trans]
theorem Prod.mk.arg_snd.preimage_rule_prod (b : β) (g : α → γ) :
    preimage (fun x => (b, g x))
    =
    fun s => g ⁻¹' {z | (b,z) ∈ s} := sorry_proof

open SciLean
variable {R} [RealScalar R] -- probably generalize following to LinearlyOrderedAddCommGroup

@[fun_trans]
theorem HAdd.hAdd.arg_a0.preimage_rule_Ioo (x' a b : R)  :
    preimage (fun x : R => x + x') (Ioo a b)
    =
    Ioo (a - x') (b - x') := by ext; simp; sorry_proof

@[fun_trans]
theorem HAdd.hAdd.arg_a1.preimage_rule_Ioo (x' a b : R)  :
    preimage (fun x : R => x' + x) (Ioo a b)
    =
    Ioo (a - x') (b - x') := by ext; simp; sorry_proof

@[fun_trans]
theorem HSub.hSub.arg_a0.preimage_rule_Ioo (x' a b : R)  :
    preimage (fun x : R => x - x') (Ioo a b)
    =
    Ioo (a + x') (b + x') := by ext; simp; sorry_proof

@[fun_trans]
theorem HSub.hSub.arg_a1.preimage_rule_Ioo (x' a b : R)  :
    preimage (fun x : R => x' - x) (Ioo a b)
    =
    Ioo (x' - b) (x' - a) := by ext; simp; sorry_proof

@[fun_trans]
theorem Neg.neg.arg_a1.preimage_rule_Ioo (a b : R)  :
    preimage (fun x : R => - x) (Ioo a b)
    =
    Ioo (-b) (-a) := by ext; simp; sorry_proof

----------------------------------------------------------------------------------------------------
-- Preimage1 ---------------------------------------------------------------------------------------
----------------------------------------------------------------------------------------------------

-- todo: turn into function transformation once we have `fun_trans` supporting two argument functions
def preimage1 {α β γ} (f : α → β → γ) (C : Set γ) : Set α := ⋃ b, (f · b) ⁻¹' C

-- @[simp, ftrans_simp]
-- theorem preimage1_id {α β} (s : Set (α×β)) :
--      s.preimage1 (fun (a : α) (b : β) => (a,b)) = ⋃ b, {a | (a, b) ∈ s} := sorry_proof

-- this probably needs non-empty `β`

@[simp, ftrans_simp]
theorem preimage1_id1 {α β} (A : Set α) :
     A.preimage1 (fun (a : α) (_ : β) => a) = A


import Mathlib.Data.Set.Defs
import Mathlib.Data.Set.Image

import Mathlib.Tactic.FunTrans.Attr
import Mathlib.Tactic.FunTrans.Elab

import SciLean.Core.Objects.Scalar
import SciLean.Util.SorryProof

variable {α β γ : Type _}

attribute [fun_trans] Set.preimage
attribute [fun_trans] Set.image

attribute [fun_trans] Set.preimage_id Set.preimage_id'

namespace Set

open Classical in
@[fun_trans]
theorem preimage_const' (b : β) (s : Set β) :
    (fun _ : α => b) ⁻¹' s = if b ∈ s then univ else ∅ := by apply preimage_const

@[fun_trans]
theorem preimage_comp' (f : β → γ) (g : α → β) :
    preimage (fun x => f (g x))
    =
    fun s => g ⁻¹' (f ⁻¹' s) := rfl

----------------------------------------------------------------------------------------------------

@[fun_trans]
theorem Prod.mk.arg_fstsnd.preimage_rule_prod (f : α → β) (g : α → γ) (B : Set β) (C : Set γ) :
    preimage (fun x => (f x, g x)) (B.prod C)
    =
    f ⁻¹' B ∩ g ⁻¹' C := sorry_proof

@[fun_trans]
theorem Prod.mk.arg_fst.preimage_rule_prod (f : α → β) (c : γ) :
    preimage (fun x => (f x, c))
    =
    fun s => f ⁻¹' {y | (y,c) ∈ s} := sorry_proof

@[fun_trans]
theorem Prod.mk.arg_snd.preimage_rule_prod (b : β) (g : α → γ) :
    preimage (fun x => (b, g x))
    =
    fun s => g ⁻¹' {z | (b,z) ∈ s} := sorry_proof

open SciLean
variable {R} [RealScalar R] -- probably generalize following to LinearlyOrderedAddCommGroup

@[fun_trans]
theorem HAdd.hAdd.arg_a0.preimage_rule_Ioo (x' a b : R)  :
    preimage (fun x : R => x + x') (Ioo a b)
    =
    Ioo (a - x') (b - x') := by ext; simp; sorry_proof

@[fun_trans]
theorem HAdd.hAdd.arg_a1.preimage_rule_Ioo (x' a b : R)  :
    preimage (fun x : R => x' + x) (Ioo a b)
    =
    Ioo (a - x') (b - x') := by ext; simp; sorry_proof

@[fun_trans]
theorem HSub.hSub.arg_a0.preimage_rule_Ioo (x' a b : R)  :
    preimage (fun x : R => x - x') (Ioo a b)
    =
    Ioo (a + x') (b + x') := by ext; simp; sorry_proof

@[fun_trans]
theorem HSub.hSub.arg_a1.preimage_rule_Ioo (x' a b : R)  :
    preimage (fun x : R => x' - x) (Ioo a b)
    =
    Ioo (x' - b) (x' - a) := by ext; simp; sorry_proof

@[fun_trans]
theorem Neg.neg.arg_a1.preimage_rule_Ioo (a b : R)  :
    preimage (fun x : R => - x) (Ioo a b)
    =
    Ioo (-b) (-a) := by ext; simp; sorry_proof

----------------------------------------------------------------------------------------------------
-- Preimage1 ---------------------------------------------------------------------------------------
----------------------------------------------------------------------------------------------------

-- todo: turn into function transformation once we have `fun_trans` supporting two argument functions
def preimage1 {α β γ} (f : α → β → γ) (C : Set γ) : Set α := ⋃ b, (f · b) ⁻¹' C

-- @[simp, ftrans_simp]
-- theorem preimage1_id {α β} (s : Set (α×β)) :
--      s.preimage1 (fun (a : α) (b : β) => (a,b)) = ⋃ b, {a | (a, b) ∈ s} := sorry_proof

-- this probably needs non-empty `β`
@[simp, ftrans_simp]
theorem preimage1_id1 {α β} (A : Set α) :
     A.preimage1 (fun (a : α) (_ : β) => a) = A := sorry_proof

-- this probably needs non-empty `B`

@[simp, ftrans_simp]
theorem preimage1_id2 {α β} (B : Set β) :
     B.preimage1 (fun (_ : α) (b : β) => b) = Set.univ


import Mathlib.Data.Set.Defs
import Mathlib.Data.Set.Image

import Mathlib.Tactic.FunTrans.Attr
import Mathlib.Tactic.FunTrans.Elab

import SciLean.Core.Objects.Scalar
import SciLean.Util.SorryProof

variable {α β γ : Type _}

attribute [fun_trans] Set.preimage
attribute [fun_trans] Set.image

attribute [fun_trans] Set.preimage_id Set.preimage_id'

namespace Set

open Classical in
@[fun_trans]
theorem preimage_const' (b : β) (s : Set β) :
    (fun _ : α => b) ⁻¹' s = if b ∈ s then univ else ∅ := by apply preimage_const

@[fun_trans]
theorem preimage_comp' (f : β → γ) (g : α → β) :
    preimage (fun x => f (g x))
    =
    fun s => g ⁻¹' (f ⁻¹' s) := rfl

----------------------------------------------------------------------------------------------------

@[fun_trans]
theorem Prod.mk.arg_fstsnd.preimage_rule_prod (f : α → β) (g : α → γ) (B : Set β) (C : Set γ) :
    preimage (fun x => (f x, g x)) (B.prod C)
    =
    f ⁻¹' B ∩ g ⁻¹' C := sorry_proof

@[fun_trans]
theorem Prod.mk.arg_fst.preimage_rule_prod (f : α → β) (c : γ) :
    preimage (fun x => (f x, c))
    =
    fun s => f ⁻¹' {y | (y,c) ∈ s} := sorry_proof

@[fun_trans]
theorem Prod.mk.arg_snd.preimage_rule_prod (b : β) (g : α → γ) :
    preimage (fun x => (b, g x))
    =
    fun s => g ⁻¹' {z | (b,z) ∈ s} := sorry_proof

open SciLean
variable {R} [RealScalar R] -- probably generalize following to LinearlyOrderedAddCommGroup

@[fun_trans]
theorem HAdd.hAdd.arg_a0.preimage_rule_Ioo (x' a b : R)  :
    preimage (fun x : R => x + x') (Ioo a b)
    =
    Ioo (a - x') (b - x') := by ext; simp; sorry_proof

@[fun_trans]
theorem HAdd.hAdd.arg_a1.preimage_rule_Ioo (x' a b : R)  :
    preimage (fun x : R => x' + x) (Ioo a b)
    =
    Ioo (a - x') (b - x') := by ext; simp; sorry_proof

@[fun_trans]
theorem HSub.hSub.arg_a0.preimage_rule_Ioo (x' a b : R)  :
    preimage (fun x : R => x - x') (Ioo a b)
    =
    Ioo (a + x') (b + x') := by ext; simp; sorry_proof

@[fun_trans]
theorem HSub.hSub.arg_a1.preimage_rule_Ioo (x' a b : R)  :
    preimage (fun x : R => x' - x) (Ioo a b)
    =
    Ioo (x' - b) (x' - a) := by ext; simp; sorry_proof

@[fun_trans]
theorem Neg.neg.arg_a1.preimage_rule_Ioo (a b : R)  :
    preimage (fun x : R => - x) (Ioo a b)
    =
    Ioo (-b) (-a) := by ext; simp; sorry_proof

----------------------------------------------------------------------------------------------------
-- Preimage1 ---------------------------------------------------------------------------------------
----------------------------------------------------------------------------------------------------

-- todo: turn into function transformation once we have `fun_trans` supporting two argument functions
def preimage1 {α β γ} (f : α → β → γ) (C : Set γ) : Set α := ⋃ b, (f · b) ⁻¹' C

-- @[simp, ftrans_simp]
-- theorem preimage1_id {α β} (s : Set (α×β)) :
--      s.preimage1 (fun (a : α) (b : β) => (a,b)) = ⋃ b, {a | (a, b) ∈ s} := sorry_proof

-- this probably needs non-empty `β`
@[simp, ftrans_simp]
theorem preimage1_id1 {α β} (A : Set α) :
     A.preimage1 (fun (a : α) (_ : β) => a) = A := sorry_proof

-- this probably needs non-empty `B`
@[simp, ftrans_simp]
theorem preimage1_id2 {α β} (B : Set β) :
     B.preimage1 (fun (_ : α) (b : β) => b) = Set.univ := sorry_proof

open Classical in

@[simp, ftrans_simp]
theorem preimage1_const {α β γ} (c : γ) (C : Set γ) :
     C.preimage1 (fun (_ : α) (_ : β) => c) = if c ∈ C then Set.univ else ∅


import SciLean.Core.Rand.Rand
import SciLean.Core.FunctionPropositions
import SciLean.Core.FloatAsReal
-- import SciLean.Modules.Prob.DRand
-- import SciLean.Modules.Prob.FDRand

namespace SciLean.Rand

variable
  {R} [RealScalar R]
  -- {X} [NormedAddCommGroup X] [NormedSpace ℝ X] [NormedSpace R X] [CompleteSpace X] [MeasurableSpace X]
  -- {Y} [NormedAddCommGroup Y] [NormedSpace ℝ Y] [NormedSpace R Y] [CompleteSpace Y] [MeasurableSpace Y]
  -- {Z} [NormedAddCommGroup Z] [NormedSpace ℝ Z] [NormedSpace R Z] [CompleteSpace Z] [MeasurableSpace Z]
  {X : Type _} [MeasurableSpace X] [AddCommGroup X] [Module ℝ X]
  {Y : Type _} [AddCommGroup Y] [Module ℝ Y] [MeasurableSpace Y]
  {Z : Type _} [AddCommGroup Z] [Module ℝ Z]

open Rand

@[rand_pull_E]
theorem bind_pull_mean (x : Rand X) (f : X → Rand Y) :
    (x >>= (fun x' => pure (f x').mean)).mean
    =
    (x >>= f).mean := by simp[rand_push_E]

@[rand_push_E]
theorem ite_push_E {c} [Decidable c] (t e : Rand X) (φ : X → Y):
    (if c then t else e).𝔼 φ = (if c then t.𝔼 φ else e.𝔼 φ) := by
  if h : c then simp[h] else simp[h]

-- I don't think this is a desirable `rand_pull_E` theorem as it duplicates the if statement
-- @[rand_pull_E]
theorem ite_pull_E {c} [Decidable c] (t e : Rand X) (φ ψ : X → Y):
    (if c then t.𝔼 φ else e.𝔼 ψ) = (if c then t else e).𝔼 (if c then φ else ψ) := by
  if h : c then simp[h] else simp[h]

@[rand_push_E]
theorem ite_push_mean {c} [Decidable c] (t e : Rand X) :
    (if c then t else e).mean = (if c then t.mean else e.mean) := by
  if h : c then simp[h] else simp[h]

@[rand_pull_E]
theorem ite_pull_mean {c} [Decidable c] (t e : Rand X) :
    (if c then t.mean else e.mean) = (if c then t else e).mean := by
  if h : c then simp[h] else simp[h]

@[rand_pull_E mid-1]
theorem ite_pull_mean_t {c} [Decidable c] (t : Rand X) (e : X) :
    (if c then t.mean else e) = (if c then t else pure e).mean := by
  if h : c then simp[h] else simp[h]

@[rand_pull_E mid-1]
theorem ite_pull_mean_f {c} [Decidable c] (t : X) (e : Rand X) :
    (if c then t else e.mean) = (if c then pure t else e).mean := by
  if h : c then simp[h] else simp[h]

-- this has messed up universes
@[rand_pull_E]
theorem pull_E_lambda (r : Rand Y) (f : X → Y → Z) :
    (fun x => r.𝔼 (fun y => f x y))
    =
    r.𝔼 (fun y x => f x y) := sorry_proof

-- this has messed up universes
@[rand_push_E]
theorem push_E_lambda (r : Rand Y) (f : X → Y → Z) :
    r.𝔼 (fun y x => f x y)
    =
    (fun x => r.𝔼 (fun y => f x y)) := sorry_proof

-- can't be simp as it has variable head
theorem pull_E_affine (r : Rand X) (φ : X → Y)
    (f : Y → Z) (hf : IsAffineMap ℝ f := by fun_prop) :
    (f (r.𝔼 φ)) = r.𝔼 (fun x => f (φ x)) := by sorry_proof -- have := hf; sorry_proof

@[rand_push_E]
theorem push_E_affine (r : Rand X) (φ : X → Y)
    (f : Y → Z) (hf : IsAffineMap ℝ f := by fun_prop) :
    r.𝔼 (fun x => f (φ x)) = (f (r.𝔼 φ)) := by rw[pull_E_affine (hf:=hf)]

@[rand_pull_E]
theorem pull_mean_add (x y : Rand X) :
    x.mean + y.mean
    =
    Rand.mean do
      let x' ← x
      let y' ← y
      return x' + y'


import SciLean.Core.Rand.Rand
import SciLean.Core.FunctionPropositions
import SciLean.Core.FloatAsReal
-- import SciLean.Modules.Prob.DRand
-- import SciLean.Modules.Prob.FDRand

namespace SciLean.Rand

variable
  {R} [RealScalar R]
  -- {X} [NormedAddCommGroup X] [NormedSpace ℝ X] [NormedSpace R X] [CompleteSpace X] [MeasurableSpace X]
  -- {Y} [NormedAddCommGroup Y] [NormedSpace ℝ Y] [NormedSpace R Y] [CompleteSpace Y] [MeasurableSpace Y]
  -- {Z} [NormedAddCommGroup Z] [NormedSpace ℝ Z] [NormedSpace R Z] [CompleteSpace Z] [MeasurableSpace Z]
  {X : Type _} [MeasurableSpace X] [AddCommGroup X] [Module ℝ X]
  {Y : Type _} [AddCommGroup Y] [Module ℝ Y] [MeasurableSpace Y]
  {Z : Type _} [AddCommGroup Z] [Module ℝ Z]

open Rand

@[rand_pull_E]
theorem bind_pull_mean (x : Rand X) (f : X → Rand Y) :
    (x >>= (fun x' => pure (f x').mean)).mean
    =
    (x >>= f).mean := by simp[rand_push_E]

@[rand_push_E]
theorem ite_push_E {c} [Decidable c] (t e : Rand X) (φ : X → Y):
    (if c then t else e).𝔼 φ = (if c then t.𝔼 φ else e.𝔼 φ) := by
  if h : c then simp[h] else simp[h]

-- I don't think this is a desirable `rand_pull_E` theorem as it duplicates the if statement
-- @[rand_pull_E]
theorem ite_pull_E {c} [Decidable c] (t e : Rand X) (φ ψ : X → Y):
    (if c then t.𝔼 φ else e.𝔼 ψ) = (if c then t else e).𝔼 (if c then φ else ψ) := by
  if h : c then simp[h] else simp[h]

@[rand_push_E]
theorem ite_push_mean {c} [Decidable c] (t e : Rand X) :
    (if c then t else e).mean = (if c then t.mean else e.mean) := by
  if h : c then simp[h] else simp[h]

@[rand_pull_E]
theorem ite_pull_mean {c} [Decidable c] (t e : Rand X) :
    (if c then t.mean else e.mean) = (if c then t else e).mean := by
  if h : c then simp[h] else simp[h]

@[rand_pull_E mid-1]
theorem ite_pull_mean_t {c} [Decidable c] (t : Rand X) (e : X) :
    (if c then t.mean else e) = (if c then t else pure e).mean := by
  if h : c then simp[h] else simp[h]

@[rand_pull_E mid-1]
theorem ite_pull_mean_f {c} [Decidable c] (t : X) (e : Rand X) :
    (if c then t else e.mean) = (if c then pure t else e).mean := by
  if h : c then simp[h] else simp[h]

-- this has messed up universes
@[rand_pull_E]
theorem pull_E_lambda (r : Rand Y) (f : X → Y → Z) :
    (fun x => r.𝔼 (fun y => f x y))
    =
    r.𝔼 (fun y x => f x y) := sorry_proof

-- this has messed up universes
@[rand_push_E]
theorem push_E_lambda (r : Rand Y) (f : X → Y → Z) :
    r.𝔼 (fun y x => f x y)
    =
    (fun x => r.𝔼 (fun y => f x y)) := sorry_proof

-- can't be simp as it has variable head
theorem pull_E_affine (r : Rand X) (φ : X → Y)
    (f : Y → Z) (hf : IsAffineMap ℝ f := by fun_prop) :
    (f (r.𝔼 φ)) = r.𝔼 (fun x => f (φ x)) := by sorry_proof -- have := hf; sorry_proof

@[rand_push_E]
theorem push_E_affine (r : Rand X) (φ : X → Y)
    (f : Y → Z) (hf : IsAffineMap ℝ f := by fun_prop) :
    r.𝔼 (fun x => f (φ x)) = (f (r.𝔼 φ)) := by rw[pull_E_affine (hf:=hf)]

@[rand_pull_E]
theorem pull_mean_add (x y : Rand X) :
    x.mean + y.mean
    =
    Rand.mean do
      let x' ← x
      let y' ← y
      return x' + y' := sorry_proof

@[rand_pull_E]
theorem pull_mean_add_1 (x : Rand X) (y : X) :
    x.mean + y
    =
    Rand.mean do
      let x' ← x
      return x' + y


import SciLean.Core.Rand.Rand
import SciLean.Core.FunctionPropositions
import SciLean.Core.FloatAsReal
-- import SciLean.Modules.Prob.DRand
-- import SciLean.Modules.Prob.FDRand

namespace SciLean.Rand

variable
  {R} [RealScalar R]
  -- {X} [NormedAddCommGroup X] [NormedSpace ℝ X] [NormedSpace R X] [CompleteSpace X] [MeasurableSpace X]
  -- {Y} [NormedAddCommGroup Y] [NormedSpace ℝ Y] [NormedSpace R Y] [CompleteSpace Y] [MeasurableSpace Y]
  -- {Z} [NormedAddCommGroup Z] [NormedSpace ℝ Z] [NormedSpace R Z] [CompleteSpace Z] [MeasurableSpace Z]
  {X : Type _} [MeasurableSpace X] [AddCommGroup X] [Module ℝ X]
  {Y : Type _} [AddCommGroup Y] [Module ℝ Y] [MeasurableSpace Y]
  {Z : Type _} [AddCommGroup Z] [Module ℝ Z]

open Rand

@[rand_pull_E]
theorem bind_pull_mean (x : Rand X) (f : X → Rand Y) :
    (x >>= (fun x' => pure (f x').mean)).mean
    =
    (x >>= f).mean := by simp[rand_push_E]

@[rand_push_E]
theorem ite_push_E {c} [Decidable c] (t e : Rand X) (φ : X → Y):
    (if c then t else e).𝔼 φ = (if c then t.𝔼 φ else e.𝔼 φ) := by
  if h : c then simp[h] else simp[h]

-- I don't think this is a desirable `rand_pull_E` theorem as it duplicates the if statement
-- @[rand_pull_E]
theorem ite_pull_E {c} [Decidable c] (t e : Rand X) (φ ψ : X → Y):
    (if c then t.𝔼 φ else e.𝔼 ψ) = (if c then t else e).𝔼 (if c then φ else ψ) := by
  if h : c then simp[h] else simp[h]

@[rand_push_E]
theorem ite_push_mean {c} [Decidable c] (t e : Rand X) :
    (if c then t else e).mean = (if c then t.mean else e.mean) := by
  if h : c then simp[h] else simp[h]

@[rand_pull_E]
theorem ite_pull_mean {c} [Decidable c] (t e : Rand X) :
    (if c then t.mean else e.mean) = (if c then t else e).mean := by
  if h : c then simp[h] else simp[h]

@[rand_pull_E mid-1]
theorem ite_pull_mean_t {c} [Decidable c] (t : Rand X) (e : X) :
    (if c then t.mean else e) = (if c then t else pure e).mean := by
  if h : c then simp[h] else simp[h]

@[rand_pull_E mid-1]
theorem ite_pull_mean_f {c} [Decidable c] (t : X) (e : Rand X) :
    (if c then t else e.mean) = (if c then pure t else e).mean := by
  if h : c then simp[h] else simp[h]

-- this has messed up universes
@[rand_pull_E]
theorem pull_E_lambda (r : Rand Y) (f : X → Y → Z) :
    (fun x => r.𝔼 (fun y => f x y))
    =
    r.𝔼 (fun y x => f x y) := sorry_proof

-- this has messed up universes
@[rand_push_E]
theorem push_E_lambda (r : Rand Y) (f : X → Y → Z) :
    r.𝔼 (fun y x => f x y)
    =
    (fun x => r.𝔼 (fun y => f x y)) := sorry_proof

-- can't be simp as it has variable head
theorem pull_E_affine (r : Rand X) (φ : X → Y)
    (f : Y → Z) (hf : IsAffineMap ℝ f := by fun_prop) :
    (f (r.𝔼 φ)) = r.𝔼 (fun x => f (φ x)) := by sorry_proof -- have := hf; sorry_proof

@[rand_push_E]
theorem push_E_affine (r : Rand X) (φ : X → Y)
    (f : Y → Z) (hf : IsAffineMap ℝ f := by fun_prop) :
    r.𝔼 (fun x => f (φ x)) = (f (r.𝔼 φ)) := by rw[pull_E_affine (hf:=hf)]

@[rand_pull_E]
theorem pull_mean_add (x y : Rand X) :
    x.mean + y.mean
    =
    Rand.mean do
      let x' ← x
      let y' ← y
      return x' + y' := sorry_proof

@[rand_pull_E]
theorem pull_mean_add_1 (x : Rand X) (y : X) :
    x.mean + y
    =
    Rand.mean do
      let x' ← x
      return x' + y := sorry_proof

@[rand_pull_E]
theorem pull_mean_add_2 (x : X) (y : Rand X) :
    x + y.mean
    =
    Rand.mean do
      let y' ← y
      return x + y'


import SciLean.Core.Rand.Rand
import SciLean.Core.FunctionPropositions
import SciLean.Core.FloatAsReal
-- import SciLean.Modules.Prob.DRand
-- import SciLean.Modules.Prob.FDRand

namespace SciLean.Rand

variable
  {R} [RealScalar R]
  -- {X} [NormedAddCommGroup X] [NormedSpace ℝ X] [NormedSpace R X] [CompleteSpace X] [MeasurableSpace X]
  -- {Y} [NormedAddCommGroup Y] [NormedSpace ℝ Y] [NormedSpace R Y] [CompleteSpace Y] [MeasurableSpace Y]
  -- {Z} [NormedAddCommGroup Z] [NormedSpace ℝ Z] [NormedSpace R Z] [CompleteSpace Z] [MeasurableSpace Z]
  {X : Type _} [MeasurableSpace X] [AddCommGroup X] [Module ℝ X]
  {Y : Type _} [AddCommGroup Y] [Module ℝ Y] [MeasurableSpace Y]
  {Z : Type _} [AddCommGroup Z] [Module ℝ Z]

open Rand

@[rand_pull_E]
theorem bind_pull_mean (x : Rand X) (f : X → Rand Y) :
    (x >>= (fun x' => pure (f x').mean)).mean
    =
    (x >>= f).mean := by simp[rand_push_E]

@[rand_push_E]
theorem ite_push_E {c} [Decidable c] (t e : Rand X) (φ : X → Y):
    (if c then t else e).𝔼 φ = (if c then t.𝔼 φ else e.𝔼 φ) := by
  if h : c then simp[h] else simp[h]

-- I don't think this is a desirable `rand_pull_E` theorem as it duplicates the if statement
-- @[rand_pull_E]
theorem ite_pull_E {c} [Decidable c] (t e : Rand X) (φ ψ : X → Y):
    (if c then t.𝔼 φ else e.𝔼 ψ) = (if c then t else e).𝔼 (if c then φ else ψ) := by
  if h : c then simp[h] else simp[h]

@[rand_push_E]
theorem ite_push_mean {c} [Decidable c] (t e : Rand X) :
    (if c then t else e).mean = (if c then t.mean else e.mean) := by
  if h : c then simp[h] else simp[h]

@[rand_pull_E]
theorem ite_pull_mean {c} [Decidable c] (t e : Rand X) :
    (if c then t.mean else e.mean) = (if c then t else e).mean := by
  if h : c then simp[h] else simp[h]

@[rand_pull_E mid-1]
theorem ite_pull_mean_t {c} [Decidable c] (t : Rand X) (e : X) :
    (if c then t.mean else e) = (if c then t else pure e).mean := by
  if h : c then simp[h] else simp[h]

@[rand_pull_E mid-1]
theorem ite_pull_mean_f {c} [Decidable c] (t : X) (e : Rand X) :
    (if c then t else e.mean) = (if c then pure t else e).mean := by
  if h : c then simp[h] else simp[h]

-- this has messed up universes
@[rand_pull_E]
theorem pull_E_lambda (r : Rand Y) (f : X → Y → Z) :
    (fun x => r.𝔼 (fun y => f x y))
    =
    r.𝔼 (fun y x => f x y) := sorry_proof

-- this has messed up universes
@[rand_push_E]
theorem push_E_lambda (r : Rand Y) (f : X → Y → Z) :
    r.𝔼 (fun y x => f x y)
    =
    (fun x => r.𝔼 (fun y => f x y)) := sorry_proof

-- can't be simp as it has variable head
theorem pull_E_affine (r : Rand X) (φ : X → Y)
    (f : Y → Z) (hf : IsAffineMap ℝ f := by fun_prop) :
    (f (r.𝔼 φ)) = r.𝔼 (fun x => f (φ x)) := by sorry_proof -- have := hf; sorry_proof

@[rand_push_E]
theorem push_E_affine (r : Rand X) (φ : X → Y)
    (f : Y → Z) (hf : IsAffineMap ℝ f := by fun_prop) :
    r.𝔼 (fun x => f (φ x)) = (f (r.𝔼 φ)) := by rw[pull_E_affine (hf:=hf)]

@[rand_pull_E]
theorem pull_mean_add (x y : Rand X) :
    x.mean + y.mean
    =
    Rand.mean do
      let x' ← x
      let y' ← y
      return x' + y' := sorry_proof

@[rand_pull_E]
theorem pull_mean_add_1 (x : Rand X) (y : X) :
    x.mean + y
    =
    Rand.mean do
      let x' ← x
      return x' + y := sorry_proof

@[rand_pull_E]
theorem pull_mean_add_2 (x : X) (y : Rand X) :
    x + y.mean
    =
    Rand.mean do
      let y' ← y
      return x + y' := sorry_proof

@[rand_pull_E]
theorem pull_mean_sub (x y : Rand X) :
    x.mean - y.mean
    =
    Rand.mean do
      let x' ← x
      let y' ← y
      return x' - y'


import SciLean.Core.FunctionPropositions.CDifferentiable
import SciLean.Core.FunctionPropositions.IsSmoothLinearMap

import SciLean.Core.Meta.GenerateLinearMapSimp

import Mathlib.Tactic.FunTrans.Attr
import Mathlib.Tactic.FunTrans.Elab

set_option linter.unusedVariables false

open LeanColls

namespace SciLean

variable
  (K : Type _) [RCLike K]
  {X : Type _} [Vec K X]
  {Y : Type _} [Vec K Y]
  {Z : Type _} [Vec K Z]
  {W : Type _} [Vec K W]
  {ι : Type _} [IndexType ι] [LawfulIndexType ι] [DecidableEq ι]
  {E : ι → Type _} [∀ i, Vec K (E i)]

@[fun_trans]
noncomputable
def cderiv (f : X → Y) (x dx : X) : Y := Curve.deriv (fun t : K => f (x + t•dx)) 0

@[ftrans_simp]
noncomputable
def scalarCDeriv (f : K → X) (t : K) : X := cderiv K f t 1

-- Basic identities ------------------------------------------------------------
--------------------------------------------------------------------------------

variable {K}
@[fun_trans]
theorem cderiv_of_linear (f : X → Y) (hf : IsSmoothLinearMap K f)
  : cderiv K f = fun x dx => f dx := sorry_proof

@[simp, ftrans_simp]
theorem cderiv_apply
  (f : X → Y → Z) (x dx : X) (y : Y)
  : cderiv K f x dx y
    =
    cderiv K (fun x' => f x' y) x dx := sorry_proof

@[simp, ftrans_simp]
theorem cderiv_apply_zero
  (f : X → Y) (x : X)
  : cderiv K f x 0
    =
    0 := sorry_proof

@[fun_prop]
theorem cderiv.arg_dx.IsLinearMap_rule_simple
  (f : X → Y) (x : X) (hf : CDifferentiableAt K f x)
  : IsLinearMap K (fun dx => cderiv K f x dx) := sorry_proof

#generate_linear_map_simps SciLean.cderiv.arg_dx.IsLinearMap_rule_simple

@[fun_prop]
theorem cderiv.arg_dx.IsLinearMap_rule
  (f : X → Y) (x : X) (dx : W → X) (hf : CDifferentiableAt K f x) (hdx : IsLinearMap K dx)
  : IsLinearMap K (fun w => cderiv K f x (dx w)) := by fun_prop

@[fun_prop]
theorem cderiv.arg_f.IsLinearMap_rule
    (f : X → Y → Z) (hf : ∀ x, CDifferentiable K (fun y => f x y)) (hf' : ∀ y, IsLinearMap K (fun x => f x y)) :
    IsLinearMap K (fun x => cderiv K (f x ·)) := sorry_proof

@[fun_prop]
theorem cderiv.arg_f.IsSmoothLinearMap_rule
    (f : X → Y → Z) (hf : CDifferentiable K (fun (x,y) => f x y)) (hf' : ∀ y, IsLinearMap K (fun x => f x y)) :
    IsSmoothLinearMap K (fun x => cderiv K (f x ·)) := by constructor; fun_prop; sorry_proof /- differentiable and linear implies smooth -/

variable (K)

-- Basic lambda calculus rules -------------------------------------------------
--------------------------------------------------------------------------------
@[fun_trans]
theorem cderiv.id_rule :
    (cderiv K fun x : X => x) = fun _ => fun dx => dx := by sorry_proof

@[fun_trans]
theorem cderiv.const_rule (x : X) :
    (cderiv K fun _ : Y => x) = fun _ => fun dx => 0 := by sorry_proof

@[fun_trans]
theorem cderiv.comp_rule_at
  (f : Y → Z) (g : X → Y) (x : X)
  (hf : CDifferentiableAt K f (g x)) (hg : CDifferentiableAt K g x)
  : (cderiv K fun x : X => f (g x)) x
    =
    let y := g x
    fun dx =>
      let dy := cderiv K g x dx
      let dz := cderiv K f y dy
      dz :=
by sorry_proof

@[fun_trans]
theorem cderiv.comp_rule
  (f : Y → Z) (g : X → Y)
  (hf : CDifferentiable K f) (hg : CDifferentiable K g)
  : (cderiv K fun x : X => f (g x))
    =
    fun x =>
      let y := g x
      fun dx =>
        let dy := cderiv K g x dx
        let dz := cderiv K f y dy
        dz :=
by sorry_proof

@[fun_trans]
theorem cderiv.let_rule_at
  (f : X → Y → Z) (g : X → Y) (x : X)
  (hf : CDifferentiableAt K ↿f (x, g x))
  (hg : CDifferentiableAt K g x)
  : (cderiv K
      fun x : X =>
        let y := g x
        f x y) x
    =
    let y  := g x
    fun dx =>
      let dy := cderiv K g x dx
      let dz := cderiv K (fun xy : X×Y => f xy.1 xy.2) (x,y) (dx, dy)
      dz :=
by sorry_proof

@[fun_trans]
theorem cderiv.let_rule
  (f : X → Y → Z) (g : X → Y)
  (hf : CDifferentiable K fun xy : X×Y => f xy.1 xy.2) (hg : CDifferentiable K g)
  : (cderiv K fun x : X =>
       let y := g x
       f x y)
    =
    fun x =>
      let y  := g x
      fun dx =>
        let dy := cderiv K g x dx
        let dz := cderiv K (fun xy : X×Y => f xy.1 xy.2) (x,y) (dx, dy)
        dz :=
by sorry_proof

@[fun_trans]
theorem cderiv.apply_rule (i : ι) :
    (cderiv K fun (x : (i : ι) → E i) => x i)
    =
    fun _ => fun dx => dx i := by sorry_proof

@[fun_trans]
theorem cderiv.pi_rule_at
  (f : X → (i : ι) → E i) (x : X) (hf : ∀ i, CDifferentiableAt K (f · i) x)
  : (cderiv K fun (x : X) (i : ι) => f x i) x
    =
    fun dx => fun i =>
      cderiv K (f · i) x dx
  := by sorry_proof

@[fun_trans]
theorem cderiv.pi_rule
  (f : X → (i : ι) → E i) (hf : ∀ i, CDifferentiable K (f · i))
  : (cderiv K fun (x : X) (i : ι) => f x i)
    =
    fun x => fun dx => fun i =>
      cderiv K (f · i) x dx
  := by sorry_proof

--------------------------------------------------------------------------------
-- Function Rules --------------------------------------------------------------
--------------------------------------------------------------------------------

open SciLean
open LeanColls

variable
  {K : Type _} [RCLike K]
  {X : Type _} [Vec K X]
  {Y : Type _} [Vec K Y]
  {Z : Type _} [Vec K Z]
  {W : Type _} [Vec K W]
  {ι : Type _} [IndexType ι] [LawfulIndexType ι] [DecidableEq ι]
  {E : ι → Type _} [∀ i, Vec K (E i)]

-- Prod.mk -----------------------------------v---------------------------------
--------------------------------------------------------------------------------

@[fun_trans]
theorem Prod.mk.arg_fstsnd.cderiv_rule_at (x : X)
    (g : X → Y) (hg : CDifferentiableAt K g x)
    (f : X → Z) (hf : CDifferentiableAt K f x) :
    cderiv K (fun x => (g x, f x)) x
    =
    fun dx =>
      (cderiv K g x dx, cderiv K f x dx) := by
  sorry_proof

@[fun_trans]
  theorem Prod.mk.arg_fstsnd.cderiv_rule
    (g : X → Y) (hg : CDifferentiable K g)
    (f : X → Z) (hf : CDifferentiable K f) :
    cderiv K (fun x => (g x, f x))
    =
    fun x => fun dx =>
      (cderiv K g x dx, cderiv K f x dx) := by funext x; fun_trans

-- Prod.fst --------------------------------------------------------------------
--------------------------------------------------------------------------------

@[fun_trans]
theorem Prod.fst.arg_self.cderiv_rule_at (x : X)
    (f : X → Y×Z) (hf : CDifferentiableAt K f x) :
    cderiv K (fun x => (f x).1) x
    =
    fun dx => (cderiv K f x dx).1 := by fun_trans

@[fun_trans]
theorem Prod.fst.arg_self.cderiv_rule
    (f : X → Y×Z) (hf : CDifferentiable K f) :
    cderiv K (fun x => (f x).1)
    =
    fun x dx => (cderiv K f x dx).1 := by funext x; fun_trans

-- Prod.snd --------------------------------------------------------------------
--------------------------------------------------------------------------------

@[fun_trans]
theorem Prod.snd.arg_self.cderiv_rule_at (x : X)
    (f : X → Y×Z) (hf : CDifferentiableAt K f x) :
    cderiv K (fun x => (f x).2) x
    =
    fun dx => (cderiv K f x dx).2 := by fun_trans

@[fun_trans]
theorem Prod.snd.arg_self.cderiv_rule
    (f : X → Y×Z) (hf : CDifferentiable K f) :
    cderiv K (fun x => (f x).2)
    =
    fun x => fun dx => (cderiv K f x dx).2 := by funext x; fun_trans

-- HAdd.hAdd -------------------------------------------------------------------
--------------------------------------------------------------------------------

@[fun_trans]
theorem HAdd.hAdd.arg_a0a1.cderiv_rule_at (x : X)
    (f g : X → Y) (hf : CDifferentiableAt K f x) (hg : CDifferentiableAt K g x) :
    (cderiv K fun x => f x + g x) x
    =
    fun dx =>
      cderiv K f x dx + cderiv K g x dx := by fun_trans

@[fun_trans]
theorem HAdd.hAdd.arg_a0a1.cderiv_rule
    (f g : X → Y) (hf : CDifferentiable K f) (hg : CDifferentiable K g) :
    (cderiv K fun x => f x + g x)
    =
    fun x => fun dx =>
      cderiv K f x dx + cderiv K g x dx := by funext x; fun_trans

@[fun_trans]
theorem HAdd.hAdd.arg_a0.cderiv_rule
    (f : X → Y) (y : Y) :
    (cderiv K fun x => f x + y)
    =
    fun x dx =>
      cderiv K f x dx


import SciLean.Core.FunctionPropositions.CDifferentiable
import SciLean.Core.FunctionPropositions.IsSmoothLinearMap

import SciLean.Core.Meta.GenerateLinearMapSimp

import Mathlib.Tactic.FunTrans.Attr
import Mathlib.Tactic.FunTrans.Elab

set_option linter.unusedVariables false

open LeanColls

namespace SciLean

variable
  (K : Type _) [RCLike K]
  {X : Type _} [Vec K X]
  {Y : Type _} [Vec K Y]
  {Z : Type _} [Vec K Z]
  {W : Type _} [Vec K W]
  {ι : Type _} [IndexType ι] [LawfulIndexType ι] [DecidableEq ι]
  {E : ι → Type _} [∀ i, Vec K (E i)]

@[fun_trans]
noncomputable
def cderiv (f : X → Y) (x dx : X) : Y := Curve.deriv (fun t : K => f (x + t•dx)) 0

@[ftrans_simp]
noncomputable
def scalarCDeriv (f : K → X) (t : K) : X := cderiv K f t 1

-- Basic identities ------------------------------------------------------------
--------------------------------------------------------------------------------

variable {K}
@[fun_trans]
theorem cderiv_of_linear (f : X → Y) (hf : IsSmoothLinearMap K f)
  : cderiv K f = fun x dx => f dx := sorry_proof

@[simp, ftrans_simp]
theorem cderiv_apply
  (f : X → Y → Z) (x dx : X) (y : Y)
  : cderiv K f x dx y
    =
    cderiv K (fun x' => f x' y) x dx := sorry_proof

@[simp, ftrans_simp]
theorem cderiv_apply_zero
  (f : X → Y) (x : X)
  : cderiv K f x 0
    =
    0 := sorry_proof

@[fun_prop]
theorem cderiv.arg_dx.IsLinearMap_rule_simple
  (f : X → Y) (x : X) (hf : CDifferentiableAt K f x)
  : IsLinearMap K (fun dx => cderiv K f x dx) := sorry_proof

#generate_linear_map_simps SciLean.cderiv.arg_dx.IsLinearMap_rule_simple

@[fun_prop]
theorem cderiv.arg_dx.IsLinearMap_rule
  (f : X → Y) (x : X) (dx : W → X) (hf : CDifferentiableAt K f x) (hdx : IsLinearMap K dx)
  : IsLinearMap K (fun w => cderiv K f x (dx w)) := by fun_prop

@[fun_prop]
theorem cderiv.arg_f.IsLinearMap_rule
    (f : X → Y → Z) (hf : ∀ x, CDifferentiable K (fun y => f x y)) (hf' : ∀ y, IsLinearMap K (fun x => f x y)) :
    IsLinearMap K (fun x => cderiv K (f x ·)) := sorry_proof

@[fun_prop]
theorem cderiv.arg_f.IsSmoothLinearMap_rule
    (f : X → Y → Z) (hf : CDifferentiable K (fun (x,y) => f x y)) (hf' : ∀ y, IsLinearMap K (fun x => f x y)) :
    IsSmoothLinearMap K (fun x => cderiv K (f x ·)) := by constructor; fun_prop; sorry_proof /- differentiable and linear implies smooth -/

variable (K)

-- Basic lambda calculus rules -------------------------------------------------
--------------------------------------------------------------------------------
@[fun_trans]
theorem cderiv.id_rule :
    (cderiv K fun x : X => x) = fun _ => fun dx => dx := by sorry_proof

@[fun_trans]
theorem cderiv.const_rule (x : X) :
    (cderiv K fun _ : Y => x) = fun _ => fun dx => 0 := by sorry_proof

@[fun_trans]
theorem cderiv.comp_rule_at
  (f : Y → Z) (g : X → Y) (x : X)
  (hf : CDifferentiableAt K f (g x)) (hg : CDifferentiableAt K g x)
  : (cderiv K fun x : X => f (g x)) x
    =
    let y := g x
    fun dx =>
      let dy := cderiv K g x dx
      let dz := cderiv K f y dy
      dz :=
by sorry_proof

@[fun_trans]
theorem cderiv.comp_rule
  (f : Y → Z) (g : X → Y)
  (hf : CDifferentiable K f) (hg : CDifferentiable K g)
  : (cderiv K fun x : X => f (g x))
    =
    fun x =>
      let y := g x
      fun dx =>
        let dy := cderiv K g x dx
        let dz := cderiv K f y dy
        dz :=
by sorry_proof

@[fun_trans]
theorem cderiv.let_rule_at
  (f : X → Y → Z) (g : X → Y) (x : X)
  (hf : CDifferentiableAt K ↿f (x, g x))
  (hg : CDifferentiableAt K g x)
  : (cderiv K
      fun x : X =>
        let y := g x
        f x y) x
    =
    let y  := g x
    fun dx =>
      let dy := cderiv K g x dx
      let dz := cderiv K (fun xy : X×Y => f xy.1 xy.2) (x,y) (dx, dy)
      dz :=
by sorry_proof

@[fun_trans]
theorem cderiv.let_rule
  (f : X → Y → Z) (g : X → Y)
  (hf : CDifferentiable K fun xy : X×Y => f xy.1 xy.2) (hg : CDifferentiable K g)
  : (cderiv K fun x : X =>
       let y := g x
       f x y)
    =
    fun x =>
      let y  := g x
      fun dx =>
        let dy := cderiv K g x dx
        let dz := cderiv K (fun xy : X×Y => f xy.1 xy.2) (x,y) (dx, dy)
        dz :=
by sorry_proof

@[fun_trans]
theorem cderiv.apply_rule (i : ι) :
    (cderiv K fun (x : (i : ι) → E i) => x i)
    =
    fun _ => fun dx => dx i := by sorry_proof

@[fun_trans]
theorem cderiv.pi_rule_at
  (f : X → (i : ι) → E i) (x : X) (hf : ∀ i, CDifferentiableAt K (f · i) x)
  : (cderiv K fun (x : X) (i : ι) => f x i) x
    =
    fun dx => fun i =>
      cderiv K (f · i) x dx
  := by sorry_proof

@[fun_trans]
theorem cderiv.pi_rule
  (f : X → (i : ι) → E i) (hf : ∀ i, CDifferentiable K (f · i))
  : (cderiv K fun (x : X) (i : ι) => f x i)
    =
    fun x => fun dx => fun i =>
      cderiv K (f · i) x dx
  := by sorry_proof

--------------------------------------------------------------------------------
-- Function Rules --------------------------------------------------------------
--------------------------------------------------------------------------------

open SciLean
open LeanColls

variable
  {K : Type _} [RCLike K]
  {X : Type _} [Vec K X]
  {Y : Type _} [Vec K Y]
  {Z : Type _} [Vec K Z]
  {W : Type _} [Vec K W]
  {ι : Type _} [IndexType ι] [LawfulIndexType ι] [DecidableEq ι]
  {E : ι → Type _} [∀ i, Vec K (E i)]

-- Prod.mk -----------------------------------v---------------------------------
--------------------------------------------------------------------------------

@[fun_trans]
theorem Prod.mk.arg_fstsnd.cderiv_rule_at (x : X)
    (g : X → Y) (hg : CDifferentiableAt K g x)
    (f : X → Z) (hf : CDifferentiableAt K f x) :
    cderiv K (fun x => (g x, f x)) x
    =
    fun dx =>
      (cderiv K g x dx, cderiv K f x dx) := by
  sorry_proof

@[fun_trans]
  theorem Prod.mk.arg_fstsnd.cderiv_rule
    (g : X → Y) (hg : CDifferentiable K g)
    (f : X → Z) (hf : CDifferentiable K f) :
    cderiv K (fun x => (g x, f x))
    =
    fun x => fun dx =>
      (cderiv K g x dx, cderiv K f x dx) := by funext x; fun_trans

-- Prod.fst --------------------------------------------------------------------
--------------------------------------------------------------------------------

@[fun_trans]
theorem Prod.fst.arg_self.cderiv_rule_at (x : X)
    (f : X → Y×Z) (hf : CDifferentiableAt K f x) :
    cderiv K (fun x => (f x).1) x
    =
    fun dx => (cderiv K f x dx).1 := by fun_trans

@[fun_trans]
theorem Prod.fst.arg_self.cderiv_rule
    (f : X → Y×Z) (hf : CDifferentiable K f) :
    cderiv K (fun x => (f x).1)
    =
    fun x dx => (cderiv K f x dx).1 := by funext x; fun_trans

-- Prod.snd --------------------------------------------------------------------
--------------------------------------------------------------------------------

@[fun_trans]
theorem Prod.snd.arg_self.cderiv_rule_at (x : X)
    (f : X → Y×Z) (hf : CDifferentiableAt K f x) :
    cderiv K (fun x => (f x).2) x
    =
    fun dx => (cderiv K f x dx).2 := by fun_trans

@[fun_trans]
theorem Prod.snd.arg_self.cderiv_rule
    (f : X → Y×Z) (hf : CDifferentiable K f) :
    cderiv K (fun x => (f x).2)
    =
    fun x => fun dx => (cderiv K f x dx).2 := by funext x; fun_trans

-- HAdd.hAdd -------------------------------------------------------------------
--------------------------------------------------------------------------------

@[fun_trans]
theorem HAdd.hAdd.arg_a0a1.cderiv_rule_at (x : X)
    (f g : X → Y) (hf : CDifferentiableAt K f x) (hg : CDifferentiableAt K g x) :
    (cderiv K fun x => f x + g x) x
    =
    fun dx =>
      cderiv K f x dx + cderiv K g x dx := by fun_trans

@[fun_trans]
theorem HAdd.hAdd.arg_a0a1.cderiv_rule
    (f g : X → Y) (hf : CDifferentiable K f) (hg : CDifferentiable K g) :
    (cderiv K fun x => f x + g x)
    =
    fun x => fun dx =>
      cderiv K f x dx + cderiv K g x dx := by funext x; fun_trans

@[fun_trans]
theorem HAdd.hAdd.arg_a0.cderiv_rule
    (f : X → Y) (y : Y) :
    (cderiv K fun x => f x + y)
    =
    fun x dx =>
      cderiv K f x dx := by funext x; sorry_proof

@[fun_trans]
theorem HAdd.hAdd.arg_a1.cderiv_rule
    (f : X → Y) (y : Y) :
    (cderiv K fun x => y + f x)
    =
    fun x dx =>
      cderiv K f x dx


import SciLean.Core.FunctionPropositions.CDifferentiable
import SciLean.Core.FunctionPropositions.IsSmoothLinearMap

import SciLean.Core.Meta.GenerateLinearMapSimp

import Mathlib.Tactic.FunTrans.Attr
import Mathlib.Tactic.FunTrans.Elab

set_option linter.unusedVariables false

open LeanColls

namespace SciLean

variable
  (K : Type _) [RCLike K]
  {X : Type _} [Vec K X]
  {Y : Type _} [Vec K Y]
  {Z : Type _} [Vec K Z]
  {W : Type _} [Vec K W]
  {ι : Type _} [IndexType ι] [LawfulIndexType ι] [DecidableEq ι]
  {E : ι → Type _} [∀ i, Vec K (E i)]

@[fun_trans]
noncomputable
def cderiv (f : X → Y) (x dx : X) : Y := Curve.deriv (fun t : K => f (x + t•dx)) 0

@[ftrans_simp]
noncomputable
def scalarCDeriv (f : K → X) (t : K) : X := cderiv K f t 1

-- Basic identities ------------------------------------------------------------
--------------------------------------------------------------------------------

variable {K}
@[fun_trans]
theorem cderiv_of_linear (f : X → Y) (hf : IsSmoothLinearMap K f)
  : cderiv K f = fun x dx => f dx := sorry_proof

@[simp, ftrans_simp]
theorem cderiv_apply
  (f : X → Y → Z) (x dx : X) (y : Y)
  : cderiv K f x dx y
    =
    cderiv K (fun x' => f x' y) x dx := sorry_proof

@[simp, ftrans_simp]
theorem cderiv_apply_zero
  (f : X → Y) (x : X)
  : cderiv K f x 0
    =
    0 := sorry_proof

@[fun_prop]
theorem cderiv.arg_dx.IsLinearMap_rule_simple
  (f : X → Y) (x : X) (hf : CDifferentiableAt K f x)
  : IsLinearMap K (fun dx => cderiv K f x dx) := sorry_proof

#generate_linear_map_simps SciLean.cderiv.arg_dx.IsLinearMap_rule_simple

@[fun_prop]
theorem cderiv.arg_dx.IsLinearMap_rule
  (f : X → Y) (x : X) (dx : W → X) (hf : CDifferentiableAt K f x) (hdx : IsLinearMap K dx)
  : IsLinearMap K (fun w => cderiv K f x (dx w)) := by fun_prop

@[fun_prop]
theorem cderiv.arg_f.IsLinearMap_rule
    (f : X → Y → Z) (hf : ∀ x, CDifferentiable K (fun y => f x y)) (hf' : ∀ y, IsLinearMap K (fun x => f x y)) :
    IsLinearMap K (fun x => cderiv K (f x ·)) := sorry_proof

@[fun_prop]
theorem cderiv.arg_f.IsSmoothLinearMap_rule
    (f : X → Y → Z) (hf : CDifferentiable K (fun (x,y) => f x y)) (hf' : ∀ y, IsLinearMap K (fun x => f x y)) :
    IsSmoothLinearMap K (fun x => cderiv K (f x ·)) := by constructor; fun_prop; sorry_proof /- differentiable and linear implies smooth -/

variable (K)

-- Basic lambda calculus rules -------------------------------------------------
--------------------------------------------------------------------------------
@[fun_trans]
theorem cderiv.id_rule :
    (cderiv K fun x : X => x) = fun _ => fun dx => dx := by sorry_proof

@[fun_trans]
theorem cderiv.const_rule (x : X) :
    (cderiv K fun _ : Y => x) = fun _ => fun dx => 0 := by sorry_proof

@[fun_trans]
theorem cderiv.comp_rule_at
  (f : Y → Z) (g : X → Y) (x : X)
  (hf : CDifferentiableAt K f (g x)) (hg : CDifferentiableAt K g x)
  : (cderiv K fun x : X => f (g x)) x
    =
    let y := g x
    fun dx =>
      let dy := cderiv K g x dx
      let dz := cderiv K f y dy
      dz :=
by sorry_proof

@[fun_trans]
theorem cderiv.comp_rule
  (f : Y → Z) (g : X → Y)
  (hf : CDifferentiable K f) (hg : CDifferentiable K g)
  : (cderiv K fun x : X => f (g x))
    =
    fun x =>
      let y := g x
      fun dx =>
        let dy := cderiv K g x dx
        let dz := cderiv K f y dy
        dz :=
by sorry_proof

@[fun_trans]
theorem cderiv.let_rule_at
  (f : X → Y → Z) (g : X → Y) (x : X)
  (hf : CDifferentiableAt K ↿f (x, g x))
  (hg : CDifferentiableAt K g x)
  : (cderiv K
      fun x : X =>
        let y := g x
        f x y) x
    =
    let y  := g x
    fun dx =>
      let dy := cderiv K g x dx
      let dz := cderiv K (fun xy : X×Y => f xy.1 xy.2) (x,y) (dx, dy)
      dz :=
by sorry_proof

@[fun_trans]
theorem cderiv.let_rule
  (f : X → Y → Z) (g : X → Y)
  (hf : CDifferentiable K fun xy : X×Y => f xy.1 xy.2) (hg : CDifferentiable K g)
  : (cderiv K fun x : X =>
       let y := g x
       f x y)
    =
    fun x =>
      let y  := g x
      fun dx =>
        let dy := cderiv K g x dx
        let dz := cderiv K (fun xy : X×Y => f xy.1 xy.2) (x,y) (dx, dy)
        dz :=
by sorry_proof

@[fun_trans]
theorem cderiv.apply_rule (i : ι) :
    (cderiv K fun (x : (i : ι) → E i) => x i)
    =
    fun _ => fun dx => dx i := by sorry_proof

@[fun_trans]
theorem cderiv.pi_rule_at
  (f : X → (i : ι) → E i) (x : X) (hf : ∀ i, CDifferentiableAt K (f · i) x)
  : (cderiv K fun (x : X) (i : ι) => f x i) x
    =
    fun dx => fun i =>
      cderiv K (f · i) x dx
  := by sorry_proof

@[fun_trans]
theorem cderiv.pi_rule
  (f : X → (i : ι) → E i) (hf : ∀ i, CDifferentiable K (f · i))
  : (cderiv K fun (x : X) (i : ι) => f x i)
    =
    fun x => fun dx => fun i =>
      cderiv K (f · i) x dx
  := by sorry_proof

--------------------------------------------------------------------------------
-- Function Rules --------------------------------------------------------------
--------------------------------------------------------------------------------

open SciLean
open LeanColls

variable
  {K : Type _} [RCLike K]
  {X : Type _} [Vec K X]
  {Y : Type _} [Vec K Y]
  {Z : Type _} [Vec K Z]
  {W : Type _} [Vec K W]
  {ι : Type _} [IndexType ι] [LawfulIndexType ι] [DecidableEq ι]
  {E : ι → Type _} [∀ i, Vec K (E i)]

-- Prod.mk -----------------------------------v---------------------------------
--------------------------------------------------------------------------------

@[fun_trans]
theorem Prod.mk.arg_fstsnd.cderiv_rule_at (x : X)
    (g : X → Y) (hg : CDifferentiableAt K g x)
    (f : X → Z) (hf : CDifferentiableAt K f x) :
    cderiv K (fun x => (g x, f x)) x
    =
    fun dx =>
      (cderiv K g x dx, cderiv K f x dx) := by
  sorry_proof

@[fun_trans]
  theorem Prod.mk.arg_fstsnd.cderiv_rule
    (g : X → Y) (hg : CDifferentiable K g)
    (f : X → Z) (hf : CDifferentiable K f) :
    cderiv K (fun x => (g x, f x))
    =
    fun x => fun dx =>
      (cderiv K g x dx, cderiv K f x dx) := by funext x; fun_trans

-- Prod.fst --------------------------------------------------------------------
--------------------------------------------------------------------------------

@[fun_trans]
theorem Prod.fst.arg_self.cderiv_rule_at (x : X)
    (f : X → Y×Z) (hf : CDifferentiableAt K f x) :
    cderiv K (fun x => (f x).1) x
    =
    fun dx => (cderiv K f x dx).1 := by fun_trans

@[fun_trans]
theorem Prod.fst.arg_self.cderiv_rule
    (f : X → Y×Z) (hf : CDifferentiable K f) :
    cderiv K (fun x => (f x).1)
    =
    fun x dx => (cderiv K f x dx).1 := by funext x; fun_trans

-- Prod.snd --------------------------------------------------------------------
--------------------------------------------------------------------------------

@[fun_trans]
theorem Prod.snd.arg_self.cderiv_rule_at (x : X)
    (f : X → Y×Z) (hf : CDifferentiableAt K f x) :
    cderiv K (fun x => (f x).2) x
    =
    fun dx => (cderiv K f x dx).2 := by fun_trans

@[fun_trans]
theorem Prod.snd.arg_self.cderiv_rule
    (f : X → Y×Z) (hf : CDifferentiable K f) :
    cderiv K (fun x => (f x).2)
    =
    fun x => fun dx => (cderiv K f x dx).2 := by funext x; fun_trans

-- HAdd.hAdd -------------------------------------------------------------------
--------------------------------------------------------------------------------

@[fun_trans]
theorem HAdd.hAdd.arg_a0a1.cderiv_rule_at (x : X)
    (f g : X → Y) (hf : CDifferentiableAt K f x) (hg : CDifferentiableAt K g x) :
    (cderiv K fun x => f x + g x) x
    =
    fun dx =>
      cderiv K f x dx + cderiv K g x dx := by fun_trans

@[fun_trans]
theorem HAdd.hAdd.arg_a0a1.cderiv_rule
    (f g : X → Y) (hf : CDifferentiable K f) (hg : CDifferentiable K g) :
    (cderiv K fun x => f x + g x)
    =
    fun x => fun dx =>
      cderiv K f x dx + cderiv K g x dx := by funext x; fun_trans

@[fun_trans]
theorem HAdd.hAdd.arg_a0.cderiv_rule
    (f : X → Y) (y : Y) :
    (cderiv K fun x => f x + y)
    =
    fun x dx =>
      cderiv K f x dx := by funext x; sorry_proof

@[fun_trans]
theorem HAdd.hAdd.arg_a1.cderiv_rule
    (f : X → Y) (y : Y) :
    (cderiv K fun x => y + f x)
    =
    fun x dx =>
      cderiv K f x dx := by funext x; sorry_proof

-- HSub.hSub -------------------------------------------------------------------
--------------------------------------------------------------------------------

@[fun_trans]
theorem HSub.hSub.arg_a0a1.cderiv_rule_at (x : X)
    (f g : X → Y) (hf : CDifferentiableAt K f x) (hg : CDifferentiableAt K g x) :
    (cderiv K fun x => f x - g x) x
    =
    fun dx =>
      cderiv K f x dx - cderiv K g x dx := by fun_trans

@[fun_trans]
theorem HSub.hSub.arg_a0a1.cderiv_rule
    (f g : X → Y) (hf : CDifferentiable K f) (hg : CDifferentiable K g) :
    (cderiv K fun x => f x - g x)
    =
    fun x => fun dx =>
      cderiv K f x dx - cderiv K g x dx := by funext x; fun_trans

@[fun_trans]
theorem HSub.hSub.arg_a0.cderiv_rule
    (f : X → Y) (y : Y) :
    (cderiv K fun x => f x - y)
    =
    fun x dx =>
      cderiv K f x dx


import SciLean.Core.FunctionPropositions.CDifferentiable
import SciLean.Core.FunctionPropositions.IsSmoothLinearMap

import SciLean.Core.Meta.GenerateLinearMapSimp

import Mathlib.Tactic.FunTrans.Attr
import Mathlib.Tactic.FunTrans.Elab

set_option linter.unusedVariables false

open LeanColls

namespace SciLean

variable
  (K : Type _) [RCLike K]
  {X : Type _} [Vec K X]
  {Y : Type _} [Vec K Y]
  {Z : Type _} [Vec K Z]
  {W : Type _} [Vec K W]
  {ι : Type _} [IndexType ι] [LawfulIndexType ι] [DecidableEq ι]
  {E : ι → Type _} [∀ i, Vec K (E i)]

@[fun_trans]
noncomputable
def cderiv (f : X → Y) (x dx : X) : Y := Curve.deriv (fun t : K => f (x + t•dx)) 0

@[ftrans_simp]
noncomputable
def scalarCDeriv (f : K → X) (t : K) : X := cderiv K f t 1

-- Basic identities ------------------------------------------------------------
--------------------------------------------------------------------------------

variable {K}
@[fun_trans]
theorem cderiv_of_linear (f : X → Y) (hf : IsSmoothLinearMap K f)
  : cderiv K f = fun x dx => f dx := sorry_proof

@[simp, ftrans_simp]
theorem cderiv_apply
  (f : X → Y → Z) (x dx : X) (y : Y)
  : cderiv K f x dx y
    =
    cderiv K (fun x' => f x' y) x dx := sorry_proof

@[simp, ftrans_simp]
theorem cderiv_apply_zero
  (f : X → Y) (x : X)
  : cderiv K f x 0
    =
    0 := sorry_proof

@[fun_prop]
theorem cderiv.arg_dx.IsLinearMap_rule_simple
  (f : X → Y) (x : X) (hf : CDifferentiableAt K f x)
  : IsLinearMap K (fun dx => cderiv K f x dx) := sorry_proof

#generate_linear_map_simps SciLean.cderiv.arg_dx.IsLinearMap_rule_simple

@[fun_prop]
theorem cderiv.arg_dx.IsLinearMap_rule
  (f : X → Y) (x : X) (dx : W → X) (hf : CDifferentiableAt K f x) (hdx : IsLinearMap K dx)
  : IsLinearMap K (fun w => cderiv K f x (dx w)) := by fun_prop

@[fun_prop]
theorem cderiv.arg_f.IsLinearMap_rule
    (f : X → Y → Z) (hf : ∀ x, CDifferentiable K (fun y => f x y)) (hf' : ∀ y, IsLinearMap K (fun x => f x y)) :
    IsLinearMap K (fun x => cderiv K (f x ·)) := sorry_proof

@[fun_prop]
theorem cderiv.arg_f.IsSmoothLinearMap_rule
    (f : X → Y → Z) (hf : CDifferentiable K (fun (x,y) => f x y)) (hf' : ∀ y, IsLinearMap K (fun x => f x y)) :
    IsSmoothLinearMap K (fun x => cderiv K (f x ·)) := by constructor; fun_prop; sorry_proof /- differentiable and linear implies smooth -/

variable (K)

-- Basic lambda calculus rules -------------------------------------------------
--------------------------------------------------------------------------------
@[fun_trans]
theorem cderiv.id_rule :
    (cderiv K fun x : X => x) = fun _ => fun dx => dx := by sorry_proof

@[fun_trans]
theorem cderiv.const_rule (x : X) :
    (cderiv K fun _ : Y => x) = fun _ => fun dx => 0 := by sorry_proof

@[fun_trans]
theorem cderiv.comp_rule_at
  (f : Y → Z) (g : X → Y) (x : X)
  (hf : CDifferentiableAt K f (g x)) (hg : CDifferentiableAt K g x)
  : (cderiv K fun x : X => f (g x)) x
    =
    let y := g x
    fun dx =>
      let dy := cderiv K g x dx
      let dz := cderiv K f y dy
      dz :=
by sorry_proof

@[fun_trans]
theorem cderiv.comp_rule
  (f : Y → Z) (g : X → Y)
  (hf : CDifferentiable K f) (hg : CDifferentiable K g)
  : (cderiv K fun x : X => f (g x))
    =
    fun x =>
      let y := g x
      fun dx =>
        let dy := cderiv K g x dx
        let dz := cderiv K f y dy
        dz :=
by sorry_proof

@[fun_trans]
theorem cderiv.let_rule_at
  (f : X → Y → Z) (g : X → Y) (x : X)
  (hf : CDifferentiableAt K ↿f (x, g x))
  (hg : CDifferentiableAt K g x)
  : (cderiv K
      fun x : X =>
        let y := g x
        f x y) x
    =
    let y  := g x
    fun dx =>
      let dy := cderiv K g x dx
      let dz := cderiv K (fun xy : X×Y => f xy.1 xy.2) (x,y) (dx, dy)
      dz :=
by sorry_proof

@[fun_trans]
theorem cderiv.let_rule
  (f : X → Y → Z) (g : X → Y)
  (hf : CDifferentiable K fun xy : X×Y => f xy.1 xy.2) (hg : CDifferentiable K g)
  : (cderiv K fun x : X =>
       let y := g x
       f x y)
    =
    fun x =>
      let y  := g x
      fun dx =>
        let dy := cderiv K g x dx
        let dz := cderiv K (fun xy : X×Y => f xy.1 xy.2) (x,y) (dx, dy)
        dz :=
by sorry_proof

@[fun_trans]
theorem cderiv.apply_rule (i : ι) :
    (cderiv K fun (x : (i : ι) → E i) => x i)
    =
    fun _ => fun dx => dx i := by sorry_proof

@[fun_trans]
theorem cderiv.pi_rule_at
  (f : X → (i : ι) → E i) (x : X) (hf : ∀ i, CDifferentiableAt K (f · i) x)
  : (cderiv K fun (x : X) (i : ι) => f x i) x
    =
    fun dx => fun i =>
      cderiv K (f · i) x dx
  := by sorry_proof

@[fun_trans]
theorem cderiv.pi_rule
  (f : X → (i : ι) → E i) (hf : ∀ i, CDifferentiable K (f · i))
  : (cderiv K fun (x : X) (i : ι) => f x i)
    =
    fun x => fun dx => fun i =>
      cderiv K (f · i) x dx
  := by sorry_proof

--------------------------------------------------------------------------------
-- Function Rules --------------------------------------------------------------
--------------------------------------------------------------------------------

open SciLean
open LeanColls

variable
  {K : Type _} [RCLike K]
  {X : Type _} [Vec K X]
  {Y : Type _} [Vec K Y]
  {Z : Type _} [Vec K Z]
  {W : Type _} [Vec K W]
  {ι : Type _} [IndexType ι] [LawfulIndexType ι] [DecidableEq ι]
  {E : ι → Type _} [∀ i, Vec K (E i)]

-- Prod.mk -----------------------------------v---------------------------------
--------------------------------------------------------------------------------

@[fun_trans]
theorem Prod.mk.arg_fstsnd.cderiv_rule_at (x : X)
    (g : X → Y) (hg : CDifferentiableAt K g x)
    (f : X → Z) (hf : CDifferentiableAt K f x) :
    cderiv K (fun x => (g x, f x)) x
    =
    fun dx =>
      (cderiv K g x dx, cderiv K f x dx) := by
  sorry_proof

@[fun_trans]
  theorem Prod.mk.arg_fstsnd.cderiv_rule
    (g : X → Y) (hg : CDifferentiable K g)
    (f : X → Z) (hf : CDifferentiable K f) :
    cderiv K (fun x => (g x, f x))
    =
    fun x => fun dx =>
      (cderiv K g x dx, cderiv K f x dx) := by funext x; fun_trans

-- Prod.fst --------------------------------------------------------------------
--------------------------------------------------------------------------------

@[fun_trans]
theorem Prod.fst.arg_self.cderiv_rule_at (x : X)
    (f : X → Y×Z) (hf : CDifferentiableAt K f x) :
    cderiv K (fun x => (f x).1) x
    =
    fun dx => (cderiv K f x dx).1 := by fun_trans

@[fun_trans]
theorem Prod.fst.arg_self.cderiv_rule
    (f : X → Y×Z) (hf : CDifferentiable K f) :
    cderiv K (fun x => (f x).1)
    =
    fun x dx => (cderiv K f x dx).1 := by funext x; fun_trans

-- Prod.snd --------------------------------------------------------------------
--------------------------------------------------------------------------------

@[fun_trans]
theorem Prod.snd.arg_self.cderiv_rule_at (x : X)
    (f : X → Y×Z) (hf : CDifferentiableAt K f x) :
    cderiv K (fun x => (f x).2) x
    =
    fun dx => (cderiv K f x dx).2 := by fun_trans

@[fun_trans]
theorem Prod.snd.arg_self.cderiv_rule
    (f : X → Y×Z) (hf : CDifferentiable K f) :
    cderiv K (fun x => (f x).2)
    =
    fun x => fun dx => (cderiv K f x dx).2 := by funext x; fun_trans

-- HAdd.hAdd -------------------------------------------------------------------
--------------------------------------------------------------------------------

@[fun_trans]
theorem HAdd.hAdd.arg_a0a1.cderiv_rule_at (x : X)
    (f g : X → Y) (hf : CDifferentiableAt K f x) (hg : CDifferentiableAt K g x) :
    (cderiv K fun x => f x + g x) x
    =
    fun dx =>
      cderiv K f x dx + cderiv K g x dx := by fun_trans

@[fun_trans]
theorem HAdd.hAdd.arg_a0a1.cderiv_rule
    (f g : X → Y) (hf : CDifferentiable K f) (hg : CDifferentiable K g) :
    (cderiv K fun x => f x + g x)
    =
    fun x => fun dx =>
      cderiv K f x dx + cderiv K g x dx := by funext x; fun_trans

@[fun_trans]
theorem HAdd.hAdd.arg_a0.cderiv_rule
    (f : X → Y) (y : Y) :
    (cderiv K fun x => f x + y)
    =
    fun x dx =>
      cderiv K f x dx := by funext x; sorry_proof

@[fun_trans]
theorem HAdd.hAdd.arg_a1.cderiv_rule
    (f : X → Y) (y : Y) :
    (cderiv K fun x => y + f x)
    =
    fun x dx =>
      cderiv K f x dx := by funext x; sorry_proof

-- HSub.hSub -------------------------------------------------------------------
--------------------------------------------------------------------------------

@[fun_trans]
theorem HSub.hSub.arg_a0a1.cderiv_rule_at (x : X)
    (f g : X → Y) (hf : CDifferentiableAt K f x) (hg : CDifferentiableAt K g x) :
    (cderiv K fun x => f x - g x) x
    =
    fun dx =>
      cderiv K f x dx - cderiv K g x dx := by fun_trans

@[fun_trans]
theorem HSub.hSub.arg_a0a1.cderiv_rule
    (f g : X → Y) (hf : CDifferentiable K f) (hg : CDifferentiable K g) :
    (cderiv K fun x => f x - g x)
    =
    fun x => fun dx =>
      cderiv K f x dx - cderiv K g x dx := by funext x; fun_trans

@[fun_trans]
theorem HSub.hSub.arg_a0.cderiv_rule
    (f : X → Y) (y : Y) :
    (cderiv K fun x => f x - y)
    =
    fun x dx =>
      cderiv K f x dx := by funext x; sorry_proof

@[fun_trans]
theorem HSub.hSub.arg_a1.cderiv_rule
    (f : X → Y) (y : Y) :
    (cderiv K fun x => y - f x)
    =
    fun x dx =>
      - cderiv K f x dx


import SciLean.Core.Rand.Rand
import SciLean.Core.Rand.Condition

import SciLean.Core.Distribution.ParametricDistribDeriv

import Mathlib.MeasureTheory.Constructions.Prod.Basic

namespace SciLean

open MeasureTheory
variable
  {R} [RealScalar R]
  {X Z} [MeasurableSpace X] [MeasurableSpace Z]

/-- Kullback–Leibler divergence of `Dₖₗ(P‖Q)` -/
noncomputable
def KLDiv (P Q : Rand X) : R := P.𝔼 (fun x => Scalar.log (P.pdf R Q.ℙ x))

/-- Evidence Lower Bound

```
𝔼_Q [log Q(Z) - log P(Z,X)]
```
reference: https://en.wikipedia.org/wiki/Variational_Bayesian_methods#Evidence_lower_bound
 -/
noncomputable
def ELBO {X Z} [MeasureSpace Z] [MeasureSpace X]
    (P : Rand (Z×X)) (Q : Rand Z) (x : X) : R :=
  - Q.𝔼 (fun z => Scalar.log (Q.pdf R volume z) - Scalar.log (P.pdf R volume (z,x)))

/-- Express `Kullback–Leibler divergence` as log evidence + ELBO -/
theorem kldiv_elbo
    {X Z} [MeasureSpace Z] [MeasureSpace X]
    (P : Rand (Z×X)) (Q : Rand Z) (x : X) :
    KLDiv Q (P.conditionSnd x)
    =
    (Scalar.log (P.snd.pdf R volume x)) - ELBO P Q x


import SciLean.Core.Rand.Rand
import SciLean.Core.Rand.Condition

import SciLean.Core.Distribution.ParametricDistribDeriv

import Mathlib.MeasureTheory.Constructions.Prod.Basic

namespace SciLean

open MeasureTheory
variable
  {R} [RealScalar R]
  {X Z} [MeasurableSpace X] [MeasurableSpace Z]

/-- Kullback–Leibler divergence of `Dₖₗ(P‖Q)` -/
noncomputable
def KLDiv (P Q : Rand X) : R := P.𝔼 (fun x => Scalar.log (P.pdf R Q.ℙ x))

/-- Evidence Lower Bound

```
𝔼_Q [log Q(Z) - log P(Z,X)]
```
reference: https://en.wikipedia.org/wiki/Variational_Bayesian_methods#Evidence_lower_bound
 -/
noncomputable
def ELBO {X Z} [MeasureSpace Z] [MeasureSpace X]
    (P : Rand (Z×X)) (Q : Rand Z) (x : X) : R :=
  - Q.𝔼 (fun z => Scalar.log (Q.pdf R volume z) - Scalar.log (P.pdf R volume (z,x)))

/-- Express `Kullback–Leibler divergence` as log evidence + ELBO -/
theorem kldiv_elbo
    {X Z} [MeasureSpace Z] [MeasureSpace X]
    (P : Rand (Z×X)) (Q : Rand Z) (x : X) :
    KLDiv Q (P.conditionSnd x)
    =
    (Scalar.log (P.snd.pdf R volume x)) - ELBO P Q x := sorry_proof

----------------------------------------------------------------------------------------------------
-- KLDiv properties --------------------------------------------------------------------------------
----------------------------------------------------------------------------------------------------

variable
  {W} [Vec R W]
  [Vec R X]

@[fun_trans]
theorem KLDiv.arg_P.cderiv_rule (P : W → Rand X) (Q : Rand X) :
    cderiv R (fun w => KLDiv (R:=R) (P w) Q)
    =
    fun w dw =>
      let dP := parDistribDeriv (fun w => (P w).ℙ.toDistribution (R:=R)) w dw
      dP.extAction' (fun x => Scalar.log ((P w).pdf R Q.ℙ x) - 1)


import SciLean.Core.Rand.Rand
import SciLean.Core.Distribution.ParametricDistribFwdDeriv

namespace SciLean

open MeasureTheory

variable
  {R} [RealScalar R]
  {W} [Vec R W]
  {X} [MeasurableSpace X] [Vec R X]
  {Y} [Vec R Y] [Module ℝ Y]

set_default_scalar R

@[fun_trans]
theorem Rand.𝔼.arg_r.cderiv_rule (r : W → Rand X) (f : X → Y) :
  cderiv R (fun w => (r w).𝔼 f)
  =
  fun w dw =>
    let d := parDistribDeriv (fun w => (r w).ℙ.toDistribution (R:=R)) w dw
    d.extAction f (fun r ⊸ fun y ⊸ ((r • y) : Y))


import Mathlib.MeasureTheory.Measure.MeasureSpaceDef
import Mathlib.MeasureTheory.Measure.Restrict
import Mathlib.MeasureTheory.Measure.Dirac
-- import Mathlib.Data.RCLike.Basic
-- import Mathlib.Analysis.LocallyConvex.Basic
-- import Mathlib.Topology.Algebra.Module.LocallyConvex

import SciLean.Core.FunctionTransformations
import SciLean.Util.SorryProof

open MeasureTheory

namespace SciLean

/-- Convenient integral - the integral I need :)
It should be Bochner integral but it should integrate function valued function point wise i.e.
```
∫ x, fun y => f x y = fun y => ∫ x, f x y
```
where rhs can be understoods as Bochenr integral and lhs defined thie `cintegral`. -/
noncomputable
opaque cintegral {α} [MeasurableSpace α] {X} [AddCommGroup X] [Module ℝ X]
  -- dragging along all of these typeclasses is really really annoying
  -- [AddCommGroup X] [TopologicalSpace X] [TopologicalAddGroup X] [Module ℝ X] [LocallyConvexSpace ℝ X]
    (f : α → X) (μ : Measure α) : X := 0

opaque CIntegrable {α} [MeasurableSpace α] {X} [AddCommGroup X] [Module ℝ X]
    (f : α → X) (μ : Measure α) : Prop

open Lean Parser  Term
syntax "∫' " funBinder (" in " term)? ", " term:60 (" ∂" term:70)? : term

macro_rules
| `(∫' $x:funBinder, $b) => `(cintegral (fun $x => $b) (by volume_tac))
| `(∫' $x:funBinder, $b ∂$μ) => `(cintegral (fun $x => $b) $μ)
| `(∫' $x:funBinder in $set, $b) => `(cintegral (fun $x => $b) (Measure.restrict (by volume_tac) $set))
| `(∫' $x:funBinder in $set, $b ∂$μ) => `(cintegral (fun $x => $b) (Measure.restrict $μ $set))

@[app_unexpander cintegral] def unexpandCIntegral : Lean.PrettyPrinter.Unexpander

  | `($(_) $f:term volume) =>
    match f with
    | `(fun $x => $b) => `(∫' $x, $b)
    | `(fun $x $xs* => $b) => `(∫' $x, fun $xs* => $b)
    | _ => throw ()

  | `($(_) $f:term $μ) =>
    match f, μ with
    -- |      `(fun $x => $b), `(Measure.restrict volume $A) => `(∫' $x in $A, $b)
    -- | `(fun $x $xs* => $b), `(Measure.restrict volume $A) => `(∫' $x in $A, (fun $xs* => $b))
    |      `(fun $x => $b), `(Measure.restrict $μ $A) => `(∫' $x in $A, $b ∂$μ)
    | `(fun $x $xs* => $b), `(Measure.restrict $μ $A) => `(∫' $x in $A, (fun $xs* => $b) ∂$μ)
    -- |      `(fun $x => $b), `(volume) => `(∫' $x, $b)
    -- | `(fun $x $xs* => $b), `(volume) => `(∫' $x, (fun $xs* => $b))
    |      `(fun $x => $b), _ => `(∫' $x, $b ∂$μ)
    | `(fun $x $xs* => $b), _ => `(∫' $x, (fun $xs* => $b) ∂$μ)
    | _, _ => throw ()

  | _ => throw ()

----------------------------------------------------------------------------------------------------
-- Algebra -----------------------------------------------------------------------------------------
----------------------------------------------------------------------------------------------------
section Algebra

variable
  {α} [MeasurableSpace α]
  {β} [MeasurableSpace β]
  {X} [AddCommGroup X] [Module ℝ X]
  {Y} [AddCommGroup Y] [Module ℝ Y]

@[simp, ftrans_simp]
theorem cintegral_zero (μ : Measure α): ∫' _, (0 : X) ∂μ = 0 := sorry_proof

@[add_pull] -- @[integral_push]
theorem cintegral_add {f g : α → X} {μ : Measure α}
    (hf : CIntegrable f μ) (hg : CIntegrable g μ) :
    ∫' x, f x + g x ∂μ = ∫' x, f x ∂μ + ∫' x, g x ∂μ := sorry_proof

@[add_push] -- @[integral_push]
theorem cintegral_add' {f g : α → X} {μ : Measure α}
    (hf : CIntegrable f μ) (hg : CIntegrable g μ) :
    ∫' x, f x ∂μ + ∫' x, g x ∂μ = ∫' x, f x + g x ∂μ  := sorry_proof

@[smul_pull]-- @[integral_push]
theorem cintegral_smul {R} [Semiring R] [Module R X] {f : α → X} (r : R) :
    ∫' x, r • f x ∂μ = r • ∫' x, f x ∂μ := sorry_proof

@[smul_push] -- @[integral_pull]
theorem cintegral_smul' {R} [Semiring R] [Module R X] {f : α → X} (r : R) :
    r • ∫' x, f x ∂μ = ∫' x, r • f x ∂μ := sorry_proof

@[add_pull] -- @[integral_push]
theorem cintegral_add_measures {f : α → X} {μ ν : Measure α}
    (hμ : CIntegrable f μ) (hν : CIntegrable f ν) :
    ∫' x, f x ∂(μ+ν) = ∫' x, f x ∂μ + ∫' x, f x ∂ν := sorry_proof

@[simp, ftrans_simp]
theorem cintegral_dirac {f : α → X} (p : α) :
    ∫' x, f x ∂(Measure.dirac p) = f p := sorry_proof

@[simp,ftrans_simp]
theorem cintegral_unit {f : Unit → X} :
    ∫' x, f x  = f () := sorry_proof

-- @[integral_push]
theorem cintegral_prod_mk {f : α → X} {g : α → Y} {μ : Measure α} :
    ∫' x, (f x, g x) ∂μ = (∫' x, f x ∂μ, ∫' x, g x ∂μ) := sorry_proof

-- @[integral_pull]
theorem cintegral_prod_mk' {f : α → X} {g : α → Y} {μ : Measure α} :
    (∫' x, f x ∂μ, ∫' x, g x ∂μ) = ∫' x, (f x, g x) ∂μ := sorry_proof

theorem cintegral_lambda {α} (f : α → β → X) (μ : Measure β) :
    (fun x => ∫' y, f x y ∂μ)
    =
    ∫' y, (fun x => f x y) ∂μ := sorry_proof

theorem cintegral.arg_f.push_lambda {α} (f : α → β → X) (μ : Measure β) :
    (fun x => ∫' y, f x y ∂μ)
    =
    ∫' y, (fun x => f x y) ∂μ := sorry_proof

@[simp, ftrans_simp]
theorem cintegral_measure_map (g : β → X) (f : α → β) (μ : Measure α) :
    (∫' y, g y ∂(μ.map f))
    =
    ∫' x, g (f x) ∂μ


import SciLean.Core.Distribution.Basic
import SciLean.Core.FunctionTransformations
import SciLean.Core.FunctionPropositions
import SciLean.Core.Notation

open MeasureTheory

namespace SciLean

open Distribution

variable
  {R} [RealScalar R]
  {W} [Vec R W]
  {X} [Vec R X] [MeasureSpace X]
  {Y} [Vec R Y] [Module ℝ Y]
  {Z} [Vec R Z] [Module ℝ Z]
  {U} [Vec R U] -- [Module ℝ U]
  {V} [Vec R V] -- [Module ℝ U]

set_default_scalar R

noncomputable
def diracDeriv (x dx : X) : 𝒟' X := fun φ ⊸ cderiv R φ x dx

@[fun_prop]
def DistribDifferentiableAt (f : X → 𝒟'(Y,Z)) (x : X) :=
  ∀ (φ : X → 𝒟 Y), CDifferentiableAt R φ x → CDifferentiableAt R (fun x => f x (φ x)) x

theorem distribDifferentiableAt_const_test_fun
    {f : X → 𝒟'(Y,Z)} {x : X}
    (hf : DistribDifferentiableAt f x)
    {φ : 𝒟 Y} :
    CDifferentiableAt R (fun x => f x φ) x := by
  apply hf
  fun_prop

@[fun_prop]
def DistribDifferentiable (f : X → 𝒟'(Y,Z)) :=
  ∀ x, DistribDifferentiableAt f x

-- TODO:
-- probably change the definition of `parDistribDeriv` to:
-- ⟨⟨fun φ =>
--    if h : DistribDifferentiableAt f x then
--      ∂ (x':=x;dx), ⟪f x', φ⟫
--    else
--      0 , sorry_proof⟩⟩
-- I believe in that case the function is indeed linear in φ

open Classical in
@[fun_trans]
noncomputable
def parDistribDeriv (f : X → 𝒟'(Y,Z)) (x dx : X) : 𝒟'(Y,Z) :=
  ⟨fun φ => ∂ (x':=x;dx), f x' φ, sorry_proof⟩

@[simp, ftrans_simp]
theorem action_parDistribDeriv (f : X → 𝒟'(Y,Z)) (x dx : X) (φ : 𝒟 Y) :
    parDistribDeriv f x dx φ = ∂ (x':=x;dx), f x' φ := rfl

----------------------------------------------------------------------------------------------------
-- Const rule --------------------------------------------------------------------------------------
----------------------------------------------------------------------------------------------------

@[fun_prop]
theorem DistribDiffrentiable.const_rule (T : 𝒟'(X,Y)) :
    DistribDifferentiable (fun _ : W => T) := by
  intro _ φ hφ; simp; fun_prop

@[fun_trans]
theorem parDistribDeriv.const_rule (T : 𝒟'(X,Y)) :
    parDistribDeriv (fun _ : W => T)
    =
    fun w dw =>
      0 := by
  funext w dw; ext φ
  unfold parDistribDeriv
  fun_trans

----------------------------------------------------------------------------------------------------
-- Pure --------------------------------------------------------------------------------------------
----------------------------------------------------------------------------------------------------

@[fun_prop]
theorem dirac.arg_xy.DistribDiffrentiable_rule
    (x : W → X) (hx : CDifferentiable R x) :
    DistribDifferentiable (R:=R) (fun w => dirac (x w))  := by
  intro x
  unfold DistribDifferentiableAt
  intro φ hφ
  simp [action_dirac, dirac]
  fun_prop

@[fun_trans]
theorem dirac.arg_x.parDistribDeriv_rule
    (x : W → X) (hx : CDifferentiable R x) :
    parDistribDeriv (R:=R) (fun w => dirac (x w))
    =
    fun w dw =>
      let xdx := fwdDeriv R x w dw
      diracDeriv xdx.1 xdx.2 := by --= (dpure (R:=R) ydy.1 ydy.2) := by
  funext w dw; ext φ
  unfold parDistribDeriv dirac diracDeriv
  simp [pure, fwdDeriv, DistribDifferentiableAt]
  fun_trans

----------------------------------------------------------------------------------------------------
-- Composition -------------------------------------------------------------------------------------
----------------------------------------------------------------------------------------------------

@[fun_prop]
theorem DistribDiffrentiable.comp_rule
    (f : Y → 𝒟'(Z,U)) (g : X → Y)
    (hf : DistribDifferentiable f) (hg : CDifferentiable R g) :
    DistribDifferentiable (fun x => f (g x)) := by
  intro x
  unfold DistribDifferentiableAt
  intro φ hφ
  apply CDifferentiable.comp_rule (K:=R) (f:=fun xy : X×Y => f xy.2 (φ xy.1)) (g:=fun x => (x, g x))
    (hg:=by fun_prop)
  intro x
  sorry_proof -- is this even true ?

@[fun_trans]
theorem parDistribDeriv.comp_rule
    (f : Y → 𝒟'(Z,U)) (g : X → Y)
    (hf : DistribDifferentiable f) (hg : CDifferentiable R g) :
    parDistribDeriv (fun x => f (g x))
    =
    fun x dx =>
      let ydy := fwdDeriv R g x dx
      parDistribDeriv f ydy.1 ydy.2 := by

  funext x dx; ext φ
  unfold parDistribDeriv
  simp[hg]
  sorry_proof

----------------------------------------------------------------------------------------------------
-- Bind --------------------------------------------------------------------------------------------
----------------------------------------------------------------------------------------------------

-- The assumptions here are definitely not right!!!
-- I think `f` has to be `deg`
@[fun_prop]
theorem Bind.bind.arg_fx.DistribDifferentiable_rule
    (f : X → Y → 𝒟'(Z,V)) (g : X → 𝒟'(Y,U)) (L : U ⊸ V ⊸ W)
    (hf : DistribDifferentiable (fun (x,y) => f x y)) -- `f` has to be nice enough to accomodate action of `g`
    (hg : DistribDifferentiable g) :
    DistribDifferentiable (fun x => (g x).bind (f x) L) := by

  intro x
  unfold DistribDifferentiableAt
  intro φ hφ
  simp
  sorry_proof

@[fun_trans]
theorem Bind.bind.arg_fx.parDistribDiff_rule
    (f : W → X → 𝒟'(Y,V)) (g : W → 𝒟'(X,U)) (L : U ⊸ V ⊸ W)
    (hf : DistribDifferentiable (fun (w,x) => f w x)) -- `f` has to be nice enough to accomodate action of `g`
    (hg : DistribDifferentiable g) :
    parDistribDeriv (fun w => (g w).bind (f w) L)
    =
    fun w dw =>
      ((parDistribDeriv g w dw).bind (f x · ) L)
      +
      ((g w).bind (fun x => parDistribDeriv (f · x) w dw) L) := sorry_proof

----------------------------------------------------------------------------------------------------
-- Move these around -------------------------------------------------------------------------------
----------------------------------------------------------------------------------------------------

@[fun_prop]
theorem Distribution.restrict.arg_T.IsSmoothLinearMap_rule (T : W → 𝒟'(X,Y)) (A : Set X)
    (hT : IsSmoothLinearMap R T) :
    IsSmoothLinearMap R (fun w => (T w).restrict A) := sorry_proof

@[fun_prop]
theorem Distribution.restrict.arg_T.IsSmoothLinearMap_rule_simple (A : Set X) :
    IsSmoothLinearMap R (fun (T : 𝒟'(X,Y)) => T.restrict A) := sorry_proof

@[fun_trans]
theorem Distribution.restrict.arg_T.parDistribDeriv_rule (T : W → 𝒟'(X,Y)) (A : Set X)
    (hT : DistribDifferentiable T) :
    parDistribDeriv (fun w => (T w).restrict A)
    =
    fun w dw =>
      (parDistribDeriv T w dw).restrict A := sorry_proof

@[fun_prop]
theorem Function.toDistribution.arg_f.CDifferentiable_rule (f : W → X → Y)
    (hf : ∀ x, CDifferentiable R (f · x)) :
    CDifferentiable R (fun w => (fun x => f w x).toDistribution (R:=R)) := sorry_proof

@[fun_trans]
theorem Function.toDistribution.arg_f.cderiv_rule (f : W → X → Y)
    (hf : ∀ x, CDifferentiable R (f · x)) :
    cderiv R (fun w => (fun x => f w x).toDistribution (R:=R))
    =
    fun w dw =>
      (fun x =>
        let dy := cderiv R (f · x) w dw
        dy).toDistribution := sorry_proof

@[fun_trans]
theorem toDistribution.linear_parDistribDeriv_rule (f : W → X → Y) (L : Y → Z)
    (hL : IsSmoothLinearMap R L) :
    parDistribDeriv (fun w => (fun x => L (f w x)).toDistribution)
    =
    fun w dw =>
      parDistribDeriv (fun w => (fun x => f w x).toDistribution) w dw |>.postComp (fun y ⊸ L y) := by
  funext w dw
  unfold parDistribDeriv Distribution.postComp Function.toDistribution
  ext φ
  simp [ftrans_simp] -- , Distribution.mk_extAction_simproc]
  sorry_proof

----------------------------------------------------------------------------------------------------
-- Integral ----------------------------------------------------------------------------------------
----------------------------------------------------------------------------------------------------

variable [MeasureSpace X] [MeasureSpace Y] [MeasureSpace (X×Y)]

open Notation

@[fun_trans]
theorem cintegral.arg_f.cderiv_distrib_rule (f : W → X → Y) :
    cderiv R (fun w => ∫' x, f w x)
    =
    fun w dw =>
      (parDistribDeriv (fun w => (f w ·).toDistribution) w dw).extAction (fun _ => (1:R)) (fun y ⊸ fun r ⊸ r • y) := sorry_proof

@[fun_trans]
theorem cintegral.arg_f.cderiv_distrib_rule' (f : W → X → R) (A : Set X):
    cderiv R (fun w => ∫' x in A, f w x)
    =
    fun w dw =>
       (parDistribDeriv (fun w => (f w ·).toDistribution) w dw).restrict A |>.extAction (fun _ => (1:R)) (fun y ⊸ fun r ⊸ r • y) := sorry_proof

@[fun_trans]
theorem cintegral.arg_f.parDistribDeriv_rule (f : W → X → Y → Z) :
    parDistribDeriv (fun w => (fun x => ∫' y, f w x y).toDistribution (R:=R))
    =
    fun w dw =>
      let Tf := (fun w => (fun x => (fun y => f w x y).toDistribution (R:=R)).toDistribution (R:=R))
      parDistribDeriv Tf w dw |>.postComp (fun T ⊸ T.extAction (fun _ => (1:R)) (fun z ⊸ fun r ⊸ r • z)) := by
  funext w dw
  unfold parDistribDeriv postComp Function.toDistribution
  ext φ
  simp [ftrans_simp] -- , Distribution.mk_extAction_simproc]
  sorry_proof

@[fun_trans]
theorem cintegral.arg_f.parDistribDeriv_rule' (f : W → X → Y → Z) (B : X → Set Y) :
    parDistribDeriv (fun w => (fun x => ∫' y in B x, f w x y).toDistribution)
    =
    fun w dw =>
      let Tf := (fun w => (fun x => ((fun y => f w x y).toDistribution (R:=R)).restrict (B x)).toDistribution (R:=R))
      parDistribDeriv Tf w dw |>.postComp (fun T ⊸ T.extAction (fun _ => (1:R)) (fun z ⊸ fun r ⊸ r • z)) := sorry_proof

----------------------------------------------------------------------------------------------------
-- Add ---------------------------------------------------------------------------------------------
----------------------------------------------------------------------------------------------------

@[fun_prop]
theorem HAdd.hAdd.arg_a0a1.DistribDifferentiable_rule (f g : W → 𝒟'(X,Y))
    (hf : DistribDifferentiable f) (hg : DistribDifferentiable g) :
    DistribDifferentiable (fun w => f w + g w)


import SciLean.Core.Distribution.Basic
import SciLean.Core.FunctionTransformations
import SciLean.Core.FunctionPropositions
import SciLean.Core.Notation

open MeasureTheory

namespace SciLean

open Distribution

variable
  {R} [RealScalar R]
  {W} [Vec R W]
  {X} [Vec R X] [MeasureSpace X]
  {Y} [Vec R Y] [Module ℝ Y]
  {Z} [Vec R Z] [Module ℝ Z]
  {U} [Vec R U] -- [Module ℝ U]
  {V} [Vec R V] -- [Module ℝ U]

set_default_scalar R

noncomputable
def diracDeriv (x dx : X) : 𝒟' X := fun φ ⊸ cderiv R φ x dx

@[fun_prop]
def DistribDifferentiableAt (f : X → 𝒟'(Y,Z)) (x : X) :=
  ∀ (φ : X → 𝒟 Y), CDifferentiableAt R φ x → CDifferentiableAt R (fun x => f x (φ x)) x

theorem distribDifferentiableAt_const_test_fun
    {f : X → 𝒟'(Y,Z)} {x : X}
    (hf : DistribDifferentiableAt f x)
    {φ : 𝒟 Y} :
    CDifferentiableAt R (fun x => f x φ) x := by
  apply hf
  fun_prop

@[fun_prop]
def DistribDifferentiable (f : X → 𝒟'(Y,Z)) :=
  ∀ x, DistribDifferentiableAt f x

-- TODO:
-- probably change the definition of `parDistribDeriv` to:
-- ⟨⟨fun φ =>
--    if h : DistribDifferentiableAt f x then
--      ∂ (x':=x;dx), ⟪f x', φ⟫
--    else
--      0 , sorry_proof⟩⟩
-- I believe in that case the function is indeed linear in φ

open Classical in
@[fun_trans]
noncomputable
def parDistribDeriv (f : X → 𝒟'(Y,Z)) (x dx : X) : 𝒟'(Y,Z) :=
  ⟨fun φ => ∂ (x':=x;dx), f x' φ, sorry_proof⟩

@[simp, ftrans_simp]
theorem action_parDistribDeriv (f : X → 𝒟'(Y,Z)) (x dx : X) (φ : 𝒟 Y) :
    parDistribDeriv f x dx φ = ∂ (x':=x;dx), f x' φ := rfl

----------------------------------------------------------------------------------------------------
-- Const rule --------------------------------------------------------------------------------------
----------------------------------------------------------------------------------------------------

@[fun_prop]
theorem DistribDiffrentiable.const_rule (T : 𝒟'(X,Y)) :
    DistribDifferentiable (fun _ : W => T) := by
  intro _ φ hφ; simp; fun_prop

@[fun_trans]
theorem parDistribDeriv.const_rule (T : 𝒟'(X,Y)) :
    parDistribDeriv (fun _ : W => T)
    =
    fun w dw =>
      0 := by
  funext w dw; ext φ
  unfold parDistribDeriv
  fun_trans

----------------------------------------------------------------------------------------------------
-- Pure --------------------------------------------------------------------------------------------
----------------------------------------------------------------------------------------------------

@[fun_prop]
theorem dirac.arg_xy.DistribDiffrentiable_rule
    (x : W → X) (hx : CDifferentiable R x) :
    DistribDifferentiable (R:=R) (fun w => dirac (x w))  := by
  intro x
  unfold DistribDifferentiableAt
  intro φ hφ
  simp [action_dirac, dirac]
  fun_prop

@[fun_trans]
theorem dirac.arg_x.parDistribDeriv_rule
    (x : W → X) (hx : CDifferentiable R x) :
    parDistribDeriv (R:=R) (fun w => dirac (x w))
    =
    fun w dw =>
      let xdx := fwdDeriv R x w dw
      diracDeriv xdx.1 xdx.2 := by --= (dpure (R:=R) ydy.1 ydy.2) := by
  funext w dw; ext φ
  unfold parDistribDeriv dirac diracDeriv
  simp [pure, fwdDeriv, DistribDifferentiableAt]
  fun_trans

----------------------------------------------------------------------------------------------------
-- Composition -------------------------------------------------------------------------------------
----------------------------------------------------------------------------------------------------

@[fun_prop]
theorem DistribDiffrentiable.comp_rule
    (f : Y → 𝒟'(Z,U)) (g : X → Y)
    (hf : DistribDifferentiable f) (hg : CDifferentiable R g) :
    DistribDifferentiable (fun x => f (g x)) := by
  intro x
  unfold DistribDifferentiableAt
  intro φ hφ
  apply CDifferentiable.comp_rule (K:=R) (f:=fun xy : X×Y => f xy.2 (φ xy.1)) (g:=fun x => (x, g x))
    (hg:=by fun_prop)
  intro x
  sorry_proof -- is this even true ?

@[fun_trans]
theorem parDistribDeriv.comp_rule
    (f : Y → 𝒟'(Z,U)) (g : X → Y)
    (hf : DistribDifferentiable f) (hg : CDifferentiable R g) :
    parDistribDeriv (fun x => f (g x))
    =
    fun x dx =>
      let ydy := fwdDeriv R g x dx
      parDistribDeriv f ydy.1 ydy.2 := by

  funext x dx; ext φ
  unfold parDistribDeriv
  simp[hg]
  sorry_proof

----------------------------------------------------------------------------------------------------
-- Bind --------------------------------------------------------------------------------------------
----------------------------------------------------------------------------------------------------

-- The assumptions here are definitely not right!!!
-- I think `f` has to be `deg`
@[fun_prop]
theorem Bind.bind.arg_fx.DistribDifferentiable_rule
    (f : X → Y → 𝒟'(Z,V)) (g : X → 𝒟'(Y,U)) (L : U ⊸ V ⊸ W)
    (hf : DistribDifferentiable (fun (x,y) => f x y)) -- `f` has to be nice enough to accomodate action of `g`
    (hg : DistribDifferentiable g) :
    DistribDifferentiable (fun x => (g x).bind (f x) L) := by

  intro x
  unfold DistribDifferentiableAt
  intro φ hφ
  simp
  sorry_proof

@[fun_trans]
theorem Bind.bind.arg_fx.parDistribDiff_rule
    (f : W → X → 𝒟'(Y,V)) (g : W → 𝒟'(X,U)) (L : U ⊸ V ⊸ W)
    (hf : DistribDifferentiable (fun (w,x) => f w x)) -- `f` has to be nice enough to accomodate action of `g`
    (hg : DistribDifferentiable g) :
    parDistribDeriv (fun w => (g w).bind (f w) L)
    =
    fun w dw =>
      ((parDistribDeriv g w dw).bind (f x · ) L)
      +
      ((g w).bind (fun x => parDistribDeriv (f · x) w dw) L) := sorry_proof

----------------------------------------------------------------------------------------------------
-- Move these around -------------------------------------------------------------------------------
----------------------------------------------------------------------------------------------------

@[fun_prop]
theorem Distribution.restrict.arg_T.IsSmoothLinearMap_rule (T : W → 𝒟'(X,Y)) (A : Set X)
    (hT : IsSmoothLinearMap R T) :
    IsSmoothLinearMap R (fun w => (T w).restrict A) := sorry_proof

@[fun_prop]
theorem Distribution.restrict.arg_T.IsSmoothLinearMap_rule_simple (A : Set X) :
    IsSmoothLinearMap R (fun (T : 𝒟'(X,Y)) => T.restrict A) := sorry_proof

@[fun_trans]
theorem Distribution.restrict.arg_T.parDistribDeriv_rule (T : W → 𝒟'(X,Y)) (A : Set X)
    (hT : DistribDifferentiable T) :
    parDistribDeriv (fun w => (T w).restrict A)
    =
    fun w dw =>
      (parDistribDeriv T w dw).restrict A := sorry_proof

@[fun_prop]
theorem Function.toDistribution.arg_f.CDifferentiable_rule (f : W → X → Y)
    (hf : ∀ x, CDifferentiable R (f · x)) :
    CDifferentiable R (fun w => (fun x => f w x).toDistribution (R:=R)) := sorry_proof

@[fun_trans]
theorem Function.toDistribution.arg_f.cderiv_rule (f : W → X → Y)
    (hf : ∀ x, CDifferentiable R (f · x)) :
    cderiv R (fun w => (fun x => f w x).toDistribution (R:=R))
    =
    fun w dw =>
      (fun x =>
        let dy := cderiv R (f · x) w dw
        dy).toDistribution := sorry_proof

@[fun_trans]
theorem toDistribution.linear_parDistribDeriv_rule (f : W → X → Y) (L : Y → Z)
    (hL : IsSmoothLinearMap R L) :
    parDistribDeriv (fun w => (fun x => L (f w x)).toDistribution)
    =
    fun w dw =>
      parDistribDeriv (fun w => (fun x => f w x).toDistribution) w dw |>.postComp (fun y ⊸ L y) := by
  funext w dw
  unfold parDistribDeriv Distribution.postComp Function.toDistribution
  ext φ
  simp [ftrans_simp] -- , Distribution.mk_extAction_simproc]
  sorry_proof

----------------------------------------------------------------------------------------------------
-- Integral ----------------------------------------------------------------------------------------
----------------------------------------------------------------------------------------------------

variable [MeasureSpace X] [MeasureSpace Y] [MeasureSpace (X×Y)]

open Notation

@[fun_trans]
theorem cintegral.arg_f.cderiv_distrib_rule (f : W → X → Y) :
    cderiv R (fun w => ∫' x, f w x)
    =
    fun w dw =>
      (parDistribDeriv (fun w => (f w ·).toDistribution) w dw).extAction (fun _ => (1:R)) (fun y ⊸ fun r ⊸ r • y) := sorry_proof

@[fun_trans]
theorem cintegral.arg_f.cderiv_distrib_rule' (f : W → X → R) (A : Set X):
    cderiv R (fun w => ∫' x in A, f w x)
    =
    fun w dw =>
       (parDistribDeriv (fun w => (f w ·).toDistribution) w dw).restrict A |>.extAction (fun _ => (1:R)) (fun y ⊸ fun r ⊸ r • y) := sorry_proof

@[fun_trans]
theorem cintegral.arg_f.parDistribDeriv_rule (f : W → X → Y → Z) :
    parDistribDeriv (fun w => (fun x => ∫' y, f w x y).toDistribution (R:=R))
    =
    fun w dw =>
      let Tf := (fun w => (fun x => (fun y => f w x y).toDistribution (R:=R)).toDistribution (R:=R))
      parDistribDeriv Tf w dw |>.postComp (fun T ⊸ T.extAction (fun _ => (1:R)) (fun z ⊸ fun r ⊸ r • z)) := by
  funext w dw
  unfold parDistribDeriv postComp Function.toDistribution
  ext φ
  simp [ftrans_simp] -- , Distribution.mk_extAction_simproc]
  sorry_proof

@[fun_trans]
theorem cintegral.arg_f.parDistribDeriv_rule' (f : W → X → Y → Z) (B : X → Set Y) :
    parDistribDeriv (fun w => (fun x => ∫' y in B x, f w x y).toDistribution)
    =
    fun w dw =>
      let Tf := (fun w => (fun x => ((fun y => f w x y).toDistribution (R:=R)).restrict (B x)).toDistribution (R:=R))
      parDistribDeriv Tf w dw |>.postComp (fun T ⊸ T.extAction (fun _ => (1:R)) (fun z ⊸ fun r ⊸ r • z)) := sorry_proof

----------------------------------------------------------------------------------------------------
-- Add ---------------------------------------------------------------------------------------------
----------------------------------------------------------------------------------------------------

@[fun_prop]
theorem HAdd.hAdd.arg_a0a1.DistribDifferentiable_rule (f g : W → 𝒟'(X,Y))
    (hf : DistribDifferentiable f) (hg : DistribDifferentiable g) :
    DistribDifferentiable (fun w => f w + g w) := sorry_proof

@[fun_trans]
theorem HAdd.hAdd.arg_a0a1.parDistribDeriv_rule (f g : W → 𝒟'(X,Y))
    (hf : DistribDifferentiable f) (hg : DistribDifferentiable g) :
    parDistribDeriv (fun w => f w + g w)
    =
    fun w dw =>
      let dy := parDistribDeriv f w dw
      let dz := parDistribDeriv g w dw
      dy + dz