include("../../../../std/util/src/util_module.jl") # for binary search
# include("../../datastructs/src/datastructs_module.jl") # for stack

module numint
using ..util
# using ..datastructs

struct Interval
    x_lo::Float64
    x_hi::Float64
    function Interval(x_lo, x_hi)
        (x_lo, x_hi) = (x_lo < x_hi) ? (x_lo, x_hi) : (x_hi, x_lo)
        return new(x_lo, x_hi)
    end
end
function range(interval::Interval)
    return interval.x_hi - interval.x_lo
end

function trapezoidal(fcn, interval::Interval, N::Int)

    int_range = range(interval)
    dx = int_range / N
    I  = 0.0
    x1 = interval.x_lo
    f1 = fcn(x1)
    @inbounds for kk = 1:N
        x2 = x1 + dx
        f2 = fcn(x2)
        I += 0.5 * (f1 + f2) * dx
        x1 = x2
        f1 = f2
    end
    return I
end

function trapezoidal(fcn, interval::Interval, xs::Vector)
    I  = 0.0
    x1 = interval.x_lo
    f1 = fcn(x1)
    @inbounds for kk = 2:length(xs)
        x2 = xs[kk]
        dx = x2 - x1
        f2 = fcn(x2)
        I += 0.5 * (f1 + f2) * dx
        x1 = x2
        f1 = f2
    end
    return I
end

struct RombergOptions
    N::Int
    abs_tol::Float64
    rel_tol::Float64
end
function RombergOptions()
    return RombergOptions(128, 1e-8, 1e-8)
end
function RombergOptions(;N = 128, abs_tol = 1e-8, rel_tol = 1e-8)
    return RombergOptions(N, abs_tol, rel_tol)
end

@enum StoppingCrit begin
    max_iter_reached
    rel_tol
    abs_tol
    unknown
end

struct RombergResult
    integral::Float64
    n_fcn_eval::Int
    stopping_crit::StoppingCrit
    I::Vector{Float64}
end

function romberg(
    fcn,
    interval::Interval,
    options::RombergOptions = RombergOptions()
    )

    # allocations
    int_range   = range(interval)
    Nmax        = ceil(Int, log2(options.N)) + 1
    I           = zeros(Nmax)

    # N == 2
    # stat
    n_fcn_eval    = 2
    stopping_crit = unknown::StoppingCrit

    idxi    = 1
    fi_a    = fcn(interval.x_lo)
    fi_b    = fcn(interval.x_hi)
    I[idxi] = 0.5 * (fi_a + fi_b) * int_range

    # R = datastructs.Stack{Float64}()
    R    = zeros(Nmax)
    R[1] = I[1]

    # auxilary arrays
    r   = zeros(Nmax)
    tmp = zeros(Nmax)

    # main loop
    @inbounds for N = 2:Nmax

        # fill up the trapezoidal table first column
        hn = int_range / 2^(N-1)

        # evaluate the function or steal it from the previous iteration
        I_partial = 0.0
        @inbounds for kk = 1:2^(N-2)
            x = interval.x_lo + hn * (2 * kk - 1)
            n_fcn_eval += 1
            I_partial  += fcn(x)
        end
        R[N] = 0.5 * R[N-1] + hn * I_partial

        copyto!(r, R)
        @inbounds for jj = 2:N
            @inbounds for ii = jj:N
                R_n_m1  = r[ii]
                R_n1_m1 = r[ii-1]
                tmp[ii] = R_n_m1 + (R_n_m1 - R_n1_m1) / (4^(jj-1)-1)
            end
            # swap the arrays such that the previous array is always in r
            r, tmp = tmp, r
        end

        idxi    += 1
        I[idxi] = r[N]
        if idxi >= 2
            dI = abs(I[idxi-1] - I[idxi])
            if dI <= options.abs_tol
                stopping_crit = abs_tol::StoppingCrit
                break
            end
            if  dI / abs(I[idxi]) <= options.rel_tol
                stopping_crit = rel_tol::StoppingCrit
                break
            end
        end

    end

    if idxi == Nmax
        stopping_crit = max_iter_reached::StoppingCrit
    end

    stat = RombergResult(I[idxi], n_fcn_eval, stopping_crit, I)
    return stat
end

