using DifferentialEquations
using Plots

function SIR(du,u,p,t)
    s,i,r = u
    b,k = p
    du[1] = -b*s*i
    du[2] = b*s*i - k*i
    du[3] = k*i
end

#initian conditions
p = [0.5/7900000.0,0.33]
u0 = [7900000.0,10.0,0.0]
# u0 = [1.,1.27e-6,0.]
tspan = (0.0,140.0)

#solve
sir = ODEProblem(SIR,u0,tspan,p)
sol_sir = solve(sir)

#plot
plot(sol_sir)


using Distributions

"""
Example of a computation of beta credible interval
"""
function betaCredibleInt()
    S = 47
    N = 100
    a = S+1
    b = (N-S)+1
    α = .05
    β = Beta(a, b)
    CI = quantile(β, [α/2, 1-α/2])
end


include("differentiation.jl")
using Plots

h = 2.0 .^ (-2:-1:-52)
ffderr = []
bfderr = []
cfderr = []
f(x) = exp(x)
x = 1.0
dfexact = exp(1.0)
for h0 in h
    append!(ffderr, abs(ffd(f, x, h0 )-dfexact))
    append!(bfderr, abs(bfd(f, x, h0 )-dfexact))
    append!(cfderr, abs(cfd(f, x, h0 )-dfexact))
end

plot(h, ffderr, xaxis=:log, yaxis=:log)
plot!(h, bfderr, xaxis=:log, yaxis=:log)
plot!(h, cfderr, xaxis=:log, yaxis=:log)


# # Chapter 3 - 数学のおさらい

## #src
# ## 3.1 線形代数
# ### 3.1.1 スカラーとベクトル
## 列ベクトル: 2行1列のArrayを定義
## Juliaはcolumn-major
w = [5.0, 3.0]

## #src
## 行ベクトル
#-
## 行ベクトル: 1行2列
x = [5.0 3.0]

## #src
#-
## すべての要素が0の1行5列のArrayを生成
x_zeros = zeros(1, 5)

## #src
#-
## すべての要素が1の1行5列のArrayを生成
x_ones = ones(1, 5)

## #src
#-
## 一様分布または正規分布にしたがってランダムに1行5列のArrayを生成
using Random
Random.seed!(1234)

## 一様分布
uniform = rand(1, 5)

## #src
#-
## 正規分布
normal = randn(1, 5)

## #src
#-
## 標準出力
for s in (:w, :x, :x_zeros, :x_ones, :uniform, :normal)
    y = eval(s)
    println("ベクトル $(s)) 形: $(size(y)), 要素の型:$(eltype(y))\n$(y)\n")
end

## #src
### 3.1.2 ベクトルの内積
## 列ベクトル: 2行1列のArrayを定義
w = [5.0, 3.0]

## 行ベクトル: 1行2列
x = [1.0 5.0]

## 内積の計算
## \cdot
using LinearAlgebra
xw = x ⋅ w

## #src
#-
## wのノルムの計算: wの転置(行ベクトル)とw(列ベクトル)の掛け算
ww = transpose(w) ⋅ w

## #src
#-
## ノルムを計算する関数が用意されている
norm(w)

## #src
#-
println("xとwの内積) 形: $(size(xw)), 型:$(typeof(xw))\n$(xw)\n")
println("wのノルム) 形: $(size(ww)), 型:$(typeof(ww))\n$(ww)\n")

## #src
# ### 3.1.5 行列とベクトルの関係
## 行列: 3行2列のArrayを定義
X = [1.0 2.0; 2.0 4.0; 3.0 6.0]

## #src
#-
## 1行目の行ベクトル(1行2列)の取り出し
Xrow1 = X[[1], :]

## #src
#-
## 2列目の列ベクトル(3行1列)の取り出し
Xcol2 = X[:, [2]]

## #src
#-
## 標準出力
println("行列X) 形:$(size(X)), 要素の型:$(eltype(X))\n$(X)\n")
println("1行目) 形:$(size(Xrow1)), 要素の型:$(eltype(Xrow1))\n$(Xrow1)\n")
println("2列目) 形:$(size(Xcol2)), 要素の型:$(eltype(Xcol2))\n$(Xcol2)\n")

## #src
#-
## リストなしの場合
X[1, :]

## #src
#-
## リストありの場合
X[[1], :]

## #src
# ### 3.1.6 行列の積と和
## 行列 X: 3行2列のArrayを定義
X = [1.0 2.0; 2.0 4.0; 3.0 6.0]

## 列ベクトル b: 2行1列のArrayを定義
w = [5.0, 3.0]

## 列ベクトル b: 3行1列のArrayを定義
b = [1.0, 1.0, 1.0]

## 行列とベクトルの積と和
res = X * w .+ b
println("積和の結果 \n$(res)")

## #src
# ### 3.1.8 行列の階数
## 行列: 2行2列のArrayを定義
A = [6.0 2.0; 2.0 5.0]
B = [6.0 3.0; 2.0 1.0]

## 行列のランクの計算
println("Rank of A: $(rank(A))")
println("Rank of B: $(rank(B))")

## #src
#-
function find_inv(mat)
    if rank(mat) == size(mat)[1]
        println("行列:\n$(mat)\nランク$(rank(mat))\n逆行列:\n$(inv(mat))\n")
    else
        println("行列:\n$(mat)\nランク$(rank(mat))\n特異行列")
    end
end

find_inv(A)
find_inv(B)

## #src
# ### 3.1.10 行列と行列式の関係
## 行列: 2行2列のArrayを定義
A = [6.0 2.0; 2.0 5.0]
B = [6.0 3.0; 2.0 1.0]

## 行列式の計算
println("行列Aの行列式: $(det(A))")
println("行列Bの行列式: $(det(B))")

## #src
# ### 3.1.12 固有値問題の解法
## 行列Aの定義
A = [3 2; 4 1]

## 固有値問題の解
eig_A = eigen(A)

## #src
# ### 3.1.13 固有値と固有ベクトルの性質
## 行列Aの定義
A = [3 1; 1 3]

## 固有値・固有ベクトル
eig_A = eigen(A)

println("行列Aの固有値:\n$(eig_A.values)")
println("行列Aの固有ベクトル:\n$(eig_A.vectors)")
println("固有ベクトルの内積:$(eig_A.vectors[:, 1] ⋅ eig_A.vectors[:, 2])")
println("固有値の和:$(sum(eig_A.values))")
println("行列Aのトレース(対角成分の和):$(tr(A))")
println("固有値の積:$(prod(eig_A.values))")
println("行列Aの行列式:$(det(A))")
println("固有ベクトルがなす行列の逆行列:\n$(eig_A.vectors)")
println("固有ベクトルがなす行列の転置行列:\n$(eig_A.vectors')")

## #src
# ## 3.2 最適化
## 行列Aの定義
A = [2 -3; 4 1]

## 列ベクトルbの定義
b = [5, -2]

## #src
#-
## 行列が非ゼロか否かを確認
if det(A) != 0
    ## Aの逆行列とベクトルbの積から解を求める
    w = inv(A) * b
    println("wの解\n$(w)")
else
    ## 解が存在しない
    println("解が存在しません.")
end

## #src
#-
# ### 3.2.2 微分を用いた解放
## 最急降下法による最適化
## 学習率の設定
α = 0.2

## パラメーターの初期化
w = 1

for ite in 1:50
    println("反復$(ite), w=$(round(w; digits = 2))")
    ## 最急降下法によるパラメーターの更新
    global w-= α * 3 * w^2
end

## #src
#-
### 3.3.2 ベルヌーイ分布
## p=6/13(青色のボールの確率)のベルヌーイ分布に従って10回試行を行う
using Distributions

b = Bernoulli(6/13)

X1 = rand(b, 10)
X2 = rand(b, 10)
X3 = rand(b, 10)

## 標準出力
println("データ X1:\n$(X1)")
println("データ X2:\n$(X2)")
println("データ X3:\n$(X3)")

## #src
### 3.3.9 正規分布
## 平均0, 分散1の正規分布に従って10回試行を行う
X1 = rand(Normal(0, 1), 10)
## 平均3, 分散1の正規分布に従って10回試行を行う
X2 = rand(Normal(3, 1), 10)
## 平均-3, 分散1の正規分布に従って10回試行を行う
X3 = rand(Normal(-3, 1), 10)

## 標準出力
println("データ X1:\n$(X1)")
println("データ X2:\n$(X2)")
println("データ X3:\n$(X3)")

## #src
# ## 3.4 統計
# ### 3.4.1 中心の統計量
## 3種類(身長、体重、胸囲)のデータを格納する8行3列の行列
X = [
    170 60 80;
    167 52 93;
    174 57 85;
    181 70 80;
    171 62 70;
    171 66 95;
    171 66 95;
    168 54 85;
]

## 行方向(dims=1)に対して、平均値を計算
means = mean(X; dims=1)

## 行方向(dims=1)に対して、中央値を計算
medians = median(X; dims=1)

## 標準出力
println("データ X:\n$X")
println("身長の平均値: $(means[1]), 体重の平均値:$(means[2]), 胸囲の平均値$(means[3])")
println("身長の中央値: $(medians[1]), 体重の中央値:$(medians[2]), 胸囲の中央値$(medians[3])")

## src
#-
## 分散共分散行列の計算
## 3種類(身長、体重、胸囲)のデータを格納する8行3列の行列
X = [
    170 60 80;
    167 52 93;
    174 57 85;
    181 70 80;
    171 62 70;
    171 66 95;
    171 66 95;
    168 54 85;
]
vars = var(X; corrected=false, dims=1)
stds = std(X; corrected=false, dims=1)
cov_bias = cov(X)
cov_nobias = cov(X; corrected=false)
println("データ X:\n$X")
println("分散) 身長:$(vars[1]), 体重:$(vars[2]), 胸囲:$(vars[3])")
println("標準偏差) 身長:$(stds[1]), 体重:$(stds[2]), 胸囲:$(stds[3])")
println("分散共分散行列 バイアスあり:\n$(cov_bias)")
println("分散共分散行列 バイアスなし:\n$(cov_nobias)")

## #src
# ### 3.4.7 相関係数と相関行列
## 3種類(身長、体重、胸囲)のデータを格納する8行3列の行列
X = [
    170 60 80;
    167 52 93;
    174 57 85;
    181 70 80;
    171 62 70;
    171 66 95;
    171 66 95;
    168 54 85;
]
println("データ X:\n$X")
println("相関行列:\n$(cor(X))")


# Spherical and cartesian coordinate conversions

"""
    azimuth(lon1, lat1, lon2, lat2, degrees::Bool=true) -> az

Compute the azimuth `az` from point (`lon1`, `lat1`) to (`lon2`, `lat2`) on the sphere.
Points and azimuth are read and returned in degrees by default; use `degrees=false`
for radians.
"""
function azimuth(lon1, lat1, lon2, lat2, degrees::Bool=true)
    if degrees
        lon1, lat1, lon2, lat2 = deg2rad(lon1), deg2rad(lat1), deg2rad(lon2), deg2rad(lat2)
    end
    azimuth = atan(sin(lon2-lon1)*cos(lat2),
                   cos(lat1)*sin(lat2) - sin(lat1)*cos(lat2)*cos(lon2-lon1))
    degrees ? rad2deg(azimuth) : azimuth
end

"""
    cart2geog(x, y, z, degrees::Bool=true) -> lon, lat, r

Compute the longitude, latitude and radius given the cartesian coordinates `x`,
`y` and `z`, where `x` is at (lon,lat) = (0,0), `y` is at (90°,0) and `z` is
through lat = 90°.
"""
function cart2geog(x, y, z, degrees::Bool=true)
    r = sqrt(x^2 + y^2 + z^2)
    r == 0. && return zero(x), zero(x), zero(x)
    lon = atan(y, x)
    lat = asin(z/r)
    degrees ? (rad2deg(lon), rad2deg(lat), r) : (lon, lat, r)
end

function cart2geog(x::AbstractArray, y::AbstractArray, z::AbstractArray, degrees::Bool=true)
    dims = size(x)
    dims == size(y) == size(z) || throw(ArgumentError("All arrays must have same length"))
    T = promote_type(float.(eltype.((x, y, z)))...)
    lon, lat, r = Array{T}(dims), Array{T}(dims), Array{T}(dims)
    for i in eachindex(lon)
        lon[i], lat[i], r[i] = cart2geog(x[i], y[i], z[i], degrees)
    end
    lon, lat, r
end

"""
    delta(lon1, lat1, lon2, lat2, degrees::Bool=true) -> d

Compute the angular distance `d` on the sphere between two points, (`lon1`, `lat1`)
and (`lon2`, `lat2`).  Points and distance are read and returned in degrees
by default; use `degrees=false` for radians.
"""
function delta(lon1, lat1, lon2, lat2, degrees::Bool=true)
    if degrees
        lon1, lat1, lon2, lat2 = deg2rad(lon1), deg2rad(lat1), deg2rad(lon2), deg2rad(lat2)
    end
    d = atan(sqrt(
               (cos(lat2)*sin(lon2-lon1))^2 + (cos(lat1)*sin(lat2) -
                sin(lat1)*cos(lat2)*cos(lon2-lon1))^2),
               sin(lat1)*sin(lat2) + cos(lat1)*cos(lat2)*cos(lon2-lon1)
              )
    degrees ? rad2deg(d) : d
end

"""
    geog2cart(lon, lat, r, degrees::Bool=true) -> x, y, z
    geog2cart(lon, lat, degrees::Bool=true) -> x, y, z

Return the cartesian coordinates given the geographic longitude, latitude and
radius `lon`, `lat` and `r`.

If `r` is not given, points are returned on the unit sphere.
"""
function geog2cart(lon, lat, r, degrees::Bool=true)
    points_valid(lon, lat, degrees) || error("geog2cart: Points are not on the sphere")
    degrees && begin lon, lat = deg2rad(lon), deg2rad(lat) end
    x = r*cos(lon)*cos(lat)
    y = r*sin(lon)*cos(lat)
    z = r*sin(lat)
    x, y, z
end
function geog2cart(lon::AbstractArray, lat::AbstractArray, r::AbstractArray, degrees::Bool=true)
    size(lon) == size(lat) == size(r) ||
        throw(ArgumentError("Sizes of lon, lat and r must be the same"))
    T = promote_type(float.(eltype.((lon, lat, r)))...)
    x = Array{T}(undef, size(lon))
    y, z = similar(x), similar(x)
    for i in eachindex(lon)
        x[i], y[i], z[i] = geog2cart(lon[i], lat[i], r[i], degrees)
    end
    x, y, z
end
geog2cart(lon::AbstractArray, lat::AbstractArray, degrees::Bool=true) =
    geog2cart(lon, lat, fill(one(eltype(lon)), size(lon)), degrees)

"""
    points_valid(lon, lat, degrees::Bool=true) -> ::Bool

Return `true` if all points in arrays `lon` and `lat` are on the sphere.
`lon` is ignored, but `lat` is checked to see if points are in the range
-90°–90° (-π–π).  Points are read and returned in degrees by default; use
`degrees=false` for radians.
"""
points_valid(lon, lat, degrees::Bool=true) =
    degrees ? !any(abs.(lat) .> 90.) : !any(abs.(lat) .> pi/2.)

