# Exercise 1

using LinearAlgebra
using StaticArrays
using BenchmarkTools

N = 3
A = rand(N, N)
x = rand(N)

@btime $A * $x  # the $ in front of variable names is sometimes important
@btime inv($A)

# Benchmark

A_s = @SMatrix rand(N,N)
x_s =  @SVector rand(N)

@btime  $A_s * $x_s
@btime inv($A_s);

# very fast matrix multiplication, use $ and denote as static vectors
@btime $A_s * $x_s
@show supertypes(typeof(A_s))
@show supertypes(typeof(A));

# Exercise 2

Σ = [0.4  0.3;
     0.3  0.45]
G = I
R = 0.5 * Σ

gain(Σ, G, R) = Σ * G' * inv(G * Σ * G' + R)
@btime gain($Σ, $G, $R)

# Benchmark

Σ_s = @SMatrix [0.4  0.3; 0.3  0.45]
R_s = 0.5 * Σ_s

@show typeof(R) #static array, too

@btime gain($Σ_s, $G, $R_s);

# Exercise 3

using Polynomials

p = Polynomial([2, -5, 2], :x)  # :x just gives a symbol for display

@show p
@show typeof(p)
p′ = derivative(p)   # gives the derivative of p, another polynomial
@show p(0.1), p′(0.1)  # call like a function
@show roots(p);   # find roots such that p(x) = 0

# plot p and p'

r = range(-2.0, 2.0, length = 20)

p_r = p.(r)
p′_r = p′.(r)

plot(r[1:end-1],p′_r[1:end-1],label="Dp(x)")
plot!(r,p_r,label="p(x)")

# Exercise 4

function newtonsmethod(p::Polynomial, x_0; tolerance = 1E-7, maxiter = 100)

    p′ = derivative(p)

    #initialise loop
    iter = 1
    normdiff = Inf
    x_old = x_0

    while normdiff > tolerance && iter < maxiter

        x_new = x_old - p(x_old)/ p′(x_old)
        normdiff = norm(x_new - x_old)

        iter += 1 #continuation
        x_old = x_new
    end
    return (root = x_old, normdiff = normdiff, iter = iter )
end

# Compute roots

p = Polynomial([2, -5, 2], :x)  # :x just gives a symbol for display

@show newtonsmethod(p, 5);   # find roots such that p(x) = 0
@show roots(p)
@btime newtonsmethod(p, 5);
@btime roots(p);

# Exercise 5
function trapezoidal(f::AbstractArray{Float64,1}, x::AbstractArray{Float64,1})
    sol = similar(x)
    del = diff(x)
    for i in 2:length(x)
        sol[i] = (f_x[i]+f_x[i-1])/2*del[i-1]
    end
    return sum(sol)
end

x = range(0.0,1.0,length=9)
f(x) = x^2
f_x = f.(x)

int = trapezoidal(f_x,x)

function trapezoidal2(f::AbstractArray{Float64}, x::AbstractRange{Float64})
    sol = similar(x)
    del = step(x)
    for i in 2:length(x)
        sol[i] = (f_x[i]+ f_x[i-1])/2*del
    end
    return sum(sol)
end

x = range(0.0,1.0,length=1100)
f(x) = x^2
f_x = f.(x)

integral = trapezoidal2(f_x,x)

f(x) = x^3
x̄ = 1.0
x̲ = 0.0
N = 1100;

function trapezoidal3(f, x̲::Real , x̄::Real , N::Int64)
    x = range(x̲, x̄, length = N)
    f_x = f.(x)

    #implementation the same as trapezodal2
    sol = similar(x)
    del = step(x)
    for i in 2:length(x)
        sol[i] = (f_x[i]+ f_x[i-1])/2*del
    end
    return sum(sol)
end

integral = trapezoidal3(f ,x̲ , x̄ , N)

# Exercise 6

using ForwardDiff

function f(a, b; N::Int64 = 50)
    r = range(a, b, length=N) # one
return mean(r)
end

Df(x) = ForwardDiff.derivative(y -> f(0.0, y), x)

@show f(0.0, 3.0)
@show f(0.0, 3.1)

Df(3.0)

function f(a, b; N::Int64 = 50)
    r = range(a, b, length=N) # one
return mean(r)
end


"""
    harmonic(n::Integer)

 Calculates harmonic numbers,
   e.g., see [http://mathworld.wolfram.com/HarmonicNumber.html](http://mathworld.wolfram.com/HarmonicNumber.html)

## Arguments
* ``n`` `::Integer`: non-negative index of the Harmonic number to calculate

## Examples
```jldoctest; setup = :(using Polylogarithms)
julia> harmonic(2)
1.5
```
"""
function harmonic(n::Integer)
    if n < 0
        throw(DomainError(n))
    elseif n==0
        return 0.0 # by convention (see Crandall, p.22)
    elseif n <= 10
        # perform exact sum for small n
        total = 0.0
        for k=1:n
            total +=  1.0 / k
        end
        return total
    else
        return γ + SpecialFunctions.digamma(n+1) # digamma(m) = ψ(m)
    end
end

"""
    harmonic(x::ComplexOrReal{Float64})

 Calculates harmonic numbers extended to non-integer arguments using the
 digamma form.

## Arguments
* ``x`` `::ComplexOrReal{Float64}`: index of the Harmonic number to calculate

## Examples
```jldoctest; setup = :(using Polylogarithms)
julia> harmonic(2.0)
1.5000000000000016
```
"""
function harmonic(x::ComplexOrReal{Float64})
    return γ + SpecialFunctions.digamma(x+1)
end
#        ``digamma(m,x) = (-1)^{m+1} m! \zeta(m+1,x), m>=1``, and integer

"""
    harmonic(n::Integer,r::Real)

 Calculates generalized harmonic numbers,
   e.g., see [http://mathworld.wolfram.com/HarmonicNumber.html](http://mathworld.wolfram.com/HarmonicNumber.html)

## Arguments
* ``n`` `::Integer`: non-negative index 1 of the Harmonic number to calculate
* ``r`` `::Real`: index 2 of the Harmonic number to calculate

It should be possible to extend this to complex r, but that requires more testing.

## Examples
```jldoctest; setup = :(using Polylogarithms)
julia> harmonic(2,1.5)
1.3535533905932737
```
"""
function harmonic(n::Integer, r::Real)
    if n < 0
        throw(DomainError(n))
    end
    if n == 0
        return 0.0
    end
    if r==1
        return harmonic(n)
    end
    total = 0.0
    for k=1:n
        total +=  1.0 / k^r
    end
    return total
end

"""
    harmonic(n::Integer,r::Integer)

 Calculates generalized harmonic numbers
   e.g., see [http://mathworld.wolfram.com/HarmonicNumber.html](http://mathworld.wolfram.com/HarmonicNumber.html)
 using a better approach which works when both inputs are integers
 [https://carma.newcastle.edu.au/resources/jon/Preprints/Papers/Published-InPress/Oscillatory%20(Tapas%20II)/Papers/coffey-zeta.pdf](https://carma.newcastle.edu.au/resources/jon/Preprints/Papers/Published-InPress/Oscillatory%20(Tapas%20II)/Papers/coffey-zeta.pdf), p.341

## Arguments
* ``n`` `::Integer`: non-negative index 1 of the Harmonic number to calculate
* ``r`` `::Integer`: index 2 of the Harmonic number to calculate

## Examples
```jldoctest; setup = :(using Polylogarithms)
julia> harmonic(2,1)
1.5000000000000002
```
"""
function harmonic(n::Integer, r::Integer)
    if r<1
        throw(DomainError(r))
    end
    return (-1)^(r-1) * ( SpecialFunctions.polygamma(r-1,n+1) - SpecialFunctions.polygamma(r-1,1) ) / SpecialFunctions.gamma(r)
end

# # better approach to calculation, but doesn't seem to work
# #  using
# function harmonic3(x::Real, r::Integer)
#     if (r < 1)
#         throw(DomainError(r))
#     end
#     return SpecialFunctions.zeta(r,1) - SpecialFunctions.zeta(r,x+1)
# end


using OrdinaryDiffEq
using DiffEqFlux, Flux
using Random
using Plots;

Random.seed!(123);

function sir_ode(u,p,t)
    (s,i,r) = u
    (β,γ) = p
    ds = -β*s*i
    di = β*s*i - γ*i
    dr = γ*i
    [ds,di,dr]
end;

p = [0.5,0.25]
u0 = [0.99, 0.01, 0.0]
tspan = (0.0, 40.0)
δt = 1;

solver = Rodas5();

sir_prob = ODEProblem(sir_ode, u0, tspan, p)
sir_sol = solve(sir_prob, solver, saveat = δt);

train_time = 30.0
train_data = Array(sir_sol(0:δt:train_time));

nhidden = 8
sir_node = FastChain(FastDense(3, nhidden, tanh),
                     FastDense(nhidden, nhidden, tanh),
                     FastDense(nhidden, nhidden, tanh),
                     FastDense(nhidden, 3));

p_ = Float64.(initial_params(sir_node));
function dudt_sir_node(u,p,t)
    s,i,r = u
    ds,di,dr = ann_node([s,i,r],p)
    [ds,di,dr]
end
prob_node = ODEProblem(dudt_sir_node, u0, tspan, p_);

tspan_train = (0,train_time)
prob_node = NeuralODE(sir_node,
                      tspan_train,
                      solver,
                      saveat=δt,
                      sensealg = InterpolatingAdjoint(autojacvec=ReverseDiffVJP()))
nump = length(prob_node.p)

function loss(p)
    sol = prob_node(u0,p)
    pred = Array(sol)
    sum(abs2, (train_data .- pred)), pred
end;

const losses = []
callback = function (p, l, pred)
    push!(losses, l)
    numloss = length(losses)
    if numloss % 50 == 0
        display("Epoch: " * string(numloss) * " Loss: " * string(l))
    end
    return false
end;

res_node = DiffEqFlux.sciml_train(loss,
                                   prob_node.p,
                                   cb = callback);

plot(losses, yaxis = :log, xaxis = :log, xlabel = "Iterations", ylabel = "Loss", legend = false)

prob_node = NeuralODE(sir_node,
                      tspan_train,
                      solver,
                      saveat=δt,
                      sensealg = InterpolatingAdjoint(autojacvec=ReverseDiffVJP()),
                      p = res_node)
sol_node = prob_node(u0);

scatter(sir_sol, label=["True Susceptible" "True Infected" "True Recovered"])
plot!(sol_node, label=["Estimated Susceptible" "Estimated Infected" "Estimated Recovered"])

tspan_test = (0.0, 40.0)
prob_node_test = NeuralODE(sir_node,
                      tspan_test,
                      solver,
                      saveat=δt,
                      sensealg = InterpolatingAdjoint(autojacvec=ReverseDiffVJP()),
                      p = res_node)
sol_node_test = prob_node_test(u0);

p_node = scatter(sol_node_test, legend = :topright, label=["True Susceptible" "True Infected" "True Recovered"], title="Neural ODE Extrapolation: training until t=30")
plot!(p_node,sol_node_test, lw=5, label=["Estimated Susceptible" "Estimated Infected" "Estimated Recovered"])

newu0 = [0.95, 0.05, 0.0]
sir_prob_u0 = remake(sir_prob,u0=newu0)
sir_sol_u0 = solve(sir_prob_u0, solver, saveat =  δt)
node_sol_u0 = prob_node(newu0)
p_node = scatter(sir_sol_u0, legend = :topright, label=["True Susceptible" "True Infected" "True Recovered"], title="Neural ODE with different initial conditions")
plot!(p_node,node_sol_u0, lw=5, label=["Estimated Susceptible" "Estimated Infected" "Estimated Recovered"])

tspan_train2 = (0.0,20.0)
prob2 = ODEProblem(sir_ode, u0, tspan_train2, p)
sol2 = solve(prob2, solver, saveat = δt)
data2 = Array(sol2)
solver2 = ExplicitRK()
prob_node2 = NeuralODE(sir_node,
                      tspan_train2,
                      solver2,
                      saveat=δt,
                      sensealg = InterpolatingAdjoint(autojacvec=ReverseDiffVJP()))
function loss2(p)
    sol = prob_node2(u0,p)
    pred = Array(sol)
    sum(abs2, (data2 .- sol)), pred
end
const losses2 = []
callback2 = function (p, l, pred)
    push!(losses2, l)
    numloss = length(losses2)
    if numloss % 50 == 0
        display("Epoch: " * string(numloss) * " Loss: " * string(l))
    end
    return false
end
res_node2 = DiffEqFlux.sciml_train(loss2,
                                   prob_node2.p,
                                   cb = callback2);

prob_node2_test = NeuralODE(sir_node,
                      tspan_test,
                      solver,
                      saveat=δt,
                      sensealg = InterpolatingAdjoint(autojacvec=ReverseDiffVJP()),
                      p = res_node2)
sol_node2_test = prob_node2_test(u0);

p_node2 = scatter(sir_sol, legend = :topright, label=["True Susceptible" "True Infected" "True Recovered"], title="Neural ODE Extrapolation: training until t=20")
plot!(p_node2, sol_node2_test, lw=5, label=["Estimated Susceptible" "Estimated Infected" "Estimated Recovered"])

function loss3(p)
    sol = prob_node(u0,p)
    pred = Array(sol)
    sum(abs2, (train_data[2,:] .- sol[2,:])), pred
end
const losses3 = []
callback3 = function (p, l, pred)
    push!(losses3, l)
    numloss = length(losses3)
    if numloss % 50 ==0
        display("Epoch: " * string(numloss) * " Loss: " * string(l))
    end
    return false
end
res_node3 = DiffEqFlux.sciml_train(loss3,
                                   prob_node.p,
                                   cb = callback3);

prob_node3_test = NeuralODE(sir_node,
                      tspan_test,
                      solver,
                      saveat=δt,
                      sensealg = InterpolatingAdjoint(autojacvec=ReverseDiffVJP()),
                      p = res_node3)
sol_node3_test = prob_node3_test(u0);

p_node3 = scatter(sir_sol, legend = :topright, label=["True Susceptible" "True Infected" "True Recovered"], title="Neural ODE Extrapolation: training on I(t) until t=30")
plot!(p_node3, sol_node3_test, lw=5, label=["Estimated Susceptible" "Estimated Infected" "Estimated Recovered"])


"""

```
shellradius(i,step)
```

Compute the point in which the radius comprises half of the
volume of the shell.
"""
function shellradius(i,step)
  rmin = (i-1)*step
  return ( 0.5*( (rmin+step)^3 + rmin^3 ) )^(1/3)
end


using QuadGK

sup = (-1,1)
f1(x) = 3/4*(1-x^2)
f2(x) = x < 0 ? x+1 : 1-x

expect(f,support) = quadgk((x) -> x*f(x),support...)[1]