struct GaussKronrodOptions
    abs_tol::Float64
    rel_tol::Float64
    min_interval::Float64
end
function GaussKronrodOptions(;
    abs_tol = 1e-8,
    rel_tol = 1e-8,
    min_interval = 1e-10
    )
    return GaussKronrodOptions(abs_tol, rel_tol, min_interval)
end

struct GaussKronrodResult
    integral::Float64
    n_fcn_eval::Int
end

function gauss_kronrod(
    fcn,
    interval::Interval,
    options::GaussKronrodOptions = GaussKronrodOptions()
    )

    (gi_idx,
    kronrod_weights,
    kronrod_nodes,
    gauss_weights) = __gk_nodes_weights()

    n_fcn_eval = 0
    x_lo = interval.x_lo
    x_hi = interval.x_hi

    (K15, n_fcn_eval) = __rec_gk(
        fcn,
        x_lo,
        x_hi,
        options,
        gi_idx,
        kronrod_weights,
        kronrod_nodes,
        gauss_weights,
        n_fcn_eval)

    stat = GaussKronrodResult(K15, n_fcn_eval)
    return stat
end

function __rec_gk(
    fcn,
    x_lo,
    x_hi,
    options,
    gi_idx,
    kronrod_weights,
    kronrod_nodes,
    gauss_weights,
    n_fcn_eval)

    # constants
    Nk = 15
    Ng = 7

    int_range = x_hi - x_lo
    mid_point = (x_lo + x_hi) / 2

    if int_range < options.min_interval
        return (0.0, n_fcn_eval)
    end

    fi = zeros(Nk)
    @inbounds for ii = 1:Nk
        xi = int_range/2 * kronrod_nodes[ii] + mid_point
        fi[ii] = fcn(xi)
    end
    n_fcn_eval += 15

    gi = fi[gi_idx]

    K15 = 0.0
    @inbounds for ii = 1:Nk
        K15 += kronrod_weights[ii] * fi[ii]
    end
    K15 *= int_range/2

    G7 = 0.0
    @inbounds for ii = 1:Ng
        G7 += gauss_weights[ii] * gi[ii]
    end
    G7 *= int_range/2

    error_est = abs(G7 - K15)
    if error_est <= options.abs_tol
        return (K15, n_fcn_eval)
    end
    if error_est / abs(K15) <= options.rel_tol
        return (K15, n_fcn_eval)
    end

    # subdivide the interval
    x_mid = mid_point

    (K15_left, n_fcn_eval) = __rec_gk(
        fcn,
        x_lo,
        x_mid,
        options,
        gi_idx,
        kronrod_weights,
        kronrod_nodes,
        gauss_weights,
        n_fcn_eval)

    (K15_right, n_fcn_eval) = __rec_gk(
        fcn,
        x_mid,
        x_hi,
        options,
        gi_idx,
        kronrod_weights,
        kronrod_nodes,
        gauss_weights,
        n_fcn_eval)

    return (K15_left + K15_right, n_fcn_eval)
end

function __gk_nodes_weights()
# gauss_nodes = [
    # -0.949107912342759  # 1
    # -0.741531185599394  # 2
    # -0.405845151377397  # 3
    # 0.0                 # 4
    # +0.405845151377397
    # +0.741531185599394
    # +0.949107912342759
    # ]

    gauss_weights = [
        0.129484966168870
        0.279705391489277
        0.381830050505119

        0.417959183673469

        0.381830050505119
        0.279705391489277
        0.129484966168870
    ]

    kronrod_nodes = [
        -0.991455371120813
        -0.949107912342759  # 1
        -0.864864423359769
        -0.741531185599394  # 2
        -0.586087235467691
        -0.405845151377397  # 3
        -0.207784955007898

        0.0                 # 4

        0.207784955007898
        0.405845151377397
        0.586087235467691
        0.741531185599394
        0.864864423359769
        0.949107912342759
        0.991455371120813
    ]

    kronrod_weights = [
        0.022935322010529
        0.063092092629979
        0.104790010322250
        0.140653259715525
        0.169004726639267
        0.190350578064785
        0.204432940075298

        0.209482141084728

        0.204432940075298
        0.190350578064785
        0.169004726639267
        0.140653259715525
        0.104790010322250
        0.063092092629979
        0.022935322010529
    ]

    gi_idx = [  2, 4, 6,
                8,
                10, 12, 14
            ]

    return (gi_idx, kronrod_weights, kronrod_nodes, gauss_weights)