"""
    geodesic_endpoint(lon, lat, az, delta, degrees::Bool=true) -> lon1, lat1

Compute the end point (`lon1`, `lat1`) reached by travelling on the sphere along
azimuth `az` for `delta` angular distance.  Points, angles and distance are read
and returned in degrees by default; use `degrees=false` for radians.
"""
function geodesic_endpoint(lon, lat, az, delta, degrees::Bool=true)
    if degrees
        lon, lat, az, delta = deg2rad(lon), deg2rad(lat), deg2rad(az), deg2rad(delta)
    end
    lat2 = asin(sin(lat)*cos(delta) + cos(lat)*sin(delta)*cos(az))
    lon2 = lon + atan(sin(az)*sin(delta)*cos(lat),
                      cos(delta)-sin(lat)*sin(lat2))
    if degrees
        lon2, lat2 = rad2deg(lon2), rad2deg(lat2)
    end
    lon2, lat2
end


# https://www.juliabox.com/notebook/notebooks/tutorials/intro-to-julia/AutoDiff.ipynb
# Dual Number

using InteractiveUtils
versioninfo()

struct D <: Number
    f::Tuple{Float64,Float64}
end

Base.show(io::IO, x::D) = print(io,x.f[1]," + ", x.f[2],"ϵ")

import Base: +, -, *, /, convert, promote_rule
+(x::D, y::D) = D(x.f .+ y.f)
-(x::D, y::D) = D(x.f .- y.f)
*(x::D, y::D) = D((x.f[1] .* y.f[1], x.f[2] * y.f[1] + x.f[1] * y.f[2]))
/(x::D, y::D) = D((x.f[1] / y.f[1], (y.f[1] * x.f[2] - x.f[1] * y.f[2]) / y.f[1]^2))
convert(::Type{D}, x::Real) = D((x,zero(x)))
promote_rule(::Type{D}, ::Type{<:Number}) = D

ϵ = D((0,1))

@show ϵ * ϵ, ϵ^2
@show 1/(1+ϵ)
@show (1+ϵ)^10

function nthroot(x, n=2; t=1, N = 10)
    for i in 1:N
        t += (x/t^(n-1)-t)/n
    end
    t
end

@show nthroot(2,3), ∛2
@show nthroot(7,12), 7^(1/12)
x=17.0; @show nthroot(x+ϵ, 3), ∛x, 1/x^(2/3)/3

#= output
Julia Version 0.7.0
Commit a4cb80f3ed (2018-08-08 06:46 UTC)
Platform Info:
  OS: macOS (x86_64-apple-darwin14.5.0)
  CPU: Intel(R) Core(TM) M-5Y51 CPU @ 1.10GHz
  WORD_SIZE: 64
  LIBM: libopenlibm
  LLVM: libLLVM-6.0.0 (ORCJIT, broadwell)
(ϵ * ϵ, ϵ ^ 2) = (0.0 + 0.0ϵ, 0.0 + 0.0ϵ)
1 / (1 + ϵ) = 1.0 + -1.0ϵ
(1 + ϵ) ^ 10 = 1.0 + 10.0ϵ
(nthroot(2, 3), ∛2) = (1.2599210498948732, 1.2599210498948732)
(nthroot(7, 12), 7 ^ (1 / 12)) = (1.1760474285795146, 1.1760474285795146)
(nthroot(x + ϵ, 3), ∛x, (1 / x ^ (2 / 3)) / 3) = (2.571281590658235 + 0.05041728609133794ϵ, 2.571281590658235, 0.05041728609133795)
=#


abstract type lossFun end

export lossFun

### binaryCrossentropy

"""
    return the average cross entropy loss over vector of labels and predictions

    input:
        a := (?1, c,m) matrix of predicted values, where c is the number of classes
        y := (?1, c,m) matrix of predicted values, where c is the number of classes

        Note: in case the number of classes is one (1) it is okay to have
              a scaler values for a and y

    output:
        J := scaler value of the cross entropy loss
"""
abstract type binaryCrossentropy <: lossFun end

function binaryCrossentropy(a, y)

    aNew = prevnextfloat.(a)
    J = .-(y .* log.(aNew) .+ (1 .- y) .* log.(1 .- aNew))
    return J
end #binaryCrossentropy

export binaryCrossentropy

"""
    compute the drivative of cross-entropy loss function to the input of the
    layer dZ
"""
function dbinaryCrossentropy(a, y)
    dJ = a .- y
    return dJ
end #dbinaryCrossentropy

export dbinaryCrossentropy

### categoricalCrossentropy

abstract type categoricalCrossentropy <: lossFun end

function categoricalCrossentropy(a, y)
    aNew = prevnextfloat.(a)
    J = .-(y .* log.(aNew))
    return J
end

function dcategoricalCrossentropy(a, y)
    dJ = a .- y
    return dJ
end

export categoricalCrossentropy, dcategoricalCrossentropy

"""
    return previous float if x == 1 and nextfloat if x == 0
"""
prevnextfloat(x) = x==0 ? nextfloat(x) : x==1 ? prevfloat(x) : x

export prevnextfloat

### cost function

"""
    function cost(
        loss::Type{categoricalCrossentropy},
        A::AbstractArray{T1,N},
        Y::AbstractArray{T2,N},
    ) where {T1, T2, N}

Compute the cost for `categoricalCrossentropy` loss function

"""
function cost(
    loss::Type{categoricalCrossentropy},
    A::AbstractArray{T1,N},
    Y::AbstractArray{T2,N},
) where {T1, T2, N}

    c, m = size(A)[N-1:N]
    costs = sum(loss(Float64.(A),Y)) / m
    return Float64.(costs)

end #function cost

"""
    function cost(
        loss::Type{binaryCrossentropy},
        A::AbstractArray{T1,N},
        Y::AbstractArray{T2,N},
    ) where {T1, T2, N}

Compute the cost for `binaryCrossentropy` loss function

"""
function cost(
    loss::Type{binaryCrossentropy},
    A::AbstractArray{T1,N},
    Y::AbstractArray{T2,N},
) where {T1, T2, N}

    c, m = size(A)[N-1:N]
    costs = sum(loss(Float64.(A),Y)) / (c*m)
    return Float64.(costs)

end #function cost


## Exercise 4-5
## Make a copy of square and change the name to polygon. Add another parameter named n and modify the body so it draws an n-sided regular polygon.

## TIP
## The exterior angles of an n-sided regular polygon are 360/n degrees.
using ThinkJulia

println("Ans: ")

function polygon(turtle::Turtle, nsides::Int, len::Int = 100)
    angle = Int(360 / nsides)
    println("Angle: " * string(angle))

    for i in 1:nsides
        forward(turtle, len)
        turn(turtle, -angle)
    end
end

# (a) Try 1
# @svg begin
#     polygon(Turtle(), 5)
# end

# (b) Try 2
@svg begin
    polygon(Turtle(), 9)
end

println("End.")


using ForwardDiff
using Plots
using ReverseDiff
using DifferentialEquations
using Flux
using Statistics
# gr()
# pyplot()
plotly()

# Define SHO Hamiltonian Function
H(q, p) = 0.5*q^2 + 0.5*p^2

function dH(q, p)
    dHdq, dHdp = ForwardDiff.gradient(x->H(x[1], x[2]), [q, p])
    return dHdq, dHdp
end

function dH_analytic(q, p)
    return q, p
end

function H_vec_field(q, p)
    dHdq, dHdp = dH(q, p)
    q̇ = dHdp
    ṗ = - dHdq

    return q̇, ṗ
end

# Visualize the H, it's level curves, and the flow as
# as tangent vectors to level curves
# i.e. governing idea is H = const
qs = -1.5:0.1:1.5
ps = -1.5:0.1:1.5

Q = []
P = []
Q̇ = []
Ṗ = []

for q ∈ qs[1:2:end], p ∈ ps[1:2:end]
    q̇, ṗ = H_vec_field(q, p)
    push!(Q, q)
    push!(P, p)
    push!(Q̇, q̇)
    push!(Ṗ, ṗ)
end

p1 = contour(qs, ps, H, xlabel="q", ylabel="p", aspect_ratio=1.0, fill=true, colorbar_title="H", size=(800,800), c=:viridis)
α=0.1
quiver!(p1, Q, P, quiver=(α .* Q̇, α .* Ṗ), aspect_ratio=1.0, color=:white)
xlims!(p1, -1.5, 1.5)
ylims!(p1, -1.5, 1.5)

Ham(q,p, param) = q^2 + p^2
p0=1.0
q0=0.0
tspan=(0.0, 10.0)
prob = HamiltonianProblem(Ham, p0, q0, tspan, dt=0.01)
sol = solve(prob, SymplecticEuler())

plot!(p1, sol[1,:], sol[2, :], color=:green, label="q₀=$(q0), p₀=$(p0)")

p0=1.5
q0=0.0
tspan=(0.0, 10.0)
prob = HamiltonianProblem(Ham, p0, q0, tspan, dt=0.01)
sol = solve(prob, SymplecticEuler())

plot!(p1, sol[1,:], sol[2, :], color=:cyan, label="q₀=$(q0), p₀=$(p0)")

# Let's set up the Neural Network to approximate the Hamiltonian
# generate training data with H(q,p) <= 1
Npoints = 500
r_train = sqrt.(rand(1, Npoints))
φ_train = 2π.*rand(1, Npoints)
data = vcat(r_train .* cos.(φ_train), r_train .* sin.(φ_train))

# set up target data: here we use analytic result for time derivatives
target = zero(data)
# i.e. q̇= dHdp, ṗ = -dHdq
target[1,:] .= data[2, :]
target[2,:] .= -data[1, :]

p2 =  scatter!(p1, data[1,:], data[2,:], c=:green, label="training points")
p3 = plot(Q, P, H.(Q, P), st = :surface, xlabel = "q", ylabel = "p", zlabel = "H")

# -------------------------------------------------------------------------------------------------------
# create model for HNN

# define the struct
struct HNN{M, R, P}
    model::M  # the internal NN
    re::R   # for recreating internal NN
    p::P  # for holding current params of internal NN

    # define the constructor
    function HNN(model)
        p, re = Flux.destructure(model)
        return new{typeof(model), typeof(re), typeof(p)}(model, re, p)
    end
end

# define the trainable paramaters
Flux.trainable(hnn::HNN) = (hnn.p,)

function _hamiltonian_flow(re, p, x)
    dHdX = Flux.gradient(x->sum(re(p)(x)), x)[1]  # note: re(p)(x) == Model(x)
    n = size(x,1) ÷ 2  # i.e. how many p's and q's
    return cat(dHdX[(n+1):2n, :], dHdX[1:n, :], dims=1)
end

# define how to call the HNN on data
(hnn::HNN)(X, p=hnn.p) = _hamiltonian_flow(hnn.re, p, X)

hnn = HNN(
    Chain(Dense(2, 200, relu), Dense(200, 200, relu), Dense(200, 1))
)

dataloader = Flux.Data.DataLoader((data, target), batchsize=500, shuffle=true)

p = hnn.p

opt = ADAM(10^(-3))

loss(x, y, p) = mean((hnn(x, p) .- y) .^ 2)
loss(data, target, p)

callback() = println("Loss Neural Hamiltonian DE = $(loss(data, target, p))")
callback()

test_gs = ReverseDiff.gradient(p -> loss(data, target, p), p)

epochs = 2000
for epoch in 1:epochs
    for (x, y) in dataloader
        gs = ReverseDiff.gradient(p -> loss(x, y, p), p)
        Flux.Optimise.update!(opt, p, gs)
    end
    if epoch % 100 == 1
        callback()
    end
end
callback()

#p3 = plot(Q, P, H.(Q, P), st = :surface, xlabel = "q", ylabel = "p", zlabel = "H")
#p3 = plot(Q, P, HNN.model.(Q, P), st = :surface, xlabel = "q", ylabel = "p", zlabel = "H")


# Using Base modules.
using Random

# Load a plotting library.
using Plots

# Load the distributions library.
using Distributions

# Set the true probability of heads in a coin.
p_true = 0.5

# Iterate from having seen 0 observations to 100 observations.
Ns = 0:100;

# Draw data from a Bernoulli distribution, i.e. draw heads or tails.
Random.seed!(12)
data = rand(Bernoulli(p_true), last(Ns))

# Here's what the first five coin flips look like:
data[1:5]

# Our prior belief about the probability of heads in a coin toss.
prior_belief = Beta(1, 1);

# Import StatsPlots for animating purposes.
using StatsPlots

# Make an animation.
animation = @gif for (i, N) in enumerate(Ns)

    # Count the number of heads and tails.
    heads = sum(data[1:i-1])
    tails = N - heads

    # Update our prior belief in closed form (this is possible because we use a conjugate prior).
    updated_belief = Beta(prior_belief.α + heads, prior_belief.β + tails)

    # Plotting
    plot(updated_belief,
        size = (500, 250),
        title = "Updated belief after $N observations",
        xlabel = "probability of heads",
        ylabel = "",
        legend = nothing,
        xlim = (0,1),
        fill=0, α=0.3, w=3)
    vline!([p_true])
end

# Load Turing and MCMCChains.
using Turing, MCMCChains

# Load the distributions library.
using Distributions

# Load StatsPlots for density plots.
using StatsPlots

@model coinflip(y) = begin

    # Our prior belief about the probability of heads in a coin.
    p ~ Beta(1, 1)

    # The number of observations.
    N = length(y)
    for n in 1:N
        # Heads or tails of a coin are drawn from a Bernoulli distribution.
        y[n] ~ Bernoulli(p)
    end
end;

# Settings of the Hamiltonian Monte Carlo (HMC) sampler.
iterations = 1000
ϵ = 0.05
τ = 10

# Start sampling.
chain = sample(coinflip(data), HMC(ϵ, τ), iterations, progress=false);

# Construct summary of the sampling process for the parameter p, i.e. the probability of heads in a coin.
p_summary = chain[:p]
plot(p_summary, seriestype = :histogram)

# Compute the posterior distribution in closed-form.
N = length(data)
heads = sum(data)
updated_belief = Beta(prior_belief.α + heads, prior_belief.β + N - heads)

# Visualize a blue density plot of the approximate posterior distribution using HMC (see Chain 1 in the legend).
p = plot(p_summary, seriestype = :density, xlim = (0,1), legend = :best, w = 2, c = :blue)

# Visualize a green density plot of posterior distribution in closed-form.
plot!(p, range(0, stop = 1, length = 100), pdf.(Ref(updated_belief), range(0, stop = 1, length = 100)),
        xlabel = "probability of heads", ylabel = "", title = "", xlim = (0,1), label = "Closed-form",
        fill=0, α=0.3, w=3, c = :lightgreen)

# Visualize the true probability of heads in red.
vline!(p, [p_true], label = "True probability", c = :red)

isdefined(Main, :TuringTutorials) && Main.TuringTutorials.tutorial_footer(WEAVE_ARGS[:folder],WEAVE_ARGS[:file])


####################
# Helper Constants #
####################

# Intermidiate calculations calculations
const ECC2 = WGS84_f * (2.0 - WGS84_f) # Square of eccentricisty

##############
# Geocentric #
##############

