import SciLean.Core.FunctionTransformations
-- import SciLean.Core.Meta.GenerateRevDeriv

open ComplexConjugate

namespace SciLean.Scalar

variable
  {R C} [Scalar R C]
  {W} [Vec C W]
  {U} [SemiInnerProductSpace C U]

--------------------------------------------------------------------------------
-- Sin -------------------------------------------------------------------------
--------------------------------------------------------------------------------

@[fun_prop]
theorem sin.arg_x.CDifferentiable_rule
    (x : W → C) (hx : CDifferentiable C x) :
    CDifferentiable C fun w => sin (x w) := sorry_proof

@[fun_trans]
theorem sin.arg_x.ceriv_rule
    (x : W → C) (hx : CDifferentiable C x) :
    cderiv C (fun w => sin (x w))
    =
    fun w dw =>
      let xdx := fwdDeriv C x w dw
      let c := cos xdx.1
      xdx.2 * c := sorry_proof

@[fun_trans]
theorem sin.arg_x.fwdDeriv_rule
    (x : W → C) (hx : CDifferentiable C x) :
    fwdDeriv C (fun w => sin (x w))
    =
    fun w dw =>
      let xdx := fwdDeriv C x w dw
      let s := sin xdx.1
      let c := cos xdx.1
      (s, xdx.2 * c) := by
  unfold fwdDeriv; fun_trans; rfl

@[fun_prop]
theorem sin.arg_x.HasAdjDiff_rule
    (x : U → C) (hx : HasAdjDiff C x) :
    HasAdjDiff C (fun u => sin (x u)) := by
  intro x
  constructor
  fun_prop
  fun_trans [fwdDeriv]; fun_prop

open ComplexConjugate
@[fun_trans]
theorem sin.arg_x.revDeriv_rule
    (x : U → C) (hx : HasAdjDiff C x) :
    revDeriv C (fun u => sin (x u))
    =
    fun u =>
      let xdx := revDeriv C x u
      (sin xdx.1, fun dy => xdx.2 (conj (cos xdx.1) * dy)) := by
  unfold revDeriv
  fun_trans only [fwdDeriv, smul_push, ftrans_simp]

--------------------------------------------------------------------------------
-- Cos -------------------------------------------------------------------------
--------------------------------------------------------------------------------

@[fun_prop]
theorem cos.arg_x.CDifferentiable_rule
  (x : W → C) (hx : CDifferentiable C x)
  : CDifferentiable C fun w => cos (x w) := sorry_proof

@[fun_trans]
theorem cos.arg_x.ceriv_rule
  (x : W → C) (hx : CDifferentiable C x)
  : cderiv C (fun w => cos (x w))
    =
    fun w dw =>
      let xdx := fwdDeriv C x w dw
      let s := -sin xdx.1
      xdx.2 * s := sorry_proof

@[fun_trans]
theorem cos.arg_x.fwdDeriv_rule
  (x : W → C) (hx : CDifferentiable C x)
  : fwdDeriv C (fun w => cos (x w))
    =
    fun w dw =>
      let xdx := fwdDeriv C x w dw
      let s := sin xdx.1
      let c := cos xdx.1
      (c, xdx.2 * -s) :=
by
  unfold fwdDeriv; fun_trans; rfl

@[fun_prop]
theorem cos.arg_x.HasAdjDiff_rule
    (x : U → C) (hx : HasAdjDiff C x) :
    HasAdjDiff C (fun u => cos (x u)) := by
  intro x
  constructor
  fun_prop
  fun_trans [fwdDeriv]; fun_prop

open ComplexConjugate
@[fun_trans]
theorem cos.arg_x.revDeriv_rule
    (x : U → C) (hx : HasAdjDiff C x) :
    revDeriv C (fun u => cos (x u))
    =
    fun u =>
      let xdx := revDeriv C x u
      (cos xdx.1, fun dy => xdx.2 (- conj (sin xdx.1) * dy)) := by
  unfold revDeriv
  fun_trans only [fwdDeriv, smul_push, neg_push, ftrans_simp]

--------------------------------------------------------------------------------
-- Tanh -------------------------------------------------------------------------
--------------------------------------------------------------------------------

@[fun_prop]
theorem tanh.arg_x.CDifferentiable_rule
    (x : W → C) (hx : CDifferentiable C x) :
    CDifferentiable C fun w => tanh (x w) := sorry_proof

@[fun_trans]
theorem tanh.arg_x.ceriv_rule
    (x : W → C) (hx : CDifferentiable C x) :
    cderiv C (fun w => tanh (x w))
    =
    fun w dw =>
      let xdx := fwdDeriv C x w dw
      let t := tanh xdx.1
      let dt := 1 - t^2
      xdx.2 * dt := sorry_proof

@[fun_trans]
theorem tanh.arg_x.fwdDeriv_rule
    (x : W → C) (hx : CDifferentiable C x) :
    fwdDeriv C (fun w => tanh (x w))
    =
    fun w dw =>
      let xdx := fwdDeriv C x w dw
      let t := tanh xdx.1
      let dt := 1-t^2
      (t, xdx.2 * dt) :=
by
  unfold fwdDeriv; fun_trans; rfl

@[fun_prop]
theorem tanh.arg_x.HasAdjDiff_rule
    (x : U → C) (hx : HasAdjDiff C x) :
    HasAdjDiff C (fun u => tanh (x u)) := by
  intro x
  constructor
  fun_prop
  fun_trans [fwdDeriv]; fun_prop

open ComplexConjugate

@[fun_trans]
theorem tanh.arg_x.revDeriv_rule
    (x : U → C) (hx : HasAdjDiff C x) :
    revDeriv C (fun u => tanh (x u))
    =
    fun u =>
      let xdx := revDeriv C x u
      (tanh xdx.1, fun dy => xdx.2 (conj (1 - tanh (xdx.1) ^ 2) * dy))


import SciLean.Core.FunctionPropositions.ContCDiff
import SciLean.Core.FunctionTransformations

import SciLean.Core.Notation.CDeriv

set_option linter.unusedVariables false

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)]

local notation "∞" => (⊤ : ℕ∞)

-- Function space --------------------------------------------------------------
--------------------------------------------------------------------------------

structure ContCDiffMap (n : ℕ∞) (X Y : Type _)  [Vec K X] [Vec K Y] where
  toFun : X → Y
  is_cont_cdiff_map : ContCDiff K n toFun

variable (n : ℕ∞)

instance : FunLike (ContCDiffMap K n X Y) X Y where
  coe f := f.toFun
  coe_injective' := sorry_proof

macro X:term:25 " ⟿[" K:term "," n:term "] " Y:term:26 : term =>
  `(ContCDiffMap $K $n $X $Y)

macro X:term:25 " ⟿[" n:term "] " Y:term:26 : term =>
  `(ContCDiffMap defaultScalar% $n $X $Y)