end

end # module


using JuMP
using GLPK

model_auto = Model(GLPK.Optimizer)
@variable(model_auto, 0 <= x <= 1)
@variable(model_auto, 0 <= y <= 1)
@constraint(model_auto, x + y <= 1)
@objective(model_auto, Max, x + 2y)
optimize!(model_auto)
objective_value(model_auto)

model_auto_no = Model()
@variable(model_auto_no, 0 <= x <= 1)
@variable(model_auto_no, 0 <= y <= 1)
@constraint(model_auto_no, x + y <= 1)
@objective(model_auto_no, Max, x + 2y)
set_optimizer(model_auto_no, GLPK.Optimizer)
optimize!(model_auto_no)
objective_value(model_auto_no)

model_manual = Model(GLPK.Optimizer, caching_mode = MOIU.MANUAL)
@variable(model_manual, 0 <= x <= 1)
@variable(model_manual, 0 <= y <= 1)
@constraint(model_manual, x + y <= 1)
@objective(model_manual, Max, x + 2y)
MOIU.attach_optimizer(model_manual)
optimize!(model_manual)
objective_value(model_manual)

model_direct = direct_model(GLPK.Optimizer())
@variable(model_direct, 0 <= x <= 1)
@variable(model_direct, 0 <= y <= 1)
@constraint(model_direct, x + y <= 1)
@objective(model_direct, Max, x + 2y)
optimize!(model_direct)
objective_value(model_direct)

using Cbc

model = Model(optimizer_with_attributes(Cbc.Optimizer, "logLevel" => 0));

model = Model(optimizer_with_attributes(Cbc.Optimizer, "max_iters" => 10000));

model = Model(optimizer_with_attributes(Cbc.Optimizer, "seconds" => 5));

termination_status(model_auto)

display(typeof(MOI.OPTIMAL))

primal_status(model_auto)

dual_status(model_auto)

display(typeof(MOI.FEASIBLE_POINT))

@show value(x)
@show value(y)
@show objective_value(model_auto)


struct Dual{V<:Real} <: Real
    value::V
    partials::AbstractVector{V}
end
partials(d::Dual) = d.partials
value(d::Dual) = d.value

Dual(v::V,ps...) where V<:Real = Dual{V}(v,collect(ps))

function Base.show(io::IO, d::Dual{V}) where V
    print(io, "Dual(", d.value)
    for p in d.partials
        print(io,",",p)
    end
    print(io, ")")
end

function jacobian(f, x::AbstractVector{T}) where T
    duals = map(1:length(x)) do i
        ps = map(1:length(x)) do j
            i==j ? T(1) : T(0)
        end
        Dual(x[i], ps)
    end
    mapreduce(r->reshape(partials(r),1,:), vcat, f(duals))
end

##########  RULES  #############################################################

Base.:+(a::Dual, b::Dual) = Dual(a.value + b.value, a.partials + b.partials)
Base.:*(a::Dual, b::Dual) = Dual(a.value * b.value, b.value * a.partials + a.value * b.partials)

function cprod(x::AbstractVector)
    y = similar(x)
    y[1] = x[1]
    for i in 2:length(y)
        y[i] = y[i-1]*x[i]
    end
    y
end

x = collect(1:4)
#csum(x) |> display

dx = [Dual(x[1], 1, 0, 0),
      Dual(x[2], 0, 1, 0),
      Dual(x[3], 0, 0, 1)]
cprod(dx) |> display
csum(dx)

jacobian(cprod, x)


# Test for solving div-grad system both complex and real and single and multiple rhs
using MUMPSjInv
using Test
using LinearAlgebra
using SparseArrays

include("getDivGrad.jl");

A = getDivGrad(32,32,16);
n = size(A,1);

# REAL: test for single rhs
println("Test for real SPD matrix: one rhs");
rhs = randn(n);

x = solveMUMPS(A,rhs,1);
err  =  norm(A*x-rhs) / norm(rhs);
@test err < 1e-14