export sGEOCtoECEF
"""
Convert geocentric position to equivalent Earth-fixed position.

Arguments:
- `geoc::Array{<:Real, 1}`: Geocentric coordinates (lon, lat, altitude) [rad] / [deg]
- `use_degrees:Bool`: If `true` interpret input as being in degrees.

Returns:
- `ecef::Array{<:Real, 1}`: Earth-fixed coordinates [m]
"""
function sGEOCtoECEF(geoc::Array{<:Real, 1} ; use_degrees::Bool=false)
    # Extract lat and lon
    lon = geoc[1]
    lat = geoc[2]

    # Handle non-explict use-degrees
    if length(geoc) == 3
        alt = geoc[3]
    else
        alt = 0.0
    end

    # Convert input to radians
    if use_degrees
        lat = lat*pi/180.0
        lon = lon*pi/180.0
    end

    # Check validity of input
    if lat < -pi/2 || lat > pi/2
        throw(ArgumentError("Lattiude, $lat, out of range. Must be between -90 and 90 degrees."))
    end

    # Compute Earth fixed coordinates
    r       = WGS84_a + alt
    x = r*cos(lat)*cos(lon)
    y = r*cos(lat)*sin(lon)
    z = r*sin(lat)

    return [x, y, z]
end

export sECEFtoGEOC
"""
Convert Earth-fixed position to geocentric location.

Arguments:
- `ecef::Array{<:Real, 1}`: Earth-fixed coordinated [m]
- `use_degrees:Bool`: If `true` returns result in units of degrees

Returns:
- `geoc`: Geocentric coordinates (lon, lat, altitude) [rad] / [deg]
"""
function sECEFtoGEOC(ecef::Array{<:Real, 1} ; use_degrees::Bool=false)
    # Expand ECEF coordinates
    x, y, z = ecef

    # Compute geocentric coordinates
    lat = atan(z, sqrt(x*x + y*y))
    lon = atan(y, x)
    alt = sqrt(x*x + y*y + z*z) - WGS84_a

    # Convert output to degrees
    if use_degrees
        lat = lat*180.0/pi
        lon = lon*180.0/pi
    end

    return [lon, lat, alt]
end

########################
# Geodetic Convertions #
########################

export sGEODtoECEF
"""
Convert geodetic position to equivalent Earth-fixed position.

Arguments:
- `geod::Array{<:Real, 1}`: Geodetic coordinates (lon, lat, altitude) [rad] / [deg]
- `use_degrees:Bool`: If `true` interpret input as being in degrees.

Returns:
- `ecef::Array{<:Real, 1}`: Earth-fixed coordinates [m]
"""
function sGEODtoECEF(geod::Array{<:Real, 1} ; use_degrees::Bool=false)
    # Extract lat and lon
    lon = geod[1]
    lat = geod[2]

    # Handle non-explict use-degrees
    if length(geod) == 3
        alt = geod[3]
    else
        alt = 0.0
    end

    # Convert input to radians
    if use_degrees
        lat = lat*pi/180.0
        lon = lon*pi/180.0
    end

    # Check validity of input
    if lat < -pi/2 || lat > pi/2
        throw(ArgumentError("Lattiude, $lat, out of range. Must be between -90 and 90 degrees."))
    end

    # Compute Earth-fixed position vector
    N = WGS84_a / sqrt(1.0 - ECC2*sin(lat)^2)
    x =           (N+alt)*cos(lat)*cos(lon)
    y =           (N+alt)*cos(lat)*sin(lon)
    z =  ((1.0-ECC2)*N+alt)*sin(lat)

    return [x, y, z]
end

export sECEFtoGEOD
"""
Convert geodetic coordinaties to Earth-fixed position

Arguments:
- `ecef::Array{<:Real, 1}`: Earth-fixed position [m]
- `use_degrees:Bool`: If `true` returns result in units of degrees

Returns:
- `geod::Array{<:Real, 1}`: Geocentric coordinates (lon, lat, altitude) [rad] / [deg]
"""
function sECEFtoGEOD(ecef::Array{<:Real, 1} ; use_degrees::Bool=false)
    # Expand ECEF coordinates
    x, y, z = ecef

    # Compute intermediate quantities
    epsilon  = eps(Float64) * 1.0e3 * WGS84_a # Convergence requirement as function of machine precision
    rho2 = x^2 + y^2                      # Square of the distance from the z-axis
    dz   = ECC2 * z
    N    = 0.0

    # Iteratively compute refine coordinates
    while true
        zdz    = z + dz
        Nh     = sqrt(rho2 + zdz^2)
        sinphi = zdz / Nh
        N      = WGS84_a / sqrt(1.0 - ECC2 * sinphi^2)
        dz_new = N * ECC2 * sinphi

        # Check convergence requirement
        if abs(dz - dz_new) < epsilon
            break
        end

        dz = dz_new
    end

    # Extract geodetic coordinates
    zdz = z + dz
    lat = atan(zdz, sqrt(rho2))
    lon = atan(y, x)
    alt = sqrt(rho2 + zdz^2) - N

    # Convert output to degrees
    if use_degrees
        lat = lat*180.0/pi
        lon = lon*180.0/pi
    end

    return [lon, lat, alt]
end

#######
# ENZ #
#######

export rECEFtoENZ
"""
Compute the rotation matrix from the Earth-fixed to the East-North-Up
coorindate basis.

Arguments:
- `station_ecef::Array{<:Real, 1}`: Earth-fixed cartesian station coordinates
- `conversion::Bool`: Conversion type to use. Can be "geocentric" or "geodetic"

Returns:
- `E::Array{Real, 2}`: Topocentric rotation matrix
"""
function rECEFtoENZ(ecef::Array{<:Real, 1} ; conversion::String="geodetic")
    if length(ecef) < 3
        throw(ArgumentError("Input coordinates must be length 3."))
    end

    # Compute Station Lat-Lon-Altitude
    if conversion == "geodetic"
        lon, lat, = sECEFtoGEOD(ecef, use_degrees=false)
    elseif conversion == "geocentric"
        lon, lat, = sECEFtoGEOC(ecef, use_degrees=false)
    else
        throw(ArgumentError("Unknown conversion method: $conversion"))
    end

    # Compute ENZ basis vectors
    eE = [-sin(lon) ; cos(lon) ; 0]
    eN = [-sin(lat)*cos(lon) ; -sin(lat)*sin(lon) ; cos(lat)]
    eZ = [cos(lat)*cos(lon) ; cos(lat)*sin(lon) ; sin(lat)]

    # Construct Rotation matrix
    E = hcat(eE, eN, eZ)'

    # Return Result
    return E
end

export rENZtoECEF
"""
Compute the rotation matrix from the Earth-fixed to the South-East-Zenith
coorindate basis.

Arguments:
- `station_ecef::Array{<:Real, 1}`: Earth-fixed cartesian station coordinates
- `conversion::Bool`: Conversion type to use. Can be "geocentric" or "geodetic"

Returns:
- `E::Array{Float64, 2}`: Topocentric rotation matrix
"""
function rENZtoECEF(ecef::Array{<:Real, 1} ; conversion::String="geodetic")
    # Check input coordinates
    if length(ecef) < 3
        throw(ArgumentError("Input coordinates must be length 3."))
    end

    return rECEFtoENZ(ecef, conversion=conversion)'
end

export sECEFtoENZ
"""
Compute the coordinates of an object in the topocentric frame of an
Earth-fixed frame

Arguments:
- `station_ecef::Array{<:Real, 1}`: Earth-fixed cartesian station coordinates
- `ecef::Array{<:Real, 1}`: Coordinates of the object in Earth-fixed station
- `conversion::Bool`: Conversion type to use. Can be "geocentric" or "geodetic"

Returns:
- `E::Array{Float64, 2}`: Topocentric rotation matrix
"""
function sECEFtoENZ(station_ecef::Array{<:Real, 1}, ecef::Array{<:Real, 1} ; conversion::String="geodetic")
    # Check input sizes
    if length(ecef) < 3
        throw(ArgumentError("Input ecef state must be at least length 3."))
    end

    if length(station_ecef) < 3
        throw(ArgumentError("Input station coordinates must be length 3."))
    end

    # Compute ENZ Rotation matrix
    E = rECEFtoENZ(station_ecef, conversion=conversion)

    # Transform range
    range_ecef = ecef[1:3] - station_ecef
    range_enz  = E * range_ecef

    # Transform range-rate (if necessary)
    if length(ecef) == 6
        range_rate_ecef = ecef[4:6]
        range_rate_enz  = E * range_rate_ecef
    end

    # Return
    if length(ecef) == 6
        sat_enz = vcat(range_enz, range_rate_enz)
    else
        sat_enz = range_enz
    end

    return sat_enz
end

export sENZtoECEF
"""
Compute the coordinates of an object in the topocentric frame of an
Earth-fixed frame

Arguments:
- `station_ecef::Array{<:Real, 1}`: Earth-fixed cartesian station coordinates
- `sez::Array{<:Real, 1}`: SEZ coordinates of the object
- `conversion::Bool`: Conversion type to use. Can be "geocentric" or "geodetic"

Returns:
- `E::Array{Float64, 2}`: Topocentric rotation matrix
"""
function sENZtoECEF(station_ecef::Array{<:Real, 1}, enz::Array{<:Real, 1} ; conversion::String="geodetic")
    # Check input sizes
    if length(enz) < 3
        throw(ArgumentError("Input ENZ state must be at least length 3."))
    end

    if length(station_ecef) < 3
        throw(ArgumentError("Input station coordinates must be length 3."))
    end

    # Compute ENZ Rotation matrix
    E = rENZtoECEF(station_ecef, conversion=conversion)

    # Transform range
    range_enz  = enz[1:3]
    range_ecef = E * range_enz

    # Transform range-rate (if necessary)
    if length(enz) == 6
        range_rate_enz  = enz[4:6]
        range_rate_ecef = E * range_rate_enz
    end

    # Return
    if length(enz) == 6
        sat_ecef = vcat(range_ecef + station_ecef, range_rate_ecef)
    else
        sat_ecef = range_ecef + station_ecef
    end

    return sat_ecef
end

#######
# SEZ #
#######

export rECEFtoSEZ
"""
Compute the rotation matrix from the Earth-fixed to the South-East-Zenith
coorindate basis.

Arguments:
- `station_ecef::Array{<:Real, 1}`: Earth-fixed cartesian station coordinates
- `conversion::Bool`: Conversion type to use. Can be "geocentric" or "geodetic"

Returns:
- `E::Array{Float64, 2}`: Topocentric rotation matrix
"""
function rECEFtoSEZ(ecef::Array{<:Real, 1} ; conversion::String="geodetic")
    if length(ecef) < 3
        throw(ArgumentError("Input coordinates must be length 3."))
    end

    # Compute Station Lat-Lon-Altitude
    if conversion == "geodetic"
        lon, lat, = sECEFtoGEOD(ecef, use_degrees=false)
    elseif conversion == "geocentric"
        lon, lat, = sECEFtoGEOC(ecef, use_degrees=false)
    else
        throw(ArgumentError("Unknown conversion method: $conversion"))
    end

    # Compute SEZ basis vectors
    eS = [sin(lat)*cos(lon) ; sin(lat)*sin(lon) ; -cos(lat)]
    eE = [-sin(lon) ; cos(lon) ; 0]
    eZ = [cos(lat)*cos(lon) ; cos(lat)*sin(lon) ; sin(lat)]

    # Construct Rotation matrix
    E = hcat(eS, eE, eZ)'

    # Return Result
    return E
end

export rSEZtoECEF
"""
Compute the rotation matrix from the Earth-fixed to the South-East-Zenith
coorindate basis.

Arguments:
- `station_ecef::Array{<:Real, 1}`: Earth-fixed cartesian station coordinates
- `conversion::Bool`: Conversion type to use. Can be "geocentric" or "geodetic"

Returns:
- `E::Array{Float64, 2}`: Topocentric rotation matrix
"""
function rSEZtoECEF(ecef::Array{<:Real, 1} ; conversion::String="geodetic")
    # Check input coordinates
    if length(ecef) < 3
        throw(ArgumentError("Input coordinates must be length 3."))
    end

    return rECEFtoSEZ(ecef, conversion=conversion)'
end

export sECEFtoSEZ
"""
Compute the coordinates of an object in the topocentric frame of an
Earth-fixed frame

Arguments:
- `station_ecef::Array{<:Real, 1}`: Earth-fixed cartesian station coordinates
- `ecef::Array{<:Real, 1}`: Coordinates of the object in Earth-fixed station
- `conversion::Bool`: Conversion type to use. Can be "geocentric" or "geodetic"

Returns:
- `E::Array{Float64, 2}`: Topocentric rotation matrix
"""
function sECEFtoSEZ(station_ecef::Array{<:Real, 1}, ecef::Array{<:Real, 1} ; conversion::String="geodetic")
    # Check input sizes
    if length(ecef) < 3
        throw(ArgumentError("Input ecef state must be at least length 3."))
    end

    if length(station_ecef) < 3
        throw(ArgumentError("Input station coordinates must be length 3."))
    end

    # Construct SEZ Rotation matrix
    E = rECEFtoSEZ(station_ecef, conversion=conversion)

    # Transform range
    range_ecef = ecef[1:3] - station_ecef
    range_sez  = E * range_ecef

    # Transform range-rate (if necessary)
    if length(ecef) == 6
        range_rate_ecef = ecef[4:6]
        range_rate_sez  = E * range_rate_ecef
    end

    # Return
    if length(ecef) == 6
        sez = vcat(range_sez, range_rate_sez)
    else
        sez = range_sez
    end

    return sez
end

export sSEZtoECEF
"""
Compute the coordinates of an object in the topocentric frame of an
Earth-fixed frame

Arguments:
- `station_ecef::Array{<:Real, 1}`: Earth-fixed cartesian station coordinates
- `sez::Array{<:Real, 1}`: SEZ coordinates of the object
- `conversion::Bool`: Conversion type to use. Can be "geocentric" or "geodetic"

Returns:
- `E::Array{Float64, 2}`: Topocentric rotation matrix
"""
function sSEZtoECEF(station_ecef::Array{<:Real, 1}, sez::Array{<:Real, 1} ; conversion::String="geodetic")
    # Check input sizes
    if length(sez) < 3
        throw(ArgumentError("Input SEZ state must be at least length 3."))
    end

    if length(station_ecef) < 3
        throw(ArgumentError("Input station coordinates must be length 3."))
    end

    # Compute ENZ Rotation matrix
    E = rSEZtoECEF(station_ecef, conversion=conversion)

    # Transform range
    range_sez  = sez[1:3]
    range_ecef = E * range_sez

    # Transform range-rate (if necessary)
    if length(sez) >= 6
        range_rate_sez  = sez[4:6]
        range_rate_ecef = E * range_rate_sez
    end

    # Return
    if length(sez) >= 6
        sat_ecef = vcat(range_ecef + station_ecef, range_rate_ecef)
    else
        sat_ecef = range_ecef + station_ecef
    end

    return sat_ecef
end

###############
# Topocentric #
###############