println("Mean 1: ", expect(f1,sup))
println("Mean 2: ", expect(f2,sup))


#=

# 107: 1D Nonlinear Storage
([source code](SOURCE_URL))

This equation comes from the transformation of the nonlinear diffuision equation.
```math
\partial_t u^\frac{1}{m} -\Delta u = 0
```
in $\Omega=(-1,1)$ with homogeneous Neumann boundary conditions.
We can derive an exact solution from the Barenblatt solution of the previous
example.

=#

module Example107_NonlinearStorage1D
using Printf
using VoronoiFVM
using ExtendableGrids
using GridVisualize

function barenblatt(x,t,m)
    tx=t^(-1.0/(m+1.0))
    xx=x*tx
    xx=xx*xx
    xx=1- xx*(m-1)/(2.0*m*(m+1));
    if xx<0.0
        xx=0.0
    end
    return tx*xx^(1.0/(m-1.0))
end

function main(;n=20,m=2.0,Plotter=nothing,verbose=false, unknown_storage=:sparse,tend=0.01,tstep=0.0001)

    ## Create a one-dimensional discretization
    h=1.0/convert(Float64,n/2)
    X=collect(-1:h:1)
    grid=VoronoiFVM.Grid(X)

    ## Flux function which describes the flux
    ## between neigboring control volumes
    function flux!(f,u0,edge)
        u=unknowns(edge,u0)
        f[1]=u[1,1]-u[1,2]
    end

    ϵ=1.0e-10

    ## Storage term
    ## This needs to be regularized as its derivative
    ## at 0 is infinity
    function storage!(f,u,node)
        f[1]=(ϵ+u[1])^(1.0/m)
    end

    ## Create a physics structure
    physics=VoronoiFVM.Physics(
        flux=flux!,
        storage=storage!)

    ## Create a finite volume system - either
    ## in the dense or  the sparse version.
    ## The difference is in the way the solution object
    ## is stored - as dense or as sparse matrix
    sys=VoronoiFVM.System(grid,physics,unknown_storage=unknown_storage)

    ## Add species 1 to region 1
    enable_species!(sys,1,[1])

    ## Create a solution array
    inival=unknowns(sys)
    solution=unknowns(sys)
    t0=0.001

    ## Broadcast the initial value
    inival[1,:].=map(x->barenblatt(x,t0,m)^m,X)

    ## Create solver control info
    control=VoronoiFVM.NewtonControl()
    control.verbose=verbose
    control.Δu_opt=0.1
    control.force_first_step=true
    tsol=solve(inival,sys,[t0,tend],control=control)

    p=GridVisualizer(Plotter=Plotter,layout=(1,1),fast=true)
    for i=1:length(tsol)
        time=tsol.t[i]
        scalarplot!(p[1,1],grid,tsol[1,:,i],title=@sprintf("t=%.3g",time),color=:red,label="numerical")
        scalarplot!(p[1,1],grid,map(x->barenblatt(x,time,m)^m,grid),clear=false,color=:green,label="exact")
        reveal(p)
        sleep(1.0e-2)
    end

    return sum(tsol[end])
end

function test()
    testval=175.20261258406686
    main(unknown_storage=:sparse) ≈ testval && main(unknown_storage=:dense) ≈ testval
end

end