# REAL: test for multiple rhs
println("Test for real SPD matrix: multiple rhs");
nrhs = 10;
rhs = randn(n,nrhs);

x = solveMUMPS(A,rhs,1);

err = zeros(nrhs)
for i=1:nrhs
        err[i] =  norm(A*x[:,i]-rhs[:,i]) / norm(rhs[:,i]);
end
@test eltype(x) == Float64
@test maximum(err) < 1e-14

# COMPLEX: test for multiple rhs
println("Test for complex symmetric matrix: one rhs");
r = rand(n);
A = A + im*spdiagm(0=>r);

rhs = randn(n) + im*randn(n);

x = solveMUMPS(A,rhs,1);

err=  norm(A*x-rhs) / norm(rhs);
@test eltype(x) == ComplexF64
@test err < 1e-14

# COMPLEX : test for multiple rhs
println("Test for complex symmetric matrix: multiple rhs");
nrhs = 10;
rhs = randn(n,nrhs) + im*randn(n,nrhs);

x = solveMUMPS(A,rhs,2);
err = zeros(nrhs)
for i=1:nrhs
        err[i] =  norm(A*x[:,i]-rhs[:,i]) / norm(rhs[:,i]);
end
@test eltype(x) == ComplexF64
@test maximum(err) < 1e-14

println("DONE!")


using LinearAlgebra
using Random
using ForwardDiff
using DataFrames
using GLM

function marquardt(f, x₀; ε=1e-6, maxiter=1000, λ=10.0^4, α=2)
    x = x₀
    fx = f(x)
    converged = false
    i = 0
    while i < maxiter
        i += 1
        ∇f = ForwardDiff.gradient(f, x)
        if norm(∇f) ≤ ε
            converged = true
            break
        end
        ∇²f = ForwardDiff.hessian(f, x)
        x′ = x - (∇²f + λ*I) \ ∇f
        fx′ = f(x′)
        if fx′ < fx
            λ *= 0.5
            fx = fx′
            x = x′
        else
            λ *= 2.0
        end
    end
    (x=x, fx=fx, converged=converged, iters=i)
end

rosenbrock(x) = sum((1-x[i])^2 + 100(x[i+1]-x[i]^2)^2 for i in 1:length(x)-1)

Random.seed!(1234)
x = rand(20);
marquardt(rosenbrock, x)

Random.seed!(123)
N = 1000
df = DataFrame(rand(N, 4))
df.y = sum.(eachrow(df)) + randn(N);
lm(@formula(y~x1+x2+x3+x4), df)

# This is not very efficient, but is good enough in our case
function ssq(x)
    m = Matrix(df)[:, startswith.(string.(names(df)), "x")]
    norm(df.y - fill(x[1], nrow(df)) - m * x[2:end])
end
marquardt(ssq, rand(5))


using VectorPartitions
using Test

@testset "VectorPartitions.jl" begin
using Combinatorics

function count_partition_iterator(n,k)
  P = all_vector_partitions(n,k)
  x = zeros(Int,n)
  count = 0
  for partition in P
    x .= partition
    count += 1
  end
  return count
end

function count_partitions_raw(n,k)
  P = all_vector_partitions(n,k)
  x = zeros(Int,n)
  next = iterate(P)
  count = 0
  while next !== nothing
      (x, state) = next
      next = iterate(P, state)
      count += 1
  end
  return count
end

function count_partitions_stateful(n,k)
  P = all_vector_partitions(n,k)
  next = iterate(P)
  if next !== nothing
    (_,state) = next
  else
    return nothing
  end
  count = 0
  while !state.done
    increment_vector_partition_state!(P,state)
    count += 1
  end
  return count
end

# some spot checks that we get the right number of partitions
@test count_partition_iterator(5,2) == stirlings2(5,2)
@test count_partitions_raw(5,2) == stirlings2(5,2)
@test count_partitions_stateful(5,2) == stirlings2(5,2)

@test count_partition_iterator(5,3) == stirlings2(5,3)
@test count_partitions_raw(5,3) == stirlings2(5,3)
@test count_partitions_stateful(5,3) == stirlings2(5,3)

# do we produce the right list of partitions in this case?
three_two_partitions = [[1,1,2],[1,2,1],[1,2,2]]
for (p1, p2) in zip(all_vector_partitions(3,2),three_two_partitions)
  @test all(p1 .== p2)