export sENZtoAZEL
"""
Convert East-North-Zenith topocentric state to azimuth, elevation, and range.

Arguments:
- `x::Array{<:Real, 1}`: East-North-Up state
- `use_degrees:Bool`: If `true` returns result in units of degrees

Returns:
- `azel::Array{<:Real, 1}`: Azimuth, elevation and range [rad; rad; m]
"""
function sENZtoAZEL(x::Array{<:Real, 1} ; use_degrees::Bool=false)
    # Check inputs
    if !(length(x) == 3 || length(x) == 6)
        throw(ArgumentError("Input ENZ state must be length 3 or 6."))
    end

    # Expand values
    rE, rN, rZ = x[1], x[2], x[3]

    # Range
    rho = norm(x[1:3])

    # Elevation
    el = atan(rZ, sqrt(rE^2 + rN^2))

    # Azimuth
    az = 0.0
    if el != pi/2 # Non-singular azimuth
        az = atan(rE, rN)
        if az < 0
            az += 2*pi
        end
    else # Azimuth may be singular for 90 deg elevation
        if length(x) != 6
            az = 0.0
            # @warn "Could not resolve singularity calculating azimuth."
        else
            # Use rate information to get azimuth if there is a singularity
            # in the position
            az = atan(x[4], x[5])
        end
    end

    # Output
    azel = [az ; el ; rho]

    if use_degrees
        azel[1] *= 180.0/pi
        azel[2] *= 180.0/pi
    end

    # Process Rate information
    if length(x) == 6
        rdE, rdN, rdZ = x[4], x[5], x[6]

        # Range-rate
        rhod = dot(x[1:3], x[4:6])/rho

        # Elevation-rate
        eld = (rdZ - norm(x[4:6])*sin(el))/sqrt(rE^2 + rN^2)

        # Azimuth-rate
        azd = (rdE*rN - rdN*rE)/(rE^2 + rN^2)

        # Output
        azel_rate = [azd ; eld ; rhod]
        if use_degrees
            azel_rate[1] *= 180/pi
            azel_rate[2] *= 180/pi
        end
    end

    # Return
    if length(x) == 6
        return vcat(azel, azel_rate)
    else
        return azel
    end
end

export sSEZtoAZEL
"""
Convert South-East-Zenith topocentric state to azimuth, elevation, and range.

Arguments:
- `x::Array{<:Real, 1}`: South-East-Zenith state
- `use_degrees:Bool`: If `true` returns result in units of degrees

Returns:
- `azel::Array{<:Real, 1}`: Azimuth, elevation and range [rad; rad; m]
"""
function sSEZtoAZEL(x::Array{<:Real, 1} ; use_degrees::Bool=false)
    # Check inputs
    if !(length(x) == 3 || length(x) == 6)
        throw(ArgumentError("Input rECEFtoSEZ state must be length 3 or 6."))
    end

    # Expand values
    rS, rE, rZ = x[1], x[2], x[3]

    # Range
    rho = norm(x[1:3])

    # Elevation
    el = atan(rZ, sqrt(rS^2 + rE^2))

    # Azimuth
    az = 0.0
    if el != pi/2
        az = atan(rE, -rS)
        if az < 0
            az += 2*pi
        end
    else
        if length(x) != 6
            az = 0.0
            # @warn "Could not resolve singularity calculating azimuth."
        else
            # Use rate information to get azimuth if there is a singularity
            # in the position
            az = atan(x[6], -x[5])
        end
    end

    # Output
    azel = [az ; el ; rho]

    if use_degrees
        azel[1] *= 180.0/pi
        azel[2] *= 180.0/pi
    end

    # Process Rate information
    if length(x) == 6
        rdS, rdE, rdZ = x[4], x[5], x[6]

        # Range-rate
        rhod = dot(x[1:3], x[4:6])/rho

        # Elevation-rate
        eld = (rdZ - norm(x[4:6])*sin(el))/sqrt(rS^2 + rE^2)

        # Azimuth-rate
        azd = (rdS*rE - rdE*rS)/(rS^2 + rE^2)

        # Output
        azel_rate = [azd ; eld ; rhod]
        if use_degrees
            azel_rate[1] *= 180/pi
            azel_rate[2] *= 180/pi
        end
    end

    # Return
    if length(x) == 6
        return vcat(azel, azel_rate)
    else
        return azel
    end
end


function derivative(f::Function, x::Array{T, 1}) where T <: AbstractFloat

  n = length(f(x))
  m = length(x)

	e  = eps(T)^(1/3)*maximum(abs,[x;one(T)])
  dh = Matrix{Float64}(I, m, m)*e
	deriv = Array{T}(undef,n,m)

  for i = 1:m
    deriv[:,i] = (-f(x+2*dh[:,i])+8*f(x+dh[:,i])-8*f(x-dh[:,i])+f(x-2*dh[:,i]))/(12*e)
  end

  return deriv

end


function michaelis_menten_model(concentration, Vm, K)
    (Vm * concentration) / (K + concentration)
end

# MicMen data is taken from the README of the NLReg.jl package: https://github.com/dmbates/NLreg.jl
# 1st column is concentration and 2nd column is the Rate
MicMenData = [
    0.02 76;
    0.02 47;
    0.06 97;
    0.06 107;
    0.11 123;
    0.11 139;
    0.22 159;
    0.22 152;
    0.56 191;
    0.56 201;
    1.1  207;
    1.1  200
];

MicMenConcentration = MicMenData[:, 1];
MicMenRate = MicMenData[:, 2];

# Fitness function takes a vector of Vm and K and calculates the RSS
function mic_men_fitness(params)
    Vm, K = params
    yhat = Float64[michaelis_menten_model(c, Vm, K) for c in MicMenConcentration]
    sumabs2(MicMenRate .- yhat)
end

using BlackBoxOptim
result = bboptimize(mic_men_fitness;
  SearchRange = (-1000.0, 1000.0), NumDimensions = 2, MaxSteps = 1e4)
Vm, K = best_candidate(result)
RSS = best_fitness(result)

println("NLReg.jl uses specific MicMen object and fit method to find:")
println("Vm = 212.684, K = 0.0641212, RSS = 1195.45")
println("\nUsing BlackBoxOptim.jl we find:")
println("Vm = $(Vm), K = $(K), RSS = $(RSS)")


# Ridge Regression (Tikhonov regularization)

#### auxiliary

function lreg_chkdims(X::AbstractMatrix, Y::AbstractVecOrMat, trans::Bool)
    mX, nX = size(X)
    dX = ifelse(trans, mX, nX)
    dY = ifelse(trans, nX, mX)
    size(Y, 1) == dY || throw(DimensionMismatch("Dimensions of X and Y mismatch."))
    return dX
end

lrsoltype(::AbstractVector{T}) where T = Vector{T}
lrsoltype(::AbstractMatrix{T}) where T = Matrix{T}

_vaug(X::AbstractMatrix{T}) where T = vcat(X, ones(T, 1, size(X,2)))::Matrix{T}
_haug(X::AbstractMatrix{T}) where T = hcat(X, ones(T, size(X,1), 1))::Matrix{T}

## linear least square

function llsq(X::AbstractMatrix{T}, Y::AbstractVecOrMat{T};
              trans::Bool=false, bias::Bool=true) where T<:AbstractFloat
    if trans
        mX, nX = size(X)
        size(Y, 1) == nX || throw(DimensionMismatch("Dimensions of X and Y mismatch."))
        mX <= nX || error("mX <= nX is required when trans is false.")
    else
        mX, nX = size(X)
        size(Y, 1) == mX || throw(DimensionMismatch("Dimensions of X and Y mismatch."))
        mX >= nX || error("mX >= nX is required when trans is false.")
    end
    _ridge(X, Y, zero(T), trans, bias)
end

## ridge regression

function ridge(X::AbstractMatrix{T}, Y::AbstractVecOrMat{T}, r::Real;
               trans::Bool=false, bias::Bool=true) where T<:AbstractFloat
    lreg_chkdims(X, Y, trans)
    r >= zero(r) || error("r must be non-negative.")
    _ridge(X, Y, convert(T, r), trans, bias)
end

function ridge(X::AbstractMatrix{T}, Y::AbstractVecOrMat{T}, r::AbstractVector{T};
               trans::Bool=false, bias::Bool=true) where T<:AbstractFloat
    d = lreg_chkdims(X, Y, trans)
    length(r) == d || throw(DimensionMismatch("Incorrect length of r."))
    _ridge(X, Y, r, trans, bias)
end

function ridge(X::AbstractMatrix{T}, Y::AbstractVecOrMat{T}, r::AbstractMatrix{T};
               trans::Bool=false, bias::Bool=true) where T<:AbstractFloat
    d = lreg_chkdims(X, Y, trans)
    size(r) == (d, d) || throw(DimensionMismatch("Incorrect size of r."))
    _ridge(X, Y, r, trans, bias)
end

## implementation