#=
## Integration with `NetworkDynamics.jl`
In this example we model a Network based on the [Kuroamoto model](https://en.wikipedia.org/wiki/Kuramoto_model) using the [`NetworkDynamics.jl`](https://github.com/FHell/NetworkDynamics.jl) package.
=#

using NetworkDynamics
using Graphs
using BlockSystems
using ModelingToolkit
using OrdinaryDiffEq
using Plots

#=
In the Kuramoto model each vertex $i$ has it's own intrinsic angular frequency $\omega_i$.
The angle of frequency is given by
```math
\dot\phi = \omega_i + \sum e
```
where $\sum e$ sums over all the connected edges.
=#
@parameters t ω edgesum(t)
@variables ϕ(t)
D = Differential(t)

kmvert = IOBlock([D(ϕ) ~ ω + edgesum],
                 [edgesum],
                 [ϕ],
                 name=:kmvert)

function generate_ode_vertex(iob::IOBlock, aggregator; kwargs...)
    gen = generate_io_function(iob; kwargs...)
    f = (du, u, edges, p, t) -> gen.f_ip(du, u, aggregator(edges), p, t)
    vars = ModelingToolkit.tosymbol.(gen.states)
    ODEVertex(f = f, dim = length(vars), sym = vars)
end

## allocation free oop aggregator. might be more difficult for more-dimensional systems
## unfortunately there are no vector variables in MDK and we can't model the aggregators
## as an IOSystem.
function sumedges(edges)
    r = 0.0
    for e in edges
        r += e[1]
    end
    return r
end

vertex = generate_ode_vertex(kmvert, sumedges)
nothing #hide

#=
The edges are defined as
```math
e_{i,j} = K \cdot \mathrm{sin}(\phi_i - \phi_j)
```
For the (static) edges we can't used IOSystem function building in the moment.
Well, we could but than we'd get an `ODEEdge` with zero-massmatrix.
=#
@parameters v₁(t) v₂(t) K
edge_ip = build_function([K*sin(v₁[1] - v₂[1])], [v₁], [v₂], K, t, expression=Val{false})[2]
edge = StaticEdge(f = edge_ip, dim = 1, coupling=:antisymmetric)

nothing #hide
#=
Lets runt the [same example](https://fhell.github.io/NetworkDynamics.jl/dev/heterogeneous_system/#Heterogenous-parameters) as in `NetworkDynamics.jl`
=#
N = 8
g = watts_strogatz(N,2,0) # ring network
nd = network_dynamics(vertex, edge, g)

ω = collect(1:N)./N
ω .-= sum(ω)/N
K = 3.
p = (ω, K); # p[1] vertex parameters, p[2] edge parameters

x0 = collect(1:N)./N
x0 .-= sum(x0)./N

#src using BenchmarkTools
#src @btime $nd($x0, $x0, $p, 0.0)

tspan = (0., 10.)
prob = ODEProblem(nd, x0, tspan, p)
sol = solve(prob, Tsit5())
plot(sol, ylabel="ϕ")


# SPDX-License-Identifier: MIT

"""
    struct TropicalSemiring{T<:AbstractFloat} <: Semiring
        val::T
    end

Log-semifield is defined as :
  * ``x \\oplus y \\triangleq \\text{max}(x, y)``
  * ``x \\otimes y \\triangleq x + y``
``\\forall x, y \\in \\mathbb{R}``.
"""
struct TropicalSemiring{T<:AbstractFloat} <: Semiring
    val::T
end

Base.:+(x::TropicalSemiring, y::TropicalSemiring) =
    TropicalSemiring(max(x.val, y.val))
Base.:*(x::TropicalSemiring, y::TropicalSemiring) =
    TropicalSemiring(x.val + y.val)
Base.zero(::Type{TropicalSemiring{T}}) where T = TropicalSemiring{T}(T(-Inf))
Base.one(::Type{TropicalSemiring{T}}) where T = TropicalSemiring{T}(T(0))
Base.isless(x::TropicalSemiring, y::TropicalSemiring) = isless(x.val, y.val)
Base.typemin(x::Type{TropicalSemiring{T}}) where T =
    TropicalSemiring{T}(typemin(T))
Base.typemax(x::Type{TropicalSemiring{T}}) where T =
    TropicalSemiring{T}(typemax(T))
Base.conj(x::Type{TropicalSemiring{T}}) where T = TropicalSemiring{T}(conj(x.val))


### A Pluto.jl notebook ###
# v0.16.1

using Markdown
using InteractiveUtils

# This Pluto notebook uses @bind for interactivity. When running this notebook outside of Pluto, the following 'mock version' of @bind gives bound variables a default value (instead of an error).
macro bind(def, element)
    quote
        local el = $(esc(element))
        global $(esc(def)) = Core.applicable(Base.get, el) ? Base.get(el) : missing
        el
    end
end

# ╔═╡ 3f890c7b-44ad-426a-bf26-32ff17fa8ec4
md"""
# Examples in Julia - Tutorial 1
"""

# ╔═╡ 5bc2c872-fded-4bc8-8044-9eda1822c505
md"""
## Assignment 1

Using the package `Polynomials.jl`, write the function which implements
[Graeffe's method](https://en.wikipedia.org/wiki/Graeffe%27s_method)
(see also [here](http://mathworld.wolfram.com/GraeffesMethod.html))
for computing roots of polynomials with only real roots with simple moduli.

In the function, use Julia's `BigFloat` numbers to overcome the main disadvantage of the method. What is the number of significant decimal digits, and the largest and the smallest number?

Test the function on the [Wilkinson's polynomial](https://en.wikipedia.org/wiki/Wilkinson%27s_polynomial) $\omega(x)$, and the [Chebyshev polynomial](https://en.wikipedia.org/wiki/Chebyshev_polynomials) $T_{50}(x)$
(the latter needs to be transformed in order to apply the method).

Compare your solutions with the exact solutions.
"""

# ╔═╡ 5b542ad3-8471-4adc-9c89-d647000b8ca0
md"""
## Assignment 2

Write the function which computes simple LU factorization (without pivoting) where the matrix is overwritten by the factors.

Make sure that the function also works with block-matrices.

Compare the speed on standard matrices and block-matrices with the built-in LU factorization (which also uses block algorithm AND pivoting). Check the accuracy.
"""

# ╔═╡ dbd0a150-2ad2-49e3-b585-d34bfc2ff3b8
md"""
## Assignment 3

Use the function `eigvals()` to compute the eigenvalues of $k$ random matrices (with uniform and normal distribution of elements) of order $n$.

Plot the results using the macro `@bind` from the package `PlutoUI.jl`.
Use `Plots.jl` for plotting.

Are the eigenvalues random? Can you describe their behaviour? Can random matrices be used to test numerical algorithms?
"""

# ╔═╡ 4cd45baf-67f6-43cc-a9bf-3a89006c3814

# ╔═╡ Cell order:
# ╟─3f890c7b-44ad-426a-bf26-32ff17fa8ec4
# ╟─5bc2c872-fded-4bc8-8044-9eda1822c505
# ╟─5b542ad3-8471-4adc-9c89-d647000b8ca0
# ╟─dbd0a150-2ad2-49e3-b585-d34bfc2ff3b8
# ╠═4cd45baf-67f6-43cc-a9bf-3a89006c3814

src/Module A - Short Julia Course/T1 Examples in Julia.jl


module MyInterpolations

export my_lin_interp

function my_lin_interp(grid, vals)
  function func(x)
    for i in 1:(length(grid)-1)
      if (grid[i]<=x)&&(x<=grid[i+1])
        return ((vals[i+1]-vals[i])/(grid[i+1]-grid[i]))*(x-grid[i])+vals[i]
      end
    end
  end

  return func
end

end


using AbstractPlotting
using AbstractPlotting.MakieLayout
using GLMakie

using QuadGK

function riemann(f::Function, a::Real, b::Real, n::Int; method="right")
    if method == "right"
        meth = f -> (lr -> begin l,r = lr; f(r) * (r-l) end)
    elseif method == "left"
        meth = f -> (lr -> begin l,r = lr; f(l) * (r-l) end)
    elseif method == "trapezoid"
        meth = f -> (lr -> begin l,r = lr; (1/2) * (f(l) + f(r)) * (r-l) end)
    elseif method == "simpsons"
        meth = f -> (lr -> begin l,r=lr; (1/6) * (f(l) + 4*(f((l+r)/2)) + f(r)) * (r-l) end)
    end

    xs = a .+ (0:n) * (b-a)/n

    sum(meth(f), zip(xs[1:end-1], xs[2:end]))
end

"""
    integration(f)

Show graphically the left Riemann approximation, the trapezoid approximation, and Simpson's approximation to the integral of `f` over [-1,1].

Assumes `f` is non-negative.
"""
function integration(f=nothing)
    if f == nothing
        f = x -> x^2*exp(-x/3)
    end

    a, b = -1, 1

function left_riemann_pts(n)
    xs = range(a, b, length=n+1)
    pts = Point2f0[(xs[1], 0)]
    for i in 1:n
        xᵢ, xᵢ₊₁ = xs[i], xs[i+1]
        fᵢ = f(xᵢ)
        push!(pts, (xᵢ, fᵢ))
        push!(pts, (xᵢ₊₁, fᵢ))
    end
    push!(pts, (xs[end], 0))
    pts
end

function trapezoid_pts(n)
    xs = range(a, b, length=n+1)
    pts = Point2f0[(xs[1], 0)]
    for i in 1:n
        xᵢ, xᵢ₊₁ = xs[i], xs[i+1]
        fᵢ = f(xᵢ)
        push!(pts, (xᵢ, f(xᵢ)))
    end
    push!(pts, (xs[end], f(xs[end])))
    push!(pts, (xs[end], 0))
    pts
end

function simpsons_pts(n)
    xs = range(a, b, length=n+1)
    pts = Point2f0[(xs[1], 0), (xs[1], f(xs[1]))]

    for i in 1:n
        xi, xi1 = xs[i], xs[i+1]
        m = xi/2 + xi1/2
        p = x -> f(xi)*(x-m)*(x - xi1)/(xi-m)/(xi-xi1) + f(m) * (x-xi)*(x-xi1)/(m-xi)/(m-xi1) + f(xi1) * (x-xi)*(x-m) / (xi1-xi) / (xi1-m)

        xs′ = range(xi, xi1, length=10)
        for j in 2:10
            x = xs′[j]
            push!(pts, (x, p(x)))
        end
    end
    push!(pts, (xs[end], 0))
    pts
end

# where we lay our scene:
scene, layout = layoutscene()
layout.halign = :left
layout.valign = :top

p1 = layout[1,1] = LScene(scene)
p2 = layout[1,2] = LScene(scene)
p3 = layout[1,3] = LScene(scene)

n = layout[2,1] = LSlider(scene, range=2:25, startvalue=2)
output = layout[2, 2:3] = LText(scene, "...")

lpts = lift(n.value) do n
    left_riemann_pts(n)
end

poly!(p1.scene, lpts, color=:gray75)
lines!(p1.scene, a..b, f, color=:black, strokewidth=10)
title(p1.scene, "Left Riemann")

tpts = lift(n.value) do n
    trapezoid_pts(n)
end
poly!(p2.scene, tpts, color=:gray75)
lines!(p2.scene, a..b, f, color=:black, strokewidth=10)
title(p2.scene, "Trapezoid")

spts = lift(n.value) do n
    simpsons_pts(n)
end
poly!(p3.scene, spts, color=:gray75)
lines!(p3.scene, a..b, f, color=:black, strokewidth=10)
title(p3.scene, "Simpson's")

on(n.value) do n
    actual,err = quadgk(f, -1, 1)
    lrr  = riemann(f, -1, 1, n, method="left")
    trap = riemann(f, -1, 1, n, method="trapezoid")
    simp = riemann(f, -1, 1, n, method="simpsons")

    Δleft = round(abs(lrr - actual), digits=8)
    Δtrap = round(abs(trap - actual), digits=8)
    Δsimp = round(abs(simp - actual), digits=8)

    txt = "Riemann: $Δleft, Trapezoid: $Δtrap, Simpson's: $Δsimp"
    output.text[] = txt
end
n.value[] = n.value[]

scene
end


@doc raw"""
    Ellipsoid(a, b, f, finv, e2, ϵ2, name)
    Ellipsoid(a, finv, name)

Earth ellipsoid model.

# Fields

- `a::Real`: semi-major axis (m).
- `b::Real`: semi-minor axis (m).
- `f::Real`: flattening (``f``).
- `finv::Real`: inverse of flattening (``f^{-1}``).
- `e2::Real`: eccentricity squared.
- `ϵ2::Real`: second eccentricity squared.
- `name::String`: ellipsoid name.

# Notes

``a`` is the semi-major axis, ``b`` is the semi-minor axis. Ellipsoids are normally defined
by ``a`` and ``f^{-1}``. The rest of parameters are derived.

- Flattening:
```math
f = 1 - \frac{b}{a}
```
- Eccentricity (first eccentricity):
```math
e^2 = \frac{a^2 - b^2}{a^2} = f (2 - f)
```
- Second eccentricity:
```math
ϵ^2 = \frac{a^2 - b^2}{b^2} = \frac{f (2 - f)}{(1 - f)^2}
```

# References

1. Stevens, B. L., Lewis, F. L., (1992). Aircraft control and simulation: dynamics, controls
   design, and autonomous systems. John Wiley & Sons. Section 1.6 (page 23)
2. Rogers, R. M. (2007). Applied mathematics in integrated navigation systems. American
   Institute of Aeronautics and Astronautics. Chapter 4 (Nomenclature differs from 1).
3. Bowring, B. R. (1976). Transformation from spatial to geographical coordinates. Survey
   review, 23(181), 323-327.
"""
struct Ellipsoid
    "Semi-major axis (m)."
    a::Real
    "Semi-minor axis (m)."
    b::Real
    "Flattening ``f = \frac{a-b}{a} = a - \frac{b}{a}``."
    f::Real
    "Inverse of flattening."
    finv::Real
    "Eccentricity. ``e^2 = \frac{a^2 - b^2}{a^2} = f (2 - f)``"
    e2::Real
    "Second eccentricity. ``ϵ^2 = \frac{a^2 - b^2}{b^2} = \frac{f (2 - f)}{(1 - f)^2}``."
    ϵ2::Real
    "Ellipsoid name."
    name::String
end

function Ellipsoid(a, finv, name)
    f = 1 / finv
    b = a * (1 - f)
    e2 = f * (2 - f)
    ϵ2 = e2 / (1 - f)^2
    return Ellipsoid(a, b, f, finv, e2, ϵ2, name)
end

# Rogers, R. M. (2007). Applied mathematics in integrated navigation systems.
# American Institute of Aeronautics and Astronautics. (Page 76, table 4.1)
const Clarke1866 = Ellipsoid(6378206.4, 294.9786982, "Clarke1866")
const Clarke1880 = Ellipsoid(6378249.145, 294.465, "Clarke1880")
const International = Ellipsoid(6378388.0, 297.0, "International")
const Bessel = Ellipsoid(6377397.155, 299.1528128, "Bessel")
const Everest = Ellipsoid(6377276.345, 300.8017, "Everest")
const ModifiedEverest = Ellipsoid(6377304.063, 300.8017, "ModifiedEverest")
const AustralianNational = Ellipsoid(6378160.0, 298.25, "AustralianNational")
const SouthAmerican1969 = Ellipsoid(6378160.0, 298.25, "SouthAmerican1969")
const Airy = Ellipsoid(6377564.396, 299.3249646, "Airy")
const ModifiedAiry = Ellipsoid(6377340.189, 299.3249646, "ModifiedAiry")
const Hough = Ellipsoid(6378270.0, 297.0, "Hough")
const Fischer1960SouthAsia = Ellipsoid(6378155.0, 298.3, "Fischer1960SouthAsia")
const Fischer1960Mercury = Ellipsoid(6378166.0, 298.3, "Fischer1960Mercury")
const Fischer1968 = Ellipsoid(6378150.0, 298.3, "Fischer1968")
const WGS60 = Ellipsoid(6378165.0, 298.3, "WGS60")
const WGS66 = Ellipsoid(6378145.0, 298.25, "WGS66")
const WGS72 = Ellipsoid(6378135.0, 298.26, "WGS72")
const WGS84 = Ellipsoid(6378137.0, 298.257223563, "WGS84")

"""
    llh2ecef(lat, lon, height; ellipsoid=WGS84)

Transform geodetic latitude, longitude (rad) and ellipsoidal height (m) to ECEF for the
given ellipsoid (default ellipsoid is WGS84).

# References

1. Rogers, R. M. (2007). Applied mathematics in integrated navigation systems. American
   Institute of Aeronautics and Astronautics. (Page 75, equations 4.20, 4.21, 4.22)
"""
function llh2ecef(lat, lon, height; ellipsoid=WGS84)
    f = ellipsoid.f
    a = ellipsoid.a
    e2 = ellipsoid.e2

    var = a / sqrt(1.0 - e2 * sin(lat)*sin(lat))

    px = (var + height) * cos(lat)*cos(lon)
    py = (var + height) * cos(lat)*sin(lon)
    pz = (var * (1.0 - e2) + height)* sin(lat)

    return [px, py, pz]
end

"""
    ecef2llh(xecef, yecef, zecef; ellipsoid=WGS84)

Transform ECEF coordinates to geodetic latitude, longitude (rad) and ellipsoidal height (m)
for the given ellipsoid (default ellipsoid is WGS84).

# Notes

* The transformation is direct without iterations as [1] introduced the need to iterate for
  near Earth positions.
* [2] is an update of increased accuracy of [1]. The former is used in this implementation
  although the latter implementation is commented in the code.
* Model becomes unstable if latitude is close to 90º. An alternative equation
  can be found in [2] equation (16) but has not been implemented.

 # References

1. Bowring, B. R. (1976). Transformation from spatial to geographical coordinates. Survey
   review, 23(181), 323-327.
2. Bowring, B. R. (1985). The accuracy of geodetic latitude and height equations. Survey
   Review, 28(218), 202-206.
"""
function ecef2llh(xecef, yecef, zecef; ellipsoid=WGS84)

    x, y, z = xecef, yecef, zecef
    e = sqrt(ellipsoid.e2);
    e2 = ellipsoid.e2
    ϵ2 = ellipsoid.ϵ2
    a = ellipsoid.a
    b = ellipsoid.b

    p = sqrt(x*x + y*y)
    R = sqrt(p*p + z*z)
    θ = atan(z, p)
    # [1] equation (1) does not change in [2]
    lon = atan(y, x)
    # u -> geographical latitude
    # [1] below equation (4) and [2] equation (6)
    # This lead to errors in latitud with maximum value 0.0018"
    #u = atan(a/b * z/x)

    # [2] equation (17) If the latitude is also required to be very accurate
    # for outer-space positions then the value of u for (6) should be obtained
    # from:
    u = atan(b*z / (a*p) * (1 + ϵ2 * b / R))
    # [1] equation (4)
    #lat = atan((z + ϵ2 * b * sin(u)^3) / (x - e2 * a * cos(u)^3))
    # [2] equation (18)
    lat = atan((z + ϵ2 * b * sin(u)^3) / (p - e2 * a * cos(u)^3))
    # [2] after equation (1)
    v = a / sqrt(1.0 - e2*sin(lat)^2)
    # [2] equation (7)
    # height = p*cos(lat) + z*sin(lat) - a*a / v
    # equivalent to [2] equation (8)
    height = R * cos(lat - θ) - a*a / v
    # which is insensitive to the error in latitude calculation (The influence
    # is of order 2 [2] equation (12))
    return [lat, lon, height]
end


"""
    mat_center(X; mat = "her")

This function centers input.
If X is a Hermitian matrix, then ``Xcen = X - Tr(X)/N*I```.
If X is a rectangular matrix, then ``Xcen = X - X*ones(M,1)*ones(1,M)/M``.

# Arguments
-   `X`: a Hermitian or rectangular matrix
-   `mat`: the type of `X`, valid options are "her" and "rec"

# Outputs:
-   `Xc`: centered version of `X`
"""
function mat_center(X::Array{T,2}; mat = "her") where T <: Number
    N, M = size(X)

    if mat == "her"
        # for Hermitian matrix, X = X - I*tr(X)/N
        if M != N
            error("X should be a symmetric matrix")
        end
        Xc = X - I*tr(X)/M
    elseif mat == "rec"
        # for rectangular matrix, A = A - mean(A)*ones(N, M)
        Xc = X - mean(X)*ones(N, M)
    end

    return Xc
end

"""
    free_whiten(Z; mat = "her")

This function whitens input array of matrices in the free probability sense

# Arguments
-   `Z`: an array of matrices of mat and of the same dimensions
-   `mat`: the type of Z[i], valid options are "her" and "rec"

# Outputs
-   `Y`: an array of centered matrices, whitened in the free probability sense,
       that is, Tr(Y[i]*Y[j]')/size(Y[1],1) = (i == j)
-   `U`: a matrix of size s-by-s, where s = size(Z, 1)
-   `Σ`: a matrix of size s-by-s,
"""
function free_whiten(Z::Array{Array{T, 2}}; mat = "her") where T <: Number
    Zc = mat_center.(Z; mat = mat)
    N = size(Zc[1], 1)

    # compute the cov matrix such that
    # cov[i,j] = tr(Zc[i]*Zc[j]')/N
    eigcov = eigen(real(tr.(Zc*Zc'/N)));
    Λ, U = eigcov.values, eigcov.vectors # only take the real part

    Λ[Λ.<0] .= 0 # keep non-negative eigs
    Σ = Diagonal(Λ).^0.5
    Y = U*pinv(Σ)*U'*Zc

    return Y, U, Σ
end

# if Z is just a single matrix, we first convert it to [Y].
function free_whiten(Z::Array{T, 2}; mat = "her") where T <: Number
    return free_whiten([Z]; mat = mat)
end


# Rational Quadratic Isotropic Covariance Function

"""
    RQIso <: Isotropic{SqEuclidean}

Isotropic Rational Quadratic kernel (covariance)
```math
k(x,x') = σ²(1 + (x - x')ᵀ(x - x')/(2αℓ²))^{-α}
```
with length scale ``ℓ``, signal standard deviation ``σ``, and shape parameter ``α``.
"""
mutable struct RQIso{T<:Real} <: Isotropic{SqEuclidean}
    "Squared length scale"
    ℓ2::T
    "Signal variance"
    σ2::T
    "Shape parameter"
    α::T
    "Priors for kernel parameters"
    priors::Array
end

"""
Rational Quadratic Isotropic Covariance Function

    RQIso(ll:T, lσ::T, lα::T)

# Arguments
  - `ll::Real`: length scale (given on log scale)
  - `lσ::Real`: signal standard deviation (given on log scale)
  - `lα::Real`: shape parameter (given on log scale)
"""
RQIso(ll::T, lσ::T, lα::T) where T = RQIso{T}(exp(2 * ll), exp(2 * lσ), exp(lα), [])

function set_params!(rq::RQIso, hyp::AbstractVector)
    length(hyp) == 3 || throw(ArgumentError("Rational Quadratic function has three parameters"))
    rq.ℓ2, rq.σ2, rq.α = exp(2 * hyp[1]), exp(2 * hyp[2]), exp(hyp[3])
end

get_params(rq::RQIso{T}) where T = T[log(rq.ℓ2) / 2, log(rq.σ2) / 2, log(rq.α)]
get_param_names(rq::RQIso) = [:ll, :lσ, :lα]
num_params(rq::RQIso) = 3

cov(rq::RQIso, r::Number) = rq.σ2 * (1 + r / (2 * rq.α * rq.ℓ2))^(-rq.α)

@inline dk_dll(rq::RQIso, r::Real) =
    (s = r / rq.ℓ2; rq.σ2 * s * (1 + s / (2 * rq.α))^(-rq.α - 1)) # dK_d(log ℓ)dK_dℓ
@inline function dk_dlα(rq::RQIso, r::Real)
    s = r / rq.ℓ2
    part = 1 + s / (2 * rq.α)
    rq.σ2 * part^(-rq.α) * (s / (2 * part) - rq.α * log(part))  # dK_d(log α)
end
@inline function dk_dθp(rq::RQIso, r::Real, p::Int)
    if p==1
        return dk_dll(rq, r)
    elseif p==2
        return dk_dlσ(rq, r)
    elseif p==3
        return dk_dlα(rq, r)
    else
        return NaN
    end
end


# Generates exponential mesh of N elements on [r_mn, r_max]
#
# The domain [r_min, r_max], the mesh will contain both endpoints:
function gen_rmesh_exp(
    r_min::Float64,
    r_max::Float64,
    a::Float64,
    Nr::Int64
)
    N = Nr - 1

    rmesh = zeros(Float64,N+1)
    SMALL = eps()

    if a < 0.0
        error("gen_rmesh_exp: a > 0 required")
    #
    elseif abs(a - 1.0) < SMALL
        α = (r_max - r_min) / N
        for i in 1:N+1
            rmesh[i] = α * (i - 1) + r_min
        end
    #
    else
        if N > 1
            β = log(a)/(N-1)
            α = (r_max - r_min) / (exp(β*N) - 1)
            for i in 1:N+1
                rmesh[i] = α * ( exp(β*(i-1)) - 1 ) + r_min
            end
        elseif N == 1
            rmesh[1] = r_min
            rmesh[2] = r_max
        else
            error("gen_rmesh_exp: N >= 1 required")
        end
    end
    return rmesh
end


# glcPointsLocal(N::Int) = [0.5(1-cos((i-1)pi/(N-1))) for i in 1:N]
# glcPointsGlobal(left::Number,right::Number,localPnts::Vector) = (right-left).*localPnts.-left

function glcPointsGlobal(left::Number,right::Number,N::Int)
    glcPointsLocal= [0.5(1-cos((i-1)pi/(N-1))) for i in 1:N]
    (right-left).*glcPointsLocal.-left
end

function numstring(x::Vector)::Float64
    pop!(cumprod(x))
end

function CofDQM(x::Vector,order::Int)
    #x: vector of sampling points;
    #order: largest derivative order;
    #return: matrix with dimension of (x,x,m)

    n = length(x)
    C = fill(0.0,(n,n,order))
    (n>order) || error("need more sampling points")

    # 1st order

    for i in 1:n, j in 1:n
        if i != j
            C[i,j,1] = numstring([((m!=i && m!=j) ? x[i]-x[m] : 1) for  m in 1:n])/numstring([((m!=j) ? x[j]-x[m] : 1) for  m in 1:n])

            # m_range_denominater = filter!(xx->xx!=i,[m for m in 1:n])
            # m_range_numerator = filter!(xx->xx!=j,m_range_denominater)
            #C[i,j,1] = numstring(filter(xx->xx!=0,[x[i]-x[m] for m in m_range_numerator]))/numstring(filter(xx->xx!=0,[x[j]-x[m] for m in m_range_denominater]))
            # C[i,j,1] = pop!(cumprod(filter(xx->xx!=0,[x[i]-x[m] for m in m_range_numerator])))/pop!(cumprod(filter(xx->xx!=0,[x[j]-x[m] for m in 1:n])))
        else

            # C[i,i,1] = sum(filter(xx->xx!=Inf,[1/(x[i]-x[m]) for m in 1:n]))
            C[i,i,1] = sum([((m!=i) ? 1/(x[i]-x[m]) : 0) for m in 1:n]) #这个慢了4个微秒

        end
    end

    # higher order
    for i in 1:n, j in 1:n, k in 2:order
        if i != j
            C[i,j,k] = k*(C[i,i,k-1]*C[i,j,1] - C[i,j,k-1]/(x[i]-x[j]))
        end
    end
    for i in 1:n, j in 1:n, k in 2:order
        if i == j
            C[i,j,k] = -sum([((m!=i) ? C[i,m,k] : 0) for m in 1:n])
        end
    end
    C
end

`
demoooo: test 1st-order approximation accuracy
`
pnts = glcPointsGlobal(-10,1,100)
# pnts = collect(0:2:10)

Coef = CofDQM(pnts,2)
f4(x::Number) = 4x^3+3x-2
f4d(x) = 12x^2+3

f4d.(pnts)-Coef[:,:,1]*f4.(pnts)


"""
    FourierV <: AbstractVApproximator{Int}

    struct FourierV <: AbstractVApproximator{Int}
        weights::Vector{Float64}
    end

Using Fourier cosine basis to approximate the state value.
`weights` is the featur vector.

See more details at Section (9.5.2) on Page 211 of the book *Sutton, Richard S., and Andrew G. Barto. Reinforcement learning: An introduction. MIT press, 2018.*
"""
struct FourierV <: AbstractVApproximator{Int}
    weights::Vector{Float64}
end

"""
    FourierV(order::Int)

By specifying the `order`, feature vector will be initialized with 0.
"""
FourierV(order::Int) = FourierV(zeros(Float64, order+1))

function (f::FourierV)(s)
    sum(w * cos((i-1) * π * s) for (i, w) in enumerate(f.weights))
end

function update!(f::FourierV, s, e)
    for i in 1:length(f.weights)
        f.weights[i] += e * cos((i-1) * π * s)
    end
end


# Code from previous section
using PyPlot
using Statistics
A = imread("sample_photo.png")
B = mean(A, dims=3)[:,:,1]
function imshow_scale(A)
    # Like imshow(A) but scales the values to [0,1] and supports grayscale

    A .-= minimum(A)            # Scale and shift to [0,1]
    A ./= maximum(A)
    if ndims(A) < 3
        A = reshape(A, size(A,1), size(A,2), 1)
    end
    if size(A,3) == 1
        A = repeat(A, 1, 1, 3)  # Set R=G=B for grayscale
    end
    imshow(A)
    end;

# Uncomment below if the package is not already installed
#using Pkg; Pkg.add("FFTW")

using FFTW

function imagefft_demo(A)
    AF = fftshift(fft(A))
    subplot(1,2,1); imshow_scale(A);
    subplot(1,2,2); imshow_scale(log.(1 .+ abs.(AF)));
    return
end

G = [zeros(256,128) ones(256,128)]
imagefft_demo(G)

G = zeros(256, 256)
G[78:178, 78:178] .= 1.0
imagefft_demo(G)

G = Float32[ (i+j<329) && (i+j>182) && (i-j>-67) && (i-j<73) for i = 1:256, j = 1:256 ]
imagefft_demo(G)

G = Float32[ sin(2pi*(10i + 20j)/256) for i = 1:256, j = 1:256 ]
imagefft_demo(G)

imagefft_demo(B)

Bpernoise = copy(B)
Bpernoise = B + 0.5*Float32[sin(2π*10j / size(B,2)) for i = 1:size(B,1), j = 1:size(B,2) ]
imagefft_demo(Bpernoise)

# Filter
BF = fftshift(fft(Bpernoise))
mid = size(B) .÷ 2 .+ 1
BF[mid[1], mid[2] + 10] = 0
BF[mid[1], mid[2] - 10] = 0
Bfiltered = real.(ifft(ifftshift(BF)))

imagefft_demo(Bfiltered)


# # [Ross-Macdonald Malaria Modeling](@id ross-macdonald)
#
#md # [![](https://img.shields.io/badge/show-nbviewer-579ACA.svg)](@__NBVIEWER_ROOT_URL__/examples/Ross-Macdonald.ipynb)

# Authors: Sean Wu and Sophie Libkind

using AlgebraicDynamics.DWDDynam
using Catlab.WiringDiagrams
using LabelledArrays

using Catlab.Graphics
using DelayDiffEq, DifferentialEquations
using Plots

# The Ross-Macdonald (RM) model is the canonical model of malaria transmission, first expressed by
# Sir Ronald Ross in 1911. The simplest form of the model is given by a system of 2 nonlinear ordinary
# differential equations describing $Z$, parasite prevalence in the mosquito vector population, and $X$,
# parasite prevalence in the vertebrate host population:
# $$ \dot X = mazbZ(1 - X) - rX, \quad \dot Z = acX (e^{-gn} - Z) - gZ  \quad (1)$$

#-

# The RM model is of supreme importance to modeling of malaria and other mosquito transmitted diseases
# because it is the simplest model that retains the fundamental processes involved in transmission: recovery
# of infected hosts, death of infected mosquitoes, and parasite transmission between discordant pairs of
# vectors and hosts. In this sense it represents the irreducible complexity to modeling mosquito borne
# diseases, such that any simpler model would be insufficient.

# The model assumes that mosquitoes take bloodmeals (bites) on vertebrate hosts with rate $a$.
# Therefore, the per-capita rate at which susceptible mosquitoes become infected is $acX$, where $c$
# is the transmission efficiency from infectious humans to susceptible mosquitoes (probability of
# parasites successfully invading the mosquito given a bite occurs) and $X$ is the probability the
#  bite lands on an infectious human. $1-Z$ is the proportion of mosquitoes who are susceptible but
# given an incubation period of $n$ days during which mosquitoes suffer per-capita mortality at rate
# $g$, only a proportion $e^{-gn}$ survive incubation to become infectious. Therefore the classic
# RM model will have statics that are consistent with the incorporation of an incubation period of
# fixed duration, but dynamics will not, as the delay is not incorporated into the equations.

#-

# Susceptible humans acquire infection at a rate $mabZ$, where $m = M/H$, the ratio of mosquitoes to
# humans, $a$ is as before, and $b$ is the transmission efficiency from infectious mosquitoes to
# susceptible humans. Therefore $mabZ$ gives the per-capita rate at which susceptible humans recieve
# parasites from the mosquito population. Infectious humans recover at a rate $r$ (the rate at which the
# immune system clears parasites). Because $m$ is considered a constant parameter, the model ignores
# seasonal fluctuations in mosquito (or host) population sizes. However, because the timescale of host
# dynamcis is typically slow in relation to mosquito dynamics, the approximation is often suitable
# over a transmission season.

#-

# The model's threshold criterion between the trivial equilibrium where both $Z,X$ are equal to zero
# (no disease) and the endemic equilibrium can be expressed by the basic reproductive number, giving
# the expected number of secondary infections arising from the introduction of a single infectious human,
# $$R_0 = \frac{ma^2bce^{-gn}}{rg}.$$ The endemic equilibrium is therefore:
# $$\bar{X}= \frac{R_{0}-1}{R_{0} + \frac{ac}{g}}, \quad \bar{Z} = \frac{ac\bar{X}}{g + ac\bar{X}} e^{-gn}.$$
#-
# From these equations, the nonlinear relationships that characterize observed malaria transmission can be
# understood. We implement them as follows:

# ## Diagram of Systems
# First we must construct a diagram of systems which describes the interaction between the mosquito and
# host populations. The arrows between the two subsystems represents the bidirectional infection during bloodmeals.

bloodmeal = WiringDiagram([], [:mosquitos, :humans])
mosq_box   = add_box!(bloodmeal, Box(:mosquitos, [:x], [:z]))
human_box  = add_box!(bloodmeal, Box(:humans, [:z], [:x]))
output_box = output_id(bloodmeal)

add_wires!(bloodmeal, Pair[
    (mosq_box, 1)  => (human_box, 1),
    (human_box, 1) => (mosq_box, 1),
    (mosq_box, 1)  => (output_box, 1),
    (human_box, 1) => (output_box, 2)]
)

to_graphviz(bloodmeal)

# ## ODE Model
# Next we implement the concrete mosquito and host dynamics given in Equation (1), and apply them to the diagram
# of systems. This composition is the complete Ross-Macdonald model which we can solve and plot.

dZdt = function(u,x,p,t)
    Z = u[1]
    X = x[1]
    [p.a*p.c*X*(exp(-p.g*p.n) - Z) - p.g*Z]
end

dXdt = function(u,x,p,t)
    X = u[1]
    Z = x[1]
    [p.m*p.a*p.b*Z*(1 - X) - p.r*X]
end

#-

mosquito_model = ContinuousMachine{Float64}(1, 1, 1, dZdt, (u,p,t) -> u)
human_model    = ContinuousMachine{Float64}(1, 1, 1, dXdt, (u,p,t) ->  u)

malaria_model = oapply(bloodmeal,
    Dict(:humans => human_model, :mosquitos => mosquito_model)
)

#-
params = LVector(a = 0.3, b = 0.55, c = 0.15,
    g = 0.1, n = 10, r = 1.0/200, m = 0.5)

u0 = [0.1, 0.3]
tspan = (0.0, 365.0*2)

prob = ODEProblem(malaria_model, u0, tspan, params)
sol = solve(prob, Tsit5());
#-

plot(sol, malaria_model,
    lw=2, title = "Ross-Macdonald Malaria model",
    xlabel = "time", ylabel = "proportion infectious",
    color = ["magenta" "blue"]
)

## Plot the equilibrium behavior as well
a, b, c, g, n, r, m = params
R0 = (m*a^2*b*c*exp(-g*n))/(r*g)
X̄ = (R0 - 1)/(R0 + (a*c)/g)
Z̄ = (a*c*X̄)/(g + a*c*X̄)*exp(-g*n)

N = length(sol)
plot!(sol.t, fill(X̄, N), label = "human equilibrium", ls = :dash, lw = 2, color = "blue")
plot!(sol.t, fill(Z̄, N), label = "mosquito equilibrium", ls = :dash, lw = 2, color = "magenta")

# ## Delay Model
# The previous models did not capture the incubation period for the disease in the
# mosquito population. To do so we can replace the models with delay differential equations
# and apply them to the same diagram of systems representing the bloodmeal.

dzdt_delay = function(u,x,h,p,t)
    Y, Z = u
    Y_delay, Z_delay = h(p, t - p.n)
    X, X_delay = x[1]

    [p.a*p.c*X*(1 - Y - Z) -
        p.a*p.c*X_delay*(1 - Y_delay - Z_delay)*exp(-p.g*p.n) -
        p.g*Y,
    p.a*p.c*X_delay*(1 - Y_delay - Z_delay)*exp(-p.g*p.n) -
        p.g*Z]
end

dxdt_delay = function(u,x,h,p,t)
    X, = u
    Z, _ = x[1]
    [p.m*p.a*p.b*Z*(1 - X) - p.r*X]
end

#-

mosquito_delay_model = DelayMachine{Float64, 2}(
    1, 2, 1, dzdt_delay, (u,h,p,t) -> [[u[2], h(p,t - p.n)[2]]])

human_delay_model = DelayMachine{Float64, 2}(
    1, 1, 1, dxdt_delay, (u,h,p,t) -> [[u[1], h(p, t - p.n)[1]]])

malaria_delay_model = oapply(bloodmeal,
    Dict(:humans => human_delay_model, :mosquitos => mosquito_delay_model)
)
#-
params = LVector(a = 0.3, b = 0.55, c = 0.15,
    g = 0.1, n = 10, r = 1.0/200, m = 0.5)

u0_delay = [0.09, .01, 0.3]
tspan = (0.0, 365.0*3)
hist(p,t) = u0_delay;

prob = DDEProblem(malaria_delay_model, u0_delay, [], hist, tspan, params)
alg = MethodOfSteps(Tsit5())
sol = solve(prob, alg)

plot(sol, label=["non-infectious mosquito population" "infectious mosquito population" "host population"],
    lw=2, title = "Ross-Macdonald malaria model",
    xlabel = "time", ylabel = "proportion infectious",
    color = ["magenta" "red" "blue"]
)

# While the equilibrium points of the two models are identical, they exhibit different dynamical behavior
# before settling down to equilibrium. Because models are often used to examine how the system may respond
# to intervention, incorporating additional biological realism can produce more plausible results regarding
# the expected time for interventions to affect the system. In any case, such simple models are best used as
# tools to explain basic concepts in malaria epidemiology rather than descriptions of real systems, which are
# likely far from equilibrium and affected by weather, climate, policy, and other external forces.


# This is Julia code for discretizing AR(1) process by using Tauchen method (1986)
# written on 2017/10/17
# z' = rho*z+e e~N(0,(sig_e)^2)
#
# INPUTS
# N:     Number of grids
# m:     Max number of std devs from mean
# rho:   Coeffecient of AR(1) process z'=mew+rho*z+e, e~N(0,sigmasq)
# sig_e: Std of innovation, e, in AR(1) process
#
# OUTPUT
# vZ:    grids
# mPI:   transition prob. matrix

## Preliminary

###############################################################
## 0. Define CDF function (from https://www.johndcook.com/blog/cpp_phi/)
###############################################################

function phi(x::Real)
    # constants
    const a1 =  0.254829592;
    const a2 = -0.284496736;
    const a3 =  1.421413741;
    const a4 = -1.453152027;
    const a5 =  1.061405429;
    const p  =  0.3275911;

    # Save the sign of x
    sign = 1;
    if (x < 0)
        sign = -1;
    end
    x = abs(x)/sqrt(2.0);

    # A&S formula 7.1.26
    t = 1.0/(1.0 + p*x);
    y = 1.0 - (((((a5*t + a4)*t) + a3)*t + a2)*t + a1)*t*exp(-x*x);

    return 0.5*(1.0 + sign*y);
end

##############################################################
##############################################################
##                          MAIN PART
##############################################################
##############################################################

###########################################################
## 1. Define Parameters
###########################################################

# Ex. Standard quarterly calibration for productivity process
# of US economy

const N = 5;            # Number of grids, N
const m = 3;            # Max number of std devs from mean
const rho = 0.95;       # AR(1) coefficient
const sig_e = 0.007;    # std of innovation

###########################################################
## 2. Generating grids
###########################################################

sig_z = sqrt((sig_e^2)/(1-rho^2));  # std of AR(1) r.v.
z_max = m*sig_z;                    # Max grid
z_min = -z_max;                     # Min grid
d = (z_max-z_min) / (N-1);          # incriment of grids
vZ = linspace(z_min, z_max, N);     # Grid vector 1*N

###########################################################
## 3. Transition matrix
###########################################################

mPI = zeros(N,N); # Transition matrix N*N

for i in 1:N
    for j in 1:N
        if j == 1
            mPI[i,j] = phi((vZ[1]+d/2-rho*vZ[i])/sig_e);       # If z' is min grid
        elseif j == N-1
            mPI[i,j] = 1 - phi((vZ[N]-d/2-rho*vZ[i])/sig_e);   # z' is the max grid
        else
            mPI[i,j] = phi((vZ[j]+d/2-rho*vZ[i])/sig_e) - phi((vZ[j]-d/2-rho*vZ[i])/sig_e);
        end
    end
end

###########################################################
## 4. Display result with 4 decimals
###########################################################
print("Grid points are \n")
for i in 1:N
    @printf "%6.4f" vZ[i]
    print(" ")
end
print("\n")
print("Transition prob. matrix is\n")
for i in 1:N
    for j in 1:N
        @printf "%6.4f" mPI[i,j]
        print(" ")
    end
    print("\n")
end


"""
    shifted_backward_sub!(x, R, λ, k) → x

Solve the problem (R[1:k,1:k] - λI) \\ x[1:k] in-place.
"""
function shifted_backward_sub!(x, R::AbstractMatrix{Tm}, λ, k) where {Tm<:Real}
    # Real arithmetic with quasi-upper triangular R
    @inbounds while k > 0
        if k > 1 && R[k,k-1] != zero(Tm)
            # Solve 2x2 problem
            R11, R12 = R[k-1,k-1] - λ, R[k-1,k]
            R21, R22 = R[k  ,k-1]    , R[k  ,k] - λ
            det = R11 * R22 - R21 * R12
            a1 = ( R22 * x[k-1] - R12 * x[k-0]) / det
            a2 = (-R21 * x[k-1] + R11 * x[k-0]) / det
            x[k-1] = a1
            x[k-0] = a2

            # Backward substitute
            for i = 1 : k-2
                x[i] -= R[i,k-1] * x[k-1] + R[i,k-0] * x[k-0]
            end
            k -= 2
        else
            # Solve 1x1 "problem"
            x[k] /= R[k,k] - λ

            # Backward substitute
            for i = 1 : k - 1
                x[i] -= R[i,k] * x[k]
            end

            k -= 1
        end
    end

    x
end

function shifted_backward_sub!(x, R::AbstractMatrix, λ, k)
    # Generic implementation, upper triangular R
    @inbounds while k > 0
        # Solve 1x1 "problem"
        x[k] /= R[k,k] - λ

        # Backward substitute
        for i = 1 : k - 1
            x[i] -= R[i,k] * x[k]
        end

        k -= 1
    end

    x
end

"""
    collect_eigen!(x, R, j) → k

Store the `j`th eigenvector of an upper triangular matrix `R` in `x`.
In the end `norm(x[1:k]) = 1` This function leaves x[k+1:end] untouched!
"""
function collect_eigen!(x::AbstractVector{Tv}, R::AbstractMatrix{Tm}, j::Integer) where {Tm<:Real,Tv<:Number}
    n = size(R, 2)

    @inbounds begin
        # If it's a conjugate pair with the next index, then just increment j.
        if j < n && R[j+1,j] != zero(Tv)
            j += 1
        end

        # Initialize the rhs and do the first backward substitution
        # Then do the rest of the shifted backward substitution in another function,
        # cause λ is either real or complex.
        if j > 1 && R[j,j-1] != 0
            # Complex arithmetic
            R11, R21 = R[j-1,j-1], R[j-0,j-1]
            R12, R22 = R[j-1,j-0], R[j-0,j-0]
            det = R11 * R22 - R21 * R12
            tr = R11 + R22
            λ = (tr + sqrt(complex(tr * tr - 4det))) / 2
            x[j-1] = -R12 / (R11 - λ)
            x[j-0] = one(Tv)
            for i = 1 : j-2
                x[i] = -R[i,j-1] * x[j-1] - R[i,j]
            end

            shifted_backward_sub!(x, R, λ, j-2)
        else
            # Real arithmetic
            λ = R[j,j]
            x[j] = one(Tv)
            for i = 1 : j-1
                x[i] = -R[i,j]
            end
            shifted_backward_sub!(x, R, λ, j-1)
        end

        # Normalize
        nrm = zero(real(Tv))
        for k = 1:j
            nrm += abs2(x[k])
        end
        scale = inv(√nrm)
        for k = 1:j
            x[k] *= scale
        end
    end

    return j
end

function collect_eigen!(x::AbstractVector{Tv}, R::AbstractMatrix, j::Integer) where {Tv}
    n = size(R, 2)

    @inbounds begin
        λ = R[j,j]
        x[j] = one(Tv)
        for i = 1 : j-1
            x[i] = -R[i,j]
        end

        shifted_backward_sub!(x, R, λ, j-1)

        # Normalize
        nrm = zero(real(Tv))
        for k = 1:j
            nrm += abs2(x[k])
        end
        scale = inv(√nrm)
        for k = 1:j
            x[k] *= scale
        end
    end

    return j
end


load("src/init.jl")

# Given inputs, ouputs, lambda and current weights,
# Assess the ridge loss function.
# Assume first entry of w is an intercept and do not penalize.

function ridge_error(x, y, l, w)
  (1.0 / 2.0) * sum(map(z -> z^2, y - x * w)) + (l / 2.0) * sum(map(z -> z^2, w[2:length(w)]))
end

ridge_error([1.0 2; 3 4], [1, 2], 1, [1, 1])

function ridge_gradient(x, y, l, w)
	unshift(l * w[2:length(w)], 0) - x' * (y - x * w)
end

ridge_gradient([1.0 1 2; 1 3 4], [1, 2], 1, [1, 1, 1])

function ridge_regression(x, y, l)
	function f(w)
    ridge_error(x, y, l, w)
  end

  function g(w)
    ridge_gradient(x, y, l, w)
  end

	w0 = zeros(size(x, 2))

	solution = gradient_descent(f, g, w0)

	solution.minimum
end

x = [1 1 2; 1 3 3; 1 5 6]
y = [1, 2, 2]

w = ridge_regression(x, y, 0.0)
x * w
norm(w)

w = ridge_regression(x, y, 1.0)
x * w
norm(w)

w = ridge_regression(x, y, 10.0)
x * w
norm(w)

w = ridge_regression(x, y, 100.0)
x * w
norm(w)

p = 10

x = hcat([1, 1], vcat([1:p]', [(p + 1):(2p)]'))
y = [1, 2]

w = ridge_regression(x, y, 0)
x * w
norm(w)

w = ridge_regression(x, y, 1)
x * w
norm(w)

w = ridge_regression(x, y, 10)
x * w
norm(w)


export Atan

"""
`Atan([domainType=Float64::Type,] dim_in::Tuple)`

Creates an inverse tangent non-linear operator with input dimensions `dim_in`:
```math
\\text{atan} ( \\mathbf{x} ).
```

"""
struct Atan{T,N} <: NonLinearOperator
	dim::NTuple{N,Int}
end

function Atan(DomainType::Type, DomainDim::NTuple{N,Int}) where {N}
	Atan{DomainType,N}(DomainDim)
end

Atan(DomainDim::NTuple{N,Int}) where {N} = Atan{Float64,N}(DomainDim)
Atan(DomainDim::Vararg{Int}) = Atan{Float64,length(DomainDim)}(DomainDim)

function mul!(y::AbstractArray{T,N}, L::Atan{T,N}, x::AbstractArray{T,N}) where {T,N}
	y .= atan.(x)
end

function mul!(y::AbstractArray,
              J::AdjointOperator{Jacobian{A,TT}},
              b::AbstractArray) where {T, N, A <: Atan{T,N}, TT <: AbstractArray{T,N}}
    L = J.A
    y .= conj.(1.0./(1.0.+ L.x.^2)).*b
end

fun_name(L::Atan) = "atan"

size(L::Atan) = (L.dim, L.dim)

domainType(L::Atan{T,N}) where {T,N} = T
codomainType(L::Atan{T,N}) where {T,N} = T


## ND LINEAR INTERPOLATION KERNELS

"""
    lerp(t, v0, v1)

Linear interpolation between `f(x0)=v0` and `f(x1)=v1` at the normalized coordinate t
"""
@inline lerp(t, v0, v1) = fma(t, v1, fma(-t, v0, v0))

"""
    bilinear(tx, ty, v00, v10, v01, v11)

Linear interpolation between `f(x0,y0)=v0` and `f(x1,y1)=v1` at the normalized coordinates `tx`, `ty`
"""
@inline function ndlinear(tx, ty, v00, v10, v01, v11)
    return lerp(ty, lerp(tx, v00, v10), lerp(tx, v01, v11))
end

@inline function ndlinear(t::NTuple{2,T1}, v::NTuple{4,T2}) where {T1,T2}
    return ndlinear(t[1], t[2], v[1], v[2], v[3], v[4])
end

bilinear(tx, ty, v00, v10, v01, v11) = ndlinear(tx, ty, v00, v10, v01, v11)

"""
    trilinear(tx, ty, tz, v000, v100, v001, v101, v010, v110, v011, v111)

Linear interpolation between `f(x0,y0,z0)=v0` and `f(x1,y1,z1)=v1` at the normalized coordinates `tx`, `ty`, `tz`
"""
@inline function ndlinear(tx, ty, tz, v000, v100, v001, v101, v010, v110, v011, v111)
    return lerp(
        ty,
        ndlinear(tx, tz, v000, v100, v001, v101),
        ndlinear(tx, tz, v010, v110, v011, v111),
    )
end

@inline function ndlinear(t::NTuple{3,T1}, v::NTuple{8,T2}) where {T1,T2}
    return ndlinear(t[1], t[2], t[3], v[1], v[2], v[3], v[4], v[5], v[6], v[7], v[8])
end

@inline function trilinear(tx, ty, tz, v000, v100, v001, v101, v010, v110, v011, v111)
    return ndlinear(tx, ty, tz, v000, v100, v001, v101, v010, v110, v011, v111)
end


module Power
export power_method, inverse_power_method

using LinearAlgebra

"""
Calculate the dominant eigenvalue of a matrix A and its associated eigenvector.
"""
function power_method(A::Matrix{<:Real}, x::Vector{<:Real}, n::Int)::Tuple
	y = x
	m = 1

	for _ in 0:n
		y = A * x
		m = maximum(abs.(y))
		x = (m^-1)y
	end

	return (m, y)
end

"""
Calculate the least dominant eigenvalue of a matrix A and its associated
eigenvector.
"""
function inverse_power_method(A::Matrix{<:Real}, x::Vector{<:Real}, n::Int)::Tuple
	return power_method(inv(A), x, n)
end

end


using DataDrivenDiffEq
using OrdinaryDiffEq
using LinearAlgebra
using Plots
gr()

# Create a linear , time discrete system
function linear_discrete(du, u, p, t)
    du[1] = 0.9u[1]
    du[2] = 0.9u[2] + 0.1u[1]
end

u0 = [10.0; -2.0]
tspan = (0.0, 10.0)
prob = DiscreteProblem(linear_discrete, u0, tspan)
sol = solve(prob, FunctionMap())

# Create Approximation
approx = DMD(sol[:,:])

# Create the associated problem
prob_approx = DiscreteProblem(approx, u0, tspan)
approx_sol = solve(prob_approx, FunctionMap())

# Show solutions
plot(sol)
plot!(approx_sol)

# Show error
plot((sol .- approx_sol)')
norm(sol .- approx_sol) # ≈ 2.23e-14

# Eigen Decomposition
eigen(approx)
# Eigenvalues
scatter(eigvals(approx), xlim = (-1, 1), ylim = (-1,1))

# Adapt system with new measurements
function linear_discrete_2(du, u, p, t)
    du[1] = 0.9u[1]
    du[2] = 0.9u[2] + 0.2u[1]
end

# Solve the new system
prob2 = DiscreteProblem(linear_discrete_2, u0, (0.0, 20.0))
sol2 = solve(prob2, FunctionMap())

# Split the data
x = sol2[:,1:20]
y = sol2[:,2:21]

# Update our approximation with new measurements
update!(approx, x, y)

# Lets have a look at the operator, which moves near the true value
operator(approx)

# Add time continouos system
function linear(du, u, p, t)
    du[1] = -0.9*u[1] + 0.1*u[2]
    du[2] = -0.8*u[2]
end

prob_cont = ODEProblem(linear, u0, tspan)
sol_cont = solve(prob_cont, Tsit5())

plot(sol_cont)

X = sol_cont[:,:]
DX = sol_cont(sol_cont.t, Val{1})[:,:]

# Giving the method a time step (which should be sequentially sampled)
# Enables us to get the continouos representation
approx_cont = gDMD(X, DX)
generator(approx_cont)
approx_sys = ODEProblem(approx_cont, u0, tspan)
approx_sol = solve(approx_sys, Tsit5(),  saveat = sol_cont.t)
# Lets have a look at the solution
plot(sol_cont)
plot!(approx_sol)
# And the error
plot(abs.(sol_cont .- approx_sol)')
norm(sol_cont .- approx_sol) # ≈ 1e-13


__precompile__()
# ----------------------------------------------------------------------------
module FAD # Forward Automatic Differentiation demo
#=
    see also:
        - the walkthrough at http://vladium.com/tutorials/study_julia_with_me/multiple_dispatch/

    working with this module from REPL:

    julia>include("<path>/FAD.jl")
    julia>using Main.FAD
    ... REPL can now use derivative(), ∂(), root_solve(), etc ...
=#
# ............................................................................

import Base; # needed in order to add base ops overloads

export derivative, ∂, myerf, root_solve; # for convenience of 'using Main.FAD' in REPL

# ............................................................................

struct Context <: Number
    v   ::Number
    ∂   ::Number
end
# ............................................................................
# binary ops:

# note: ':()' is necessary to reference certain functions with single-character names

Base.:(+)(lhs ::Context, rhs ::Context) = Context(lhs.v + rhs.v, lhs.∂ + rhs.∂)
Base.:(-)(lhs ::Context, rhs ::Context) = Context(lhs.v - rhs.v, lhs.∂ - rhs.∂)

Base.:(*)(lhs ::Context, rhs ::Context) = Context(lhs.v * rhs.v, lhs.v * rhs.∂ + lhs.∂ * rhs.v)
Base.:(/)(lhs ::Context, rhs ::Context) = Context(lhs.v / rhs.v, (lhs.∂ * rhs.v - lhs.v * rhs.∂) / rhs.v^2)

# unary ops:

Base.:(+)(x ::Context) = x
Base.:(-)(x ::Context) = Context(- x.v, - x.∂)

# math ops:

Base.sin(d ::Context) = Context(sin(d.v),   cos(d.v) * d.∂)
Base.cos(d ::Context) = Context(cos(d.v), - sin(d.v) * d.∂)
# ... add more as needed for your own functions ...

# ............................................................................
# conversion/promotion:

Base.promote_rule(::Type{Context}, ::Type{<: Number}) = Context
Base.convert(::Type{Context}, x ::Real) = Context(x, 0.0)

# ............................................................................
"""
    Given a function 'f' of one variable, return its derivative with respect to that variable.
"""
function derivative(f ::Function)
    return x ::Number -> f(Context(x, 1.0)).∂ # discard value, return derivative
end

∂ = derivative # add a nice-looking alias

# ............................................................................

# to get "official" erf(x), do
# pkg> add SpecialFunctions
# julia> using SpecialFunctions

function myerf(x ::Number)
    Σ = 0.0
    x² = x * x
    for k in 0 : 20 # hardcoding the number of summation terms for simplicity
        Σ += x / (factorial(k) * (2k + 1))
        x *= -x²
    end
    return 2.0 / √π * Σ
end
# ............................................................................
"""
    Newton's method for finding root(s) of 'f'
"""
function root_solve(f ::Function, x₀ ::Number; ϵ = 1e-8)
    i = 1
    while true
        ctx = f(Context(x₀, 1.0))
        println("[$i]: f($x₀)\t= $(ctx.v)")
        abs(ctx.v) < ϵ && break
        x₀ -= ctx.v / ctx.∂ # ignoring issues with division by small floats, etc for simplicity
        i += 1
    end
    return x₀
end

end # of module
# ----------------------------------------------------------------------------


# arb_hypgeom.jl

module arb_hypgeom
export one_f_two_fast, one_f_two

using ArbNumerics

"""
Implementation of the hypergeometric function 1F2 using the expansion:

    1F2(a1; {b1, b2}; z) = ∑_{k=0}^∞ (a1)_k z^k / ((b1)_k ⋅ (b2)_k ⋅ k!)

where (x)_k is the Pochhammer rising factorial = Γ(x+k)/Γ(x) with Γ(x) the Gamma function.
"""

function one_f_two(a, b, x; prec = 64) # z > 0 & n >= 0

    # Initialise precision of ArbReal to prec.
    setextrabits(0)
    setprecision(ArbReal, prec + 8)

    a = ArbReal("$a")
    b = (ArbReal("$i") for i in b)
    x = ArbReal("$x")

    k = ArbReal("0")
    result = ArbReal("0.0")
    err = eps(result)  # Machine accuracy of specified precision prec.

    risingfact(w, n) = ArbReal(ArbNumerics.gamma(w + n) / ArbNumerics.gamma(w))

    while true
        term = ArbReal(risingfact(a, k) * x^k / (prod(risingfact.(b, k)) * ArbNumerics.gamma(k + 1)))

        # Break loop if term smaller than accuracy of result.
        if abs(term) < err
            break
        end

        result += term
        k += ArbReal("1")

        # Double precision if rounding error in result exceeds accuracy specified by prec.
        if ball(result)[2] > err
            setprecision(ArbReal, precision(result) * 2)
            k = ArbReal("0")
            a = ArbReal("$a")
            b = (ArbReal("$i") for i in b)
            x = ArbReal("$x")
            result = ArbReal("0.0")
        end
    end
    ArbReal(result, bits = prec + 8)
end

"""
Implementation of specific 1F2 hypergeometric function used for integrals:

    ∫_0^1 cosh(zx) x^{2m} dx = 1F2(m+1/2; {1/2, m+3/2}; z^2/4)/(2m+1)
                             = ∑_{t=0}^∞ z^{2t} / ((2t+2m+1)⋅(2t)!)

    ∫_0^1 sinh(zx) x^{2m} dx = z ⋅ 1F2(m+1; {3/2, m+2}; z^2/4)/(2m+2)
                             = ∑_{t=0}^∞ z^{2t+1} / ((2t+2m+2)⋅(2t+1)!)

which we can combine into a generic summation:

    ∑_{t=0}^∞ z^{2t+h} / ((2t+2m+1+h)⋅(2t+h)!)

which gives the cosh version for h=0 and the sinh version for h=1. This specialised 1F2 converges faster than generic 1F2 algorithm, one_f_two.
"""

function one_f_two_fast(x, m, h; prec = 64) # z > 0 & n >= 0

    # Initialise precision of ArbReal to prec.
    p = prec
    setextrabits(0)
    setprecision(ArbReal, p)

    x = ArbReal("$x")
    m = ArbReal("$m")
    h = ArbReal("$h")

    k = ArbReal("0")
    result = ArbReal("0.0")
    term = ArbReal("1.0")
    err = eps(result)  # Machine accuracy of specified precision prec.

    while abs(midpoint(term)) > err * abs(midpoint(result))
        term = ArbReal(x^(2 * k + h) / ((2 * k + 2 * m + 1 + h) * ArbNumerics.gamma(2 * k + 1 + h)))
        result += term
        # println("term: k = ", k, "\nterm value: ", ball(ArbReal(term, bits = prec)), "\ncumulant result: ", ball(ArbReal(result, bits = prec)), "\n")
        k += 1

        # Double precision if rounding error in result exceeds accuracy specified by prec.
        if radius(result) > err * abs(midpoint(result))
            p *= 2
            setprecision(ArbReal, p)
            # println("Not precise enough. Error = ", abs(radius(result)/midpoint(result)), " > ", err, ". Increasing precision to ", p, " bits.\n")
            x = ArbReal("$x")
            m = ArbReal("$m")
            h = ArbReal("$h")
            k = ArbReal("0")
            result = ArbReal("0.0")
            term = ArbReal("1.0")
        end
    end
    # println("x: ", ArbReal(x, bits = prec), ". Final result: ", ArbReal(result, bits = prec))
    ArbReal(result, bits = prec)
end

# β = ArbReal("2.0")
# α = ArbReal("7.0")
# v = ArbReal("5.8")
# w = ArbReal("1.6")
# R = ArbReal((v^2 - w^2) / (w^2 * v))
# a = ArbReal(sqrt(β^2 / 4 + R * β * coth(β * v / 2)))
# z = ArbReal("90.61")
# m = ArbReal("1.0")
# @time c = one_f_two_fast(z, m, 0; prec = 64)
# @show(c)

end # end module

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

"""
Some plots and tests of the above functions to check that they produce the correct values to the specified precision.
"""

# using Plots
# using QuadGK
# plotly()
# Plots.PlotlyBackend()
#
# cosh_integral(z, m) = quadgk(t -> (m + 1) * cosh(z * t) * t^m, BigFloat(0), BigFloat(1))[1]
#
# # x_range = 0:200
# # m = 1
# # @time I = [cosh_integral(x, m) for x in x_range]
# # @time F = abs.([arb_hypgeom_1f2(m / 2 + 1 / 2, (1 / 2, m / 2 + 3/2), x^2 / 4.0; prec = 128) for x in x_range])
# # p = plot(x_range, F, yaxis=:log)
# # plot!(x_range, I)
# # display(p)
# # @show(I, F)
# # println(m / 2 + 1 / 2, (1 / 2, m / 2 + 3/2), x_range^2 / 4.0)
#
# """
# Conclusion: The integral is a bit faster than the hypergeometric function, but not as accurate. So, probably use 1F2 since the time difference doesn't seem to be major!
# """
#
# # @time F_test = arb_hypgeom_1f2(1/2, (1/2, 1/2), 1/2; prec = 512)
# # setprecision(BigFloat, 512)
# # t_test = parse(BigFloat, "2.17818355660857086398922206782012528343129403292165693281081574094992093930206091916837394310903919052930938811320102906774905335969974706202490049810447181898779259195640840870264051023127290510625691429186995752638681515944856002645204867940492548077249544")
# # @show(F_test, t_test)
# # @show(isequal("$F_test"[end-1], "$t_test"[end-1]))
#
# """
# The 1F2 function above can accurate produce high precision values as shown in the above test. But it can be simplified and optimized mathematically initially as shown below.
# """
#
# x_range = 0:200
# m = 1
# h = 0
# @time L = abs.([arb_hypergeom_1f2_fast(x, m, h; prec = 128) for x in x_range])
# # @time F = abs.([arb_hypgeom_1f2(m + 1 / 2, (1 / 2, m + 3/2), x^2 / 4.0; prec = 128) / (2 * m + 1) for x in x_range])
# p = plot(x_range, L, yaxis=:log)
# # plot!(x_range, F)
# display(p)
# @show(L)


# We consider as an illustrative example a simple harmonic oscillator system.
# The set propagation definitions introduced so far are illustrated using a single
# degree of freedom second order problem given by the following equation:

# ```math
# u''(t) + ω^2 u(t) = 0,
# ```
# where $ω$ is a scalar and $u(t)$ is the unknown.
#
# This problem can be associated with a spring-mass system, where $u(t)$ is the
# elongation of the spring at time $t$ and the mass and stiffness set a natural
# frequency $ω$. In this case we consider $ω = 4π$.

# Let us introduce the state variable $v(t) := u'(t)$ and define the vector
# $x(t) = [u(t), v(t)]^T$. Then we can rewrite the system in the following first order form
#
#
# ```math
# x'(t) = \begin{pmatrix} 0 & 1 \\ -\omega^2 & 0 \end{pmatrix} x(t)
#
# Two cases are considered for the initial conditions, first a single initial condition
# given by $u(0)=1$ and $v(0)=0$, and a second case where the initial conditions belong to a set.

# This problem is solved using set propagation and the results obtained for the
# single initial condition are shown in the following figures.
#
# The time step-size used is $\delta = 0.025$.
#
# The set $\Omega_0$ includes the analytic trajectory within $[0, \delta]$, and
# such set is propagated to cover the solution for further time intervals as shown below.
#
# It is worth noting that even when the initial condition is a singleton,
# the set propagation technique returns a sequence of sets, and such sequence guarantees
# an enclosure of the analytic solution at \emph{any} time $t \in [0, T] \subset \mathbb{R}$.
#
# For comparison, we also plot the analytic solution (magenta solid line) and a
# numerical solution obtained using the Newmark's method with time step $\delta/5$
# (red triangle markers).

# The result obtained when a set of initial conditions, together with few dozen
# trajectories with random initial conditions drawn from the initial set
# $\mathcal{X}_0 = \mathcal{U}_0 \times \mathcal{V}_0 = [0.9, 1.1] \times [-0.1, 0.1]$,
# where $\times$ denotes the Cartesian product.

# In both cases, it is shown that a single set-based integration covers infinitely
# many trajectories \emph{in dense time}, i.e. for all intermediate times between
# $0$ and $T$, where $T > 0$ is the time horizon, there is a reach-set that covers
# all the exact solutions.


export lse, lomax, randomindex

# Log-sum exp.
function lse(x::AbstractMatrix{<:Real}; dims = 2)
    m = maximum(x, dims = dims)
    v = exp.(x .- m)
    return log.(sum(v, dims = dims)) + m
end

# see https://www.cs.ubc.ca/~murphyk/Papers/bayesGauss.pdf
function tpdf(μ::T, σ::T, ν::T, x::V) where {T<:Real,V<:Real}
    c = gamma(ν/2 + 1/2) / gamma(ν/2)
    c *= 1/(sqrt(ν*π)*σ)
    return c * (1 + 1/ν * ((x-μ)/σ)^2)^((ν+1)/2)
end

function tlogpdf(μ::T, σ::T, ν::T, x::V) where {T<:Real,V<:Real}
    c = lgamma(ν/2 + 1/2) - lgamma(ν/2)
    c -= log(sqrt(ν*π)*σ)
    return c + ((ν+1)/2)*log(1 + 1/ν * ((x-μ)/σ)^2)
end

@inline lomax(x, α, λ) = log(α) - log(λ) + -(α + 1) * log(1 + (x / λ))

function randomindex(p::AbstractVector{T}) where T<: Real
    return randomindex(p, rand(T))
end

function randomindex(p::AbstractVector{T}, rnd::T) where T<: Real

    csum = T(0.0)
    thres = rnd * sum(p)

    for i in 1:length(p)
        @inbounds begin
            csum += p[i]
            if csum >= thres
                return Int(i)
            end
        end
    end

    return Int(length(p))
end


#   Unit statistics.jl, part of PosDefManifold Package for julia language
#   v 0.3.5 - last update 6th of November 2019
#
#   MIT License
#   Copyright (c) 2019, Marco Congedo, CNRS, Grenobe, France:
#   https://sites.google.com/site/marcocongedo/home
#
#   DESCRIPTION
#   This Unit implements routines for statistics and probability.
#
#   CONTENT
#   1.  Utilities
#   2.  Probability
#   3.  Descriptive Statistocs
# __________________________________________________________________

#  ------------------------
## 1. Utilities
#  ------------------------

#  ------------------------
## 2. Probability
#  ------------------------

"""
    softmax(χ::Vector{T}) where T<:Real

 Given a real vector of ``k`` non-negative scores ``χ=c_1,...,c_k``,
 return the vector ``π=p_1,...,p_k`` of their
 [softmax](https://en.wikipedia.org/wiki/Softmax_function) probabilities,
 as per

 `` p_i=\\frac{\\textrm{e}^{c_i}}{\\sum_{i=1}^{k}\\textrm{e}^{c_i}} ``.

 ## Examples
    χ=[1.0, 2.3, 0.4, 5.0]
    π=softmax(χ)

"""
softmax(χ::Vector{T}) where T<:Real = exp.(χ) ./ 𝚺(exp.(χ))

#  -------------------------
## 3. Descriptive Statistics
#  -------------------------
"""
    mean(metric::Metric, ν::Vector{T}) where T<:RealOrComplex

 Mean of ``k`` real or complex scalars, using the specified `metric`
 of type [Metric::Enumerated type](@ref). Note that using the Fisher,
 logEuclidean and Jeffrey metric, the resulting mean
 is the scalar geometric mean. Note also that the code of this method
 is in unit *statistics.jl*, while the code for all the others is
 in unit *riemannianGeometry.jl*.

 ## Examples
    using PosDefManifold
    # Generate 10 random numbers distributed as a chi-square with 2 df.
    ν=[randχ²(2) for i=1:10]
    arithmetic_mean=mean(Euclidean, ν)
    geometric_mean=mean(Fisher, ν)
    harmonic_mean=mean(invEuclidean, ν)
    harmonic_mean<=geometric_mean<=arithmetic_mean # AGH inequality

"""
function mean(metric::Metric, ν::Vector{T}) where T<:RealOrComplex
    if      metric == Euclidean     return mean(ν)
    elseif  metric == invEuclidean  return inv(mean(inv, ν))
    elseif  metric in (Fisher,
                       logEuclidean,
                       Jeffrey)     return exp(mean(log, ν))
    elseif  metric == logdet0
        @warn "function statistics.mean (scalar mean) not implemented for metric $metric"
    elseif  metric == ChoEuclidean  return (mean(√, ν))^2
    elseif  metric == logCholesky   return exp((mean(log, sqrt.(ν))))^2
    elseif  metric == VonNeumann
        @warn "function statistics.mean and .geodesic not defined for metric $metric"
    elseif  metric == Wasserstein   return (mean(√, ν))^(-0.5)
    else
        @error "in RiemannianGeometry.mean function: the chosen 'metric' does not exist"
    end # if metric
end

"""
    std(metric::Metric, ν::Vector{T};
        corrected::Bool=true,
        mean=nothing) where T<:RealOrComplex

 Standard deviation of ``k`` real or complex scalars,
 using the specified `metric`
 of type [Metric::Enumerated type](@ref) and the
 specified `mean` if provided.

 Only the Euclidean and Fisher
 metric are supported by this function. Using the Euclidean
 metric return the output of standard Julia
 [std](https://docs.julialang.org/en/v1/stdlib/Statistics/#Statistics.std)
 function. Using the Fisher metric return the scalar geometric standard deviation,
 which is defined such as,

 ``\\sigma=\\text{exp}\\Big(\\sqrt{k^{-1}\\sum_{i=1}^{k}\\text{ln}^2(v_i/\\mu})\\Big)``.

 If `corrected` is `true`, then the sum is scaled with ``k-1``,
 whereas if it is `false` the sum is scaled with ``k``.

 ## Examples
    using PosDefManifold
    # Generate 10 random numbers distributed as a chi-square with 2 df.
    ν=[randχ²(2) for i=1:10]
    arithmetic_sd=std(Euclidean, ν) # mean not provided
    geometric_mean=mean(Fisher, ν)
    geometric_sd=std(Fisher, ν, mean=geometric_mean) # mean provided

"""
function std(metric::Metric, ν::Vector{T};
             corrected::Bool=true,
             mean=nothing) where T<:RealOrComplex

    if      metric == Euclidean     return std(ν; corrected=corrected, mean=mean)

    elseif  metric == Fisher
            mean==nothing ? μ=mean(Fisher, ν) : μ=mean
            if corrected return exp(√(𝚺(log(w/μ)^2 for w in ν)/(length(ν)-1)))
            else         return exp(√(𝛍(log(w/μ)^2 for w in ν))) end

    elseif  metric in (invEuclidean, logEuclidean, Jeffrey,
                        logdet0, ChoEuclidean, logCholesky,
                        VonNeumann, Wasserstein)
        @warn "function statistics.mean (scalar mean) not implemented for metric $metric"
    else
        @error "in RiemannianGeometry.mean function: the chosen 'metric' does not exist"
    end # if metric
end


# f(x)=elu(x)=x*[x>0] + (exp(x)-1)*[x<0]

alpha=1.0
function Felu(x)
    y=zeros(size(x))
    for i=1:length(x)
        if x[i]>0
            y[i]=x[i]
        else
            y[i]=alpha*(exp(x[i])-1.0)
        end
    end
    return (y,nothing)
end

function Felu_inplace(value::Array,auxvalue,x) # inplace
    for i=1:length(x)
        if x[i]>0
            value[i]=x[i]
        else
            value[i]=alpha*(exp(x[i])-1.0)
        end
    end
end

function Delu(derivativeIDX,f_c,faux_c,grad_c,grad_n,x)
    tmp=ones(size(x))
    for i=1:length(x)
        if x[i]<0
            tmp[i]=f_c[i]+alpha
        end
    end
    axpy!(1.0,tmp.*grad_c,grad_n)
end

if 1==0 # TODO GPU version

    function Felu_inplace(value::CudaArray,auxvalue,x::CudaArray...) # inplace
    end

    function Delu(derivativeIDX,f_c,faux_c,grad_c,grad_n,x::CudaArray...)
    end

end

Derivative[Felu]=Delu # Define dictionary lookup
Inplace[Felu]=Felu_inplace

elu(n::ADnode)=ADnode(Felu,n)

###elu(A::ADtrans)=transpose(elu(node[A.parent])) # elu(A')=(elu(A))' TODO:check

export elu


#=
  Where is my bag problem.

  From
  "Probabilistic Reasoning Under Uncertainty with Bayesian Networks"
  https://www.bayesia.com/2018-03-02-probabilistic-reasoning-under-uncertainty

  (Also see The BaysiaLabBook "Bayesian Networks BayesiaLab" v1, page 51f.)

  Travel from Singapore -> Tokyo -> Los Angeles.
  1) The departure of Singapore -> Tokyo is delayed
  2) Just made the Tokyo -> LA trip. There's a 50/50 change the luggage is on the plane to LA.
  3) Arriving at LA, waiting for the luggage in 5 minutes, the luggage has not been seen.
  What is the probability that the bag was in the Tokyo -> LA plane?

  Answer: The probability that the bag was with the LA plane is 33% (as in the talk).

  From Judea Pearl: The Book of Why.

  For p(bag_on_plane) = 0.5
    parameters      mean       std   naive_se      mcse         ess      rhat   ess_per_sec
  bag_on_plane    0.3364    0.4725     0.0047    0.0053   8953.8759    1.0001     1581.1188

  For p(bag_on_plane) = 0.8
    parameters      mean       std   naive_se      mcse         ess      rhat   ess_per_sec
  bag_on_plane    0.6641    0.4723     0.0047    0.0051   8363.8461    0.9999     1469.6619

  Cf ~/blog/where_is_my_bag.blog
     ~/webppl/where_is_my_bag.wppl

=#
using Turing, StatsPlots, Distributions
include("jl_utils.jl")

@model function where_is_my_bag(maxTime=10,p=0.5,time=5.0,bag_on_carousel=false)

    # We know that the bag was on the plane with 50% probability.
    # Variant: https://www.youtube.com/watch?v=c71pnonOMkI&t=1074s it's 80/20.
    bag_on_plane ~ flip(p)

    # We assume uniform time step
    time ~ DiscreteUniform(1,maxTime)

    # Probability that the bag is on the carousel given a specific time
    # (if it was on the plane, that is).
    # The probability that the bag is on the carousel on a given time
    # is directly proportional on the time:
    #     time / maxTime
    bag_on_carousel ~ bag_on_plane == true ? flip(1.0*time/maxTime) : Dirac(0.0)

    # He waited 5 minutes without seeing the luggage.
    # Note: It's faster if we observe via the parameters
    # true ~ Dirac(time == 5.0)
    # true ~ Dirac(bag_on_carousel == false)

end

maxTime = 10
p = 0.5
# p = 0.8
time = 5.0
bag_on_carousel = false
model = where_is_my_bag(maxTime,p,time,bag_on_carousel)

# chns = sample(model, Prior(), 10_000)
# chns = sample(model, MH(), 10_000)
chns = sample(model, PG(15), 10_000)
# chns = sample(model, SMC(), 10_000)
# chns = sample(model, IS(), 10_000)

display(chns)
# display(plot(chns))


"""
    PeriodicKernel(r::AbstractVector)
    PeriodicKernel(dims::Int)
    PeriodicKernel(T::DataType, dims::Int)

Periodic Kernel as described in http://www.inference.org.uk/mackay/gpB.pdf eq. 47.
```
    κ(x,y) = exp( - 0.5 sum_i(sin (π(x_i - y_i))/r_i))
```
"""
struct PeriodicKernel{T} <: SimpleKernel
    r::Vector{T}
    function PeriodicKernel(; r::AbstractVector{T} = ones(Float64, 1)) where {T<:Real}
        @assert all(r .> 0)
        new{T}(r)
    end
end

PeriodicKernel(dims::Int) = PeriodicKernel(Float64, dims)

PeriodicKernel(T::DataType, dims::Int = 1) = PeriodicKernel(r = ones(T, dims))

metric(κ::PeriodicKernel) = Sinus(κ.r)

kappa(κ::PeriodicKernel, d::Real) = exp(- 0.5d)

Base.show(io::IO, κ::PeriodicKernel) = print(io, "Periodic Kernel, length(r) = ", length(κ.r), ")")


using Plots
using LinearAlgebra

dx = 0.001
L = pi
x = L * (-1+dx:dx:1)
n = length(x)
nquart = Int(floor(n/4))

# Define hat function
f = zeros(n);
f[nquart:2*nquart] = 4*(1:nquart+1)/n
f[2*nquart+1:3*nquart] = (-4*(0:nquart-1).+1)/n

p1 = plot(x, f)

# Compute Fourier series

A0 = sum(f) * dx

fFS = A0/2
A = zeros(20)
B = zeros(20)

for k = 1:20
    A[k] = sum(f.*cos.(pi*k*x/L)) * dx
    B[k] = sum(f.*sin.(pi*k*x/L)) * dx
    global fFS = A[k]*cos.(k*pi*x/L) .+ B[k]*sin.(k*pi*x/L) .+ fFS
    plot!(x, fFS)
end

fFS2 = A0/2

kmax = 100
A2 = zeros(kmax+1)
B2 = zeros(kmax+1)
ERR = zeros(kmax+1)

A2[1] = A0/2
ERR[1] = norm(f-fFS)
p2 = plot(x, f)

for k = 1:kmax
    A2[k+1] = sum(f.*cos.(pi*k*x/L)) * dx
    B2[k+1] = sum(f.*sin.(pi*k*x/L)) * dx
    plot!(x, B2[k] * sin.(2k*pi*x/L))
    global fFS2 = A2[k+1]*cos.(k*pi*x/L) .+ B2[k+1]*sin.(k*pi*x/L) .+ fFS2
    ERR[k+1] = norm(f-fFS2)/norm(f)
end

p3 = plot(0:kmax, A2)
p4 = plot(0:kmax, ERR, yaxis = :log)


using Plots

function branin(x, a=1, b=5.1/(4π^2), c=5/π, r=6, s=10, t=1/(8π))
    return a * (x[2] - b * x[1]^2 + c * x[1] - r)^2 + s*(1-t) * cos(x[1]) + s;
end

# plot

contour(-50:50, -50:50, (x, y)->branin([x, y]))


# Define exit time
A = 600

# Starting amount of money (S0)
k = 500

# Define probability of winning
p = .49
q = 1 - p

# Define b vector
b = zeros(A + 1)
b[A+1] = 1

# Define M matrix
M = zeros(A+1, A+1)
M[1,1] = 1
M[A+1,A+1] = 1

for i=2:A
    global M, q, p
    M[i,i] = -1
    M[i,i-1] = q
    M[i,i+1] = p
end

# Compute exit probabilities by computing M^(-1)b
u = M\b

# Get probability of interest
using Printf
@printf("Probability of walking away with \$%g: %f", k,
 u[k+1])


"""
    euler(f,a,b,y0,M)
berisi program untuk mencari solusi persamaan differensial `y'=f(t,y)`
dengan masalah nilai awal `y(a) = y0` pada interval `[a, b]`. Program ini secara default
berisi 5 masukan, yaitu fungsi `f(t,y)`, titik ujung interval penyelesaian `[a,b]`, nilai awal
`y0`, dan jumlah sub-interval `M`.
# Example
```jl
julia> f(t,y) = (t-y)/2;

julia> sol = euler(f,0,3,1,6)
7×2 Array{Float64,2}:
 0.0  1.0
 0.5  0.75
 1.0  0.6875
 1.5  0.765625
 2.0  0.949219
 2.5  1.21191
 3.0  1.53394
```
return solusi masalah nilai awal `sol`.
"""
function euler(f,a,b,y0,M)
    M = Int(M)
    h = (b-a)/M;
    T = a:h:b;
    Y = Array{Float64}(undef,length(T),1)
    Y[1] = y0;
    # Mulai langkah iterasi euler
    for k = 1:M
        # Rumus iterasi euler
        Y[k+1] = Y[k]+h*f(T[k],Y[k])
    end
    sol = [T Y]
end


doc"""
    Skellam(μ1, μ2)

A *Skellam distribution* describes the difference between two independent [`Poisson`](:func:`Poisson`) variables, respectively with rate `μ1` and `μ2`.

$P(X = k) = e^{-(\mu_1 + \mu_2)} \left( \frac{\mu_1}{\mu_2} \right)^{k/2} I_k(2 \sqrt{\mu_1 \mu_2}) \quad \text{for integer } k$

where $I_k$ is the modified Bessel function of the first kind.

```julia
Skellam(mu1, mu2)   # Skellam distribution for the difference between two Poisson variables,
                    # respectively with expected values mu1 and mu2.

params(d)           # Get the parameters, i.e. (mu1, mu2)
```

External links:

* [Skellam distribution on Wikipedia](http://en.wikipedia.org/wiki/Skellam_distribution)
"""
immutable Skellam <: DiscreteUnivariateDistribution
    μ1::Float64
    μ2::Float64

    function Skellam(μ1::Real, μ2::Real)
        @check_args(Skellam, μ1 > zero(μ1) && μ2 > zero(μ2))
        new(μ1, μ2)
    end
    Skellam(μ::Real) = Skellam(μ, μ)
    Skellam() = new(1.0, 1.0)
end

@distr_support Skellam -Inf Inf

### Parameters

params(d::Skellam) = (d.μ1, d.μ2)

### Statistics

mean(d::Skellam) = d.μ1 - d.μ2

var(d::Skellam) = d.μ1 + d.μ2

skewness(d::Skellam) = mean(d) / (var(d)^1.5)

kurtosis(d::Skellam) = 1.0 / var(d)

### Evaluation

function logpdf(d::Skellam, x::Int)
    μ1, μ2 = params(d)
    - (μ1 + μ2) + (x / 2.0) * log(μ1 / μ2) + log(besseli(x, 2.0 * sqrt(μ1) * sqrt(μ2)))
end

pdf(d::Skellam, x::Int) = exp(logpdf(d, x))

function mgf(d::Skellam, t::Real)
    μ1, μ2 = params(d)
    exp(-(μ1 + μ2) + μ1 * exp(t) + μ2 * exp(-t))
end

function cf(d::Skellam, t::Real)
    μ1, μ2 = params(d)
    exp(-(μ1 + μ2) + μ1 * cis(t) + μ2 * cis(-t))
end

### Sampling

rand(d::Skellam) = rand(Poisson(d.μ1)) - rand(Poisson(d.μ2))


# Parts of this code were taken / derived from NetworkX. See LICENSE for
# licensing details.

"""
    pagerank(g, α=0.85, n=100, ϵ=1.0e-6)

Calculate the [PageRank](https://en.wikipedia.org/wiki/PageRank) of the
directed graph `g` parameterized by damping factor `α`, number of
iterations `n`, and convergence threshold `ϵ`. Return a vector representing
the centrality calculated for each node in `g`, or an error if convergence
is not reached within `n` iterations.
"""
function pagerank end

@traitfn function pagerank(g::::IsDirected, α=0.85, n=100, ϵ=1.0e-6)
    A = adjacency_matrix(g,:in,Float64)
    S = vec(sum(A,1))
    S = 1./S
    S[find(S .== Inf)]=0.0
    M = A' # need a separate line due to bug #17456 in julia
    M = (Diagonal(S) * M)'
    N = Int(nv(g))
    x = fill(1.0/N, N)
    p = fill(1.0/N, N)
    dangling_weights = p
    is_dangling = find(S .== 0)

    for _ in 1:n
        xlast = x
        x = α * (M * x + sum(x[is_dangling]) * dangling_weights) + (1 - α) * p
        err = sum(abs, (x - xlast))
        if (err < N * ϵ)
            return x
        end
    end
    error("Pagerank did not converge after $n iterations.")
end


# we write a function that returns a random vector from [0,1]^d with exponential density
function randExp(d, lam)

	# we generate a uniformly random vector
	y = rand(d)

	# we transform (elementwise) to random numbers in [exp(-lambda),1]
	y = exp(-lam) .+ (1-exp(-lam))*y

	# we now get numbers with exponential density
	x = -log.(y)/lam
	return x
end


# The prime factors of 13195 are 5, 7, 13 and 29.
# What is the largest prime factor of the number 600851475143 ?

num = BigInt(600851475143)
#num = BigInt(13195)

is_prime(n) = !any(n%i == 0 for i in 2:BigInt(floor(sqrt(n))))

function largest_factor(num)
    #step_ = (num % 2 == 0) ? -1 : -2
    num_leftover = num
    for p in (BigInt(p) for p=2:num_leftover if is_prime(p))
        if p == num_leftover
            println(p)
            return p
        elseif num_leftover % p == 0
            num_leftover = BigInt(num_leftover / p)
        end
    end
end
largest_factor(num)
println()


export ad, mad, mad_core

"""
`mad_model(G)` returns the solved linear program whose optimum
value is `mad(G)`.
"""
function mad_model(G::SimpleGraph)
    EE = elist(G)
    VV = vlist(G)

    n = length(VV)
    m = length(EE)

    MOD = Model(get_solver())

    # variables
    @variable(MOD, z >= 0)
    @variable(MOD, x[i = 1:n, j = 1:m; in(VV[i], EE[j])] >= 0)

    # vertex constraints
    for i = 1:n
        v = VV[i]
        # get all edges that have v as an end point
        j_iterator = (j for j = 1:m if EE[j][1] == v || EE[j][2] == v)
        j_list = collect(j_iterator)
        @constraint(MOD, sum(x[i, j] for j in j_list) <= z)
    end

    # edge constraints
    for j = 1:m
        v, w = EE[j]
        i1 = first(findall(x -> x == v, VV))
        i2 = first(findall(x -> x == w, VV))
        @constraint(MOD, x[i1, j] + x[i2, j] == 2)
    end

    @objective(MOD, Min, z)
    optimize!(MOD)
    return MOD, x, VV, EE
end

"""
`mad(G)` computes the maximum average degree of `G`.
See also `mad_core`.
"""
function mad(G::SimpleGraph)
    if cache_check(G, :mad)
        return cache_recall(G, :mad)
    end
    MOD, x, VV, EE = mad_model(G)
    mval = objective_value(MOD)
    cache_save(G, :mad, mval)
    return mval
end

"""
`ad(G)` is the average degree of a vertex in `G`.
"""
ad(G::SimpleGraph) = 2 * NE(G) / NV(G)

"""
`mad_core(G)` returns a subgraph `H` of `G` for which
`ad(H)==mad(G)`.
"""
function mad_core(G::SimpleGraph)
    if cache_check(G, :mad_core)
        return cache_recall(G, :mad_core)
    end
    T = eltype(G)
    GG = deepcopy(G)
    while true
        if NV(GG) == 0
            error("Bad graph or something went wrong")
        end
        # solve the LP for GG
        MOD, x, VV, EE = mad_model(GG)
        n = length(VV)
        m = length(EE)
        err = 0.1 / n
        mval = objective_value(MOD)

        # if balanced, return
        if abs(ad(GG) - mval) <= err
            cache_save(G, :mad_core, GG)
            return GG
        end

        # otherwise, find and destroy defective vertices
        for i = 1:n
            v = VV[i]
            total = 0.0
            for j = 1:m
                e = EE[j]
                if in(v, e)
                    val = value(x[i, j])
                    total += val
                end
            end
            if total < mval - err   # this is a defective vertex
                delete!(GG, v)
            end
        end
    end # end while
    error("This can't happen (but it did)")
    cache_save(G, :mad_core, GG)
    return GG
end


#=

# 101: 1D Laplace equation
([source code](SOURCE_URL))

Let $\Omega=(\gamma_1,\gamma_2)$ with $\gamma_1=0$, $\gamma_2=1$.
This is the simplest second order boundary value problem (BVP)
for a partial differential equation  (PDE):

```math
-\Delta u =0\\
u(\gamma_1)=g_1\\
u(\gamma_2)=g_2.
```

We replace the Dirichlet boundary condition by a Robin boundary
condition with a penalty parameter $\frac{1}{\varepsilon}$:

```math
\nabla u(\gamma_1) + \frac{1}{\varepsilon}(u(\gamma_1)-g_1)=0  \\
-\nabla u(\gamma_2) + \frac{1}{\varepsilon}(u(\gamma_2)-g_2)
=0
```

This penalty method for the implementation of Dirichlet
boundary conditions is used throughout VoronoiFVM.

In order to discretize it, we choose collocation points
$\gamma_1=x_1 < x_2 < \dots < x_n=\gamma_2$.

For instance, we can choose 6 collocation points in $(0,1)$:
From these, we create a discretization grid structure
for working with the method.

This implicitely creates a number of control volumes $\omega_k $ around each
discretization point $x_k$: Let
$\sigma_{k,k+1}=\frac{x_k+x_{k+1}}{2}$. Then $\omega_1=(\gamma_1,\sigma_{1,2})$,
$\omega_k= (\sigma_{k-1,k}, \sigma_{k,k+1})$ for $k=2\dots n-1$, $\omega_{n}=(\sigma_{n-1,n},\gamma_2)$.

```
 x1    x2    x3    x4    x5    x6
 o-----o-----o-----o-----o-----o
 |--|-----|-----|-----|-----|--|
  ω1  ω2     ω3    ω4    ω5  ω6
```

For each $\omega_k$, we integrate the equation

```math
\begin{aligned}
0&=\int_{\omega_k} -\Delta u d\omega=  -\int_{\partial \omega_k} \nabla u ds\\
&= \begin{cases}
u'(\sigma_{1,2}) - u'(0)& k=1\\
u'(\sigma_{k,k+1}) - u'(\sigma_{k-1,k}) & 1<k<n\\
u'(1)- u'(\sigma_{n,n+1})&k=n
\end{cases}\\
&\approx \begin{cases}
\frac{1}{x_2-x_1} g(u_1,u_2) + \frac{1}{\varepsilon}(u_1-0)& k=1\\
\frac{1}{x_k-x_{k-1}}g(u_k,u_{k-1}) -\frac{1}{x_{k+1}-x_{k}}g(u_{k+1},u_{k}) & 1<k<n\\
\frac{1}{\varepsilon}(u_n-1)+ \frac{1}{x_n-x_{n-1}} g(u_{n},u_{n-1})&k=n
\end{cases}
\end{aligned}
```

In the last equation, we wrote $u_k=u(x_k)$ and $g(u_k,u_l)=u_k-u_l$. For
the interior interfaces between control volumes, we replaced $u'$ by a
difference quotient. In the boundary control volumes, we replaced $u'$  by the boundary
conditions.

In the example below, we fix a number of species and  write a Julia function describing $g$,
we create a physics record, and a finite volume system with one unknown species and
a dense matrix to describe it's degrees of freedom
(the matrix used  to calculate the solution is sparse).
We give the species the number 1 and enable it for grid region number one 1.
Then, we set boundary conditions for species 1 at $\gamma_1, \gamma_2$.

We create a zero initial value and a solution vector
and initialize them.

With these data, we solve the system.

We wrap this example and all later ones
into a module structure. This allows to load
all of them at once into the REPL without name
clashes. We shouldn't forget the corresponding end
statement.

=#

module Example101_Laplace1D

using VoronoiFVM

## Flux function which describes the flux
## between neigboring control volumes
function g!(f,u0,edge)
    u=unknowns(edge,u0)
    f[1]=u[1,1]-u[1,2]
end

function main()

    nspecies=1 ## Number of species
    ispec=1    ## Index of species we are working with

    ## Create a one dimensional discretization grid
    ## Each grid cell belongs to a region marked by a region number
    ## By default, there is only one region numbered with 1
    X=collect(0:0.2:1)
    grid=VoronoiFVM.Grid(X)

    ## Create a physics structure
    physics=VoronoiFVM.Physics(num_species=nspecies,flux=g!)

    ## Create a finite volume system
    sys=VoronoiFVM.System(grid,physics)

    ## Enable species 1 in region 1
    enable_species!(sys,ispec,[1])

    ## Set boundary conditions at boundary regions 1 and 2
    boundary_dirichlet!(sys,ispec,1,0.0)
    boundary_dirichlet!(sys,ispec,2,1.0)

    ## Create & initialize array for solution and initial value
    inival=unknowns(sys,inival=0)
    solution=unknowns(sys)

    ## Solve stationary problem
    solve!(solution,inival,sys)

    ## Return test value
    return solution[3]
end

## Called by unit test

function test()
    main() ≈ 0.4
end

end


function recurring_decimal_length(n)
    d = 1
    while true
        if powermod(10, d, n) == 1
            return d
        end
        d += 1
    end
end

function longest_recurring_decimal(upper_limit)
    longest = 0; arg = 0

    for i in 2:upper_limit
        if i % 2 == 0 || i % 5 == 0
            continue
        end

        len = recurring_decimal_length(i)
        if len > longest
            longest = len; arg = i
        end
    end

    return arg
end

println(longest_recurring_decimal(1000))


using Random
using LinearAlgebra
"""
dir_sample(dir_type::String, d::Int, m::Int, seed::Int)

"dir_sample"
samples "m" directions in dimension "d" of type "dir_type"

Inputs:
(1) dir_type: a string, the direction type
(2) d: an integer, the dimension in which the directions are sampled
(3) m: an integer, the number of directions to generate
(4) seed: an integer, used for random sampling

Output:
an m x d matrix, where each row is a direction in R^d
"""
function dir_sample(dir_type::String, d::Int, m::Int, seed::Int,
    theta = pi/4)
    ## set the seed for random number generation
    Random.seed!(seed)
    if dir_type=="pos"
        directions = zeros(0,d)
        num_row = 0
        while num_row<m
            direction = randn(1,d)
            if sum(direction.>0)==d
                directions = [directions; direction]
                num_row+=1
            end
        end
        return directions
    elseif dir_type=="unif"
        directions = randn(m,d)
        return directions
    elseif dir_type=="angle"
        directions = zeros(0,d)
        num_row = 0
        vec_1 = ones(d)/sqrt(d) ## the vector [1,1,...,1]\in\R^d
        num_row = 0
        while num_row<m
            direction = randn(1,d)
            direction = direction/norm(direction)
            if acos(dot(vec_1, direction'))<theta
                directions = [directions; direction]
                num_row+=1
            end
        end
        return directions
    end
end

# ## sample usage code
# dir_type = "pos" ## "unif"
# d = 2
# m = 100
# seed = 1
# A = dir_sample(dir_type::String, d::Int, m::Int, seed::Int)

# using Plots
# pyplot()

# plot()
# for i = 1:size(A,1)
#     plot!([0, A[i,1]], [0,A[i,2]], legend = false)
# end
# plot!()

"""
"UnifRandBall"
generates a d x N matrix, whose columns represent points,
randomly uniformly sampled in the unit ball of R^d
"""
function UnifRandBall(d::Int,N::Int)
    Collect = zeros(d,0) ## Set an (d x 0) empty matrix for storing
    i = 0 ## number of collected points
    while i<N
        pt = 2*rand(d).-1 ## (uniformly) generate a random point inside [-1,1]^d
        if norm(pt)<1 ## if the point is inside the unit ball, collect it
            i = i+1
            Collect = [Collect pt]
        end
    end
    return Collect
end

"""
pt_sample(pt_type::String, d::Int, n::Int, seed::Int)

"pt_sample"
samples "n" points in R^d of type "pt_type"

Inputs:
(1) pt_type: a string, the point type
(2) d: an integer, the dimension in which the directions are sampled
(3) n: an integer, the number of directions to generate
(4) seed: an integer, used for random sampling
(5) r(optional): the reciprocal of the curvature of the de-sitter space, between 0 and 1, exclusive

Output:
a d x n matrix, where each row is a direction in R^d
"""
function pt_sample(pt_type::String, d::Int, n::Int, seed::Int,
        r=0.1::Float64)
    ## set the seed for random number generation
    Random.seed!(seed)
    if pt_type=="uniform"
        pts = UnifRandBall(d,n)
        return pts
    elseif pt_type=="sphere"
        pts = UnifRandBall(d,n)
        for a = 1:n
            pts[:,a] = pts[:,a]/norm(pts[:,a])
        end
        return pts
    elseif pt_type=="de-sitter"
        c = 1/r ## the constant/curvature c in x_1^2+...+x_(d-1)^2-x_d^2 = 1/c^2
        pts = zeros(d,0)
        a = 0
        while a<n
            new_pt = UnifRandBall(d-1,1)
            x_d_a2 = sum(new_pt.^2)-1/c^2
            if x_d_a2>0
                new_pt = [new_pt;(-1)^a*sqrt(x_d_a2)] ## (-1)^a is used to cover both the upper and lower piece
                pts = [pts new_pt]
                a+=1
            end
        end
        return pts
    end
end

# ## sample usage code
# pt_type = "sphere" ## "uniform"
# d = 2
# n = 100
# seed = 1
# X = pt_sample(pt_type::String, d::Int, n::Int, seed::Int)

# using Plots
# pyplot()

# scatter(X[1,:], X[2,:], legend = false)