end

# test increment_vector_partition_state!
I = all_vector_partitions(7,3)
_, state = iterate(I)
state.partition .= [1, 1, 2, 3, 3, 3, 3]
state.upper_bound .= [1, 1, 2, 3, 3, 3, 3]
state.first_index .= [1, 3, 4]
increment_vector_partition_state!(I,state)
@test all(state.partition .== [1,2,1,3,1,1,1])
@test all(state.upper_bound .== [1,2,2,3,3,3,3])
@test all(state.first_index .== [1,2,4])

increment_vector_partition_state!(I,state)
@test all(state.partition .== [1,2,2,3,1,1,1])
@test all(state.upper_bound .== [1,2,2,3,3,3,3])
@test all(state.first_index .== [1,2,4])

# test increment_surjective_monotone_sequence!
seq = [1, 2, 2, 3, 4]
step_locations = [1,2,4,5]
stop = VectorPartitions.increment_surjective_monotone_sequence!(seq,4,step_locations)
@test all(seq .== [1, 2, 3, 3, 4])
@test all(step_locations .== [1,2,3,5])
@test stop == false

seq = [1,2,3,4,4]
step_locations = [1,2,3,4]
stop = VectorPartitions.increment_surjective_monotone_sequence!(seq,4,step_locations)
@test all(seq .== [1, 2, 3, 4, 4])
@test all(step_locations .== [1,2,3,4])
@test stop == true

end


"""
    continuedfraction(x::arb)
Compute a continued fraction representation of `x`.
"""
function continuedfraction(x::arb)
    i = floor(x)

    if isint(i)
        return [i; continuedfraction(inv(x - i))]
    else
        return [i]
    end
end

"""
    convergents(expansion::Vector{T}) where {T <: Integer}
Compute the convergents given a continued faction expansion `as`.
"""
function convergents(as::Vector{T}) where {T<:Integer}
    hs = zeros(T, length(as))
    ks = zeros(T, length(as))

    hs[1] = as[1]
    ks[1] = one(T)

    hs[2] = as[1] * as[2] + 1
    ks[2] = as[2]

    for n = 3:length(as)
        hs[n] = as[n] * hs[n-1] + hs[n-2]
        ks[n] = as[n] * ks[n-1] + ks[n-2]
    end

    return hs .// ks
end

"""
    convergents(expansion::Vector{arb})
Compute the convergents given a continued faction expansion `as`. The
computations are done using ball arithmetic and it returns a list of
pairs corresponding to the nominator and denominator of the
convergents.
"""
function convergents(as::Vector{arb})
    hs = similar(as)
    ks = similar(as)

    hs[1] = as[1]
    ks[1] = one(as[1])

    hs[2] = as[1] * as[2] + 1
    ks[2] = as[2]

    for n = 3:length(as)
        hs[n] = as[n] * hs[n-1] + hs[n-2]
        ks[n] = as[n] * ks[n-1] + ks[n-2]
    end

    return collect(zip(hs, ks))
end


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


include("prime.jl")
function count_consecutive_primes(a::Int,b::Int)
	index=0
	while true
		n=index^2+index*a+b
		if !is_prime(n)
			return index
		end
		index+=1
	end
end

function compute()
	len=0;
	ans=0;
	for i=-999:1000
		for j=2:1000
			if( count_consecutive_primes(i,j)>len)
				len=count_consecutive_primes(i,j)
				ans=i*j
			end
		end
	end
	return ans

end

println(compute())


# 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)
=#


using HypothesisTests
using Test
using StatsBase
using Statistics
using DelimitedFiles