function _ridge(X::AbstractMatrix{T}, Y::AbstractVecOrMat{T},
                r::Union{Real, AbstractVecOrMat}, trans::Bool, bias::Bool) where T<:AbstractFloat
    if bias
        if trans
            X_ = _vaug(X)
            A = cholesky!(Hermitian(_ridge_reg!(X_ * transpose(X_), r, bias))) \ (X_ * Y)
        else
            X_ = _haug(X)
            A = cholesky!(Hermitian(_ridge_reg!(X_'X_, r, bias))) \ (X_'Y)
        end
    else
        if trans
            A = cholesky!(Hermitian(_ridge_reg!(X * X', r, bias))) \ (X * Y)
        else
            A = cholesky!(Hermitian(_ridge_reg!(X'X, r, bias))) \ (X'Y)
        end
    end
    return A::lrsoltype(Y)
end

function _ridge_reg!(Q::Matrix, r::Real, bias::Bool)
    if r > zero(r)
        n = size(Q, 1) - Int(bias)
        for i = 1:n
            @inbounds Q[i,i] += r
        end
    end
    return Q
end

function _ridge_reg!(Q::AbstractMatrix, r::AbstractVector, bias::Bool)
    n = size(Q, 1) - Int(bias)
    @assert length(r) == n
    for i = 1:n
        @inbounds Q[i,i] += r[i]
    end
    return Q
end

function _ridge_reg!(Q::AbstractMatrix, r::AbstractMatrix, bias::Bool)
    n = size(Q, 1) - Int(bias)
    @assert size(r) == (n, n)
    for j = 1:n, i = 1:n
        @inbounds Q[i,j] += r[i,j]
    end
    return Q
end


using Distributions, Plots, LaTeXStrings; pyplot()

prior(lam) = pdf(TriangularDist(0, 10, 3), lam)
data = [2,1,0,0,1,0,2,2,5,2,4,0,3,2,5,0]

like(lam) = *([pdf(Poisson(lam),x) for x in data]...)
posteriorUpToK(lam) = like(lam)*prior(lam)

delta = 10^-4.
lamRange = 0:delta:10
K = sum([posteriorUpToK(lam)*delta for lam in lamRange])
posterior(lam) = posteriorUpToK(lam)/K

bayesEstimate = sum([lam*posterior(lam)*delta for lam in lamRange])
println("Bayes estimate: ",bayesEstimate)

plot(lamRange, prior.(lamRange),
	c=:blue, label="Prior distribution")
plot!(lamRange, posterior.(lamRange),
	c=:red, label="Posterior distribution",
	xlims=(0, 10), ylims=(0, 1.2),
	xlabel=L"\lambda",ylabel="Density")


### A Pluto.jl notebook ###
# v0.12.7

using Markdown
using InteractiveUtils

# ╔═╡ 6fd616f6-1ea1-11eb-3814-8bfb4a096c49
using Pkg, DrWatson

# ╔═╡ 6ffe2628-1ea1-11eb-24ea-57f985146a72
begin
	@quickactivate "StatisticsWithJuliaPlutoNotebooks"
	using Distributions, Random, StatsBase, DataFrames, Plots
	Random.seed!(1)
end;

# ╔═╡ ed174bc4-1ea0-11eb-1e2f-a32874cec549
md"## Listing 3.4"

# ╔═╡ 70179b62-1ea1-11eb-13ce-7fb4d54bad4b
begin
	dist = TriangularDist(4,6,5)
	N = 10^6
	data = rand(dist,N)
	yData=(data .- 5).^2
end

# ╔═╡ 0b18657a-1f96-11eb-0dbe-614426f96594
("Mean: $(mean(yData)), Variance: $(var(data))")

# ╔═╡ 0b18a3f0-1f96-11eb-1f2b-0f12d42055fa
begin
	p1 = histogram(data, xlabel="x", bins=80, normed=true, ylims=(0,1.1))
	p2 = histogram(yData, xlabel="y", bins=80, normed=true, ylims=(0,15))
	plot(p1, p2, ylabel="Proportion", size=(800, 400), legend=:none)
end

# ╔═╡ 70182cd8-1ea1-11eb-094d-8d0d49cf15f3
md"## End of listing 3.4"

# ╔═╡ Cell order:
# ╟─ed174bc4-1ea0-11eb-1e2f-a32874cec549
# ╠═6fd616f6-1ea1-11eb-3814-8bfb4a096c49
# ╠═6ffe2628-1ea1-11eb-24ea-57f985146a72
# ╠═70179b62-1ea1-11eb-13ce-7fb4d54bad4b
# ╠═0b18657a-1f96-11eb-0dbe-614426f96594
# ╠═0b18a3f0-1f96-11eb-1f2b-0f12d42055fa
# ╟─70182cd8-1ea1-11eb-094d-8d0d49cf15f3


module Example1

export Point
struct Point{T}
    x::T
    y::T
end

function Base. +(p::Point, q::Point)
    Point(p.x+ q.x, p.y + q.y)
end

export MPoint
mutable struct MPoint{T}
    x::T
    y::T
end

greet() = println("Hello World!")

################################################################

export Polynomial
struct Polynomial{T}
    coeffs::Vector{T}
end

export evaluate
function evaluate(p::Polynomial,x)
    sum(c * x^(i-1) for (i,c) in enumerate(p.coeffs))
end

function Base. +(p::Polynomial, q::Polynomial)
    l = max(length(p.coeffs), length(q.coeffs))
    rcoeffs = zeros(l)
    for (i,c) in enumerate(p.coeffs)
        rcoeffs[i] +=c
    end
    Polynomial(rcoeffs)
end

function Base. *(a, p::Polynomial)
    Polynomial(map(x->a*x, p.coeffs))
end

#Base .*(p::Polynomial,a) = a*p

export deriv
function deriv(p::Polynomial)
    Polynomial([(i-1) * c for (i,c) in enumerate(p.coeffs)][2::end])
end

end # module


#=
# Based on the file JacobiP.m from the book:
# Nodal Discontinuous Galerkin Method - Hesthaven, Warburton
#  https://link.springer.com/book/10.1007/978-0-387-72067-8
=#

# Evaluate Nth order Jacobi Polynomial of type (alpha,beta) at points x.
# Note   : They are normalized to be orthonormal.
function jacobi_polynomial(x, alpha::Int, beta::Int, N::Int)
    PL = zeros(N+1,length(x));

    # Initial values P_0(x) and P_1(x)
    gamma0 = 2^(alpha+beta+1)/(alpha+beta+1)*factorial(alpha)*
                factorial(beta)/factorial(alpha+beta);
    PL[1,:] = 1.0/sqrt(gamma0) .* ones(length(x));
    if N==0
        return PL[1,:];
    end
    gamma1 = (alpha+1)*(beta+1)/(alpha+beta+3) * gamma0;
    PL[2,:] = ((alpha+beta+2)/(2*sqrt(gamma1))) .* x .* ones(length(x)) .+ (alpha-beta)/(2*sqrt(gamma1));
    if N==1
        return PL[2,:];
    end

    # Higher orders are generated by this recurrence.
    # P(i+1) = -(a(i-1)/a(i))P(i-1) + ((x-b(i))/a(i))P(i)
    aold = 2/(2+alpha+beta)*sqrt((alpha+1)*(beta+1)/(alpha+beta+3));
    for i=1:N-1
      h1 = 2*i+alpha+beta;
      anew = 2/(h1+2)*sqrt( (i+1)*(i+1+alpha+beta)*(i+1+alpha)*
          (i+1+beta)/((h1+1)*(h1+3)));
      bnew = - (alpha^2-beta^2)/(h1*(h1+2));
      PL[i+2,:] = (-aold/anew) .* PL[i,:] + ((x.-bnew)./anew) .* PL[i+1,:];
      aold =anew;
    end;

    return PL[N+1,:];
end

# sqrt(N*(N+alpha+beta+1))*JacobiP(r(:),alpha+1,beta+1, N-1);
function grad_jacobi_polynomial(x, alpha::Int, beta::Int, N::Int)
    if N < 1
        return zeros(size(x));
    end

    return sqrt(N*(N+alpha+beta+1)) .* jacobi_polynomial(x, alpha+1, beta+1, N-1);
end


#=
    For this particular positions
        r = [1 2; 3 5]
    the expression below are exact solutions.

    Check sagemath notebook for the proof.
=#
N = 2
r = [1 2; 3 5]

using SpecialFunctions, LinearAlgebra
H_0 = [1 besselh.(0,sqrt(13)); besselh.(0,sqrt(13)) 1]
H_m = [0 exp(-2im*atan(3,2))*besselh.(2,sqrt(13));
            exp(-2im*atan(3,2))*besselh.(2,sqrt(13)) 0]

H_p = [0 exp(+2im*atan(3,2))*besselh.(2,sqrt(13));
    exp(+2im*atan(3,2))*besselh.(2,sqrt(13)) 0]

K_exact = [H_0 H_p; H_m H_0]

#=
        Creating matrices individually
=#
let
    N = 2
    r = [1 2; 3 5]
    Dx = [r[j,1] - r[l,1] for j=1:N, l=1:N]
    Dy = [r[j,2] - r[l,2] for j=1:N, l=1:N]
    r_jl = sqrt.( Dx.^2 + Dy.^2  )
    φ_jl = atan.(Dy, Dx)

    #Hankel functions
    H0 = besselh.(0,r_jl)
    H0[findall(  isnan.(H0) )] .= 0.0
    H0[diagind(H0)] .= 1

    Hp = exp.(+2*1im*φ_jl).*besselh.(2,r_jl)
    Hp[diagind(Hp)] .= 0
    Hm = exp.(-2*1im*φ_jl).*besselh.(2,r_jl)
    Hm[diagind(Hm)] .= 0

    K = [H0 Hp; Hm H0]
    K .≈ K_exact
end

#=
        Creating matrices generically
=#
let
    N = 2
    r = [1 2; 3 5]
    δ(x,y) = Int(==(x,y))
    Dx = [r[j,1] - r[l,1] for j=1:N, l=1:N]
    Dy = [r[j,2] - r[l,2] for j=1:N, l=1:N]
    r_jl = sqrt.( Dx.^2 + Dy.^2  )
    φ_jl = atan.(Dy, Dx)

    H0 = besselh.(0,r_jl)
    H2 = besselh.(2,r_jl)

    A = []
    α_range = β_range = [-1, +1]
    for (α_idx, α) ∈ enumerate(α_range)
        B = []
        for (β_idx, β) ∈ enumerate(β_range)
            term1 = I(N)*δ(α,β)
            term2 = ones(N, N) - I(N)

            term3 = δ(α,β)*H0
            term3[findall(isnan.(term3))] .= 0
            term3[findall(isinf.(term3))] .= 0

            term4 = (1 - δ(α,β))*exp.(2*α*im*φ_jl).*H2
            term4[findall(isnan.(term4))] .= 0

            temp =  term1 + term2.*(term3 + term4 )

            push!(B, temp)
        end
        push!(A, vcat(B[1:length(β_range)]...)  )
    end

    K = hcat(A[1:length(α_range)]...)
    K .≈ K_exact
end


using FdeSolver
using Plots

## inputs
tSpan = [0, 25]        # [intial time, final time]
y0 = [34, 6]           # initial values
β = [0.98, 0.99]       # order of derivatives

## ODE model

par = [0.55, 0.028, 0.84, 0.026] # model parameters

function F(t, y, par)

    α1 = par[1]      # growth rate of the prey population
    β1 = par[2]      # rate of shrinkage relative to the product of the population sizes
    γ = par[3]       # shrinkage rate of the predator population
    δ = par[4]       # growth rate of the predator population as a factor of the product
                     # of the population sizes

    u = y[1]         # population size of the prey species at time t[n]
    v = y[2]         # population size of the predator species at time t[n]

    F1 = α1 .* u .- β1 .* u .* v
    F2 = - γ .* v .+ δ .* u .* v

    [F1, F2]

end

## Solution
t, Yapp = FDEsolver(F, tSpan, y0, β, par)

# plotting
plot(t, Yapp, linewidth = 5, title = "Solution to LV model with 2 FDEs",
     xaxis = "Time (t)", yaxis = "y(t)", label = ["Prey" "Predator"])
plot!(legendtitle = "Population of")


module Parameters

export FreeParameters, lognormal, ScaledLogitNormal

using Oceananigans.Architectures: CPU, arch_array, architecture
using Oceananigans.Utils: prettysummary
using Oceananigans.TurbulenceClosures: AbstractTurbulenceClosure
using Oceananigans.TurbulenceClosures: AbstractTimeDiscretization, ExplicitTimeDiscretization

using Printf
using Distributions
using DocStringExtensions
using LinearAlgebra

using SpecialFunctions: erfinv
using Distributions: AbstractRNG, ContinuousUnivariateDistribution

#####
##### Priors
#####

"""
    lognormal(; mean, std)

Return `Lognormal` distribution parameterized by
the distribution `mean` and standard deviation `std`.

Notes
=====

A variate `X` is `LogNormal` distributed if

```math
\\log(X) ∼ 𝒩(μ, σ²) ,
```

where ``𝒩(μ, σ²)`` is the `Normal` distribution with mean ``μ``
and variance ``σ²``.

The `mean` and variance ``s²`` (where ``s`` is the standard
deviation or `std`) are related to the parameters ``μ``
and ``σ²`` via

```math
 m = \\exp(μ + σ² / 2),
```
```math
s² = [\\exp(σ²) - 1] m².
```

These formula allow us to calculate ``μ`` and ``σ`` given
``m`` and ``s²``, since rearranging the formula for ``s²``
gives

```math
\\exp(σ²) = m² / s² + 1
```

which then yields

```math
σ = \\sqrt{\\log(m² / s² + 1)}.
```

We then find that

```math
μ = \\log(m) - σ² / 2 .
```

See also
[wikipedia](https://en.wikipedia.org/wiki/Log-normal_distribution#Generation_and_parameters).
"""
function lognormal(; mean, std)
    k = std^2 / mean^2 + 1 # intermediate variable
    σ = sqrt(log(k))
    μ = log(mean) - σ^2 / 2
    return LogNormal(μ, σ)
end

struct ScaledLogitNormal{T} <: ContinuousUnivariateDistribution
    μ :: T
    σ :: T
    lower_bound :: T
    upper_bound :: T

    ScaledLogitNormal{T}(μ, σ, L, U) where T = new{T}(T(μ), T(σ), T(L), T(U))
end

"""Return a logit-normally distributed variate given the normally-distributed variate `X`."""
normal_to_scaled_logit_normal(L, U, X) = L + (U - L) / (1 + exp(X))

"""Return a normally-distributed variate given the logit-normally distributed variate `Y`."""
scaled_logit_normal_to_normal(L, U, Y) = log((U - Y) / (Y - L))

Base.rand(rng::AbstractRNG, d::ScaledLogitNormal) =
    normal_to_scaled_logit_normal(d.lower_bound, d.upper_bound, rand(rng, Normal(d.μ, d.σ)))

unit_normal_std(mass) = 1 / (2 * √2 * erfinv(mass))

"""
    ScaledLogitNormal([FT=Float64;] bounds=(0, 1), mass=0.5, interval=nothing, μ=nothing, σ=nothing)

Return a `ScaledLogitNormal` distribution with compact support within `bounds`.

`interval` is an optional 2-element tuple or Array. When specified,
the parameters `μ` and `σ` of the underlying `Normal` distribution
are calculated so that `mass` fraction of the probability density
lies within `interval`.

If `interval` is not specified, then `μ=0` and `σ=1` by default.

Notes
=====

`ScaledLogitNormal` is a four-parameter distribution
generated by the transformation

```math
Y = L + (U - L) / [1 + \\exp(X)],
```

of the normally-distributed variate ``X ∼ 𝒩(μ, σ)``. The four parameters
governing the distribution of ``Y`` are thus

- ``L``:  lower bound (0 for the `LogitNormal` distribution)
- ``U``:  upper bound (1 for the `LogitNormal` distribution)
- ``μ``:  mean of the underlying `Normal` distribution
- ``σ²``: variance of the underlying `Normal` distribution
"""
function ScaledLogitNormal(FT=Float64; bounds=(0, 1), mass=0.5, interval=nothing, μ=nothing, σ=nothing)
    L, U = bounds

    if isnothing(interval) # use default μ=0 and σ=1 if not set
        isnothing(μ) && (μ = 0)
        isnothing(σ) && (σ = 1)

    elseif !isnothing(interval) # try to compute μ and σ

        Li, Ui = interval

        # User friendliness
        (!isnothing(μ) || !isnothing(σ)) && @warn "Using interval and mass to determine μ and σ."
        0 < mass < 1 || throw(ArgumentError("Mass must lie between 0 and 1."))
        Li > L && Ui < U || throw(ArgumentError("Interval limits must lie between `bounds`."))

        # Compute lower and upper limits of midspread in unconstrained space
        #
        # Note that the _lower_ bound in unconstrained space is associated with the
        # _upper_ bound in constrained space, and vice versa.
        L̃i = scaled_logit_normal_to_normal(L, U, Ui)
        Ũi = scaled_logit_normal_to_normal(L, U, Li)

        μ = (Ũi + L̃i) / 2

        # Note that the mass beneath a half-width `δ` of the
        # standard Normal distribution is
        #
        # mass = 2 / √(2π) ∫₀ᵟ exp(-x^2 / 2) dx
        #      = erf(δ / √2)
        #
        # For an `interval = (Ũi, L̃i)` of the normal distribution,
        # the non-dimensional half-width is
        #
        # δ = (Ũi - L̃i) / 2σ
        #
        # where σ is the distribution's standard deviation.
        # We then find
        #
        # erfinv(mass) = (Ũi - L̃i) / (2 * √2 * σ) ,
        #
        # and rearranging to solve for σ yields
        #
        σ = (Ũi - L̃i) / (2 * √2 * erfinv(mass))
    end

    return ScaledLogitNormal{FT}(μ, σ, L, U)
end

# Calculate the prior in unconstrained space given a prior in constrained space
unconstrained_prior(Π::LogNormal)         = Normal(Π.μ / abs(Π.μ), Π.σ / abs(Π.μ))
unconstrained_prior(Π::Normal)            = Normal(0, 1)
unconstrained_prior(Π::ScaledLogitNormal) = Normal(Π.μ, Π.σ)

"""
    transform_to_unconstrained(Π, Y)

Transform the "constrained" (physical) variate `Y` into it's
unconstrained (normally-distributed) counterpart `X` through the
forward map associated with `Π`.

If some mapping between ``Y`` and the normally-distributed ``X`` is
defined via

```math
Y = g(X).
```

Then `transform_to_unconstrained` is the inverse ``X = g^{-1}(Y)``.
The change of variables ``g(X)`` determines the distribution `Π` of `Y`.

Example
=======

The logarithm of a `LogNormal(μ, σ)` distributed variate is normally-distributed,
such that the forward trasform ``f ≡ \\exp``,

```math
Y = \\exp(X),
```

and the inverse trasnform is the natural logarithm ``f^{-1} ≡ \\log``,

```math
\\log(Y) = X ∼ 𝒩(μ, σ).
```
"""
transform_to_unconstrained(Π::Normal,    Y) = (Y - Π.μ) / Π.σ
transform_to_unconstrained(Π::LogNormal, Y) = log(Y^(1 / abs(Π.μ))) # log(Y) / abs(Π.μ)

transform_to_unconstrained(Π::ScaledLogitNormal, Y) =
    scaled_logit_normal_to_normal(Π.lower_bound, Π.upper_bound, Y)

"""
    transform_to_constrained(Π, X)

Transform an "unconstrained", normally-distributed variate `X`
to "constrained" (physical) space via the map associated with
the distribution `Π` of `Y`.
"""
transform_to_constrained(Π::Normal, X)    = X * Π.σ + Π.μ
transform_to_constrained(Π::LogNormal, X) = exp(X * abs(Π.μ))
transform_to_constrained(Π::ScaledLogitNormal, X) =
    normal_to_scaled_logit_normal(Π.lower_bound, Π.upper_bound, X)

# Convenience vectorized version
transform_to_constrained(priors::NamedTuple, X::AbstractVector) =
    NamedTuple(name => transform_to_constrained(priors[name], X[i])
               for (i, name) in enumerate(keys(priors)))

# Convenience matrixized version assuming particles vary on 2nd dimension
transform_to_constrained(priors::NamedTuple, X::AbstractMatrix) =
    [transform_to_constrained(priors, X[:, k]) for k = 1:size(X, 2)]

function inverse_covariance_transform(Π, X, covariance)
    diag = [covariance_transform_diagonal(Π[i], X[i]) for i=1:length(Π)]
    dT = Diagonal(diag)
    return dT * covariance * dT'
end

covariance_transform_diagonal(::LogNormal, X) = exp(X)
covariance_transform_diagonal(::Normal, X)    = 1
covariance_transform_diagonal(Π::ScaledLogitNormal, X) = - (Π.upper_bound - Π.lower_bound) * exp(X) / (1 + exp(X))^2

#####
##### Free parameters
#####

"""
    struct FreeParameters{N, P, D}

A container for free parameters that includes the parameter names and their
corresponding prior distributions.

$(FIELDS)
"""
struct FreeParameters{N, P, D}
    "free parameters"
    names :: N
    "prior distributions for free parameters"
    priors :: P
    "dependent parameters"
    dependent_parameters :: D
end

"""
    FreeParameters(priors; names = Symbol.(keys(priors)), dependent_parameters=NamedTuple())

Return named `FreeParameters` with priors. Free parameter `names` are inferred from
the keys of `priors` if not provided. Optionally, `dependent_parameters` are prescribed
as a `NamedTuple` whose keys are the names of "additional" parameters, and whose values
are functions that return those parameters given a vector of free parameters in `names`.

Example
=======

```jldoctest
julia> using Distributions, ParameterEstimocean

julia> priors = (ν = Normal(1e-4, 1e-5), κ = Normal(1e-3, 1e-5))
(ν = Normal{Float64}(μ=0.0001, σ=1.0e-5), κ = Normal{Float64}(μ=0.001, σ=1.0e-5))

julia> free_parameters = FreeParameters(priors)
FreeParameters with 2 parameters
├── names: (:ν, :κ)
├── priors: Dict{Symbol, Any}
│   ├── ν => Normal{Float64}(μ=0.0001, σ=1.0e-5)
│   └── κ => Normal{Float64}(μ=0.001, σ=1.0e-5)
└── dependent parameters: Dict{Symbol, Any}

julia> c(p) = p.ν + p.κ # compute a third dependent parameter `c` as a function of `ν` and `κ`
c (generic function with 1 method)

julia> free_parameters_with_a_dependent = FreeParameters(priors, dependent_parameters=(; c))
FreeParameters with 2 parameters and 1 dependent parameter
├── names: (:ν, :κ)
├── priors: Dict{Symbol, Any}
│   ├── ν => Normal{Float64}(μ=0.0001, σ=1.0e-5)
│   └── κ => Normal{Float64}(μ=0.001, σ=1.0e-5)
└── dependent parameters: Dict{Symbol, Any}
    └── c => c (generic function with 1 method)
```
"""
function FreeParameters(priors; names = Symbol.(keys(priors)), dependent_parameters=NamedTuple())
    priors = NamedTuple(name => priors[name] for name in names)
    return FreeParameters(Tuple(names), priors, dependent_parameters)
end

Base.summary(fp::FreeParameters) = "$(fp.names)"

function prior_show(io, priors, name, prefix, width)
    print(io, @sprintf("%s %s => ", prefix, lpad(name, width, " ")))
    show(io, priors[name])
    return nothing
end

function dependent_parameter_show(io, dependent_parameters, name, prefix, width)
    print(io, @sprintf("%s %s => ", prefix, lpad(name, width, " ")))
    print(io, prettysummary(dependent_parameters[Symbol(name)]))
    return nothing
end

parameter_str(N) = N>1 ? "parameters" : "parameter"

function Base.show(io::IO, p::FreeParameters)
    Np, Nd = length(p), length(p.dependent_parameters)

    free_parameters_summary = "FreeParameters with $Np " * parameter_str(Np)

    title = Nd > 0 ?
            free_parameters_summary * " and $Nd dependent " * parameter_str(Nd) :
            free_parameters_summary

    print(io, title, '\n',
              "├── names: $(p.names)", '\n',
              "├── priors: Dict{Symbol, Any}")

    maximum_name_length = maximum([length(string(name)) for name in p.names])

    for (i, name) in enumerate(p.names)
        prefix = i == length(p.names) ? "│   └──" : "│   ├──"
        print(io, '\n')
        prior_show(io, p.priors, name, prefix, maximum_name_length)
    end

    print(io, '\n')

    print(io, "└── dependent parameters: Dict{Symbol, Any}")

    if !isempty(p.dependent_parameters)
        maximum_name_length = maximum([length(string(name)) for name in p.dependent_parameters])

        for (i, name) in enumerate(p.dependent_parameters)
            prefix = i == length(p.dependent_parameters) ? "    └──" : "    ├──"
            print(io, '\n')
            dependent_parameter_show(io, p.dependent_parameters, name, prefix, maximum_name_length)
        end
    end

    return nothing
end

Base.length(p::FreeParameters) = length(p.names)

function build_parameters_named_tuple(p::FreeParameters, free_θ)
    if free_θ isa Dict # convert to NamedTuple with
        free_θ = NamedTuple(name => free_θ[name] for name in p.names)
    elseif !(free_θ isa NamedTuple) # mostly likely a Vector: convert to NamedTuple with
        free_θ = NamedTuple{p.names}(Tuple(free_θ))
    end

    # Compute dependent parameters
    dependent_names = keys(p.dependent_parameters)
    maps = values(p.dependent_parameters)
    dependent_θ = NamedTuple(name => maps[name](free_θ) for name in dependent_names)

    return merge(dependent_θ, free_θ) # prioritize free_θ
end

#####
##### Setting parameters
#####

const ParameterValue = Union{Number, AbstractArray}

"""
    construct_object(specification_dict, parameters; name=nothing, type_parameter=nothing)

    construct_object(d::ParameterValue, parameters; name=nothing)

Return a composite type object whose properties are prescribed by the `specification_dict`
dictionary. All parameter values are given the values in `specification_dict` *unless* they
are included as a parameter name-value pair in the named tuple `parameters`, in which case
the value in `parameters` is asigned.

The `construct_object` is recursively called upon every property that is included in `specification_dict`
until a property with a numerical value is reached. The object's constructor name must be
included in `specification_dict` under key `:type`.

Example
=======

```jldoctest; filter = [r".*Dict{Symbol.*", r".*:type       => Closure.*", r".*:c          => 3.*", r".*:subclosure => Dict{Symbol.*"]
julia> using ParameterEstimocean.Parameters: construct_object, dict_properties, closure_with_parameters

julia> struct Closure; subclosure; c end

julia> struct ClosureSubModel; a; b end

julia> sub_closure = ClosureSubModel(1, 2)
ClosureSubModel(1, 2)

julia> closure = Closure(sub_closure, 3)
Closure(ClosureSubModel(1, 2), 3)

julia> specification_dict = dict_properties(closure)
Dict{Symbol, Any} with 3 entries:
  :type       => Closure
  :c          => 3
  :subclosure => Dict{Symbol, Any}(:a=>1, :b=>2, :type=>ClosureSubModel)

julia> new_closure = construct_object(specification_dict, (a=2.1,))
Closure(ClosureSubModel(2.1, 2), 3)

julia> another_new_closure = construct_object(specification_dict, (b=π, c=2π))
Closure(ClosureSubModel(1, π), 6.283185307179586)
```
"""
construct_object(d::ParameterValue, parameters; name=nothing) =
    name ∈ keys(parameters) ? getproperty(parameters, name) : d

function construct_object(specification_dict, parameters;
                          name=nothing, type_parameter=nothing)

    type = Constructor = specification_dict[:type]
    kwargs_vector = [construct_object(specification_dict[name], parameters; name)
                        for name in fieldnames(type) if name != :type]

    return isnothing(type_parameter) ? Constructor(kwargs_vector...) :
                                       Constructor{type_parameter}(kwargs_vector...)
end

"""
    dict_properties(object)

Return a dictionary with all properties of an `object` and their values, including the
`object`'s type name. If any of the `object`'s properties is not a numerical value but
instead a composite type, then `dict_properties` is called recursively on that `object`'s
property returning a dictionary with all properties of that composite type. Recursion
ends when properties of type `ParameterValue` are found.
"""
function dict_properties(object)
    p = Dict{Symbol, Any}(n => dict_properties(getproperty(object, n)) for n in propertynames(object))
    p[:type] = typeof(object).name.wrapper

    return p
end

dict_properties(object::ParameterValue) = object

"""
    closure_with_parameters(closure, parameters)

Return a new object where for each (`parameter_name`, `parameter_value`) pair
in `parameters`, the value corresponding to the key in `closure` that matches
`parameter_name` is replaced with `parameter_value`.

Example
=======

Create a placeholder `Closure` type that includes a parameter `c` and a sub-closure
with two parameters: `a` and `b`. Then construct a closure with values `a, b, c = 1, 2, 3`.

```jldoctest closure_with_parameters
julia> struct Closure; subclosure; c end

julia> struct ClosureSubModel; a; b end

julia> sub_closure = ClosureSubModel(1, 2)
ClosureSubModel(1, 2)

julia> closure = Closure(sub_closure, 3)
Closure(ClosureSubModel(1, 2), 3)
```

Providing `closure_with_parameters` with a named tuple of parameter names and values,
and a recursive search in all types and subtypes within `closure` is done and whenever
a parameter is found whose name exists in the named tuple we provided, its value is
then replaced with the value provided.

```jldoctest closure_with_parameters
julia> new_parameters = (a = 12, d = 7)
(a = 12, d = 7)

julia> using ParameterEstimocean.Parameters: closure_with_parameters

julia> closure_with_parameters(closure, new_parameters)
Closure(ClosureSubModel(12, 2), 3)
```
"""
closure_with_parameters(closure, parameters) = construct_object(dict_properties(closure), parameters)

closure_with_parameters(closure::AbstractTurbulenceClosure{ExplicitTimeDiscretization}, parameters) =
    construct_object(dict_properties(closure), parameters, type_parameter=nothing)

closure_with_parameters(closure::AbstractTurbulenceClosure{TD}, parameters) where {TD <: AbstractTimeDiscretization} =
    construct_object(dict_properties(closure), parameters; type_parameter=TD)

closure_with_parameters(closures::Tuple, parameters) =
    Tuple(closure_with_parameters(closure, parameters) for closure in closures)

"""
    update_closure_ensemble_member!(closures, p_ensemble, parameters)

Use `parameters` to update the `p_ensemble`-th closure from and array of `closures`.
The `p_ensemble`-th closure corresponds to ensemble member `p_ensemble`.
"""
update_closure_ensemble_member!(closure, p_ensemble, parameters) = nothing

update_closure_ensemble_member!(closures::AbstractVector, p_ensemble, parameters) =
    closures[p_ensemble] = closure_with_parameters(closures[p_ensemble], parameters)

function update_closure_ensemble_member!(closures::AbstractMatrix, p_ensemble, parameters)
    for j in 1:size(closures, 2) # Assume that ensemble varies along first dimension
        closures[p_ensemble, j] = closure_with_parameters(closures[p_ensemble, j], parameters)
    end

    return nothing
end

function update_closure_ensemble_member!(closure_tuple::Tuple, p_ensemble, parameters)
    for closure in closure_tuple
        update_closure_ensemble_member!(closure, p_ensemble, parameters)
    end
    return nothing
end

"""
    new_closure_ensemble(closures, θ, arch=CPU())

Return a new set of `closures` in which all closures that have free parameters are updated.
Closures with free parameters are expected as `AbstractArray` of `TurbulenceClosures`, and
this allows `new_closure_ensemble` to go through all closures in `closures` and only update
the parameters for the any closure that is of type `AbstractArray`. The `arch`itecture
(`CPU()` or `GPU()`) defines whethere `Array` or `CuArray` is returned.
"""
function new_closure_ensemble(closures::AbstractArray, θ, arch)
    cpu_closures = arch_array(CPU(), closures)

    for (p, θp) in enumerate(θ)
        update_closure_ensemble_member!(cpu_closures, p, θp)
    end

    return arch_array(arch, cpu_closures)
end
new_closure_ensemble(closures::Tuple, θ, arch) =
    Tuple(new_closure_ensemble(closure, θ, arch) for closure in closures)

new_closure_ensemble(closure, θ, arch) = closure

end # module


function makePlots()

println("Generating plots")

plotsDir = (pwd()[end-3:end] == "docs") ? "build/plots" : "docs/build/plots"
mkdir(plotsDir)
Plots.pyplot()

# PID design functions
P = tf(1,[1,1])^4
gangoffourplot(P,tf(1))
Plots.savefig(plotsDir*"/pidgofplot.svg")

ωp = 0.8
kp,ki,C = loopshapingPI(P,ωp,phasemargin=60, doplot=false)
gangoffourplot(P, [tf(1), C])
Plots.savefig(plotsDir*"/pidgofplot2.svg")
nyquistplot([P, P*C])
Plots.savefig(plotsDir*"/pidnyquistplot.svg")

ωp = 2
kp,ki,C60 = loopshapingPI(P,ωp,rl=1,phasemargin=60, doplot=true)
gangoffourplot(P, [tf(1), C60])
Plots.savefig(plotsDir*"/pidgofplot3.svg")
nyquistplot([P, P*C60])
Plots.savefig(plotsDir*"/pidnyquistplot2.svg")

# Advanced pole placement
ζ = 0.2
ω = 1
B = [1]
A   = [1, 2ζ*ω, ω^2]
P  = tf(B,A)
# Control design
ζ0 = 0.7
ω0 = 2
Am = [1, 2ζ0*ω0, ω0^2]
Ao = conv(2Am, [1/2, 1]) # Observer polynomial
AR = [1,0] # Force the controller to contain an integrator

B⁺  = [1] # The process numerator polynomial can be facored as B = B⁺B⁻ where B⁻ contains the zeros we do not want to cancel (non-minimum phase and poorly damped zeros)
B⁻  = [1]
Bm  = conv(B⁺, B⁻) # In this case, keep the entire numerator polynomial of the process

R,S,T = rstc(B⁺,B⁻,A,Bm,Am,Ao,AR) # Calculate the 2-DOF controller polynomials

Gcl = tf(conv(B,T),zpconv(A,R,B,S)) # Form the closed loop polynomial from reference to output

stepplot([P,Gcl]) # Visualize the open and closed loop responses.
Plots.savefig(plotsDir*"/ppstepplot.svg")
gangoffourplot(P, tf(-S,R)) # Plot the gang of four to check that all tranfer functions are OK
Plots.savefig(plotsDir*"/ppgofplot.svg")

P1 = "exp(-sqrt(s))"
f1 = stabregionPID(P1,logspace(-5,1,1000)); Plots.savefig(plotsDir*"/stab1.svg")
P2 = "100*(s+6).^2./(s.*(s+1).^2.*(s+50).^2)"
f2 = stabregionPID(P2,logspace(-5,2,1000)); Plots.savefig(plotsDir*"/stab2.svg")
P3 = tf(1,[1,1])^4
f3 = stabregionPID(P3,logspace(-5,0,1000)); Plots.savefig(plotsDir*"/stab3.svg")

# PID plots

P = tf([1.],[1., 1])
ζ = 0.5 # Desired damping
ws = logspace(-1,2,8) # A vector of closed-loop bandwidths
kp = 2*ζ*ws-1 # Simple pole placement with PI given the closed-loop bandwidth, the poles are placed in a butterworth pattern
ki = ws.^2
pidplots(P,:nyquist,;kps=kp,kis=ki, ω= logspace(-2,2,500))
Plots.savefig(plotsDir*"/pidplotsnyquist1.svg")
pidplots(P,:gof,;kps=kp,kis=ki, ω= logspace(-2,2,500))
Plots.savefig(plotsDir*"/pidplotgof1.svg")

kp = linspace(-1,1,8) # Now try a different strategy, where we have specified a gain crossover frequency of 0.1 rad/s
ki = sqrt(1-kp.^2)/10
pidplots(P,:nyquist,;kps=kp,kis=ki)
Plots.savefig(plotsDir*"/pidplotsnyquist2.svg")
pidplots(P,:gof,;kps=kp,kis=ki)
Plots.savefig(plotsDir*"/pidplotsgof2.svg")

#AR = [1,0]
#B⁺  = [1]
#B⁻  = [1]
#Bm  = conv(B⁺, B⁻)
#R,S,T = rstc(B⁺,B⁻,A,bm,Am,Ao,AR)
#Gcl = tf(conv(B,T),zpconv(A,R,B,S))
#f1 = stepplot([P,Gcl])
#f2 = gangoffourplot(P, tf(-S,R))
#Plots.savefig(f1, "$plotsDir/rstcstepplot.svg")
#Plots.savefig(f2, "$plotsDir/rstcgofplot.svg")
end


"""
    expectation_step(X, z, dists)

Computes the expectation step with:
- `X`: (m×n) multi-variate observations (rows are observations, columns are variables)
- `dists`: distributions, one for each category (unique label)
- `prior`: prior probability of each label
Returns `Z`: (m×n) matrix of probability distributions of each label for each observation
"""
function expectation_step(X, dists, prior)
    n = size(X, 1)
    Z = Matrix{Float64}(undef, n, length(prior))
    for k in eachindex(prior)
        for i in 1:n
            Z[i,k] = prior[k] * pdf(dists[k], X[i,:]) /
                     sum(
                        prior[j] * pdf(dists[j], X[i,:])
                        for j in eachindex(prior)
                     )
        end
    end
    return Z
end

"""
    maximization_step(::Type{<:D}, X, Z)

Maximization step, parameterized by the type of distribution used.
Returns a tuple (dists, π, loglike) with:
- `dists`: the created vector of distributions
- `prior` vector of updated prior probabilities of each label
"""
function maximization_step(::Type{<:MvNormal}, X, Z)
    n = size(X, 1)
    Nk = size(Z, 2)
    µ = map(1:Nk) do k
        num = sum(Z[i,k] .* X[i,:] for i in 1:n)
        den = sum(Z[i,k] for i in 1:n)
        num / den
    end

    Σ = map(1:Nk) do k
        num = zeros(size(X, 2), size(X, 2))
        for i in 1:n
            r = X[i,:] .- µ[k]
            num .= num .+ Z[i,k] .* (r * r')
        end
        den = sum(Z[i,k] for i in 1:n)
        num ./ den
    end
    prior = [inv(n) * sum(Z[i,k] for i in 1:n)
            for k in 1:Nk
        ]
    dists = map(1:Nk) do k
        MvNormal(µ[k], Σ[k] + 10e-7I)
    end
    return (dists, prior)
end

"""
    expectation_maximization(D::Type{<:Distribution}, X, Nk, maxiter::Integer = 500, loglike_diff = 10e-5; all_dists = nothing)

Complete EM algorithm, takes a type of distribution `D` used to
determine the maximization step, observations `X` and number of classes `Nk`.
"""
function expectation_maximization(D::Type{<:Distribution}, X, Nk, maxiter::Integer = 500, loglike_diff = 10e-5; all_dists = nothing)
    # initialize classes
    n = size(X,1)
    Z = zeros(n, Nk)
    for i in 1:n
        j0 = mod(i,Nk)+1
        j1 = j0 > 1 ? j0-1 : 2
        Z[i,j0] = 0.75
        Z[i,j1] = 0.25
    end
    (dists, prior) = maximization_step(D, X, Z)
    all_dists isa Vector && push!(all_dists, (dists, prior))
    l = loglike_mixture(X, dists, prior)
    lprev = 0.0
    iter = 0
    # EM iterations
    while iter < maxiter && abs(lprev-l) > loglike_diff
        Z = expectation_step(X, dists, prior)
        (dists, prior) = maximization_step(D, X, Z)
        all_dists isa Vector && push!(all_dists, (dists, prior))
        lprev = l
        l = loglike_mixture(X, dists, prior)
        iter += 1
    end
    return (dists, prior, Z, l, iter)
end

function loglike_mixture(X, dists, prior)
    l = zero(eltype(X))
    n = size(X,1)
    for i in 1:n
        l += log(
            sum(prior[k] * pdf(dists[k], X[i,:]) for k in eachindex(prior))
        )
    end
    return l
end

let
    Random.seed!(42)
    X = [randn(1000, 2); 3.0 .+ randn(1000, 2)]
    n = size(X, 1)
    for maxiter in (1, 5, 10, 100, 150, 250)
        (dists, prior, Z, l, iter) = expectation_maximization(MvNormal, X, 2, maxiter)
        contour_mixture(x1, x2) = maximum([pdf(dists[k], [x1,x2]) for k in 1:2])
        colors = map(1:n) do i
            (_, j) = findmax(Z[i,:])
            j == 1 ? :red : :blue
        end
        plt = Plots.contour(-2.8:0.05:6.0, -2.8:0.05:6.0, contour_mixture, contours = false, legend = nothing)
        Plots.scatter!(plt, X[:,1],X[:,2], label = "", color = colors)
        Plots.savefig(joinpath(base_img, "mixture_iter$(maxiter)_$iter.pdf"))
    end
end


using DiscreteFiltering
using OrdinaryDiffEq: ODEProblem, solve, Tsit5
using LinearAlgebra
using Plots
using SparseArrays

## Domain
a = 0.0
b = 2π
domain = PeriodicIntervalDomain(a, b)
# domain = ClosedIntervalDomain(a, b)

## Discretization
N = 100
M = N
x = discretize(domain, M)
ξ = discretize(domain, N)
Δx = (b - a) / N

## Filter
h₀ = 2.1Δx
h(x) = h₀ * (1 - 1 / 2 * cos(x))
dh(x) = h₀ / 2 * sin(x)
α(x) = 1 / 3 * dh(x) * h(x)
f = TopHatFilter(h)

# h₀ = 5.1Δx
# h(x) = h₀ # * (1 - 1 / 2 * cos(x))
# dh(x) = 0.0 # h₀ / 2 * sin(x)
# α(x) = 1 / 3 * dh(x) * h(x)
# f = GaussianFilter(h, Δx / 2)

## Time
T = 1.5
t = T

## Plot filter
plot(x, h)
ylims!((0, ylims()[2]))

## Get matrices
C = advection_matrix(domain, N)
D = diffusion_matrix(domain, N)
W = filter_matrix(f, domain, M, N)
R = reconstruction_matrix(f, domain, M, N)
A = spdiagm(α.(ξ))

## Inspect matrices
spy(W)
spy(R)

## Exact solutions
u(ξ, t) = sin(ξ - t) + 3 / 5 * cos(5(ξ - t)) + 1 / 25 * sin(20(ξ - 1 - t))
u_int(ξ, t) = -cos(ξ - t) + 3 / 25 * sin(5(ξ - t)) - 1 / 25 / 20 * cos(20(ξ - 1 - t))
ū(x, t) = 1 / 2h(x) * (u_int(x + h(x), t) - u_int(x - h(x), t))

## Discrete initial conditions
uₕ = u.(ξ, 0.0)
ūₕ = ū.(x, 0.0)
uₕ_allbar = W * uₕ

plot(x, uₕ, label = "Discretized")
plot!(x, ūₕ, label = "Filtered-then-discretized")
plot!(x, uₕ_allbar, label = "Discretized-then-filtered")

## Solve discretized problem
duₕ(uₕ, p, t) = -C * uₕ
prob = ODEProblem(duₕ, uₕ, (0, T))
sol = solve(prob, Tsit5(), abstol = 1e-6, reltol = 1e-4)

plot(x, uₕ, label = "Initial conditions")
plot!(x, sol(t), label = "Discretized")
plot!(x, u.(x, t), label = "Exact")

## Solve filtered-and-then-discretized problem
dūₕ(ūₕ, p, t) = (-C + A * D) * ūₕ
prob_bar = ODEProblem(dūₕ, ūₕ, (0, T))
sol_bar = solve(prob_bar, Tsit5(), abstol = 1e-6, reltol = 1e-4)

plot(x, uₕ, label = "Initial")
plot!(x, ūₕ, label = "Initial filtered")
plot!(x, sol_bar(t), label = "Filtered-then-discretized")
plot!(x, ū.(x, t), label = "Filtered exact")
# plot!(x, 500*α.(ξ))

## Solve discretized-and-then-filtered problem
duₕ_allbar(uₕ_allbar, p, t) = -W * (C * (R * uₕ_allbar))
# duₕ_allbar(uₕ_allbar, p, t) = -W * (C * (W \ uₕ_allbar))
prob_allbar = ODEProblem(duₕ_allbar, W * uₕ, (0, T))
sol_allbar = solve(prob_allbar, Tsit5(), abstol = 1e-6, reltol = 1e-4)

plot(x, uₕ, label = "Initial")
plot!(x, uₕ_allbar, label = "Initial discretized-then-filtered")
plot!(x, sol_allbar(t), label = "Discretized-then-filtered")
plot!(x, W * u.(ξ, t), label = "Exact")

## PGFPlotsX
pgfplotsx()

## Comparison
p = plot(
    size = (400, 300),
    xlabel = "\$x\$",
    legend = :topright,
)
plot!(p, x, uₕ, label = "\$u(x, t = 0.0)\$")
# plot!(p, x, ūₕ, label = "Initial filtered")
# plot!(p, x, sol(t), label = "Discretized")
# plot!(p, x, sol_bar(t), label = "Filtered-then-discretized")
# plot!(p, x, sol_allbar(t), label = "Discretized-then-filtered")
# plot!(p, x, [u.(x, t), ū.(x, t)], label = "Exact")
plot!(p, x, u.(x, t), label = "\$u(x, t = $t)\$")
plot!(p, x, ū.(x, t), label = "\$u(x, t = $t)\$")
# ylims!(p, minimum(uₕ), maximum(uₕ))
display(p)
savefig(p, "output/advection/solution.tikz")

## Relative error
u_exact = u.(ξ, t)
ū_exact = ū.(x, t)
err = abs.(sol(t) - u_exact) ./ maximum(abs.(u_exact))
err_bar = abs.(sol_bar(t) - u_exact) ./ maximum(abs.(u_exact))
err_allbar = abs.(sol_allbar(t) - u_exact) ./ maximum(abs.(u_exact))

plot()
plot!(ξ, err, label = "Unfiltered discretized")
plot!(x, err_bar, label = "Filtered-then-discretized")
plot!(x, err_allbar, label = "Discretized-then-filtered")


# # Variational problems
#
# In this example, we will numerically simulate an entropy-regularised Wasserstein gradient flow
# approximating the Fokker-Planck and porous medium equations.
#
# The connection between Wasserstein gradient flows and (non)-linear PDEs is due to Jordan, Kinderlehrer and Otto [^JKO98], and
# an easy-to-read overview of the topic is provided in Section 9.3 [^PC19]
#
# [^JKO98]: Jordan, Richard, David Kinderlehrer, and Felix Otto. "The variational formulation of the Fokker--Planck equation." SIAM journal on mathematical analysis 29.1 (1998): 1-17.
# [^PC19]: Peyré, Gabriel, and Marco Cuturi. "Computational optimal transport: With applications to data science." Foundations and Trends® in Machine Learning 11.5-6 (2019): 355-607.
#
# ## Fokker-Planck equation as a $W_2$ gradient flow
# For a potential function $\Psi$ and noise level $\sigma^2$, the Fokker-Planck equation (FPE) is
# ```math
# \partial_t \rho_t = \nabla \cdot (\rho_t \nabla \Psi) + \frac{\sigma^2}{2} \Delta \rho_t,
# ```
# and we take no-flux (Neumann) boundary conditions.
#
# This describes the evolution of a massless particle undergoing both diffusion (with diffusivity $\sigma^2$) and drift (along potential $\Psi$) according to the stochastic differential equation
# ```math
# dX_t = -\nabla \Psi(X_t) dt + \sigma dB_t.
# ```
# The result of Jordan, Kinderlehrer and Otto (commonly referred to as the JKO theorem) states that
# $\rho_t$ evolves following the 2-Wasserstein gradient flow of the Gibbs free energy functional
# ```math
#   F(\rho) = \int \Psi d\rho + \int \log(\rho) d\rho.
# ```
#
# ## Implicit schemes for gradient flows
# In an Euclidean space, the gradient flow of a functional $F$ is simply the solution of an ordinary differential equation
# ```math
#  \dfrac{dx(t)}{dt} = -\nabla F(x(t)).
# ```
# Of course, there is a requirement that $F$ is smooth. A more general formulation of a gradient flow that allows
# $F$ to be non-smooth is the implicit scheme
# ```math
#   x_{t+\tau} = \operatorname{argmin}_x \frac{1}{2} \| x - x_t \|_2^2 + \tau F(x).
# ```
# As the timestep $\tau$ shrinks, $x_t$ becomes a better and better approximation to the gradient flow of $F$.
#
# ## Wasserstein gradient flow
# In the context of the JKO theorem, we seek $\rho_t$ that is the gradient flow of $F$ with
# respect to the 2-Wasserstein distance. This can be achieved by choosing the $W_2$ metric in the implicit step:
# ```math
#   \rho_{t + \tau} = \operatorname{argmin}_{\rho} d_{W_2}^2(\rho_{t}, \rho) + \tau F(\rho).
# ```
# Finally, a numerical scheme for computing this gradient flow can be developed by using the entropic regularisation
# of optimal transport on a discretised domain
# ```math
#   \rho_{t + \tau} = \operatorname{argmin}_{\rho} \operatorname{OT}_\varepsilon(\rho_{t}, \rho) + \tau F(\rho),
# ```
# where
# ```math
#   \operatorname{OT}_\varepsilon(\alpha, \beta) = \min_{\gamma \in \Pi(\alpha, \beta)} \sum_{i,j} \frac{1}{2} \| x_i - x_j \|_2^2 \gamma_{ij} + \varepsilon \sum_{i, j} \gamma_{ij} \log(\gamma_{ij}).
# ```
# Each step of this problem is a minimisation problem with respect to $\rho$.
# Since we use entropic optimal transport which is differentiable, this can be solved using gradient-based methods.

# ## Problem setup
#
using OptimalTransport
using Distances
using LogExpFunctions
using Optim
using Plots
using StatsBase

using LinearAlgebra
using Logging

# Here, we set up the computational domain that we work on - we discretize the interval $[-1, 1]$.
# The natural boundary conditions to use will be Neumann (zero flux), see e.g. [^Santam2017]
#
# [^Santam2017]: Santambrogio, Filippo. "{Euclidean, metric, and Wasserstein} gradient flows: an overview." Bulletin of Mathematical Sciences 7.1 (2017): 87-154.

support = range(-1, 1; length=64)
C = pairwise(SqEuclidean(), support');

# Now we set up various functionals that we will use.
#
# We define the generalised entropy (Equation (4.4) of [^Peyre2015]) as follows. For $m = 1$ this is just the "regular" entropy, and $m = 2$ this is squared $L_2$.
#
# [^Peyre2015]: Peyré, Gabriel. "Entropic approximation of Wasserstein gradient flows." SIAM Journal on Imaging Sciences 8.4 (2015): 2323-2351.
function E(ρ; m=1)
    if m == 1
        return sum(xlogx.(ρ)) - sum(ρ)
    elseif m > 1
        return dot(ρ, @. (ρ^(m - 1) - m) / (m - 1))
    end
end;

# Now define $\psi(x)$ to be a potential energy function that has two potential wells at $x = ± 0.5$.
ψ(x) = 10 * (x - 0.5)^2 * (x + 0.5)^2;
plot(support, ψ.(support); color="black", label="Scalar potential")

# Having defined $\psi$, this induces a potential energy functional $\Psi$ on probability distributions $\rho$:
# ```math
#    \Psi(\rho) = \int \psi(x) \rho(x) dx = \langle \psi, \rho \rangle .
# ```
Ψ = ψ.(support);

# Define the time step $\tau$ and entropic regularisation level $\varepsilon$, and form the associated Gibbs kernel $K = e^{-C/\varepsilon}$.
τ = 0.05
ε = 0.01
K = @. exp(-C / ε)

# We define the (non-smooth) initial condition $\rho_0$ in terms of step functions.
H(x) = x > 0
ρ0 = @. H(support + 0.25) - H(support - 0.25)
ρ0 = ρ0 / sum(ρ0)
plot(support, ρ0; label="Initial condition ρ0", color="blue")

# `G_fpe` is the objective function for the implicit step scheme
# ```math
# G_\mathrm{fpe}(\rho) = \operatorname{OT}_\varepsilon(\rho_{t}, \rho) + \tau F(\rho),
# ```
# and we seek to minimise in $\rho$.
function G_fpe(ρ, ρ0, τ, ε, C)
    return sinkhorn2(ρ, ρ0, C, ε; regularization=true, maxiter=250) + τ * (dot(Ψ, ρ) + E(ρ))
end;

# `step` solves the implicit step problem to produce $\rho_{t + \tau}$ from $\rho_t$.
function step(ρ0, τ, ε, C, G)
    ## only print error messages
    opt = with_logger(SimpleLogger(stderr, Logging.Error)) do
        optimize(
            u -> G(softmax(u), ρ0, τ, ε, C),
            ones(size(ρ0)),
            LBFGS(),
            Optim.Options(; iterations=50, g_tol=1e-6);
            autodiff=:forward,
        )
    end
    return softmax(Optim.minimizer(opt))
end
# Now we simulate `N = 10` iterates of the gradient flow and plot the result.
N = 10
ρ = similar(ρ0, size(ρ0, 1), N)
ρ[:, 1] = ρ0
for i in 2:N
    ρ[:, i] = step(ρ[:, i - 1], τ, ε, C, G_fpe)
end
colors = range(colorant"red"; stop=colorant"blue", length=N)
plot(
    support,
    ρ;
    title=raw"$F(\rho) = \langle \psi, \rho \rangle + \langle \rho, \log(\rho) \rangle$",
    palette=colors,
    legend=nothing,
)

# ## Porous medium equation
#
# The porous medium equation (PME) is the nonlinear PDE
# ```math
# \partial_t \rho = \nabla \cdot (\rho \nabla \Psi) + \Delta \rho^m,
# ```
# again with Neumann boundary conditions. The value of $m$ in the PME corresponds to picking $m$ in the generalised entropy functional.
# Now, we will solve the PME with $m = 2$ as a Wasserstein gradient flow.
#
function G_pme(ρ, ρ0, τ, ε, C)
    return sinkhorn2(ρ, ρ0, C, ε; regularization=true, maxiter=250) +
           τ * (dot(Ψ, ρ) + E(ρ; m=2))
end;

# set up as previously
N = 10
ρ = similar(ρ0, size(ρ0, 1), N)
ρ[:, 1] = ρ0
for i in 2:N
    ρ[:, i] = step(ρ[:, i - 1], τ, ε, C, G_pme)
end
plot(
    support,
    ρ;
    title=raw"$F(\rho) = \langle \psi, \rho \rangle + \langle \rho, \rho - 1\rangle$",
    palette=colors,
    legend=nothing,
)


function parametrized_kuramoto(du,u,p,t)
    # NumberOfOscillators, CouplingConstant, and NaturalFreqs
    N,K,w = p
    # Centroid
    u[1] = 0
    for i in 1:N
        u[1] += cos(u[i+1]) + sin(u[i+1])*im
    end
    u[1] = u[1]/N
    r = abs(u[1])
    ψ = angle(u[1])
    # Oscillators
    for i in 1:N
    	du[i+1] = w[i] + K*r*sin(ψ-u[i+1])
    end
end

function getparameters_kuramoto(N,K,γ,distribution)
    distribution == "Uniform" ? nothing : error("$(distribution) is not yet supported")
    j = collect(1:N)
    w = broadcast(j->((2j-N-1)/(N-1))*((pi*γ)/2),j)
    p = (N=N, K=K, w=w)
    return p
end

function getstartpoints_kuramoto(N)
    d = Uniform(-pi,pi)
    u0::Array{ComplexF64} = rand(d,N)
    pushfirst!(u0,mean(u0))
    return u0
end

function kuramoto(N,K,T,fs,γ,distribution)
    tspan = (0.0,T)
    u0 = getstartpoints_kuramoto(N)
    p = getparameters_kuramoto(N,K,γ,distribution)
    prob = ODEProblem(parametrized_kuramoto,u0,tspan,p)
    sol = solve(prob,RK4(),saveat=collect(0:1/fs:T),progress=true)
end


# ---
# title: 96. Unique Binary Search Trees
# id: problem96
# author: zhwang
# date: 2022-02-23
# difficulty: Medium
# categories: Dynamic Programming, Tree
# link: <https://leetcode.com/problems/unique-binary-search-trees/description/>
# hidden: true
# ---
#
# Given _n_ , how many structurally unique **BST 's** (binary search trees) that
# store values 1 ...  _n_?
#
# **Example:**
#
#
#
#     Input: 3
#     Output: 5
#     Explanation: Given _n_ = 3, there are a total of 5 unique BST's:
#
#        1         3     3      2      1
#         \       /     /      / \      \
#          3     2     1      1   3      2
#         /     /       \                 \
#        2     1         2                 3
#
#
#
#
# **Constraints:**
#
#   * `1 <= n <= 19`
#
#
## @lc code=start
using LeetCode

num_of_trees(n::Int) = binomial(2 * n, n) - binomial(2 * n, n + 1)
## @lc code=end


module TF

import ..d; import ..dxasy

sigmoid{F <: AbstractFloat}(x::F) = one(F)/(one(F) + exp(-x))
dasysigmoid{F <: AbstractFloat}(y::F) = y * (one(F) - y)
#explicitly reinterpret the dasysigmoid.
dsigmoid{F <: AbstractFloat}(x::F) = dasysigmoid(sigmoid(x))

id{F <: AbstractFloat}(x::F) = x

softplus{F <: AbstractFloat}(x::F) = log(one(F) + exp(x))
dasysoftplus{F <: AbstractFloat}(y::F) = (exp(y) - one(F)) / exp(y)

#broadcast functions over arrays.
sigmoid{F <: AbstractFloat}(a::Array{F})      = sigmoid.(a)
dasysigmoid{F <: AbstractFloat}(a::Array{F})  = dasysigmoid.(a)
dsigmoid{F <: AbstractFloat}(a::Array{F})     = dsigmoid.(a)
id{F <: AbstractFloat}(a::Array{F})           = a
softplus{F <: AbstractFloat}(a::Array{F})     = softplus.(a)
dasysoftplus{F <: AbstractFloat}(a::Array{F}) = dasysoftplus.(a)

#non-simple transfer functions.
function softmax{F <: AbstractFloat}(a::Array{F})
  #go over the entire array and add up the values.
  totalexpsum = zero(F)
  for idx = 1:length(a)
    totalexpsum += exp(a[idx])
  end
  exp.(a) / totalexpsum
end
dsoftmax{F <: AbstractFloat}(a::Array{F}) = dasysigmoid.(softmax(a))

dxasy(::typeof(sigmoid)) = dasysigmoid
dxasy(::typeof(id)) = one
dxasy(::typeof(softplus)) = dasysoftplus
dxasy(::typeof(softmax)) = dasysigmoid

d(::typeof(sigmoid)) = dsigmoid
d(::typeof(id)) = one
d(::typeof(softplus)) = sigmoid
d(::typeof(softmax)) = dsoftmax

end


module supply_allocation

using JuMP
using Cbc
using CSV

function solve_supply_allocation()

    #------
    # DATA
    #------

    n = 4;
    m = 3;
    Q = [30 15 20 35];
    d = [40 10 35];

    c = [
        40	39	40;
        35	36	31;
        33	43	42;
        37	35	38
    ];

    #------
    # MODEL
    #------

    model = Model(with_optimizer(Cbc.Optimizer));

    @variable(model, x[1:n,1:m] >= 0)

    @objective(model, Min, sum(c[i,j] * x[i,j] for i in 1:n, j in 1:m));

    @constraint(model,[j = 1:m],sum(x[i,j] for i in 1:n) == d[j]);

    @constraint(model,[i = 1:n],sum(x[i,j] for j in 1:m) <= Q[i]);

    f = open("supply_allocation.lp", "w")
    print(f, model)
    close(f)

    #-------
    # SOLVE
    #-------

    optimize!(model)

    #------------------------
    # WRITE SOLUTION TO FILE
    #------------------------

    f = open("supply_allocation.csv","w");

    print(f,",");
    print(f,"Munich,Salzburg,Prague");
    println(f,"");

    ports = ["Rotterdam" "Wilhelmshaven" "Bremerhaven" "Hamburg" ];

    for i = 1:n
        print(f,ports[i],",");
        for j = 1:m
            print(f,value(x[i,j]),",");
        end
        println(f,"");
    end

    close(f);

end

solve_supply_allocation();

end


#The key to the non-parametric method describedhere
#is that the sum of squared distances between pointsand
#their centroid is equal to (and can be calculateddirectly from)
#the sum of squared interpoint distancesdivided by the number of points
# https://onlinelibrary.wiley.com/doi/epdf/10.1111/j.1442-9993.2001.01070.pp.x
SSₜ(D,N) = sum(D .^2)/N # given a triangular dist matrix
function SS_W(D,N,n)
    return  SSₜ(D,N)/n
end

function F(D,W,N,n,a)
    SST = SSₜ(D,N)
    SSW = SS_W(W,N,n)
    A = SST - SSW
    F = (A/(a-1))/(W/(N-a))
end

permutest(D,group,N,n,a, n_perm = 1000)
    N =  length(group)
    a = length(unique(group))
    n = N/a
    W_inds = [CartesianIndex(i,j) for i in 1:(N-1) for j in i+1:N if group[i] == group[j]]
    W = D[W_inds]
    notW = D[.!W_inds]
    D = vcat(W,notW)

    f = F(D,W,N,n,a)
    Fs = Vector{Float64}(undef,n_perm)
    inds = 1:length(W)

    for i in 1:n_perm
        shuffle!.(D)
        Fs[i] = Fknown(view(D,inds),N,n,a)
    end
    P = sum(Fs .>= f)/n_perm
    return P
end


#=
def sample_actual_demand(price):
    avg_demand = 65 + (-0.8) * price
    theta = 0.1/4
    k = avg_demand / theta
    return np.random.gamma(k*theta, k*theta**2, 1)[0]
=#

function sample_actual_demand(price)
    avg_demand = 65 + (-0.8) * price
    theta = 0.1/4
    k = avg_demand / theta
    return rand(Gamma(k*theta, k*theta^2), 1)
end

#=
def emperical_mean(sampler, n):
    mean = 0
    for i in range(1, n):
        mean = mean + sampler()
    return mean/n
=#

function emperical_mean(sampler, n)
    mean = 0
    for i in 1:n-1
        mean += sampler()[1]
    end
    return mean / n
end

#=
def emperical_demand_curve(min_price, max_price, n):
    prices = np.linspace(min_price, max_price, n)
    sampling = 5000
    demands = map(lambda p: emperical_mean(functools.partial(sample_actual_demand, p), sampling), prices)
    return np.dstack((prices, list(demands)))[0]
=#

function emperical_demand_curve(min_price, max_price, n)
    prices = range(min_price, max_price, length=n)
    sampling = 5000
    demands = []
    for p in prices
        sampler() = sample_actual_demand(p)
        demands = append!(demands, emperical_mean(sampler, sampling))
    end
    return DataFrame(
        :prices => prices,
        :demands => demands
        )
end


#############################
##  Numerical experiments associated with section 5.3
#############################
## This file contains the numerical experiments used for the highly oscillatory Bessel kernel example,
## which was originally featured by Nick Hale in https://doi.org/10.1093/imanum/dry042
##

using ApproxFun, MultivariateOrthogonalPolynomials, BandedMatrices, BlockBandedMatrices, SpecialFunctions, Plots
using SparseVolterraExamples

####
## First we define the Kernel and g(x) as in section 5.3.
## As mentioned in the paper, it is sensible to use pre-existing implementations to compute the Bessel functions
Kfun(x,y)=besselj(μ,ω*(x-y))
g(x) = besselj(μ+ν,ω*x)+1/(2*x^2)*((ν-1)*(ν-2)*besselj(ν-1,ω*x)+(ν+1)*(ν+2)*besselj(ν+1,ω*x))

#####
## Solver function with step-by-step explanation for given polynomial degree 'n'
## Note the factor 10^(-3) placed as discussed in the paper.
function solveSec54(n,g,Kfun,ν,μ,ω)
    gF = Fun(x->g(x),Jacobi(2,3, 0..1))
    V = triVolterraFullKernelOpP01(Kfun,n+2,true,339)
        V = reflectPabtoPba(n+2)*WLoweringP01P00(n+2)*V[1:n+2,1:n+2]
        V = Conversion(Jacobi(0,0,0..1),Jacobi(2,3,0..1))[1:n+2,1:n+2]*V[1:n+2,1:n+2]
        V = 10^(-3)*Derivative(Jacobi(0,1,0..1),2)[1:n+2,1:n+2]+ω^2*Conversion(Jacobi(0,1,0..1),Jacobi(2,3,0..1))[1:n+2,1:n+2]+ω*V
    coeff = [DirectEvalLHSP10at0(n+2);DirectEvalLHSP10atPrime0(n+2)*Derivative(Jacobi(0,1,0..1),1)[1:n+2,1:n+2];V[1:n,1:n+2]] \ [0.0;0.0;pad(gF.coefficients,n)]
    return Fun(Jacobi(0,1, 0..1), coeff)
end

####
## Compute the solution with same parameters as discussed by Hale and in our paper
## This correponds to computing the solution in Figure 7(a).
ν = 3
μ = 2
ω = 20
u = solveSec54(2000,g,Kfun,ν,μ,ω)
plot(u,color=:black, legend=:false, xlabel = "x", ylabel = "u(x)", legendfontsize=12, tickfontsize=10, thickness_scaling = 1 , grid=:none)

####
## A basic loop to compute errors between the above high order solution and increasingly accurate approximations
function errorvals(max,step,u)
    errorvec = []
    for n=10:step:max
        diff = maximum(abs.([maximum(solveSec54(n,g,Kfun,ν,μ,ω)-u),minimum(solveSec54(n,g,Kfun,ν,μ,ω)-u)]))
        errorvec = push!(errorvec,diff)
    end
    return errorvec
end
####
## Plot successive errors to produce something like Figure 7(b). The exact shape of the plot depends on how dense we make the steps in n.
## Note that for low polynomial orders, since the order is too low to resolve the oscillations, the error varies a lot.
## Only once the polynomial order can resolve the oscillations do we see sensible convergence.
errorvec = errorvals(450,40,u)
plot(Array(10:40:450),errorvec,color=:black, legend=:bottomleft , yscale=:log10 , xlabel = "n", label="sparse method" , ylabel = "error", legendfontsize=12, tickfontsize=10, thickness_scaling = 1.2 , grid=:none)

examples/Section 5-3 - Bessel kernels with highly oscillatory solutions.jl