macro X:term:25 " ⟿ " Y:term:26 : term =>
  `(ContCDiffMap defaultScalar% ∞ $X $Y)

@[app_unexpander ContCDiffMap] def unexpandContCDiffMap : Lean.PrettyPrinter.Unexpander
  | `($(_) $R $n $X $Y) => `($X ⟿[$R,$n] $Y)
  | _  => throw ()

@[fun_prop]
theorem ContCDiffMap_apply_right (f : X ⟿[K,n] Y) : ContCDiff K n (fun x => f x) := f.2

-- Lambda function notation ----------------------------------------------------
--------------------------------------------------------------------------------

variable {K}

-- def ContCDiffMap.mk' (f : X → Y) (hf : ContCDiff K n f) :
--     ContCDiffMap K n X Y := ⟨f, hf⟩

@[simp, ftrans_simp]
theorem ContCDiffMap.mk'_eval (x : X) (f : X → Y) (hf : ContCDiff K n f) :
    mk f hf x = f x := by rfl

@[simp, ftrans_simp]
theorem ContCDiffMap.eta_reduce (f : ContCDiffMap K n X Y) :
    (mk f.1 f.2) = f := by rfl

variable (K)
def ContCDiffMap.mk' (f : X → Y) (hf : ContCDiff K n f) : X ⟿[K,n] Y := ⟨f,hf⟩

open Lean Parser Term
macro "fun " x:funBinder " ⟿[" K:term "," n:term "] " b:term : term =>
  `(ContCDiffMap.mk' $K $n (fun $x => $b) (by fun_prop (disch:=norm_num; linarith)))

open Lean Parser Term
macro "fun " x:funBinder " ⟿[" n:term "] " b:term : term =>
  `(ContCDiffMap.mk' defaultScalar% $n (fun $x => $b) (by fun_prop (disch:=norm_num; linarith)))

open Lean Parser Term
macro "fun " x:funBinder " ⟿ " b:term : term =>
  `(ContCDiffMap.mk' defaultScalar% ∞ (fun $x => $b) (by fun_prop (disch:=norm_num; linarith)))

@[app_unexpander ContCDiffMap.mk'] def unexpandContCDiffMapMk : Lean.PrettyPrinter.Unexpander

  | `($(_) $R $n $f:term $_:term) =>
    match f with
    | `(fun $x':ident => $b:term) => `(fun $x' ⟿[$R,$n] $b)
    | `(fun ($x':ident : $ty) => $b:term) => `(fun ($x' : $ty) ⟿[$R,$n] $b)
    | `(fun $x':ident : $ty => $b:term) => `(fun ($x' : $ty) ⟿[$R,$n] $b)
    | _  => throw ()
  | _  => throw ()

@[simp, ftrans_simp]
theorem ContCDiffMap_eta (f : X ⟿[K,n] Y) : (fun x ⟿[K,n] f x) = f


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)


import SciLean.Core.FunctionPropositions
import SciLean.Core.FunctionTransformations
import SciLean.Data.ArrayType.Algebra

open SciLean

set_option linter.unusedVariables false

open LeanColls

section GenericArrayType

variable
  {K : Type} [RCLike K]
  {Cont : Type} {Idx : Type |> outParam} {Elem : Type |> outParam}
  [ArrayType Cont Idx Elem] [IndexType Idx] [LawfulIndexType Idx] [DecidableEq Idx]

-- Indexed.get -----------------------------------------------------------------
--------------------------------------------------------------------------------

section OnVectorSpaces

variable {R : Type} [CommSemiring R] [AddCommGroup Elem] [Module R Elem]

@[fun_prop]
theorem LeanColls.Indexed.get.arg_cont.IsLinearMap_rule_simple (idx : Idx) :
    IsLinearMap R (fun xs : Cont => xs[idx]) := by
  constructor
  . intros; simp
  . intros; simp

#generate_linear_map_simps LeanColls.Indexed.get.arg_cont.IsLinearMap_rule_simple

end OnVectorSpaces

section OnNormedSpaces

variable
  {X : Type _} [NormedAddCommGroup X] [NormedSpace K X]
  [NormedAddCommGroup Elem] [NormedSpace K Elem]

@[fun_prop]
theorem LeanColls.Indexed.get.arg_cont.IsContinuousLinearMap_rule_simple (idx : Idx) :
    IsContinuousLinearMap K (λ cont : Cont => cont[idx]) := sorry_proof

-- automatically infer: Differentiable, DifferentiableAt, ContDiff, fderiv, fwdFDeriv

end OnNormedSpaces

section OnVec

variable
  {X : Type} [Vec K X]
  [Vec K Elem]

@[fun_prop]
theorem GetElem.getElem.arg_cont.IsSmoothLinearMap_rule_simple
  (idx : Idx)
  : IsSmoothLinearMap K (fun xs : Cont => xs[idx]) := sorry_proof

-- automatically infer: CDifferentiable, CDifferentiableAt, CContDiff, cderiv, fwdDeriv

end OnVec

section OnSemiInnerProductSpace

variable
  {X : Type} [SemiInnerProductSpace K X]
  [SemiInnerProductSpace K Elem]

@[fun_prop]
theorem GetElem.getElem.arg_cont.HasSemiAdjoint_rule_simple (idx : Idx) :
    HasSemiAdjoint K (fun cont : Cont => cont[idx]) := sorry_proof

@[fun_trans]
theorem GetElem.getElem.arg_cont.semiAdjoint_rule_simple (idx : Idx) :
    semiAdjoint K (fun cont : Cont => cont[idx])
    =
    fun elem => oneHot (X:=Cont) idx elem :=
by
  sorry_proof

@[fun_trans]
theorem GetElem.getElem.arg_cont.revCDeriv_rule
  (f : X → Cont) (idx : Idx)
  (hf : HasAdjDiff K f)
  : revDeriv K (fun x => (f x)[idx])
    =
    fun x =>
      let ydf := revDerivProj K Idx f x
      (ydf.1[idx],
       fun delem => ydf.2 idx delem) :=
by
  unfold revDeriv;
  conv =>
    lhs; fun_trans; fun_trans; simp

@[fun_trans]
theorem GetElem.getElem.arg_cont.revCDerivUpdate_rule
  (f : X → Cont) (idx : Idx)
  (hf : HasAdjDiff K f)
  : revDerivUpdate K (fun x => (f x)[idx])
    =
    fun x =>
      let ydf := revDerivProjUpdate K Idx f x
      (ydf.1[idx],
       fun delem dx => ydf.2 idx delem dx)

:=
by
  unfold revDerivUpdate;
  conv =>
    lhs; fun_trans; fun_trans; simp


import SciLean.Core.FunctionTransformations
-- import SciLean.Core.Meta.GenerateRevDeriv

open ComplexConjugate

namespace SciLean.Scalar

variable
  {R C} [Scalar R C] [RealScalar R]
  {W} [Vec C W]
  {U} [SemiInnerProductSpace C U]

--------------------------------------------------------------------------------
-- Exp -------------------------------------------------------------------------
--------------------------------------------------------------------------------

@[fun_prop]
theorem exp.arg_x.CDifferentiable_rule
  (x : W → C) (hx : CDifferentiable C x)
  : CDifferentiable C fun w => exp (x w) := sorry_proof

@[fun_trans]
theorem exp.arg_x.ceriv_rule
  (x : W → C) (hx : CDifferentiable C x)
  : cderiv C (fun w => exp (x w))
    =
    fun w dw =>
      let xdx := fwdDeriv C x w dw
      let e := exp xdx.1
      xdx.2 * e := sorry_proof

@[fun_trans]
theorem exp.arg_x.fwdDeriv_rule
    (x : W → C) (hx : CDifferentiable C x) :
    fwdDeriv C (fun w => exp (x w))
    =
    fun w dw =>
      let xdx := fwdDeriv C x w dw
      let e := exp xdx.1
      (e, xdx.2 * e) := by
  unfold fwdDeriv; fun_trans; rfl

@[fun_prop]
theorem exp.arg_x.HasAdjDiff_rule
    (x : U → C) (hx : HasAdjDiff C x) :
    HasAdjDiff C (fun u => exp (x u)) := by
  intro x
  constructor
  fun_prop
  fun_trans [fwdDeriv]; fun_prop

open ComplexConjugate
@[fun_trans]
theorem exp.arg_x.revDeriv_rule
    (x : U → C) (hx : HasAdjDiff C x) :
    revDeriv C (fun u => exp (x u))
    =
    fun u =>
      let xdx := revDeriv C x u
      (exp xdx.1, fun dy => xdx.2 (conj (exp xdx.1) * dy)) := by
  unfold revDeriv
  fun_trans only [fwdDeriv, smul_push, ftrans_simp]

--------------------------------------------------------------------------------
-- Log -------------------------------------------------------------------------
--------------------------------------------------------------------------------

section Log

variable
  {R : Type _} [RealScalar R]
  {W : Type _} [Vec R W]
  {U : Type _} [SemiInnerProductSpace R U]

@[fun_prop]
theorem log.arg_x.CDifferentiableAt_rule
    (w : W) (x : W → R) (hx : CDifferentiableAt R x w) (hw : x w ≠ 0) :
    CDifferentiableAt R (fun w => log (x w)) w := sorry_proof

@[fun_prop]
theorem log.arg_x.CDifferentiable_rule
    (x : W → R) (hx : CDifferentiable R x) (hw : ∀ w, x w ≠ 0) :
    CDifferentiable R (fun w => log (x w))

:= by
  intro x; fun_prop (disch:=aesop)


import SciLean.Core.FunctionTransformations
import SciLean.Core.FloatAsReal
-- import SciLean.Core.Meta.GenerateRevDeriv

open ComplexConjugate

namespace SciLean.Scalar

variable
  {R} [RealScalar R]
  {W} [Vec R W]
  {U} [SemiInnerProductSpace R U]

----------------------------------------------------------------------------------------------------
-- Sqrt --------------------------------------------------------------------------------------------
----------------------------------------------------------------------------------------------------

@[fun_prop]
theorem sqrt.arg_x.CDifferentiableAt_rule
    (w : W) (x : W → R) (hx : CDifferentiableAt R x w) (hw : x w ≠ 0) :
    CDifferentiableAt R (fun w => sqrt (x w)) w := sorry_proof

@[fun_prop]
theorem sqrt.arg_x.CDifferentiable_rule
    (x : W → R) (hx : CDifferentiable R x) (hw : ∀ w, x w ≠ 0) :
    CDifferentiable R (fun w => sqrt (x w)) := by
  intro x; fun_prop (disch:=aesop)

@[fun_trans]
theorem sqrt.arg_x.cderiv_rule_at
    (w : W) (x : W → R) (hx : CDifferentiableAt R x w) (hw : 0 < x w) :
    cderiv R (fun w => sqrt (x w)) w
    =
    fun dw =>
      let xdx := fwdDeriv R x w dw
      xdx.2 / (2 * sqrt xdx.1) := sorry_proof

@[fun_trans]
theorem sqrt.arg_x.cderiv_rule
    (x : W → R) (hx : CDifferentiable R x) (hw : ∀ w, 0 < x w) :
    cderiv R (fun w => sqrt (x w))
    =
    fun w dw =>
      let xdx := fwdDeriv R x w dw
      xdx.2 / (2 * sqrt xdx.1) := by
  funext x
  fun_trans (disch:=aesop)

@[fun_trans]
theorem sqrt.arg_x.fwdDeriv_rule_at
    (w : W) (x : W → R) (hx : CDifferentiableAt R x w) (hw : 0 < x w) :
    fwdDeriv R (fun w => sqrt (x w)) w
    =
    fun dw =>
      let xdx := fwdDeriv R x w dw
      let x' := sqrt xdx.1
      (x', xdx.2 / (2 * x'))

:=
by
  unfold fwdDeriv
  fun_trans (disch:=assumption) [fwdDeriv]


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)


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


import SciLean.Core.FunctionTransformations
-- import SciLean.Core.Meta.GenerateRevDeriv

open ComplexConjugate

namespace SciLean.Scalar

variable
  {R C} [Scalar R C] [RealScalar R]
  {W} [Vec C W]
  {U} [SemiInnerProductSpace C U]

--------------------------------------------------------------------------------
-- Exp -------------------------------------------------------------------------
--------------------------------------------------------------------------------

@[fun_prop]
theorem exp.arg_x.CDifferentiable_rule
  (x : W → C) (hx : CDifferentiable C x)
  : CDifferentiable C fun w => exp (x w) := sorry_proof

@[fun_trans]
theorem exp.arg_x.ceriv_rule
  (x : W → C) (hx : CDifferentiable C x)
  : cderiv C (fun w => exp (x w))
    =
    fun w dw =>
      let xdx := fwdDeriv C x w dw
      let e := exp xdx.1
      xdx.2 * e := sorry_proof

@[fun_trans]
theorem exp.arg_x.fwdDeriv_rule
    (x : W → C) (hx : CDifferentiable C x) :
    fwdDeriv C (fun w => exp (x w))
    =
    fun w dw =>
      let xdx := fwdDeriv C x w dw
      let e := exp xdx.1
      (e, xdx.2 * e) := by
  unfold fwdDeriv; fun_trans; rfl

@[fun_prop]
theorem exp.arg_x.HasAdjDiff_rule
    (x : U → C) (hx : HasAdjDiff C x) :
    HasAdjDiff C (fun u => exp (x u)) := by
  intro x
  constructor
  fun_prop
  fun_trans [fwdDeriv]; fun_prop

open ComplexConjugate
@[fun_trans]
theorem exp.arg_x.revDeriv_rule
    (x : U → C) (hx : HasAdjDiff C x) :
    revDeriv C (fun u => exp (x u))
    =
    fun u =>
      let xdx := revDeriv C x u
      (exp xdx.1, fun dy => xdx.2 (conj (exp xdx.1) * dy)) := by
  unfold revDeriv
  fun_trans only [fwdDeriv, smul_push, ftrans_simp]

--------------------------------------------------------------------------------
-- Log -------------------------------------------------------------------------
--------------------------------------------------------------------------------

section Log

variable
  {R : Type _} [RealScalar R]
  {W : Type _} [Vec R W]
  {U : Type _} [SemiInnerProductSpace R U]

@[fun_prop]
theorem log.arg_x.CDifferentiableAt_rule
    (w : W) (x : W → R) (hx : CDifferentiableAt R x w) (hw : x w ≠ 0) :
    CDifferentiableAt R (fun w => log (x w)) w := sorry_proof

@[fun_prop]
theorem log.arg_x.CDifferentiable_rule
    (x : W → R) (hx : CDifferentiable R x) (hw : ∀ w, x w ≠ 0) :
    CDifferentiable R (fun w => log (x w)) := by
  intro x; fun_prop (disch:=aesop)

@[fun_trans]
theorem log.arg_x.cderiv_rule_at
    (w : W) (x : W → R) (hx : CDifferentiableAt R x w) (hw : x w ≠ 0) :
    cderiv R (fun w => log (x w)) w
    =
    fun dw =>
      let xdx := fwdDeriv R x w dw
      xdx.2 / abs xdx.1 := sorry_proof

@[fun_trans]
theorem log.arg_x.cderiv_rule
    (x : W → R) (hx : CDifferentiable R x) (hw : ∀ w, x w ≠ 0) :
    cderiv R (fun w => log (x w))
    =
    fun w dw =>
      let xdx := fwdDeriv R x w dw
      xdx.2 / abs xdx.1 := by
  funext x
  fun_trans (disch:=aesop)

@[fun_trans]
theorem log.arg_x.fwdDeriv_rule_at
    (w : W) (x : W → R) (hx : CDifferentiableAt R x w) (hw : x w ≠ 0) :
    fwdDeriv R (fun w => log (x w)) w
    =
    fun dw =>
      let xdx := fwdDeriv R x w dw
      let l := log xdx.1
      (l, xdx.2 / abs xdx.1) :=
by
  unfold fwdDeriv; fun_trans (disch:=assumption); simp[fwdDeriv]

@[fun_trans]
theorem log.arg_x.fwdDeriv_rule
    (x : W → R) (hx : CDifferentiable R x) (hw : ∀ w, x w ≠ 0) :
    fwdDeriv R (fun w => log (x w))
    =
    fun w dw =>
      let xdx := fwdDeriv R x w dw
      let l := log xdx.1
      (l, xdx.2 / abs xdx.1) :=
by
  unfold fwdDeriv; fun_trans (disch:=assumption); simp[fwdDeriv]

@[fun_prop]
theorem log.arg_x.HasAdjDiffAt_rule
    (u : U) (x : U → R) (hx : HasAdjDiffAt R x u) (hu : x u ≠ 0) :
    HasAdjDiffAt R (fun u => log (x u)) u := by
  constructor
  fun_prop (disch:=aesop)
  fun_trans (disch:=aesop) [fwdDeriv]; fun_prop

@[fun_prop]
theorem log.arg_x.HasAdjDiff_rule
    (x : U → R) (hx : HasAdjDiff R x) (hu : ∀ u, x u ≠ 0) :
    HasAdjDiff R (fun u => log (x u))

:= by
  intro u;
  fun_prop (disch:=aesop)


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 ∅


import SciLean.Core.FunctionTransformations
-- import SciLean.Core.Meta.GenerateRevDeriv

open ComplexConjugate

namespace SciLean.Scalar

variable
  {R C} [Scalar R C]
  {W} [Vec C W]
  {U} [SemiInnerProductSpace C U]

--------------------------------------------------------------------------------
-- Sin -------------------------------------------------------------------------
--------------------------------------------------------------------------------

@[fun_prop]
theorem sin.arg_x.CDifferentiable_rule
    (x : W → C) (hx : CDifferentiable C x) :
    CDifferentiable C fun w => sin (x w) := sorry_proof

@[fun_trans]
theorem sin.arg_x.ceriv_rule
    (x : W → C) (hx : CDifferentiable C x) :
    cderiv C (fun w => sin (x w))
    =
    fun w dw =>
      let xdx := fwdDeriv C x w dw
      let c := cos xdx.1
      xdx.2 * c := sorry_proof

@[fun_trans]
theorem sin.arg_x.fwdDeriv_rule
    (x : W → C) (hx : CDifferentiable C x) :
    fwdDeriv C (fun w => sin (x w))
    =
    fun w dw =>
      let xdx := fwdDeriv C x w dw
      let s := sin xdx.1
      let c := cos xdx.1
      (s, xdx.2 * c) := by
  unfold fwdDeriv; fun_trans; rfl

@[fun_prop]
theorem sin.arg_x.HasAdjDiff_rule
    (x : U → C) (hx : HasAdjDiff C x) :
    HasAdjDiff C (fun u => sin (x u)) := by
  intro x
  constructor
  fun_prop
  fun_trans [fwdDeriv]; fun_prop

open ComplexConjugate
@[fun_trans]
theorem sin.arg_x.revDeriv_rule
    (x : U → C) (hx : HasAdjDiff C x) :
    revDeriv C (fun u => sin (x u))
    =
    fun u =>
      let xdx := revDeriv C x u
      (sin xdx.1, fun dy => xdx.2 (conj (cos xdx.1) * dy)) := by
  unfold revDeriv
  fun_trans only [fwdDeriv, smul_push, ftrans_simp]

--------------------------------------------------------------------------------
-- Cos -------------------------------------------------------------------------
--------------------------------------------------------------------------------

@[fun_prop]
theorem cos.arg_x.CDifferentiable_rule
  (x : W → C) (hx : CDifferentiable C x)
  : CDifferentiable C fun w => cos (x w) := sorry_proof

@[fun_trans]
theorem cos.arg_x.ceriv_rule
  (x : W → C) (hx : CDifferentiable C x)
  : cderiv C (fun w => cos (x w))
    =
    fun w dw =>
      let xdx := fwdDeriv C x w dw
      let s := -sin xdx.1
      xdx.2 * s := sorry_proof

@[fun_trans]
theorem cos.arg_x.fwdDeriv_rule
  (x : W → C) (hx : CDifferentiable C x)
  : fwdDeriv C (fun w => cos (x w))
    =
    fun w dw =>
      let xdx := fwdDeriv C x w dw
      let s := sin xdx.1
      let c := cos xdx.1
      (c, xdx.2 * -s) :=
by
  unfold fwdDeriv; fun_trans; rfl

@[fun_prop]
theorem cos.arg_x.HasAdjDiff_rule
    (x : U → C) (hx : HasAdjDiff C x) :
    HasAdjDiff C (fun u => cos (x u))


import SciLean.Core.Objects.IsomorphicType

open Lean

namespace SciLean

variable
  {α β γ : Type _}
  {α' β' γ' : outParam (Type _)}
  (tag : Name)
  [IsomorphicType tag α α']
  [IsomorphicType tag β β']
  [IsomorphicType tag γ γ']

namespace isomorph

@[fun_trans]
theorem id_rule :
    isomorph tag (fun (x : α) => x)
    =
    fun (x : α') => x := by
  funext _; simp[isomorph]

@[fun_trans]
theorem const_rule (y : β) :
    isomorph tag (fun (_ : α) => y)
    =
    fun (_ : α') => (IsomorphicType.equiv tag) y := by
  funext _; simp[isomorph]

@[fun_trans]
theorem comp_rule (f : β → γ) (g : α → β) :
    isomorph tag (fun x => f (g x))
    =
    fun x => isomorph tag f (isomorph tag g x) := by
  funext _; unfold isomorph; simp

@[fun_trans]
theorem let_rule (f : α → β → γ) (g : α → β) :
    isomorph tag (fun x => let y := g x; f x y)
    =
    fun x' =>
      let y' := isomorph tag g x'
      isomorph tag (fun (xy : α×β) => f xy.1 xy.2) (x', y') := by
  funext x'; simp[isomorph, IsomorphicType.equiv, -isomorph.app]

@[fun_trans]
theorem apply_rule (x : α) :
    isomorph tag (fun (f : α → β) => f x)
    =
    fun (f : α' → β') => f ((IsomorphicType.equiv tag) x) := by
  funext _; simp[isomorph, invIsomorph, IsomorphicType.equiv]

@[fun_trans]
theorem pi_rule (f : α → β → γ) :
    isomorph tag (fun x y => f x y)
    =
    fun x' => isomorph tag (f ((IsomorphicType.equiv tag).symm x')) := by
  funext x'; simp[isomorph, IsomorphicType.equiv]

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

@[fun_trans]
theorem _root_.Prod.mk.isomorph_rule (f : α → β) (g : α → γ) :
    isomorph tag (fun x => (f x, g x))
    =
    fun x' =>
      let y' := isomorph tag f x'
      let z' := isomorph tag g x'
      (y', z')


import SciLean.Core.FunctionPropositions.CDifferentiable
import SciLean.Core.FunctionPropositions.HasSemiAdjoint

import SciLean.Core.FunctionTransformations.CDeriv

import Lean.Meta.Tactic.Assumption

set_option linter.unusedVariables false

open LeanColls

namespace SciLean

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

@[fun_prop]
def HasAdjDiffAt (f : X → Y) (x : X) : Prop :=
    CDifferentiableAt K f x ∧ HasSemiAdjoint K (cderiv K f x)

@[fun_prop]
def HasAdjDiff (f : X → Y) : Prop := ∀ x, HasAdjDiffAt K f x

@[fun_prop]
theorem HasAdjDiff.hasAdjDiffAt (f : X → Y) (x) (hf : HasAdjDiff K f) :
    HasAdjDiffAt K f x := hf x

@[fun_prop]
theorem HasAdjDiffAt.cdifferentiableAt (f : X → Y) (x) (hf : HasAdjDiffAt K f x) :
    CDifferentiableAt K f x


import SciLean.Core.FunctionTransformations
import SciLean.Core.FloatAsReal
-- import SciLean.Core.Meta.GenerateRevDeriv

open ComplexConjugate

namespace SciLean.Scalar

variable
  {R} [RealScalar R]
  {W} [Vec R W]
  {U} [SemiInnerProductSpace R U]

----------------------------------------------------------------------------------------------------
-- Sqrt --------------------------------------------------------------------------------------------
----------------------------------------------------------------------------------------------------

@[fun_prop]
theorem sqrt.arg_x.CDifferentiableAt_rule
    (w : W) (x : W → R) (hx : CDifferentiableAt R x w) (hw : x w ≠ 0) :
    CDifferentiableAt R (fun w => sqrt (x w)) w := sorry_proof

@[fun_prop]
theorem sqrt.arg_x.CDifferentiable_rule
    (x : W → R) (hx : CDifferentiable R x) (hw : ∀ w, x w ≠ 0) :
    CDifferentiable R (fun w => sqrt (x w)) := by
  intro x; fun_prop (disch:=aesop)

@[fun_trans]
theorem sqrt.arg_x.cderiv_rule_at
    (w : W) (x : W → R) (hx : CDifferentiableAt R x w) (hw : 0 < x w) :
    cderiv R (fun w => sqrt (x w)) w
    =
    fun dw =>
      let xdx := fwdDeriv R x w dw
      xdx.2 / (2 * sqrt xdx.1) := sorry_proof

@[fun_trans]
theorem sqrt.arg_x.cderiv_rule
    (x : W → R) (hx : CDifferentiable R x) (hw : ∀ w, 0 < x w) :
    cderiv R (fun w => sqrt (x w))
    =
    fun w dw =>
      let xdx := fwdDeriv R x w dw
      xdx.2 / (2 * sqrt xdx.1)

:= by
  funext x
  fun_trans (disch:=aesop)


import SciLean.Util.SorryProof
import SciLean.Data.ListN
import SciLean.Data.StructType.Basic
import SciLean.Data.Function
import LeanColls

namespace SciLean

open LeanColls

/-- This class says that `Cont` behaves like an array with `Elem` values indexed by `Idx`

Examples for `Idx = Fin n` and `Elem = ℝ` are: `ArrayN ℝ n` or `ℝ^{n}`

For `array : Cont` you can:
  1. get values: `Indexed.get array i : Elem` for `i : Idx`
  2. set values: `Indexed.set array i x : Cont` for `i : Idx` and `x : Elem`
  3. make new a array: `Indexed.ofFn f : Cont` for `f : Idx → Elem`

Alternative notation:
  1. `array[x]`
  2. in `do` block: `array[x] := y`, `array[x] += y`, ...
  3. `λ [x] => f x` this notation works only if the type `Cont` can be infered from the context
     Common use: `let array : Cont := λ [x] => f x` where the type asscription `: Cont` is important.
-/
class ArrayType (Cont : Type u) (Idx : Type v |> outParam) (Elem : Type w |> outParam)
    extends Indexed.{u,v,w,w} Cont Idx Elem, LawfulIndexed.{u,v,w,w} Cont Idx Elem where
  get_injective : Function.Injective (fun (c : Cont) i => c[i])

instance {Cont Idx Elem} [ArrayType Cont Idx Elem] : StructType Cont Idx (fun _ => Elem) where
  structProj := fun c i => c[i]
  structMake := Indexed.ofFn
  structModify := fun i f c => Indexed.update c i f
  left_inv := sorry_proof
  right_inv := sorry_proof
  structProj_structModify := sorry_proof
  structProj_structModify' := sorry_proof

namespace ArrayType

variable
  {Cont : Type _} {Idx : Type _ |> outParam} {Elem : Type _ |> outParam}
  [IndexType Idx] [outParam <| DecidableEq Idx] [LawfulIndexType Idx]
  [ArrayType Cont Idx Elem]

@[ext]
theorem ext (x y : Cont) : (∀ i, x[i] = y[i]) → x = y := by sorry_proof

@[simp]
theorem eta (cont : Cont) : (Indexed.ofFn fun i => cont[i]) = cont := sorry_proof

def mapMono (f : Elem → Elem) (cont : Cont) : Cont :=
  Fold.fold (IndexType.univ Idx) (fun c i => Indexed.update c i f) cont

def mapIdxMono (f : Idx → Elem → Elem) (cont : Cont) : Cont :=
  Fold.fold (IndexType.univ Idx) (fun c i => Indexed.update c i (f i)) cont

@[simp]
theorem getElem_mapMono (f : Elem → Elem) (cont : Cont) (i : Idx) :
    (mapMono f cont)[i] = f cont[i] := sorry_proof

@[simp]
theorem getElem_mapIdxMono (f : Idx → Elem → Elem) (cont : Cont) (i : Idx) :
    (mapIdxMono f cont)[i] = f i cont[i] := sorry_proof

@[simp]
theorem getElem_map (f : Elem → Elem) (cont : Cont) (i : Idx) :
    (mapMono f cont)[i] = f cont[i] := sorry_proof

instance (priority:=low) [ArrayType Cont Idx Elem] [ToString Elem] [IndexType Idx] :
    ToString (Cont) := ⟨λ x => Id.run do
  let mut fst := true
  let mut s := "⊞["
  for i in IndexType.univ Idx do
    if fst then
      s := s ++ toString x[i]
      fst := false
    else
      s := s ++ ", " ++ toString x[i]
  s ++ "]"⟩

/-- Converts array to ArrayType -/
def _root_.Array.toArrayType {Elem} (Cont : Type u) (Idx : Type v) [IndexType Idx] [Indexed Cont Idx Elem]
  (a : Array Elem) (h : IndexType.card Idx = a.size) : Cont :=
  Indexed.ofFn fun (i : Idx) => a[h ▸ IndexType.toFin i]

-- /-- Converts ListN to ArrayType

--   WARNING: Does not do what expected for lists of size bigger or equal then USize.size
--     For example, array of size USize.size is converted to an array of size zero
--   -/
-- def _root_.ListN.toArrayType {n Elem} (Cont : Type) [ArrayType Cont (SciLean.Idx (n.toUSize)) Elem]
--   (l : ListN Elem n) : Cont :=
--   introElem fun i => l.toArray[i.1.toNat]'sorry_proof

instance [StructType Elem I ElemI] : StructType Cont (Idx×I) (fun (_,i) => ElemI i) where
  structProj := fun x (i,j) => structProj x[i] j
  structMake := fun f => Indexed.ofFn fun i => structMake fun j => f (i,j)
  structModify := fun (i,j) f x => Indexed.update x i (fun xi => structModify j f xi)
  left_inv := by intro x; simp
  right_inv := by intro x; simp
  structProj_structModify := by intro x; simp
  structProj_structModify' := by intro (i,j) (i',j') _ _ _; sorry_proof

section Operations

  instance (priority:=low) [Add Elem] : Add Cont := ⟨λ f g => mapIdxMono (λ x fx => fx + g[x]) f⟩
  instance (priority:=low) [Sub Elem] : Sub Cont := ⟨λ f g => mapIdxMono (λ x fx => fx - g[x]) f⟩
  instance (priority:=low) [Mul Elem] : Mul Cont := ⟨λ f g => mapIdxMono (λ x fx => fx * g[x]) f⟩
  instance (priority:=low) [Div Elem] : Div Cont := ⟨λ f g => mapIdxMono (λ x fx => fx / g[x]) f⟩

  -- instance (priority:=low) {R} [HMul R Elem Elem] : HMul R Cont Cont := ⟨λ r f => map (λ fx => r*(fx : Elem)) f⟩
  instance (priority:=low) {R} [SMul R Elem] : SMul R Cont := ⟨λ r f => mapMono (λ fx => r•(fx : Elem)) f⟩

  instance (priority:=low) [Neg Elem] : Neg Cont := ⟨λ f => mapMono (λ fx => -(fx : Elem)) f⟩
  instance (priority:=low) [Inv Elem] : Inv Cont := ⟨λ f => mapMono (λ fx => (fx : Elem)⁻¹) f⟩

  instance (priority:=low) [One Elem]  : One Cont  := ⟨Indexed.ofFn fun (_ : Idx) => 1⟩
  instance (priority:=low) [Zero Elem] : Zero Cont := ⟨Indexed.ofFn fun (_ : Idx) => 0⟩

  instance (priority:=low) [LT Elem] : LT Cont := ⟨λ f g => ∀ x, f[x] < g[x]⟩
  instance (priority:=low) [LE Elem] : LE Cont := ⟨λ f g => ∀ x, f[x] ≤ g[x]⟩

  instance (priority:=low) [DecidableEq Elem] : DecidableEq Cont :=
    λ f g => Id.run do
      let mut eq : Bool := true
      for x in IndexType.univ Idx do
        if f[x] ≠ g[x] then
          eq := false
          break
      if eq then isTrue sorry_proof else isFalse sorry_proof

  instance (priority:=low) [LT Elem] [∀ x y : Elem, Decidable (x < y)] (f g : Cont) :
      Decidable (f < g) := Id.run do
    let mut lt : Bool := true
    for x in IndexType.univ Idx do
      if ¬(f[x] < g[x]) then
        lt := false
        break
    if lt then isTrue sorry_proof else isFalse sorry_proof

  instance (priority:=low) [LE Elem] [∀ x y : Elem, Decidable (x ≤ y)] (f g : Cont) :
      Decidable (f ≤ g) := Id.run do
    let mut le : Bool := true
    for x in IndexType.univ Idx do
      if ¬(f[x] ≤ g[x]) then
        le := false
        break
    if le then isTrue sorry_proof else isFalse sorry_proof

  @[simp, ftrans_simp]
  theorem add_ofFn [Add Elem] (f g: Idx → Elem)
    : Indexed.ofFn (C:=Cont) f + Indexed.ofFn (C:=Cont) g
      =
      Indexed.ofFn fun i => f i + g i  := by apply ArrayType.ext (Idx:=Idx); simp[HAdd.hAdd, Add.add]

  @[simp, ftrans_simp]
  theorem sub_ofFn [Sub Elem] (f g: Idx → Elem)
    : Indexed.ofFn (C:=Cont) f - Indexed.ofFn (C:=Cont) g
      =
      Indexed.ofFn fun i => f i - g i  := by apply ArrayType.ext (Idx:=Idx); simp[HSub.hSub, Sub.sub]

  @[simp, ftrans_simp]
  theorem neg_ofFn [Neg Elem] (f: Idx → Elem)
    : - Indexed.ofFn (C:=Cont) f
      =
      Indexed.ofFn fun i => - f i  := by apply ArrayType.ext (Idx:=Idx); simp[Neg.neg]

  @[simp, ftrans_simp]
  theorem smul_ofFn [SMul K Elem] (f : Idx → Elem) (a : K)
    : a • Indexed.ofFn (C:=Cont) f
      =
      Indexed.ofFn fun i => a • f i := by apply ArrayType.ext (Idx:=Idx); simp[HSMul.hSMul, SMul.smul];

  @[simp, ftrans_simp]
  theorem sum_ofFn [AddCommMonoid Elem] {ι} [IndexType ι] (f : ι → Idx → Elem)
    : ∑ j, Indexed.ofFn (C:=Cont) (fun i => f j i)
      =
      Indexed.ofFn fun i => ∑ j, f j i
    := sorry_proof

  @[simp, ftrans_simp]
  theorem add_get [Add Elem] (x y : Cont) (i : Idx) :
      (x + y)[i] = x[i] + y[i] := by simp[HAdd.hAdd,Add.add]

  @[simp, ftrans_simp]
  theorem sub_get [Sub Elem] (x y : Cont) (i : Idx) :
      (x - y)[i] = x[i] - y[i] := by simp[HSub.hSub,Sub.sub]

  @[simp, ftrans_simp]
  theorem mul_get [Mul Elem] (x y : Cont) (i : Idx) :
      (x * y)[i] = x[i] * y[i] := by simp[HMul.hMul,Mul.mul]

@[simp, ftrans_simp]
  theorem div_get [Div Elem] (x y : Cont) (i : Idx) :
      (x / y)[i] = x[i] / y[i]


import SciLean.Core.FunctionTransformations
-- import SciLean.Core.Meta.GenerateRevDeriv

open ComplexConjugate

namespace SciLean.Scalar

variable
  {R C} [Scalar R C] [RealScalar R]
  {W} [Vec C W]
  {U} [SemiInnerProductSpace C U]

--------------------------------------------------------------------------------
-- Exp -------------------------------------------------------------------------
--------------------------------------------------------------------------------

@[fun_prop]
theorem exp.arg_x.CDifferentiable_rule
  (x : W → C) (hx : CDifferentiable C x)
  : CDifferentiable C fun w => exp (x w) := sorry_proof

@[fun_trans]
theorem exp.arg_x.ceriv_rule
  (x : W → C) (hx : CDifferentiable C x)
  : cderiv C (fun w => exp (x w))
    =
    fun w dw =>
      let xdx := fwdDeriv C x w dw
      let e := exp xdx.1
      xdx.2 * e := sorry_proof

@[fun_trans]
theorem exp.arg_x.fwdDeriv_rule
    (x : W → C) (hx : CDifferentiable C x) :
    fwdDeriv C (fun w => exp (x w))
    =
    fun w dw =>
      let xdx := fwdDeriv C x w dw
      let e := exp xdx.1
      (e, xdx.2 * e) := by
  unfold fwdDeriv; fun_trans; rfl

@[fun_prop]
theorem exp.arg_x.HasAdjDiff_rule
    (x : U → C) (hx : HasAdjDiff C x) :
    HasAdjDiff C (fun u => exp (x u)) := by
  intro x
  constructor
  fun_prop
  fun_trans [fwdDeriv]; fun_prop

open ComplexConjugate
@[fun_trans]
theorem exp.arg_x.revDeriv_rule
    (x : U → C) (hx : HasAdjDiff C x) :
    revDeriv C (fun u => exp (x u))
    =
    fun u =>
      let xdx := revDeriv C x u
      (exp xdx.1, fun dy => xdx.2 (conj (exp xdx.1) * dy)) := by
  unfold revDeriv
  fun_trans only [fwdDeriv, smul_push, ftrans_simp]

--------------------------------------------------------------------------------
-- Log -------------------------------------------------------------------------
--------------------------------------------------------------------------------

section Log

variable
  {R : Type _} [RealScalar R]
  {W : Type _} [Vec R W]
  {U : Type _} [SemiInnerProductSpace R U]

@[fun_prop]
theorem log.arg_x.CDifferentiableAt_rule
    (w : W) (x : W → R) (hx : CDifferentiableAt R x w) (hw : x w ≠ 0) :
    CDifferentiableAt R (fun w => log (x w)) w := sorry_proof

@[fun_prop]
theorem log.arg_x.CDifferentiable_rule
    (x : W → R) (hx : CDifferentiable R x) (hw : ∀ w, x w ≠ 0) :
    CDifferentiable R (fun w => log (x w)) := by
  intro x; fun_prop (disch:=aesop)

@[fun_trans]
theorem log.arg_x.cderiv_rule_at
    (w : W) (x : W → R) (hx : CDifferentiableAt R x w) (hw : x w ≠ 0) :
    cderiv R (fun w => log (x w)) w
    =
    fun dw =>
      let xdx := fwdDeriv R x w dw
      xdx.2 / abs xdx.1 := sorry_proof

@[fun_trans]
theorem log.arg_x.cderiv_rule
    (x : W → R) (hx : CDifferentiable R x) (hw : ∀ w, x w ≠ 0) :
    cderiv R (fun w => log (x w))
    =
    fun w dw =>
      let xdx := fwdDeriv R x w dw
      xdx.2 / abs xdx.1 := by
  funext x
  fun_trans (disch:=aesop)

@[fun_trans]
theorem log.arg_x.fwdDeriv_rule_at
    (w : W) (x : W → R) (hx : CDifferentiableAt R x w) (hw : x w ≠ 0) :
    fwdDeriv R (fun w => log (x w)) w
    =
    fun dw =>
      let xdx := fwdDeriv R x w dw
      let l := log xdx.1
      (l, xdx.2 / abs xdx.1) :=
by
  unfold fwdDeriv; fun_trans (disch:=assumption); simp[fwdDeriv]

@[fun_trans]
theorem log.arg_x.fwdDeriv_rule
    (x : W → R) (hx : CDifferentiable R x) (hw : ∀ w, x w ≠ 0) :
    fwdDeriv R (fun w => log (x w))
    =
    fun w dw =>
      let xdx := fwdDeriv R x w dw
      let l := log xdx.1
      (l, xdx.2 / abs xdx.1)


import SciLean.Core.FunctionPropositions.ContCDiff
import SciLean.Core.FunctionTransformations

import SciLean.Core.Notation.CDeriv

set_option linter.unusedVariables false

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)]

local notation "∞" => (⊤ : ℕ∞)

-- Function space --------------------------------------------------------------
--------------------------------------------------------------------------------

structure ContCDiffMapFD (n : ℕ∞) (X Y : Type _)  [Vec K X] [Vec K Y] where
  toFun : X → X → Y×Y
  is_cont_cdiff_map : ContCDiff K n (fun x => (toFun x 0).1)
  cderiv_snd : cderiv K (fun x => (toFun x 0).1) = fun x dx => (toFun x dx).2
  dir_independence : ∀ x dx, toFun x dx = toFun x 0

attribute [simp, ftrans_simp] ContCDiffMapFD.cderiv_snd ContCDiffMapFD.dir_independence

variable (n : ℕ∞)

instance : FunLike (ContCDiffMapFD K n X Y) X Y where
  coe f := fun x => (f.toFun x 0).1
  coe_injective' := sorry_proof

macro X:term:25 " ⟿FD[" K:term "," n:term "] " Y:term:26 : term =>
  `(ContCDiffMapFD $K $n $X $Y)

macro X:term:25 " ⟿FD[" n:term "] " Y:term:26 : term =>
  `(ContCDiffMapFD defaultScalar% $n $X $Y)

macro X:term:25 " ⟿FD " Y:term:26 : term =>
  `(ContCDiffMapFD defaultScalar% ∞ $X $Y)

@[app_unexpander ContCDiffMapFD] def unexpandContCDiffMapFD : Lean.PrettyPrinter.Unexpander
  | `($(_) $R $n $X $Y) => `($X ⟿FD[$R,$n] $Y)
  | _  => throw ()

@[fun_prop]
theorem ContCDiffMapFD_apply_right (f : X ⟿FD[K,n] Y) : ContCDiff K n (fun x => f x) := f.2

-- Lambda function notation ----------------------------------------------------
--------------------------------------------------------------------------------

def ContCDiffMapFD.mk' (f : X → Y) (f' : X → X → Y×Y) (h : fwdDeriv K f = f') (hf : ContCDiff K n f) : X ⟿FD[K,n] Y :=
  ⟨f', sorry_proof, sorry_proof, sorry_proof⟩

open Lean Parser Term
macro "fun " x:funBinder " ⟿FD[" K:term "," n:term "] " b:term : term =>
  `(ContCDiffMapFD.mk' $K $n (fun $x => $b) ((fwdDeriv $K fun $x => $b) rewrite_by autodiff /- check that derivative has been eliminated -/) (sorry_proof /- todo: add proof -/) sorry_proof)

macro "fun " x:funBinder " ⟿FD[" n:term "] " b:term : term => `(fun $x ⟿FD[defaultScalar%, $n] $b)
macro "fun " x:funBinder " ⟿FD "             b:term : term => `(fun $x ⟿FD[defaultScalar%,  ∞] $b)

variable {K n}

@[app_unexpander ContCDiffMapFD.mk'] def unexpandContCDiffMapFDMk : Lean.PrettyPrinter.Unexpander

  | `($(_) $R $n $f:term $_ $_ $_) =>
    match f with
    | `(fun $x':ident => $b:term) => `(fun $x' ⟿FD[$R,$n] $b)
    | `(fun ($x':ident : $ty) => $b:term) => `(fun ($x' : $ty) ⟿FD[$R,$n] $b)
    | `(fun $x':ident : $ty => $b:term) => `(fun ($x' : $ty) ⟿FD[$R,$n] $b)
    | _  => throw ()
  | _  => throw ()

@[simp, ftrans_simp]
theorem ContCDiffMapFD_eta (f : X ⟿FD[K,n] Y) : (fun x ⟿FD[K,n] f x) = f := by sorry_proof

@[pp_dot]
def ContCDiffMapFD.FD (f : X ⟿FD[K,n] Y) (x dx : X) : Y×Y := f.toFun x dx

@[fun_trans]
theorem ContCDiffMapFD_eval_fwdDeriv (f : X ⟿FD[K,n] Y) :
    fwdDeriv K (fun x => f x) = f.FD := by
  unfold ContCDiffMapFD.FD fwdDeriv
  simp[DFunLike.coe]

@[fun_prop]
theorem ContCDiffMapFD_eval_cdifferentiable (f : X ⟿FD[K,n] Y) (h : 0 < n) :
    CDifferentiable K (fun x => f x) := by
  simp[DFunLike.coe]
  apply CDifferentaible.ContCDiff_rule
  apply ContCDiffMapFD.is_cont_cdiff_map
  assumption

@[fun_prop]
theorem ContCDiffMapFD_eval_cdifferentiable' (f : X ⟿FD[K,∞] Y) :
    CDifferentiable K (fun x => f x) := by
  fun_prop (disch:=apply ENat.zero_lt_top)

@[simp, ftrans_simp]
theorem ContCDiffMapFD.FD_fst (f : X ⟿FD[K,n] Y) (x dx : X) :
    (f.FD x dx).1 = f x := by rw[← ContCDiffMapFD_eval_fwdDeriv]; unfold fwdDeriv; simp

-- Algebra ---------------------------------------------------------------------
--------------------------------------------------------------------------------

instance : Add (X ⟿FD[K,n] Y) := ⟨fun f g => fun x ⟿FD[K,n] f x + g x⟩
instance : Sub (X ⟿FD[K,n] Y) := ⟨fun f g => fun x ⟿FD[K,n] f x - g x⟩
instance : Neg (X ⟿FD[K,n] Y) := ⟨fun f => fun x ⟿FD[K,n] - f x⟩
instance : SMul K (X ⟿FD[K,n] Y) := ⟨fun r f => fun x ⟿FD[K,n] r • f x⟩
instance : Zero (X ⟿FD[K,n] Y) := ⟨fun (x : X) ⟿FD[K,n] (0:Y)⟩

section AlgebraSimps

variable (f g : X ⟿FD[K,n] Y) (x : X) (r : K)

@[simp, ftrans_simp]
theorem ContCDiffMapFD.add_apply : (f + g) x = f x + g x := by rfl

@[simp, ftrans_simp]
theorem ContCDiffMapFD.sub_apply : (f - g) x = f x - g x


import SciLean.Core.FunctionPropositions
import SciLean.Core.FunctionTransformations
import SciLean.Data.ArrayType.Algebra

open SciLean

set_option linter.unusedVariables false

open LeanColls

section GenericArrayType

variable
  {K : Type} [RCLike K]
  {Cont : Type} {Idx : Type |> outParam} {Elem : Type |> outParam}
  [ArrayType Cont Idx Elem] [IndexType Idx] [LawfulIndexType Idx] [DecidableEq Idx]

-- Indexed.get -----------------------------------------------------------------
--------------------------------------------------------------------------------

section OnVectorSpaces

variable {R : Type} [CommSemiring R] [AddCommGroup Elem] [Module R Elem]

@[fun_prop]
theorem LeanColls.Indexed.get.arg_cont.IsLinearMap_rule_simple (idx : Idx) :
    IsLinearMap R (fun xs : Cont => xs[idx]) := by
  constructor
  . intros; simp
  . intros; simp

#generate_linear_map_simps LeanColls.Indexed.get.arg_cont.IsLinearMap_rule_simple

end OnVectorSpaces

section OnNormedSpaces

variable
  {X : Type _} [NormedAddCommGroup X] [NormedSpace K X]
  [NormedAddCommGroup Elem] [NormedSpace K Elem]

@[fun_prop]
theorem LeanColls.Indexed.get.arg_cont.IsContinuousLinearMap_rule_simple (idx : Idx) :
    IsContinuousLinearMap K (λ cont : Cont => cont[idx]) := sorry_proof

-- automatically infer: Differentiable, DifferentiableAt, ContDiff, fderiv, fwdFDeriv

end OnNormedSpaces

section OnVec

variable
  {X : Type} [Vec K X]
  [Vec K Elem]

@[fun_prop]
theorem GetElem.getElem.arg_cont.IsSmoothLinearMap_rule_simple
  (idx : Idx)
  : IsSmoothLinearMap K (fun xs : Cont => xs[idx]) := sorry_proof

-- automatically infer: CDifferentiable, CDifferentiableAt, CContDiff, cderiv, fwdDeriv

end OnVec

section OnSemiInnerProductSpace

variable
  {X : Type} [SemiInnerProductSpace K X]
  [SemiInnerProductSpace K Elem]

@[fun_prop]
theorem GetElem.getElem.arg_cont.HasSemiAdjoint_rule_simple (idx : Idx) :
    HasSemiAdjoint K (fun cont : Cont => cont[idx]) := sorry_proof

@[fun_trans]
theorem GetElem.getElem.arg_cont.semiAdjoint_rule_simple (idx : Idx) :
    semiAdjoint K (fun cont : Cont => cont[idx])
    =
    fun elem => oneHot (X:=Cont) idx elem :=
by
  sorry_proof

@[fun_trans]
theorem GetElem.getElem.arg_cont.revCDeriv_rule
  (f : X → Cont) (idx : Idx)
  (hf : HasAdjDiff K f)
  : revDeriv K (fun x => (f x)[idx])
    =
    fun x =>
      let ydf := revDerivProj K Idx f x
      (ydf.1[idx],
       fun delem => ydf.2 idx delem) :=
by
  unfold revDeriv;
  conv =>
    lhs; fun_trans; fun_trans; simp

@[fun_trans]
theorem GetElem.getElem.arg_cont.revCDerivUpdate_rule
  (f : X → Cont) (idx : Idx)
  (hf : HasAdjDiff K f)
  : revDerivUpdate K (fun x => (f x)[idx])
    =
    fun x =>
      let ydf := revDerivProjUpdate K Idx f x
      (ydf.1[idx],
       fun delem dx => ydf.2 idx delem dx) :=
by
  unfold revDerivUpdate;
  conv =>
    lhs; fun_trans; fun_trans; simp

@[fun_trans]
theorem GetElem.getElem.arg_cont.revCDerivProj_rule
  {J ElemJ} [StructType Elem J ElemJ] [IndexType J] [LawfulIndexType J] [DecidableEq J]
  [∀ j, SemiInnerProductSpace K (ElemJ j)] [SemiInnerProductSpaceStruct K Elem J ElemJ]
  (f : X → Cont) (idx : Idx)
  (hf : HasAdjDiff K f)
  : revDerivProj K J (fun x => (f x)[idx])
    =
    fun x =>
      let ydf := revDerivProj K (Idx×J) f x
      (ydf.1[idx],
       fun j delem => ydf.2 (idx,j) delem) :=
by
  unfold revDerivProj
  conv =>
    lhs; fun_trans; fun_trans; simp
  funext x; simp[revDerivProj]
  funext j de
  simp [revDeriv]
  congr
  sorry_proof

end OnSemiInnerProductSpace

-- Indexed.set ----------------------------------------------------------------
--------------------------------------------------------------------------------

section OnVectorSpaces

variable {R : Type} [CommSemiring R] [AddCommGroup Elem] [Module R Elem]

@[fun_prop]
theorem LeanColls.Indexed.set.arg_contelem.IsLinearMap_rule_simple (idx : Idx) :
    IsLinearMap R (fun ((cont,elem) : Cont×Elem) => Indexed.set cont idx elem) := by
  constructor
  . intros; simp; ext; sorry_proof
  . intros; simp; ext; sorry_proof

#generate_linear_map_simps LeanColls.Indexed.set.arg_contelem.IsLinearMap_rule_simple

end OnVectorSpaces

section OnNormedSpaces

variable
  {X : Type _} [NormedAddCommGroup X] [NormedSpace K X]
  [NormedAddCommGroup Elem] [NormedSpace K Elem]

@[fun_prop]
theorem LeanColls.Indexed.set.arg_contelem.IsContinuousLinearMap_rule_simple :
    IsContinuousLinearMap K (fun ((cont,elem) : Cont×Elem) => Indexed.set cont idx elem) :=
  sorry_proof

-- automatically infer: Differentiable, DifferentiableAt, ContDiff, fderiv, fwdFDeriv

end OnNormedSpaces

section OnVec

variable
  {X : Type} [Vec K X]
  [Vec K Elem]

@[fun_prop]
theorem LeanColls.Indexed.set.arg_contelem.IsSmoothLinearMap_rule_simple (idx : Idx) :
    IsSmoothLinearMap K (fun ((cont,elem) : Cont×Elem) => Indexed.set cont idx elem) := sorry_proof

-- automatically infer: CDifferentiable, CDifferentiableAt, CContDiff, cderiv, fwdDeriv

end OnVec

section OnSemiInnerProductSpace

variable
  {X : Type} [SemiInnerProductSpace K X]
  [SemiInnerProductSpace K Elem]

@[fun_prop]
theorem LeanColls.Indexed.set.arg_cont.HasSemiAdjoint_rule_simple (idx : Idx) :
    HasSemiAdjoint K (fun ((cont,elem) : Cont×Elem) => Indexed.set cont idx elem) := sorry_proof

@[fun_trans]
theorem LeanColls.Indexed.set.arg_cont.semiAdjoint_rule_simple (idx : Idx) :
    semiAdjoint K (fun ((cont,elem) : Cont×Elem) => Indexed.set cont idx elem)
    =
    fun cont' : Cont => (Indexed.set cont' idx 0, cont'[idx]) := sorry_proof

@[fun_trans]
theorem LeanColls.Indexed.set.arg_contelem.revCDeriv_rule
  (cont : X → Cont) (idx : Idx) (elem : X → Elem)
  (hcont : HasAdjDiff K cont) (helem : HasAdjDiff K elem)
  : revDeriv K (fun x => Indexed.set (cont x) idx (elem x))
    =
    fun x =>
      let cdc := revDeriv K cont x
      let ede := revDerivUpdate K elem x
      (Indexed.set cdc.1 idx ede.1,
       fun dc : Cont =>
         let dci := dc[idx]
         let dc := Indexed.set dc idx 0
         ede.2 dci (cdc.2 dc)) := by
  unfold revDeriv;
  conv =>
    lhs; fun_trans; fun_trans; simp

@[fun_trans]
theorem LeanColls.Indexed.set.arg_contelem.revCDerivUpdate_rule
  (cont : X → Cont) (idx : Idx) (elem : X → Elem)
  (hcont : HasAdjDiff K cont) (helem : HasAdjDiff K elem)
  : revDerivUpdate K (fun x => Indexed.set (cont x) idx (elem x))
    =
    fun x =>
      let cdc := revDerivUpdate K cont x
      let ede := revDerivUpdate K elem x
      (Indexed.set cdc.1 idx ede.1,
       fun (dc : Cont) dx =>
         let dci := dc[idx]
         let dc := Indexed.set dc idx 0
         ede.2 dci (cdc.2 dc dx)) := by
  unfold revDerivUpdate
  conv =>
    lhs; fun_trans; fun_trans; simp
  simp[revDerivUpdate,add_assoc]

-- Consider formulating variant for index type `Idx×J`
@[fun_trans]
theorem LeanColls.Indexed.set.arg_contelem.revCDerivProj_rule
    (cont : X → Cont) (idx : Idx) (elem : X → Elem)
    (hcont : HasAdjDiff K cont) (helem : HasAdjDiff K elem) :
    revDerivProj K Idx (fun x => Indexed.set (cont x) idx (elem x))
    =
    fun x =>
      let cdc := revDerivProj K Idx cont x
      let ede := revDeriv K elem x
      (Indexed.set cdc.1 idx ede.1,
       fun i delem =>
         if i = idx then
           ede.2 delem
         else
           cdc.2 i delem) :=
by
  unfold revDerivProj
  conv =>
    lhs; fun_trans; fun_trans; simp
  funext x; simp[revDerivProj,revDerivUpdate]
  funext i de
  if h : i = idx then
    simp [revDeriv,h,oneHot]
    sorry_proof
  else
    simp [revDeriv,h,oneHot]
    sorry_proof

-- Consider formulating variant for index type `Idx×J`
@[fun_trans]
theorem LeanColls.Indexed.set.arg_contelem.revCDerivProjUpdate_rule
    (cont : X → Cont) (idx : Idx) (elem : X → Elem)
    (hcont : HasAdjDiff K cont) (helem : HasAdjDiff K elem) :
    revDerivProjUpdate K Idx (fun x => Indexed.set (cont x) idx (elem x))
    =
    fun x =>
      let cdc := revDerivProjUpdate K Idx cont x
      let ede := revDerivUpdate K elem x
      (Indexed.set cdc.1 idx ede.1,
       fun i delem dx =>
         if i = idx then
           ede.2 delem dx
         else
           cdc.2 i delem dx) :=
by
  unfold revDerivProjUpdate
  conv =>
    lhs; fun_trans; fun_trans; simp
  funext x; simp[revDerivProj,revDerivUpdate]
  funext i de
  if h : i = idx then
    simp [h]
  else
    simp [h]

end OnSemiInnerProductSpace

-- LeanColls.Indexed.ofFn -------------------------------------------------------------
--------------------------------------------------------------------------------

section OnVectorSpaces

variable {R : Type} [CommSemiring R] [AddCommGroup Elem] [Module R Elem]

@[fun_prop]
theorem LeanColls.Indexed.ofFn.arg_f.IsLinearMap_rule_simple :
    IsLinearMap R (fun f : Idx → Elem => Indexed.ofFn (C:=Cont) f) := by
  constructor
  . intros; ext; simp
  . intros; ext; simp

#generate_linear_map_simps LeanColls.Indexed.ofFn.arg_f.IsLinearMap_rule_simple

end OnVectorSpaces

section OnNormedSpaces

variable
  {X : Type _} [NormedAddCommGroup X] [NormedSpace K X]
  [NormedAddCommGroup Elem] [NormedSpace K Elem]

@[fun_prop]
theorem LeanColls.Indexed.ofFn.arg_f.IsContinuousLinearMap_rule_simple :
    IsContinuousLinearMap K (fun f : Idx → Elem => Indexed.ofFn (C:=Cont) f) := sorry_proof

-- automatically infer: Differentiable, DifferentiableAt, ContDiff, fderiv, fwdFDeriv

end OnNormedSpaces

section OnVec

variable
  {X : Type} [Vec K X]
  [Vec K Elem]

@[fun_prop]
theorem LeanColls.Indexed.ofFn.arg_f.IsSmoothLinearMap_rule_simple
  (idx : Idx)
  : IsSmoothLinearMap K (fun f : Idx → Elem => Indexed.ofFn (C:=Cont) f) := sorry_proof

-- automatically infer: CDifferentiable, CDifferentiableAt, CContDiff, cderiv, fwdDeriv

end OnVec

section OnSemiInnerProductSpace

variable
  {X : Type} [SemiInnerProductSpace K X]
  [SemiInnerProductSpace K Elem]

@[fun_prop]
theorem LeanColls.Indexed.ofFn.arg_cont.HasSemiAdjoint_rule_simple :
    HasSemiAdjoint K (fun f : Idx → Elem => Indexed.ofFn (C:=Cont) f) := sorry_proof

@[fun_trans]
theorem LeanColls.Indexed.ofFn.arg_cont.semiAdjoint_rule_simple :
    semiAdjoint K (fun f : Idx → Elem => Indexed.ofFn (C:=Cont) f)
    =
    fun (cont : Cont) idx => cont[idx] :=
by
  sorry_proof

@[fun_trans]
theorem LeanColls.Indexed.ofFn.arg_cont.revCDeriv_rule
    (f : X → Idx → Elem) (hf : HasAdjDiff K f) :
    revDeriv K (fun x => Indexed.ofFn (C:=Cont) (f x))
    =
    fun x =>
      let fdf := revDeriv K f x
      (Indexed.ofFn fdf.1,
       fun dcont : Cont =>
         fdf.2 (fun idx => dcont[idx])) :=
by
  unfold revDeriv;
  conv => lhs; fun_trans; simp

@[fun_trans]
theorem LeanColls.Indexed.ofFn.arg_cont.revCDerivUpdate_rule
    (f : X → Idx → Elem) (hf : HasAdjDiff K f) :
    revDerivUpdate K (fun x => Indexed.ofFn (C:=Cont) (f x))
    =
    fun x =>
      let fdf := revDerivUpdate K f x
      (Indexed.ofFn fdf.1,
       fun (dcont : Cont) dx =>
         fdf.2 (fun idx => dcont[idx]) dx) :=
by
  unfold revDerivUpdate;
  conv => lhs; fun_trans; simp

@[fun_trans]
theorem LeanColls.Indexed.ofFn.arg_cont.revCDerivProj_rule
    (f : X → Idx → Elem) (hf : HasAdjDiff K f) :
    revDerivProj K Idx (fun x => Indexed.ofFn (C:=Cont) (f x))
    =
    fun x =>
      let fdf := revDerivProj K Idx f x
      (Indexed.ofFn fdf.1,
       fdf.2)

:= by
  unfold revDerivProj;
  conv => lhs; fun_trans; simp
  funext x; simp[oneHot]


import SciLean.Core.Objects.Vec
import SciLean.Core.Objects.SemiInnerProductSpace
import SciLean.Core.Objects.FinVec
import SciLean.Core.Objects.Scalar
import SciLean.Core.FunctionPropositions.CDifferentiable

import SciLean.Core.Meta.ToAnyPoint

set_option linter.unusedVariables false

open LeanColls

namespace SciLean

local notation "∞" => (⊤ : ℕ∞)

variable
  (K : Type _) [RCLike K] (n : ℕ∞)
  {X : Type _} [Vec K X]
  {Y : Type _} [Vec K Y]
  {Z : Type _} [Vec K Z]
  {ι : Type _} [IndexType ι]
  {E : ι → Type _} [∀ i, Vec K (E i)]

/-- `CDifferentiableAt f x` - conveniently differentiable function `f` at point `x`. -/
@[irreducible, fun_prop]
def ContCDiffAt (f : X → Y) (x : X) : Prop :=
  ∀ (c : K → X),
      c 0 = x
      →
      Curve.ContDiffAt c 0 n
      →
      Curve.ContDiffAt (f∘c) 0 n

/-- `ContCDiff f` - conveniently differentiable function `f`.  -/
@[fun_prop]
def ContCDiff (f : X → Y) : Prop := ∀ x, ContCDiffAt K n f x

@[fun_prop, to_any_point]
theorem ContCDiffAt.id_rule (x : X) :
    ContCDiffAt K n (fun x : X => x) x := by sorry_proof

@[fun_prop, to_any_point]
theorem ContCDiffAt.const_rule (y : Y) (x : X) :
    ContCDiffAt K n (fun _ : X => y) x := by sorry_proof

@[fun_prop, to_any_point]
theorem ContCDiffAt.comp_rule (x : X)
    (f : Y → Z) (g : X → Y)
    (hf : ContCDiffAt K n f (g x)) (hg : ContCDiffAt K n g x) :
    ContCDiffAt K n (fun x => f (g x)) x := by sorry_proof

@[fun_prop, to_any_point]
theorem ContCDiffAt.apply_rule
    (i : ι) (x) : ContCDiffAt K n (fun x : (i : ι) → E i => x i) x := by sorry_proof

@[fun_prop, to_any_point]
theorem ContCDiffAt.pi_rule (x : X)
    (f : X → (i : ι) → E i)
    (hf : ∀ i, ContCDiffAt K n (f · i) x) :
    ContCDiffAt K n (fun x i => f x i) x := by sorry_proof

@[fun_prop, to_any_point]
theorem ContCDiffAt.le_rule (x : X) (f : X → Y) (hf : ContCDiffAt K m f x) (h : n ≤ m) :
    ContCDiffAt K n f x := sorry_proof

@[fun_prop, to_any_point]
theorem ContCDiff.le_rule (f : X → Y) (hf : ContCDiff K m f) (h : n ≤ m) :
    ContCDiff K n f := sorry_proof

@[fun_prop]
theorem ContCDiff.ContCDiffAt_rule (x : X) (f : X → Y) (hf : ContCDiff K n f) :
    ContCDiffAt K n f x := hf x

@[fun_prop]
theorem ContCDiff.id_rule :
    ContCDiff K n (fun x : X => x) := by intro x; fun_prop

@[fun_prop]
theorem ContCDiff.const_rule (y : Y) :
    ContCDiff K n (fun _ : X => y) := by intro x; fun_prop

@[fun_prop]
theorem ContCDiff.comp_rule
    (f : Y → Z) (g : X → Y)
    (hf : ContCDiff K n f) (hg : ContCDiff K n g) :
    ContCDiff K n (fun x => f (g x)) := by intro x; fun_prop

@[fun_prop]
theorem ContCDiff.apply_rule
    (i : ι) : ContCDiff K n (fun x : (i : ι) → E i => x i) := by intro x; fun_prop

@[fun_prop]
theorem ContCDiff.pi_rule (x : X)
    (f : X → (i : ι) → E i)
    (hf : ∀ i, ContCDiff K n (f · i)) :
    ContCDiff K n (fun x i => f x i) := by intro x; apply ContCDiffAt.pi_rule; fun_prop

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

-- transition rules to CDifferentiable

@[fun_prop]
theorem CDifferentaibleAt.ContCDiffAt_rule (x : X) (f : X → Y) (hf : ContCDiffAt K n f x) (h : 0 < n) :
    CDifferentiableAt K f x := sorry_proof

@[fun_prop]
theorem CDifferentaibleAt.ContCDiffAt_rule' (x : X) (f : X → Y) (hf : ContCDiffAt K ∞ f x) :
    CDifferentiableAt K f x := by fun_prop (disch:=aesop)

@[fun_prop]
theorem CDifferentaible.ContCDiff_rule (f : X → Y) (hf : ContCDiff K n f) (h : 0 < n) :
    CDifferentiable K f

:= by intro x; fun_prop (disch:=assumption)


import SciLean.Core.FunctionTransformations
-- import SciLean.Core.Meta.GenerateRevDeriv

open ComplexConjugate

namespace SciLean.Scalar

variable
  {R C} [Scalar R C] [RealScalar R]
  {W} [Vec C W]
  {U} [SemiInnerProductSpace C U]

--------------------------------------------------------------------------------
-- Exp -------------------------------------------------------------------------
--------------------------------------------------------------------------------

@[fun_prop]
theorem exp.arg_x.CDifferentiable_rule
  (x : W → C) (hx : CDifferentiable C x)
  : CDifferentiable C fun w => exp (x w) := sorry_proof

@[fun_trans]
theorem exp.arg_x.ceriv_rule
  (x : W → C) (hx : CDifferentiable C x)
  : cderiv C (fun w => exp (x w))
    =
    fun w dw =>
      let xdx := fwdDeriv C x w dw
      let e := exp xdx.1
      xdx.2 * e := sorry_proof

@[fun_trans]
theorem exp.arg_x.fwdDeriv_rule
    (x : W → C) (hx : CDifferentiable C x) :
    fwdDeriv C (fun w => exp (x w))
    =
    fun w dw =>
      let xdx := fwdDeriv C x w dw
      let e := exp xdx.1
      (e, xdx.2 * e) := by
  unfold fwdDeriv; fun_trans; rfl

@[fun_prop]
theorem exp.arg_x.HasAdjDiff_rule
    (x : U → C) (hx : HasAdjDiff C x) :
    HasAdjDiff C (fun u => exp (x u)) := by
  intro x
  constructor
  fun_prop
  fun_trans [fwdDeriv]; fun_prop

open ComplexConjugate

@[fun_trans]
theorem exp.arg_x.revDeriv_rule
    (x : U → C) (hx : HasAdjDiff C x) :
    revDeriv C (fun u => exp (x u))
    =
    fun u =>
      let xdx := revDeriv C x u
      (exp xdx.1, fun dy => xdx.2 (conj (exp xdx.1) * dy))


import SciLean.Core.FunctionTransformations
import SciLean.Core.FloatAsReal
-- import SciLean.Core.Meta.GenerateRevDeriv

open ComplexConjugate

namespace SciLean.Scalar

variable
  {R} [RealScalar R]
  {W} [Vec R W]
  {U} [SemiInnerProductSpace R U]

----------------------------------------------------------------------------------------------------
-- Sqrt --------------------------------------------------------------------------------------------
----------------------------------------------------------------------------------------------------

@[fun_prop]
theorem sqrt.arg_x.CDifferentiableAt_rule
    (w : W) (x : W → R) (hx : CDifferentiableAt R x w) (hw : x w ≠ 0) :
    CDifferentiableAt R (fun w => sqrt (x w)) w := sorry_proof

@[fun_prop]
theorem sqrt.arg_x.CDifferentiable_rule
    (x : W → R) (hx : CDifferentiable R x) (hw : ∀ w, x w ≠ 0) :
    CDifferentiable R (fun w => sqrt (x w)) := by
  intro x; fun_prop (disch:=aesop)

@[fun_trans]
theorem sqrt.arg_x.cderiv_rule_at
    (w : W) (x : W → R) (hx : CDifferentiableAt R x w) (hw : 0 < x w) :
    cderiv R (fun w => sqrt (x w)) w
    =
    fun dw =>
      let xdx := fwdDeriv R x w dw
      xdx.2 / (2 * sqrt xdx.1) := sorry_proof

@[fun_trans]
theorem sqrt.arg_x.cderiv_rule
    (x : W → R) (hx : CDifferentiable R x) (hw : ∀ w, 0 < x w) :
    cderiv R (fun w => sqrt (x w))
    =
    fun w dw =>
      let xdx := fwdDeriv R x w dw
      xdx.2 / (2 * sqrt xdx.1) := by
  funext x
  fun_trans (disch:=aesop)

@[fun_trans]
theorem sqrt.arg_x.fwdDeriv_rule_at
    (w : W) (x : W → R) (hx : CDifferentiableAt R x w) (hw : 0 < x w) :
    fwdDeriv R (fun w => sqrt (x w)) w
    =
    fun dw =>
      let xdx := fwdDeriv R x w dw
      let x' := sqrt xdx.1
      (x', xdx.2 / (2 * x')) :=
by
  unfold fwdDeriv
  fun_trans (disch:=assumption) [fwdDeriv]

@[fun_trans]
theorem sqrt.arg_x.fwdDeriv_rule
    (x : W → R) (hx : CDifferentiable R x) (hw : ∀ w, 0 < x w) :
    fwdDeriv R (fun w => sqrt (x w))
    =
    fun w dw =>
      let xdx := fwdDeriv R x w dw
      let x' := sqrt xdx.1
      (x', xdx.2 / (2 * x')) :=
by
  unfold fwdDeriv; fun_trans (disch:=assumption) [fwdDeriv]

@[fun_prop]
theorem sqrt.arg_x.HasAdjDiffAt_rule
    (u : U) (x : U → R) (hx : HasAdjDiffAt R x u) (hu : 0 < x u) :
    HasAdjDiffAt R (fun u => sqrt (x u)) u := by
  constructor
  fun_prop (disch:=aesop)
  fun_trans (disch:=aesop) [fwdDeriv]; fun_prop

@[fun_prop]
theorem sqrt.arg_x.HasAdjDiff_rule
    (x : U → R) (hx : HasAdjDiff R x) (hu : ∀ u, 0 < x u) :
    HasAdjDiff R (fun u => sqrt (x u))


import SciLean.Core.Objects.Vec
import SciLean.Core.Objects.SemiInnerProductSpace
import SciLean.Core.Objects.FinVec
import SciLean.Core.Objects.Scalar
import SciLean.Core.FunctionPropositions.CDifferentiable

import SciLean.Core.Meta.ToAnyPoint

set_option linter.unusedVariables false

open LeanColls

namespace SciLean

local notation "∞" => (⊤ : ℕ∞)

variable
  (K : Type _) [RCLike K] (n : ℕ∞)
  {X : Type _} [Vec K X]
  {Y : Type _} [Vec K Y]
  {Z : Type _} [Vec K Z]
  {ι : Type _} [IndexType ι]
  {E : ι → Type _} [∀ i, Vec K (E i)]

/-- `CDifferentiableAt f x` - conveniently differentiable function `f` at point `x`. -/
@[irreducible, fun_prop]
def ContCDiffAt (f : X → Y) (x : X) : Prop :=
  ∀ (c : K → X),
      c 0 = x
      →
      Curve.ContDiffAt c 0 n
      →
      Curve.ContDiffAt (f∘c) 0 n

/-- `ContCDiff f` - conveniently differentiable function `f`.  -/
@[fun_prop]
def ContCDiff (f : X → Y) : Prop := ∀ x, ContCDiffAt K n f x

@[fun_prop, to_any_point]
theorem ContCDiffAt.id_rule (x : X) :
    ContCDiffAt K n (fun x : X => x) x := by sorry_proof

@[fun_prop, to_any_point]
theorem ContCDiffAt.const_rule (y : Y) (x : X) :
    ContCDiffAt K n (fun _ : X => y) x := by sorry_proof

@[fun_prop, to_any_point]
theorem ContCDiffAt.comp_rule (x : X)
    (f : Y → Z) (g : X → Y)
    (hf : ContCDiffAt K n f (g x)) (hg : ContCDiffAt K n g x) :
    ContCDiffAt K n (fun x => f (g x)) x := by sorry_proof

@[fun_prop, to_any_point]
theorem ContCDiffAt.apply_rule
    (i : ι) (x) : ContCDiffAt K n (fun x : (i : ι) → E i => x i) x := by sorry_proof

@[fun_prop, to_any_point]
theorem ContCDiffAt.pi_rule (x : X)
    (f : X → (i : ι) → E i)
    (hf : ∀ i, ContCDiffAt K n (f · i) x) :
    ContCDiffAt K n (fun x i => f x i) x := by sorry_proof

@[fun_prop, to_any_point]
theorem ContCDiffAt.le_rule (x : X) (f : X → Y) (hf : ContCDiffAt K m f x) (h : n ≤ m) :
    ContCDiffAt K n f x := sorry_proof

@[fun_prop, to_any_point]
theorem ContCDiff.le_rule (f : X → Y) (hf : ContCDiff K m f) (h : n ≤ m) :
    ContCDiff K n f := sorry_proof

@[fun_prop]
theorem ContCDiff.ContCDiffAt_rule (x : X) (f : X → Y) (hf : ContCDiff K n f) :
    ContCDiffAt K n f x


import SciLean.Util.SorryProof
import SciLean.Data.ListN
import SciLean.Data.StructType.Basic
import SciLean.Data.Function
import LeanColls

namespace SciLean

open LeanColls

/-- This class says that `Cont` behaves like an array with `Elem` values indexed by `Idx`

Examples for `Idx = Fin n` and `Elem = ℝ` are: `ArrayN ℝ n` or `ℝ^{n}`

For `array : Cont` you can:
  1. get values: `Indexed.get array i : Elem` for `i : Idx`
  2. set values: `Indexed.set array i x : Cont` for `i : Idx` and `x : Elem`
  3. make new a array: `Indexed.ofFn f : Cont` for `f : Idx → Elem`

Alternative notation:
  1. `array[x]`
  2. in `do` block: `array[x] := y`, `array[x] += y`, ...
  3. `λ [x] => f x` this notation works only if the type `Cont` can be infered from the context
     Common use: `let array : Cont := λ [x] => f x` where the type asscription `: Cont` is important.
-/
class ArrayType (Cont : Type u) (Idx : Type v |> outParam) (Elem : Type w |> outParam)
    extends Indexed.{u,v,w,w} Cont Idx Elem, LawfulIndexed.{u,v,w,w} Cont Idx Elem where
  get_injective : Function.Injective (fun (c : Cont) i => c[i])

instance {Cont Idx Elem} [ArrayType Cont Idx Elem] : StructType Cont Idx (fun _ => Elem) where
  structProj := fun c i => c[i]
  structMake := Indexed.ofFn
  structModify := fun i f c => Indexed.update c i f
  left_inv := sorry_proof
  right_inv := sorry_proof
  structProj_structModify := sorry_proof
  structProj_structModify' := sorry_proof

namespace ArrayType

variable
  {Cont : Type _} {Idx : Type _ |> outParam} {Elem : Type _ |> outParam}
  [IndexType Idx] [outParam <| DecidableEq Idx] [LawfulIndexType Idx]
  [ArrayType Cont Idx Elem]

@[ext]
theorem ext (x y : Cont) : (∀ i, x[i] = y[i]) → x = y := by sorry_proof

@[simp]
theorem eta (cont : Cont) : (Indexed.ofFn fun i => cont[i]) = cont := sorry_proof

def mapMono (f : Elem → Elem) (cont : Cont) : Cont :=
  Fold.fold (IndexType.univ Idx) (fun c i => Indexed.update c i f) cont

def mapIdxMono (f : Idx → Elem → Elem) (cont : Cont) : Cont :=
  Fold.fold (IndexType.univ Idx) (fun c i => Indexed.update c i (f i)) cont

@[simp]
theorem getElem_mapMono (f : Elem → Elem) (cont : Cont) (i : Idx) :
    (mapMono f cont)[i] = f cont[i] := sorry_proof

@[simp]
theorem getElem_mapIdxMono (f : Idx → Elem → Elem) (cont : Cont) (i : Idx) :
    (mapIdxMono f cont)[i] = f i cont[i] := sorry_proof

@[simp]
theorem getElem_map (f : Elem → Elem) (cont : Cont) (i : Idx) :
    (mapMono f cont)[i] = f cont[i] := sorry_proof

instance (priority:=low) [ArrayType Cont Idx Elem] [ToString Elem] [IndexType Idx] :
    ToString (Cont) := ⟨λ x => Id.run do
  let mut fst := true
  let mut s := "⊞["
  for i in IndexType.univ Idx do
    if fst then
      s := s ++ toString x[i]
      fst := false
    else
      s := s ++ ", " ++ toString x[i]
  s ++ "]"⟩

/-- Converts array to ArrayType -/
def _root_.Array.toArrayType {Elem} (Cont : Type u) (Idx : Type v) [IndexType Idx] [Indexed Cont Idx Elem]
  (a : Array Elem) (h : IndexType.card Idx = a.size) : Cont :=
  Indexed.ofFn fun (i : Idx) => a[h ▸ IndexType.toFin i]

-- /-- Converts ListN to ArrayType

--   WARNING: Does not do what expected for lists of size bigger or equal then USize.size
--     For example, array of size USize.size is converted to an array of size zero
--   -/
-- def _root_.ListN.toArrayType {n Elem} (Cont : Type) [ArrayType Cont (SciLean.Idx (n.toUSize)) Elem]
--   (l : ListN Elem n) : Cont :=
--   introElem fun i => l.toArray[i.1.toNat]'sorry_proof

instance [StructType Elem I ElemI] : StructType Cont (Idx×I) (fun (_,i) => ElemI i) where
  structProj := fun x (i,j) => structProj x[i] j
  structMake := fun f => Indexed.ofFn fun i => structMake fun j => f (i,j)
  structModify := fun (i,j) f x => Indexed.update x i (fun xi => structModify j f xi)
  left_inv := by intro x; simp
  right_inv := by intro x; simp
  structProj_structModify := by intro x; simp
  structProj_structModify' := by intro (i,j) (i',j') _ _ _; sorry_proof

section Operations

  instance (priority:=low) [Add Elem] : Add Cont := ⟨λ f g => mapIdxMono (λ x fx => fx + g[x]) f⟩
  instance (priority:=low) [Sub Elem] : Sub Cont := ⟨λ f g => mapIdxMono (λ x fx => fx - g[x]) f⟩
  instance (priority:=low) [Mul Elem] : Mul Cont := ⟨λ f g => mapIdxMono (λ x fx => fx * g[x]) f⟩
  instance (priority:=low) [Div Elem] : Div Cont := ⟨λ f g => mapIdxMono (λ x fx => fx / g[x]) f⟩

  -- instance (priority:=low) {R} [HMul R Elem Elem] : HMul R Cont Cont := ⟨λ r f => map (λ fx => r*(fx : Elem)) f⟩
  instance (priority:=low) {R} [SMul R Elem] : SMul R Cont := ⟨λ r f => mapMono (λ fx => r•(fx : Elem)) f⟩

  instance (priority:=low) [Neg Elem] : Neg Cont := ⟨λ f => mapMono (λ fx => -(fx : Elem)) f⟩
  instance (priority:=low) [Inv Elem] : Inv Cont := ⟨λ f => mapMono (λ fx => (fx : Elem)⁻¹) f⟩

  instance (priority:=low) [One Elem]  : One Cont  := ⟨Indexed.ofFn fun (_ : Idx) => 1⟩
  instance (priority:=low) [Zero Elem] : Zero Cont := ⟨Indexed.ofFn fun (_ : Idx) => 0⟩

  instance (priority:=low) [LT Elem] : LT Cont := ⟨λ f g => ∀ x, f[x] < g[x]⟩
  instance (priority:=low) [LE Elem] : LE Cont := ⟨λ f g => ∀ x, f[x] ≤ g[x]⟩

  instance (priority:=low) [DecidableEq Elem] : DecidableEq Cont :=
    λ f g => Id.run do
      let mut eq : Bool := true
      for x in IndexType.univ Idx do
        if f[x] ≠ g[x] then
          eq := false
          break
      if eq then isTrue sorry_proof else isFalse sorry_proof

  instance (priority:=low) [LT Elem] [∀ x y : Elem, Decidable (x < y)] (f g : Cont) :
      Decidable (f < g) := Id.run do
    let mut lt : Bool := true
    for x in IndexType.univ Idx do
      if ¬(f[x] < g[x]) then
        lt := false
        break
    if lt then isTrue sorry_proof else isFalse sorry_proof

  instance (priority:=low) [LE Elem] [∀ x y : Elem, Decidable (x ≤ y)] (f g : Cont) :
      Decidable (f ≤ g) := Id.run do
    let mut le : Bool := true
    for x in IndexType.univ Idx do
      if ¬(f[x] ≤ g[x]) then
        le := false
        break
    if le then isTrue sorry_proof else isFalse sorry_proof

  @[simp, ftrans_simp]
  theorem add_ofFn [Add Elem] (f g: Idx → Elem)
    : Indexed.ofFn (C:=Cont) f + Indexed.ofFn (C:=Cont) g
      =
      Indexed.ofFn fun i => f i + g i  := by apply ArrayType.ext (Idx:=Idx); simp[HAdd.hAdd, Add.add]

  @[simp, ftrans_simp]
  theorem sub_ofFn [Sub Elem] (f g: Idx → Elem)
    : Indexed.ofFn (C:=Cont) f - Indexed.ofFn (C:=Cont) g
      =
      Indexed.ofFn fun i => f i - g i  := by apply ArrayType.ext (Idx:=Idx); simp[HSub.hSub, Sub.sub]

  @[simp, ftrans_simp]
  theorem neg_ofFn [Neg Elem] (f: Idx → Elem)
    : - Indexed.ofFn (C:=Cont) f
      =
      Indexed.ofFn fun i => - f i  := by apply ArrayType.ext (Idx:=Idx); simp[Neg.neg]

  @[simp, ftrans_simp]
  theorem smul_ofFn [SMul K Elem] (f : Idx → Elem) (a : K)
    : a • Indexed.ofFn (C:=Cont) f
      =
      Indexed.ofFn fun i => a • f i := by apply ArrayType.ext (Idx:=Idx); simp[HSMul.hSMul, SMul.smul];

  @[simp, ftrans_simp]
  theorem sum_ofFn [AddCommMonoid Elem] {ι} [IndexType ι] (f : ι → Idx → Elem)
    : ∑ j, Indexed.ofFn (C:=Cont) (fun i => f j i)
      =
      Indexed.ofFn fun i => ∑ j, f j i
    := sorry_proof

  @[simp, ftrans_simp]
  theorem add_get [Add Elem] (x y : Cont) (i : Idx) :
      (x + y)[i] = x[i] + y[i] := by simp[HAdd.hAdd,Add.add]

  @[simp, ftrans_simp]
  theorem sub_get [Sub Elem] (x y : Cont) (i : Idx) :
      (x - y)[i] = x[i] - y[i] := by simp[HSub.hSub,Sub.sub]

  @[simp, ftrans_simp]
  theorem mul_get [Mul Elem] (x y : Cont) (i : Idx) :
      (x * y)[i] = x[i] * y[i] := by simp[HMul.hMul,Mul.mul]

  @[simp, ftrans_simp]
  theorem div_get [Div Elem] (x y : Cont) (i : Idx) :
      (x / y)[i] = x[i] / y[i] := by simp[HDiv.hDiv,Div.div]

@[simp, ftrans_simp]
  theorem smul_get {R} [SMul R Elem] (r : R) (x : Cont) (i : Idx) :
      (r • x)[i] = r • x[i]


import SciLean.Core.FunctionTransformations
-- import SciLean.Core.Meta.GenerateRevDeriv

open ComplexConjugate

namespace SciLean.Scalar

variable
  {R C} [Scalar R C] [RealScalar R]
  {W} [Vec C W]
  {U} [SemiInnerProductSpace C U]

--------------------------------------------------------------------------------
-- Exp -------------------------------------------------------------------------
--------------------------------------------------------------------------------

@[fun_prop]
theorem exp.arg_x.CDifferentiable_rule
  (x : W → C) (hx : CDifferentiable C x)
  : CDifferentiable C fun w => exp (x w) := sorry_proof

@[fun_trans]
theorem exp.arg_x.ceriv_rule
  (x : W → C) (hx : CDifferentiable C x)
  : cderiv C (fun w => exp (x w))
    =
    fun w dw =>
      let xdx := fwdDeriv C x w dw
      let e := exp xdx.1
      xdx.2 * e := sorry_proof

@[fun_trans]
theorem exp.arg_x.fwdDeriv_rule
    (x : W → C) (hx : CDifferentiable C x) :
    fwdDeriv C (fun w => exp (x w))
    =
    fun w dw =>
      let xdx := fwdDeriv C x w dw
      let e := exp xdx.1
      (e, xdx.2 * e) := by
  unfold fwdDeriv; fun_trans; rfl

@[fun_prop]
theorem exp.arg_x.HasAdjDiff_rule
    (x : U → C) (hx : HasAdjDiff C x) :
    HasAdjDiff C (fun u => exp (x u)) := by
  intro x
  constructor
  fun_prop
  fun_trans [fwdDeriv]; fun_prop

open ComplexConjugate
@[fun_trans]
theorem exp.arg_x.revDeriv_rule
    (x : U → C) (hx : HasAdjDiff C x) :
    revDeriv C (fun u => exp (x u))
    =
    fun u =>
      let xdx := revDeriv C x u
      (exp xdx.1, fun dy => xdx.2 (conj (exp xdx.1) * dy)) := by
  unfold revDeriv
  fun_trans only [fwdDeriv, smul_push, ftrans_simp]

--------------------------------------------------------------------------------
-- Log -------------------------------------------------------------------------
--------------------------------------------------------------------------------

section Log

variable
  {R : Type _} [RealScalar R]
  {W : Type _} [Vec R W]
  {U : Type _} [SemiInnerProductSpace R U]

@[fun_prop]
theorem log.arg_x.CDifferentiableAt_rule
    (w : W) (x : W → R) (hx : CDifferentiableAt R x w) (hw : x w ≠ 0) :
    CDifferentiableAt R (fun w => log (x w)) w := sorry_proof

@[fun_prop]
theorem log.arg_x.CDifferentiable_rule
    (x : W → R) (hx : CDifferentiable R x) (hw : ∀ w, x w ≠ 0) :
    CDifferentiable R (fun w => log (x w)) := by
  intro x; fun_prop (disch:=aesop)

@[fun_trans]
theorem log.arg_x.cderiv_rule_at
    (w : W) (x : W → R) (hx : CDifferentiableAt R x w) (hw : x w ≠ 0) :
    cderiv R (fun w => log (x w)) w
    =
    fun dw =>
      let xdx := fwdDeriv R x w dw
      xdx.2 / abs xdx.1 := sorry_proof

@[fun_trans]
theorem log.arg_x.cderiv_rule
    (x : W → R) (hx : CDifferentiable R x) (hw : ∀ w, x w ≠ 0) :
    cderiv R (fun w => log (x w))
    =
    fun w dw =>
      let xdx := fwdDeriv R x w dw
      xdx.2 / abs xdx.1 := by
  funext x
  fun_trans (disch:=aesop)

@[fun_trans]
theorem log.arg_x.fwdDeriv_rule_at
    (w : W) (x : W → R) (hx : CDifferentiableAt R x w) (hw : x w ≠ 0) :
    fwdDeriv R (fun w => log (x w)) w
    =
    fun dw =>
      let xdx := fwdDeriv R x w dw
      let l := log xdx.1
      (l, xdx.2 / abs xdx.1) :=
by
  unfold fwdDeriv; fun_trans (disch:=assumption); simp[fwdDeriv]

@[fun_trans]
theorem log.arg_x.fwdDeriv_rule
    (x : W → R) (hx : CDifferentiable R x) (hw : ∀ w, x w ≠ 0) :
    fwdDeriv R (fun w => log (x w))
    =
    fun w dw =>
      let xdx := fwdDeriv R x w dw
      let l := log xdx.1
      (l, xdx.2 / abs xdx.1) :=
by
  unfold fwdDeriv; fun_trans (disch:=assumption); simp[fwdDeriv]

@[fun_prop]
theorem log.arg_x.HasAdjDiffAt_rule
    (u : U) (x : U → R) (hx : HasAdjDiffAt R x u) (hu : x u ≠ 0) :
    HasAdjDiffAt R (fun u => log (x u)) u := by
  constructor
  fun_prop (disch:=aesop)
  fun_trans (disch:=aesop) [fwdDeriv]; fun_prop

@[fun_prop]
theorem log.arg_x.HasAdjDiff_rule
    (x : U → R) (hx : HasAdjDiff R x) (hu : ∀ u, x u ≠ 0) :
    HasAdjDiff R (fun u => log (x u)) := by
  intro u;
  fun_prop (disch:=aesop)

open ComplexConjugate

@[fun_trans]
theorem log.arg_x.revDeriv_rule_at
    (u : U) (x : U → R) (hx : HasAdjDiffAt R x u) (hu : x u ≠ 0) :
    revDeriv R (fun u => log (x u)) u
    =
    let xdx := revDeriv R x u
    (log xdx.1, fun dy => xdx.2 ((abs (x u))⁻¹ * dy))

:= by
  unfold revDeriv
  fun_trans (disch:=aesop) only [fwdDeriv, smul_push, ftrans_simp]


import SciLean.Core.FunctionPropositions.ContCDiff
import SciLean.Core.FunctionTransformations

import SciLean.Core.Notation.CDeriv

set_option linter.unusedVariables false

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)]

local notation "∞" => (⊤ : ℕ∞)

-- Function space --------------------------------------------------------------
--------------------------------------------------------------------------------

structure ContCDiffMapFD (n : ℕ∞) (X Y : Type _)  [Vec K X] [Vec K Y] where
  toFun : X → X → Y×Y
  is_cont_cdiff_map : ContCDiff K n (fun x => (toFun x 0).1)
  cderiv_snd : cderiv K (fun x => (toFun x 0).1) = fun x dx => (toFun x dx).2
  dir_independence : ∀ x dx, toFun x dx = toFun x 0

attribute [simp, ftrans_simp] ContCDiffMapFD.cderiv_snd ContCDiffMapFD.dir_independence

variable (n : ℕ∞)

instance : FunLike (ContCDiffMapFD K n X Y) X Y where
  coe f := fun x => (f.toFun x 0).1
  coe_injective' := sorry_proof

macro X:term:25 " ⟿FD[" K:term "," n:term "] " Y:term:26 : term =>
  `(ContCDiffMapFD $K $n $X $Y)

macro X:term:25 " ⟿FD[" n:term "] " Y:term:26 : term =>
  `(ContCDiffMapFD defaultScalar% $n $X $Y)

macro X:term:25 " ⟿FD " Y:term:26 : term =>
  `(ContCDiffMapFD defaultScalar% ∞ $X $Y)

@[app_unexpander ContCDiffMapFD] def unexpandContCDiffMapFD : Lean.PrettyPrinter.Unexpander
  | `($(_) $R $n $X $Y) => `($X ⟿FD[$R,$n] $Y)
  | _  => throw ()

@[fun_prop]
theorem ContCDiffMapFD_apply_right (f : X ⟿FD[K,n] Y) : ContCDiff K n (fun x => f x) := f.2

-- Lambda function notation ----------------------------------------------------
--------------------------------------------------------------------------------

def ContCDiffMapFD.mk' (f : X → Y) (f' : X → X → Y×Y) (h : fwdDeriv K f = f') (hf : ContCDiff K n f) : X ⟿FD[K,n] Y :=
  ⟨f', sorry_proof, sorry_proof, sorry_proof⟩

open Lean Parser Term
macro "fun " x:funBinder " ⟿FD[" K:term "," n:term "] " b:term : term =>
  `(ContCDiffMapFD.mk' $K $n (fun $x => $b) ((fwdDeriv $K fun $x => $b) rewrite_by autodiff /- check that derivative has been eliminated -/) (sorry_proof /- todo: add proof -/) sorry_proof)

macro "fun " x:funBinder " ⟿FD[" n:term "] " b:term : term => `(fun $x ⟿FD[defaultScalar%, $n] $b)
macro "fun " x:funBinder " ⟿FD "             b:term : term => `(fun $x ⟿FD[defaultScalar%,  ∞] $b)

variable {K n}

@[app_unexpander ContCDiffMapFD.mk'] def unexpandContCDiffMapFDMk : Lean.PrettyPrinter.Unexpander

  | `($(_) $R $n $f:term $_ $_ $_) =>
    match f with
    | `(fun $x':ident => $b:term) => `(fun $x' ⟿FD[$R,$n] $b)
    | `(fun ($x':ident : $ty) => $b:term) => `(fun ($x' : $ty) ⟿FD[$R,$n] $b)
    | `(fun $x':ident : $ty => $b:term) => `(fun ($x' : $ty) ⟿FD[$R,$n] $b)
    | _  => throw ()
  | _  => throw ()

@[simp, ftrans_simp]
theorem ContCDiffMapFD_eta (f : X ⟿FD[K,n] Y) : (fun x ⟿FD[K,n] f x) = f := by sorry_proof

@[pp_dot]
def ContCDiffMapFD.FD (f : X ⟿FD[K,n] Y) (x dx : X) : Y×Y := f.toFun x dx

@[fun_trans]
theorem ContCDiffMapFD_eval_fwdDeriv (f : X ⟿FD[K,n] Y) :
    fwdDeriv K (fun x => f x) = f.FD := by
  unfold ContCDiffMapFD.FD fwdDeriv
  simp[DFunLike.coe]

@[fun_prop]
theorem ContCDiffMapFD_eval_cdifferentiable (f : X ⟿FD[K,n] Y) (h : 0 < n) :
    CDifferentiable K (fun x => f x) := by
  simp[DFunLike.coe]
  apply CDifferentaible.ContCDiff_rule
  apply ContCDiffMapFD.is_cont_cdiff_map
  assumption

@[fun_prop]
theorem ContCDiffMapFD_eval_cdifferentiable' (f : X ⟿FD[K,∞] Y) :
    CDifferentiable K (fun x => f x)

:= by
  fun_prop (disch:=apply ENat.zero_lt_top)


import SciLean.Core.FunctionPropositions
import SciLean.Core.FunctionTransformations
import SciLean.Data.ArrayType.Algebra

open SciLean

set_option linter.unusedVariables false

open LeanColls

section GenericArrayType

variable
  {K : Type} [RCLike K]
  {Cont : Type} {Idx : Type |> outParam} {Elem : Type |> outParam}
  [ArrayType Cont Idx Elem] [IndexType Idx] [LawfulIndexType Idx] [DecidableEq Idx]

-- Indexed.get -----------------------------------------------------------------
--------------------------------------------------------------------------------

section OnVectorSpaces

variable {R : Type} [CommSemiring R] [AddCommGroup Elem] [Module R Elem]

@[fun_prop]
theorem LeanColls.Indexed.get.arg_cont.IsLinearMap_rule_simple (idx : Idx) :
    IsLinearMap R (fun xs : Cont => xs[idx]) := by
  constructor
  . intros; simp
  . intros; simp

#generate_linear_map_simps LeanColls.Indexed.get.arg_cont.IsLinearMap_rule_simple

end OnVectorSpaces

section OnNormedSpaces

variable
  {X : Type _} [NormedAddCommGroup X] [NormedSpace K X]
  [NormedAddCommGroup Elem] [NormedSpace K Elem]

@[fun_prop]
theorem LeanColls.Indexed.get.arg_cont.IsContinuousLinearMap_rule_simple (idx : Idx) :
    IsContinuousLinearMap K (λ cont : Cont => cont[idx]) := sorry_proof

-- automatically infer: Differentiable, DifferentiableAt, ContDiff, fderiv, fwdFDeriv

end OnNormedSpaces

section OnVec

variable
  {X : Type} [Vec K X]
  [Vec K Elem]

@[fun_prop]
theorem GetElem.getElem.arg_cont.IsSmoothLinearMap_rule_simple
  (idx : Idx)
  : IsSmoothLinearMap K (fun xs : Cont => xs[idx]) := sorry_proof

-- automatically infer: CDifferentiable, CDifferentiableAt, CContDiff, cderiv, fwdDeriv

end OnVec

section OnSemiInnerProductSpace

variable
  {X : Type} [SemiInnerProductSpace K X]
  [SemiInnerProductSpace K Elem]

@[fun_prop]
theorem GetElem.getElem.arg_cont.HasSemiAdjoint_rule_simple (idx : Idx) :
    HasSemiAdjoint K (fun cont : Cont => cont[idx]) := sorry_proof

@[fun_trans]
theorem GetElem.getElem.arg_cont.semiAdjoint_rule_simple (idx : Idx) :
    semiAdjoint K (fun cont : Cont => cont[idx])
    =
    fun elem => oneHot (X:=Cont) idx elem :=
by
  sorry_proof

@[fun_trans]
theorem GetElem.getElem.arg_cont.revCDeriv_rule
  (f : X → Cont) (idx : Idx)
  (hf : HasAdjDiff K f)
  : revDeriv K (fun x => (f x)[idx])
    =
    fun x =>
      let ydf := revDerivProj K Idx f x
      (ydf.1[idx],
       fun delem => ydf.2 idx delem) :=
by
  unfold revDeriv;
  conv =>
    lhs; fun_trans; fun_trans; simp

@[fun_trans]
theorem GetElem.getElem.arg_cont.revCDerivUpdate_rule
  (f : X → Cont) (idx : Idx)
  (hf : HasAdjDiff K f)
  : revDerivUpdate K (fun x => (f x)[idx])
    =
    fun x =>
      let ydf := revDerivProjUpdate K Idx f x
      (ydf.1[idx],
       fun delem dx => ydf.2 idx delem dx) :=
by
  unfold revDerivUpdate;
  conv =>
    lhs; fun_trans; fun_trans; simp

@[fun_trans]
theorem GetElem.getElem.arg_cont.revCDerivProj_rule
  {J ElemJ} [StructType Elem J ElemJ] [IndexType J] [LawfulIndexType J] [DecidableEq J]
  [∀ j, SemiInnerProductSpace K (ElemJ j)] [SemiInnerProductSpaceStruct K Elem J ElemJ]
  (f : X → Cont) (idx : Idx)
  (hf : HasAdjDiff K f)
  : revDerivProj K J (fun x => (f x)[idx])
    =
    fun x =>
      let ydf := revDerivProj K (Idx×J) f x
      (ydf.1[idx],
       fun j delem => ydf.2 (idx,j) delem) :=
by
  unfold revDerivProj
  conv =>
    lhs; fun_trans; fun_trans; simp
  funext x; simp[revDerivProj]
  funext j de
  simp [revDeriv]
  congr
  sorry_proof

end OnSemiInnerProductSpace

-- Indexed.set ----------------------------------------------------------------
--------------------------------------------------------------------------------

section OnVectorSpaces

variable {R : Type} [CommSemiring R] [AddCommGroup Elem] [Module R Elem]

@[fun_prop]
theorem LeanColls.Indexed.set.arg_contelem.IsLinearMap_rule_simple (idx : Idx) :
    IsLinearMap R (fun ((cont,elem) : Cont×Elem) => Indexed.set cont idx elem) := by
  constructor
  . intros; simp; ext; sorry_proof
  . intros; simp; ext; sorry_proof

#generate_linear_map_simps LeanColls.Indexed.set.arg_contelem.IsLinearMap_rule_simple

end OnVectorSpaces

section OnNormedSpaces

variable
  {X : Type _} [NormedAddCommGroup X] [NormedSpace K X]
  [NormedAddCommGroup Elem] [NormedSpace K Elem]

@[fun_prop]
theorem LeanColls.Indexed.set.arg_contelem.IsContinuousLinearMap_rule_simple :
    IsContinuousLinearMap K (fun ((cont,elem) : Cont×Elem) => Indexed.set cont idx elem) :=
  sorry_proof

-- automatically infer: Differentiable, DifferentiableAt, ContDiff, fderiv, fwdFDeriv

end OnNormedSpaces

section OnVec

variable
  {X : Type} [Vec K X]
  [Vec K Elem]

@[fun_prop]
theorem LeanColls.Indexed.set.arg_contelem.IsSmoothLinearMap_rule_simple (idx : Idx) :
    IsSmoothLinearMap K (fun ((cont,elem) : Cont×Elem) => Indexed.set cont idx elem) := sorry_proof

-- automatically infer: CDifferentiable, CDifferentiableAt, CContDiff, cderiv, fwdDeriv

end OnVec

section OnSemiInnerProductSpace

variable
  {X : Type} [SemiInnerProductSpace K X]
  [SemiInnerProductSpace K Elem]

@[fun_prop]
theorem LeanColls.Indexed.set.arg_cont.HasSemiAdjoint_rule_simple (idx : Idx) :
    HasSemiAdjoint K (fun ((cont,elem) : Cont×Elem) => Indexed.set cont idx elem) := sorry_proof

@[fun_trans]
theorem LeanColls.Indexed.set.arg_cont.semiAdjoint_rule_simple (idx : Idx) :
    semiAdjoint K (fun ((cont,elem) : Cont×Elem) => Indexed.set cont idx elem)
    =
    fun cont' : Cont => (Indexed.set cont' idx 0, cont'[idx]) := sorry_proof

@[fun_trans]
theorem LeanColls.Indexed.set.arg_contelem.revCDeriv_rule
  (cont : X → Cont) (idx : Idx) (elem : X → Elem)
  (hcont : HasAdjDiff K cont) (helem : HasAdjDiff K elem)
  : revDeriv K (fun x => Indexed.set (cont x) idx (elem x))
    =
    fun x =>
      let cdc := revDeriv K cont x
      let ede := revDerivUpdate K elem x
      (Indexed.set cdc.1 idx ede.1,
       fun dc : Cont =>
         let dci := dc[idx]
         let dc := Indexed.set dc idx 0
         ede.2 dci (cdc.2 dc)) := by
  unfold revDeriv;
  conv =>
    lhs; fun_trans; fun_trans; simp

@[fun_trans]
theorem LeanColls.Indexed.set.arg_contelem.revCDerivUpdate_rule
  (cont : X → Cont) (idx : Idx) (elem : X → Elem)
  (hcont : HasAdjDiff K cont) (helem : HasAdjDiff K elem)
  : revDerivUpdate K (fun x => Indexed.set (cont x) idx (elem x))
    =
    fun x =>
      let cdc := revDerivUpdate K cont x
      let ede := revDerivUpdate K elem x
      (Indexed.set cdc.1 idx ede.1,
       fun (dc : Cont) dx =>
         let dci := dc[idx]
         let dc := Indexed.set dc idx 0
         ede.2 dci (cdc.2 dc dx)) := by
  unfold revDerivUpdate
  conv =>
    lhs; fun_trans; fun_trans; simp
  simp[revDerivUpdate,add_assoc]

-- Consider formulating variant for index type `Idx×J`
@[fun_trans]
theorem LeanColls.Indexed.set.arg_contelem.revCDerivProj_rule
    (cont : X → Cont) (idx : Idx) (elem : X → Elem)
    (hcont : HasAdjDiff K cont) (helem : HasAdjDiff K elem) :
    revDerivProj K Idx (fun x => Indexed.set (cont x) idx (elem x))
    =
    fun x =>
      let cdc := revDerivProj K Idx cont x
      let ede := revDeriv K elem x
      (Indexed.set cdc.1 idx ede.1,
       fun i delem =>
         if i = idx then
           ede.2 delem
         else
           cdc.2 i delem) :=
by
  unfold revDerivProj
  conv =>
    lhs; fun_trans; fun_trans; simp
  funext x; simp[revDerivProj,revDerivUpdate]
  funext i de
  if h : i = idx then
    simp [revDeriv,h,oneHot]
    sorry_proof
  else
    simp [revDeriv,h,oneHot]
    sorry_proof

-- Consider formulating variant for index type `Idx×J`
@[fun_trans]
theorem LeanColls.Indexed.set.arg_contelem.revCDerivProjUpdate_rule
    (cont : X → Cont) (idx : Idx) (elem : X → Elem)
    (hcont : HasAdjDiff K cont) (helem : HasAdjDiff K elem) :
    revDerivProjUpdate K Idx (fun x => Indexed.set (cont x) idx (elem x))
    =
    fun x =>
      let cdc := revDerivProjUpdate K Idx cont x
      let ede := revDerivUpdate K elem x
      (Indexed.set cdc.1 idx ede.1,
       fun i delem dx =>
         if i = idx then
           ede.2 delem dx
         else
           cdc.2 i delem dx) :=
by
  unfold revDerivProjUpdate
  conv =>
    lhs; fun_trans; fun_trans; simp
  funext x; simp[revDerivProj,revDerivUpdate]
  funext i de
  if h : i = idx then
    simp [h]
  else
    simp [h]

end OnSemiInnerProductSpace

-- LeanColls.Indexed.ofFn -------------------------------------------------------------
--------------------------------------------------------------------------------

section OnVectorSpaces

variable {R : Type} [CommSemiring R] [AddCommGroup Elem] [Module R Elem]

@[fun_prop]
theorem LeanColls.Indexed.ofFn.arg_f.IsLinearMap_rule_simple :
    IsLinearMap R (fun f : Idx → Elem => Indexed.ofFn (C:=Cont) f)

:= by
  constructor
  . intros; ext; simp
  . intros; ext; simp


import SciLean.Core.FunctionTransformations
-- import SciLean.Core.Meta.GenerateRevDeriv

open ComplexConjugate

namespace SciLean.Scalar

variable
  {R C} [Scalar R C]
  {W} [Vec C W]
  {U} [SemiInnerProductSpace C U]

--------------------------------------------------------------------------------
-- Sin -------------------------------------------------------------------------
--------------------------------------------------------------------------------

@[fun_prop]
theorem sin.arg_x.CDifferentiable_rule
    (x : W → C) (hx : CDifferentiable C x) :
    CDifferentiable C fun w => sin (x w) := sorry_proof

@[fun_trans]
theorem sin.arg_x.ceriv_rule
    (x : W → C) (hx : CDifferentiable C x) :
    cderiv C (fun w => sin (x w))
    =
    fun w dw =>
      let xdx := fwdDeriv C x w dw
      let c := cos xdx.1
      xdx.2 * c := sorry_proof

@[fun_trans]
theorem sin.arg_x.fwdDeriv_rule
    (x : W → C) (hx : CDifferentiable C x) :
    fwdDeriv C (fun w => sin (x w))
    =
    fun w dw =>
      let xdx := fwdDeriv C x w dw
      let s := sin xdx.1
      let c := cos xdx.1
      (s, xdx.2 * c) := by
  unfold fwdDeriv; fun_trans; rfl

@[fun_prop]
theorem sin.arg_x.HasAdjDiff_rule
    (x : U → C) (hx : HasAdjDiff C x) :
    HasAdjDiff C (fun u => sin (x u))