@testset "Equality of Variances" begin

    # https://en.wikipedia.org/wiki/One-way_analysis_of_variance#Example
    groups = [
        [6, 8, 4, 5, 3, 4],
        [8, 12, 9, 11, 6, 8],
        [13, 9, 11, 8, 7, 12]
    ]
    t = OneWayANOVATest(groups...)
    @test nobs(t) == fill(6, 3)
    @test dof(t) == (2,15)
    @test pvalue(t) ≈ 0.002 atol=1e-3
    @test occursin("reject h_0", sprint(show, t))

    show(IOContext(IOBuffer(), :table => true), t)
    show(IOBuffer(), t)

    # http://www.real-statistics.com/one-way-analysis-of-variance-anova/confidence-interval-anova/
    groups = [
        [51, 87, 50, 48, 79, 61, 53],
        [82, 91, 92, 80, 52, 79, 73, 74],
        [79, 84, 74, 98, 63, 83, 85, 58],
        [85, 80, 65, 71, 67, 51],
    ]
    t = OneWayANOVATest(groups...)
    @test nobs(t) == [7, 8, 8, 6]
    @test dof(t) == (3,25)
    @test pvalue(t) ≈ 0.072 atol=1e-3
    @test occursin("reject h_0", sprint(show, t))

    # http://www.real-statistics.com/one-way-analysis-of-variance-anova/homogeneity-variances/levenes-test/
    groups = [
        [51, 87, 50, 48, 79, 61, 53, 54],
        [82, 91, 92, 80, 52, 79, 73, 74],
        [79, 84, 74, 98, 63, 83, 85, 58],
        [85, 80, 65, 71, 67, 51, 63, 93],
    ]
    # with means
    l = LeveneTest(groups...; statistic=mean)
    @test nobs(l) == fill(8, 4)
    @test dof(l) == (3,28)
    @test pvalue(l) ≈ 0.90357 atol=1e-4
    # with medians
    l = LeveneTest(groups...; statistic=median)
    @test pvalue(l) ≈ 0.97971 atol=1e-4
    # with 10% trimmed means
    l = LeveneTest(groups...; statistic=v->mean(trim(v, prop=0.1)))
    @test pvalue(l) ≈ 0.90357 atol=1e-4

    # Fligner-Killeen Test
    t = FlignerKilleenTest(groups...)
    @test nobs(t) == fill(8, 4)
    @test dof(t) == 3
    @test pvalue(t) ≈ 0.9878 atol=1e-4
    @test HypothesisTests.teststatistic(t) ≈ 0.1311 atol=1e-5
    @test occursin("reject h_0", sprint(show, t))

    # https://www.itl.nist.gov/div898/handbook/eda/section3/eda35a.htm

    # Columns are gear diameter and batch number
    gear = readdlm(joinpath(@__DIR__, "data", "gear.txt"))
    samples = reshape(gear[:,1],:,10)
    groups2 = tuple((samples[:,i] for i in 1:size(samples,1))...)

    # Brown-Forsythe Test
    l = BrownForsytheTest(groups2...)
    @test nobs(l) == fill(10, 10)
    @test dof(l) == (9,90)
    @test HypothesisTests.teststatistic(l) ≈ 1.705910 atol=1e-5
    @test pvalue(l) ≈ 0.0991 atol=1e-4
    @test occursin("reject h_0", sprint(show, l))
end


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


"""
## glob_to_axial

This subroutine transforms the global end reactions
into an axial force for rod fin_els (2- or 3-d).

### Function
```julia
glob_to_axial(glob, coord)
```

### Arguments
```julia
* glob::Vector{Float64}      : Globale forces and moments
* coord::Matrix{Float64}     : Nodal coordinates
```

### REturn value
```julia
* ::Float64                  : Axial force
```
"""
function glob_to_axial(glob::Vector{Float64}, coord::LinearAlgebra.Adjoint{Float64,Array{Float64,2}})
  #
  # This subroutine transforms the global end reactions
  # into an axial force for rod fin_els (2- or 3-d).
  #
  local axial = 0.0
  local ndim = size(coord,2)
  local add = 0.0
  for i in 1:ndim
    add = add+(coord[2,i]-coord[1,i])^2
  end
  ell = sqrt(add)
  for i in 1:ndim
    axial = axial+(coord[2,i]-coord[1,i])/ell*glob[ndim+i]
  end
  axial
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


using GSL, SpecialFunctions
using Test, LinearAlgebra

@testset "Special functions" begin

    @testset "bessel functions" begin

        x = rand() - 0.5 + (rand() - 0.5) * im

        @test diffbessel(besselj,2,x,2) ≈ diffbesselj(2,x,2)
        @test diffbessel(hankelh1,2,x,2) ≈ diffhankelh1(2,x,2)

        @test diffbessel(besselj,4,x,3) ≈ diffbesselj(4,x,3)
        @test diffbessel(hankelh1,4,x,3) ≈ diffhankelh1(4,x,3)

        # spherical bessel functions
        @test shankelh1(1, x) ≈ - exp(im*x) * (x + im) / (x^2)
        @test sbesselj(1, x) ≈ sin(x)/x^2 - cos(x)/x

        @test sbesselj(1, eps(zero(Float64))) ≈ 0.0
        @test sbesselj(0, eps(zero(Float64))) ≈ 1.0

        n = 1
        @test 2 * diffsbessel(shankelh1,n,x) ≈ shankelh1(n-1, x) - (shankelh1(n, x) + x*shankelh1(n+1, x))/x

        @test diffsbessel(shankelh1,n,x) ≈ diffshankelh1(n,x)
        @test diffsbessel(sbesselj,n,x) ≈ diffsbesselj(n,x)

        n = 2
        @test 2 * diffsbessel(shankelh1,n,x) ≈ shankelh1(n-1, x) - (shankelh1(n, x) + x*shankelh1(n+1, x))/x

        xs = 0.01:0.9:11.0;
        n = 15; ns = 0:n

        # Compare with spherical bessel from GSL, which can only accept real arguments..
        errs = map(xs) do x
            maximum(abs.(sf_bessel_jl_array(n,x) - [sbesselj(n,x) for n in ns]))
        end;
        @test maximum(errs) < 20*eps(Float64)

        errs = map(xs) do x
            maximum(
                abs.(
                    im .* sf_bessel_yl_array(n,x) -
                    [shankelh1(n,x) for n in ns] +
                    sf_bessel_jl_array(n,x)
                ) ./ abs.(sf_bessel_yl_array(n,x))
            )
        end;
        @test maximum(errs) < 1e4*eps(Float64)

    end

    @testset "Associated legendre functions" begin
        x = rand(1)[1] * 0.99
        l_max = 2

        # a few associated legendre functions without the Condon-Shortly factor
        Plm_arr = [1,x,sqrt(1-x^2),(3x^2-1)/2, 3x*sqrt(1-x^2),3*(1-x^2)]

        ls, ms = associated_legendre_indices(l_max)

        @test sf_legendre_array(GSL_SF_LEGENDRE_NONE, l_max, x)[1:length(ls)] ≈ Plm_arr

        sph_factors = (-1).^ms .* sqrt.((2 .* ls .+ 1) ./ (4pi) .* factorial.(ls-ms) ./ factorial.(ls+ms))

        condon_phase = (-1).^ms

        @test condon_phase .* sf_legendre_array(GSL_SF_LEGENDRE_SPHARM, l_max, x)[1:length(ls)] ≈ sph_factors .* Plm_arr

    end

    @testset "Spherical harmonics" begin
        θ = rand(1)[1] * 0.99
        φ = rand(1)[1] * 0.99

        l_max = 8 # small l_max due to factorial formula below

        ls, ms = spherical_harmonics_indices(l_max)
        sphs = spherical_harmonics(l_max, θ, φ)

        @test maximum(i - lm_to_spherical_harmonic_index(ls[i],ms[i]) for i in eachindex(ls)) == 0

        # check special case l == abs(m)
        is = findall(ls .== abs.(ms))

        for i in is
            @test sphs[i] ≈ (- sign(ms[i]))^ls[i] / (2^ls[i] * factorial(ls[i])) *
                sqrt(factorial(2*ls[i] + 1) / (4pi)) * sin(θ)^ls[i] * exp(im * ms[i] * φ)
        end

        l_max = 30
        ls, ms = spherical_harmonics_indices(l_max)
        sphs = spherical_harmonics(l_max, θ, φ)

        # special case m == 0, reduce to just Legendre polynomials
        Ps = sf_legendre_Pl_array(l_max, cos(θ))
        is = findall(ms .== 0)

        for l in 0:l_max
            @test sphs[is][l+1] ≈ sqrt((2l+1)/(4pi)) * Ps[l+1]
        end

        #sphs[inds] .≈ sqrt.((2 .* (0:l_max) .+ 1) ./ (4pi)) .* Ps

        # special case, north pole
        θ = 0.0
        sphs = spherical_harmonics(l_max, θ, φ)

        for i in eachindex(sphs)
            if ms[i] == 0
                @test sphs[i] ≈ sqrt((2*ls[i] + 1)/(4pi))
            else
                @test sphs[i] ≈ 0.0
            end
            i += 1
        end

    end

    @testset "Gaunt and Wigner symbols" begin
        l1 = rand(1:100)
        l2 = rand(1:100)
        l3 = rand(abs(l1-l2):2:(l1+l2)) # guarantees that iseven(l1+l2+l3)

        @test iseven(l1+l2+l3)

        m1 = rand(-l1:l1)
        m2 = rand(-l2:l2)
        m3 = m1-m2

        @test_throws(DomainError,gaunt_coefficient(l1,2*l1,l2,m2,l3,m3))
        @test_throws(MethodError,gaunt_coefficient(l1,m1,l2,m2,0.1,m3))

        # the spherical harmonics linearisation formula
        θ, φ = rand(2) .* pi

        l_small = 6
        l_max = 2*l_small # needs to be larger than l_small

        ls, ms = spherical_harmonics_indices(l_max);
        Ys = spherical_harmonics(l_max, θ, φ);

        cs = reshape(
            [
                gaunt_coefficient(l1,m1,l2,m2,l3,m3)

                    for l3 = 0:l_max for m3 = -l3:l3
                for l2 = 0:l_small for m2 = -l2:l2
            for l1 = 0:l_small for m1 = -l1:l1]
        , ((l_small + 1)^2,(l_small + 1)^2,(l_max+1)^2));

        for n2 in 1:(l_small + 1)^2, n3 in 1:(l_small + 1)^2
            @test 4pi * (-1)^(ms[n3]+ms[n2]) * (1.0im)^(ls[n3]-ls[n2]) * Ys[n3] * conj(Ys[n2]) ≈
                sum( (1.0im).^(-ls) .* (-1.0).^ms .* conj.(Ys) .* cs[n2,n3,:]) atol = 1e-12
        end

    end

    @testset "complex radial coordinates" begin

    # Test 3-dimensional transforms
        xs = [rand(-1.01:0.1:1.0,3) + rand(-1.01:0.1:1.0,3)*im for i = 1:100]
        rθφs = cartesian_to_radial_coordinates.(xs)
        @test maximum(norm.(xs - radial_to_cartesian_coordinates.(rθφs))) < 2e-14

        xs = [rand(-1.01:0.1:1.0,3) for i = 1:100]
        rθφs = cartesian_to_radial_coordinates.(xs)

        @test maximum(rθφ[2] for rθφ in rθφs) <= pi
        @test minimum(rθφ[2] for rθφ in rθφs) >= 0.0

        @test pi/2 < maximum(rθφ[3] for rθφ in rθφs) <= pi
        @test -pi <= minimum(rθφ[3] for rθφ in rθφs) < -pi/2

    # Test 2-dimensional transforms
        xs = [rand(-1.01:0.1:1.0,2) + rand(-1.01:0.1:1.0,2)*im for i = 1:100]
        rθs = cartesian_to_radial_coordinates.(xs)

        @test maximum(norm.(xs - radial_to_cartesian_coordinates.(rθs))) < 1e-14
    end
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")


#solve linear equation

#Example1
a=float([0 4 5 2
         1 0 2 -6;
         4 1 0 -2;
         1 7 1 0])::Array{Float64,2}
b=float([9,-3, 1, -3])::Array{Float64,1}
x=float([0, 0, 0, 0])::Array{Float64,1}

function main(a,b,x)
    ori_a=a
    a=hcat(a,b)
    row,col=size(a)
    for j in 1:row
        #search pivot
        maxidx=j-1+indmax(abs.(a[j:end,j]))
        #swap
        a[j,:],a[maxidx,:]=a[maxidx,:],a[j,:]
        #push foward
        for i in j+1:row
            p=-a[i,j]/a[j,j]
            a[i,:]+=p*a[j,:]
        end
    end
    #back foward
    for i in row:-1:1
        x[i]=(a[i,col]-dot(a[i,i+1:row],x[i+1:row]))/a[i,i]
    end
    println("answer=",x)
    println("check ax-b=",ori_a*x-b)
end

main(a,b,x)


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