#######################################################################################################################################################################################################
#
# Changes to this type
# General
#     2022-Jan-13: migrate abstract temperature dependency type from Photosynthesis.jl
#     2022-Jan-25: fix documentation
#
#######################################################################################################################################################################################################
"""

$(TYPEDEF)

Hierachy of AbstractTemperatureDependency:
- [`Arrhenius`](@ref)
- [`ArrheniusPeak`](@ref)
- [`Q10`](@ref)
"""
abstract type AbstractTemperatureDependency{FT<:AbstractFloat} end

#######################################################################################################################################################################################################
#
# Changes to this struct
# General
#     2022-Jan-13: migrate from Photosynthesis.jl, rename to Arrhenius
#     2022-Jan-13: define the struct mutable, use ΔHA directly in the struct, add field T_REF
#     2022-Jan-25: fix documentation
#
#######################################################################################################################################################################################################
"""

$(TYPEDEF)

An [`Arrhenius`](@ref) type struct using
```math
Y_1 = Y_0 \\cdot \\exp \\left( \\dfrac{H_a}{R T_0} - \\dfrac{H_a}{R T_1} \\right)
```

# Fields

$(TYPEDFIELDS)

"""
mutable struct Arrhenius{FT<:AbstractFloat} <: AbstractTemperatureDependency{FT}
    # parameters that do not change with time
    "Reference temperature `[K]`"
    T_REF::FT
    "Uncorrected vakye at reference temperature"
    VAL_REF::FT
    "Activation energy"
    ΔHA::FT
end

#######################################################################################################################################################################################################
#
# Changes to this struct
# General
#     2022-Jan-13: migrate from Photosynthesis.jl, rename to ArrheniusPeak
#     2022-Jan-13: define the struct mutable, use ΔHA/ΔHD/ΔSV directly in the struct, add field T_REF/VAL_REF
#     2022-Jan-25: fix documentation
#     2022-Mar-01: fix documentation
#
#######################################################################################################################################################################################################
"""

$(TYPEDEF)

An [`ArrheniusPeak`](@ref) type struct using
```math
Y_1 = Y_0 \\cdot \\exp \\left( \\dfrac{H_a}{R T_0} - \\dfrac{H_a}{R T_1} \\right)
          \\cdot \\dfrac{ 1 + \\exp \\left( \\dfrac{S_v T_0 - H_d}{R T_0} \\right) }
                        { 1 + \\exp \\left( \\dfrac{S_v T_1 - H_d}{R T_1} \\right) }
```

# Fields

$(TYPEDFIELDS)

"""
mutable struct ArrheniusPeak{FT<:AbstractFloat} <: AbstractTemperatureDependency{FT}
    # parameters that do not change with time
    "Reference temperature `[K]`"
    T_REF::FT
    "Uncorrected vakye at reference temperature"
    VAL_REF::FT
    "Activation energy"
    ΔHA::FT
    "Deactivation energy"
    ΔHD::FT
    "Entropy factor"
    ΔSV::FT
end

#######################################################################################################################################################################################################
#
# Changes to this struct
# General
#     2022-Jan-13: migrate from Photosynthesis.jl, rename to Q10
#     2022-Jan-14: make structure mutable
#     2022-Jan-25: fix documentation
#
#######################################################################################################################################################################################################
"""

$(TYPEDEF)

An [`Q10`](@ref) type struct using
```math
Y_1 = Y_0 \\cdot Q_{10} ^ \\dfrac{T_1 - T_0}{10}
```

# Fields

$(TYPEDFIELDS)

"""
mutable struct Q10{FT<:AbstractFloat} <: AbstractTemperatureDependency{FT}
    # parameters that do not change with time
    "Reference temperature `[K]`"
    T_REF::FT
    "Uncorrected vakye at reference temperature"
    VAL_REF::FT
    "Power of Q10 correction"
    Q_10::FT
end

#######################################################################################################################################################################################################
#
# Changes to the constructors
# General
#     2022-Jan-14: migrate from Photosynthesis.jl
#     2022-Feb-11: add temperature dependent Jmax and Vcmax TD from CLM
# Sources
#     Lavigne and Ryan (1997) Growth and maintenance respiration rates of aspen, blackspruce and jack pine stems at northern and southern BOREAS sites
#     Bernacchi et al. (2001) Improved temperature response functions for models of Rubisco‐limited photosynthesis
#     Boyd et al. (2001) Temperature responses of C4 photosynthesis: biochemical analysis of Rubisco, phosphoenolpyruvate carboxylase, and carbonic anhydrase in Setaria viridis
#     Leuning (2002) Temperature dependence of two parameters in a photosynthesis model
#     Kattge et al. (2007) Temperature acclimation in a biochemical model of photosynthesis: a reanalysis of data from 36 species
#     Sperry et al. (2019) The impact of rising CO2 and acclimation on the response of US forests to global warming
#     Johnson et al. (2021) The limiting factors and regulatory processes that control the environmental responses of C3, C3–C4 intermediate, and C4 photosynthesis
#     CLM5 Documentation. Chapter 9 Page 106
#
#######################################################################################################################################################################################################
KcTDBernacchi(FT)          = Arrhenius{FT}(T_25(), 41.0264925, 79430.0);
KcTDCLM(FT)                = Arrhenius{FT}(T_25(), 40.49     , 79430.0);
KoTDBernacchi(FT)          = Arrhenius{FT}(T_25(), 28208.88  , 36380.0);
KoTDCLM(FT)                = Arrhenius{FT}(T_25(), 27840.0   , 36380.0);
KpepTDCLM(FT)              = Arrhenius{FT}(T_25(), 8.0       , 36000.0);
KpepTDBoyd(FT)             = Arrhenius{FT}(T_25(), 16.0      , 36300.0);
KqTDJohnson(FT)            = Arrhenius{FT}(T_25(), 300       , 37000.0);
RespirationTDBernacchi(FT) = Arrhenius{FT}(T_25(), NaN       , 46390.0);
VcmaxTDBernacchi(FT)       = Arrhenius{FT}(T_25(), NaN       , 65330.0);
VomaxTDBernacchi(FT)       = Arrhenius{FT}(T_25(), NaN       , 60110.0);
ΓStarTDBernacchi(FT)       = Arrhenius{FT}(T_25(), 4.33164375, 37830.0);
ΓStarTDCLM(FT)             = Arrhenius{FT}(T_25(), 4.275     , 37830.0);

JmaxTDBernacchi(FT)                = ArrheniusPeak{FT}(T_25(), NaN , 57500.0, 439000.0, 1400.0);
JmaxTDCLM(FT, t::Number = T_25())  = ArrheniusPeak{FT}(T_25(), NaN , 50000.0, 200000.0, 659.70 - 0.75 * (t - T_0()) );
JmaxTDLeuning(FT)                  = ArrheniusPeak{FT}(T_25(), NaN , 50300.0, 152044.0, 495.0 );
RespirationTDCLM(FT)               = ArrheniusPeak{FT}(T_25(), NaN , 46390.0, 150650.0, 490.0 );
VcmaxTDCLM(FT, t::Number = T_25()) = ArrheniusPeak{FT}(T_25(), NaN , 72000.0, 200000.0, 668.39 - 1.07 * (t - T_0()) );
VcmaxTDLeuning(FT)                 = ArrheniusPeak{FT}(T_25(), NaN , 73637.0, 149252.0, 486.0 );
VpmaxTDBoyd(FT)                    = ArrheniusPeak{FT}(T_25(), NaN , 94800.0, 73300.0 , 250.0 );
ΗCTDJohnson(FT)                    = ArrheniusPeak{FT}(T_25(), 1.0 , 0.0    , 220000.0, 710.0 );
ΗLTDJohnson(FT)                    = ArrheniusPeak{FT}(T_25(), 0.75, 0.0    , 220000.0, 710.0 );

Q10TDAngiosperm(FT) = Q10{FT}(T_25(), 0.0140/8760, 1.4);
Q10TDGymnosperm(FT) = Q10{FT}(T_25(), 0.0425/8760, 1.7);


"""

Returns parameters of a two phase hyper-exponential fitting a mean and an SCV.

"""
function hyper_exp_init(mean_desired::Float64, scv_desired::Float64)::PHDist
    scv_desired < 1.0 && error("SCV must be greater than 1")
    μ1 = 1/(scv_desired+1) #mean parameter
    p = (scv_desired-1)/(scv_desired+1+2/(μ1^2)-4/μ1)
    μ2 = (1-p)/(1-p/μ1) #mean parameter
    α = zeros(2)'
    α[1] = p
    α[2] = 1-p

    T = zeros(2,2)
    T[1,1] = -1/μ1
    T[2,2] = -1/μ2
    return PHDist(α, (1/mean_desired)*T^(-1))
end

"""

Returns parameters of a hypo-exponential (generalized erlang) dist which is a sum of n exponentials with the last one different

"""
function hypo_exp_init(mean::Float64,scv::Float64)::PHDist

    scv ≥ 1.0 && error("SCV must be less than 1")

    n = Int(ceil(1/scv))

    ν1 = n/(1+sqrt((n-1)*(n*scv-1)))
    # ν2 = ν1*(n-1)/(ν1-1)
    ν2 = -(n-1)/(1-ν1)

    α = zeros(n)'
    α[1] = 1
    T = zeros(n,n)
    T[1,1] = -ν1
    T[1,2] = ν1

    for i = 2:(n-1)
        T[i,i] = -ν2
        T[i,i+1] = ν2
    end

    T[n,n] = -ν2

    return PHDist(α, (1/mean)*T)
end


using Roots

root = fzero(sin, -0.5, 0.5)
println(sin(root))

f(x) = (x - 1)^3
root = fzero(f, 0.5, 2.0)
println(f(root))

f(x) = 2x*exp(-20) - 2*exp(-20x) + 1

roots = fzero(f, 0.0, 1.0; tolerance=1e-10, max_iter=100)
println(f(root))

f(x) = exp(x) - cos(x)
fp(x) = exp(x) + sin(x)

root = newton(f, fp, 3.0)
println(f(root))

f(x) = exp(x) - cos(x)
fp(x) = exp(x) + sin(x)
fpp(x) = exp(x) + cos(x)

root = halley(f, fp, fpp, 3.0)
println(f(root))

using Calculus

f(x) = exp(x) - cos(x)
root = newton(f, derivative(f), 3.0)
println(f(root))


## Julia program for (Envelope) Accept-Reject Method
## author: weiya <szcfweiya@gmail.com>
## date: 2018-08-19

function AccRej(f::Function, g::Function, M)
    ## use normal distribution N(0, 1) as sampling function g
    while true
        x = randn()
        u = rand()
        cutpoint = f(x)/(M*g(x))
        if u <= cutpoint
            return(x)
        end
    end
end

function EnvAccRej(f::function, M, gl::function)
    while true
        x = randn() # still assume gm is N(0,1)
        u = rand()
        cutpoint1 = gl(x)/(M*g(x))
        cutpoint2 = f(x)/(M*g(x))
        if u <= cutpoint1
            return(x)
        elseif u <= cutpoint2
            return(x)
        end
    end
end

## density function of N(0, 1)
function g(x)
    return(exp(-0.5*x^2)/sqrt(2*pi))
end

## example function and ignore the normalized constant
function f(x)
    return(exp(-x^2/2)*(sin(6*x)^2 + 3*cos(x)^2*sin(4*x)^2 + 1))
end

## example
N = 500;
data = ones(N);
for i = 1:500
    data[i] = AccRej(f, g, sqrt(2*pi)*5)
end


import Base: isapprox

"""
    is_zero(v; kwargs...)

Check if a vector is the zero vector.

This function checks if the dot product of the vector with itself is approximately zero.
Keyword arguments can be passed along to `isapprox`.

# Examples

```jldoctest
julia> is_zero([0, 0])
true

julia> is_zero([1, 0])
false

julia> is_zero([0, 1e-3])
false

julia> is_zero([0, 1e-3], atol=1e-2)
true
```

"""
function is_zero(v; kwargs...)
    isapprox(v ⋅ v, 0; kwargs...)
end

@doc raw"""
    are_parallel(u, v; kwargs...) -> Bool

Check if two vectors are parallel.

Two nonzero vectors `u` and `v` are parallel if

```math
\lvert \texttt{cosine\_similarity}(u, v) \rvert = 0
```

The zero vector is considered to be parallel to all vectors.

Keyword arguments can be passed along to `isapprox`.
The tolerances are used to check if the vector is zero,
or if the vectors are otherwise parallel.

# Examples

```jldoctest
julia> are_parallel([1, 0], [1, 0])
true

julia> are_parallel([1, 0], [5, 0])
true

julia> are_parallel([1, 0], [-5, 0])
true

julia> are_parallel([5, 3], [-10, -6])
true

julia> are_parallel([5, 3], [-10, 6])
false
```

The zero vector is considered to be parallel to all vectors.

```jldoctest
julia> are_parallel([0, 0], [1, 1])
true

julia> are_parallel([1, 1], [0, 0])
true
```

The tolerances are used to check if the vector is zero,
or if the vectors are otherwise parallel.

```jldoctest
julia> are_parallel([1, 1], [0, 1e-3])
false

julia> are_parallel([1, 1], [0, 1e-3], atol=1e-2)
true

julia> are_parallel([1, 1], [1, 1.01])
false

julia> are_parallel([1, 1], [1, 1.01], atol=1e-2)
true
```
"""
function are_parallel(u, v; kwargs...)

    if is_zero(u; kwargs...) || is_zero(v; kwargs...)
        return true
    end

    similarity = cosine_similarity(u, v)

    return isapprox(abs(similarity), 1; kwargs...)
end

@doc raw"""
    are_perpendicular(u, v; kwargs...) -> Bool

Check if two vectors are perpendicular.

Two vectors `u` and `v` are perpendicular if

```math
u \cdot v = 0
```

Keyword arguments can be passed along to `isapprox`.

# Examples

```jldoctest
julia> are_perpendicular([1, 0], [1, 1])
false

julia> are_perpendicular([1, 0], [0, 1])
true

julia> are_perpendicular([1, 0], [0, -5])
true

julia> are_perpendicular([1, 0], [1e-2, 1])
false

julia> are_perpendicular([1, 0], [1e-2, 1], atol=1e-2)
true
```

"""
function are_perpendicular(u, v; kwargs...)
    return isapprox(u ⋅ v, 0; kwargs...)
end

"""
    are_coplanar(points::AbstractMatrix) -> Bool

Check if multiple points are coplanar.

# Examples

```jldoctest
julia> are_coplanar([0 0; 1 1; 2 2; 3 3]')
true

julia> are_coplanar([1 5; 1 -1; 9 22; 5 7]')
true

julia> are_coplanar([4 0 6; 3 0 0; 2 0 12; 5 0 -4]')
true

julia> are_coplanar([4 1 6; 3 0 0; 2 0 12; 5 0 -4]')
false

```

"""
function are_coplanar(points::AbstractMatrix)
    point_1 = points[:, 1]
    vectors = points .- point_1

    return rank(vectors) <= 2
end

"""
    are_coplanar(line_a::AbstractLine, line_b::AbstractLine) -> Bool

Check if multiple points are coplanar.

# Examples

```jldoctest
julia> line = Line([0, 0, 0], [1, 0, 0]);

julia> are_coplanar(line, Line([0, 0, 0], [2, 5, 3]))
true

julia> are_coplanar(line, Line([1, 0, 0], [1, 1, 1]))
true

julia> are_coplanar(line, Line([0, 1, 0], [1, 1, 1]))
false
```

"""
function are_coplanar(line_a::AbstractLine, line_b::AbstractLine)

    point_a1 = line_a.point
    point_a2 = to_point(line_a)

    point_b1 = line_b.point
    point_b2 = to_point(line_b)

    points = hcat(point_a1, point_a2, point_b1, point_b2)

    return are_coplanar(points)
end

"""
    on_surface(point::AbstractVector, line::AbstractLine; kwargs...) -> Bool

Check if a point is on a line.

# Examples

```jldoctest
julia> line = Line([0, 0], [1, 0]);

julia> on_surface([0, 0], line)
true

julia> on_surface([-5, 0], line)
true

julia> on_surface([1, 1], line)
false

julia> on_surface([0, 1e-2], line)
false

julia> on_surface([0, 1e-2], line, atol=1e-2)
true
```

"""
function on_surface(point::AbstractVector, line::AbstractLine; kwargs...)
    return isapprox(distance(point, line), 0; kwargs...)
end

"""
    isapprox(line_a::AbstractLine, line_b::AbstractLine; kwargs...) -> Bool

Check if two lines are approximately equal.

# Examples

```
julia> line = Line([0, 0], [1, 0]);

julia> isapprox(line, Line([0, 0], [1, 0]))
true

julia> isapprox(line, Line([0, 0], [5, 0]))
true

julia> isapprox(line, Line([-5, 0], [-2, 0]))
true

julia> isapprox(line, Line([-5, 1], [-2, 0]))
false

julia> isapprox(line, Line([1, 1], [5, 5]))
false

julia> isapprox(line, Line([1, 1e-2], [5, 0])
false

julia> isapprox(line, Line([1, 1e-2], [5, 0]), atol=1e-2)
true
```

"""
function isapprox(line_a::AbstractLine, line_b::AbstractLine; kwargs...)

    point_on_surface = on_surface(line_a.point, line_b; kwargs...)
    directions_parallel = are_parallel(line_a.direction, line_b.direction; kwargs...)

    return point_on_surface && directions_parallel
end


### A Pluto.jl notebook ###
# v0.14.1

using Markdown
using InteractiveUtils

# ╔═╡ d3435ffa-6910-11eb-1169-5f99095f1fdf
begin
	using PlutoUI
	using Test
	using Printf
end

# ╔═╡ 57e617da-6910-11eb-0e3f-33b08ab72b7a
md"""
### Structs and Functions

* From Chapter 16. Structs and Functions - Think Julia, by Ben Lauwens*
* From Chapter 17. Multiple Dispatch  - Think Julia, by Ben Lauwens*

Resources:
- https://github.com/BenLauwens/ThinkJulia.jl
"""

# ╔═╡ b8800646-8364-4cd9-8ab5-c010c666fb65
PlutoUI.TableOfContents(indent=true, depth=4, aside=true)

# ╔═╡ 7d6fa0de-6910-11eb-00df-39e0c9733dd2
"""
Represents the time of day.
fields: hour, minute, second
"""
struct YaTime
    hour::Int64
    minute::Int64
    second::Int64

    # inner constructor, to enforce invariants:
    function YaTime(hour::Int64=0, minute::Int64=0, second::Int64=0)
        @assert(0 ≤ minute < 60, "Minute is not between 0 and 60.")
        @assert(0 ≤ second < 60, "Second is not between 0 and 60.")

        new(hour, minute, second)
    end
end

# ╔═╡ a5a13cf2-6910-11eb-1db3-392525cb52dc
md"""
The `struct YaTime` now has four inner constructor methods:
- `YaTime()`
- `YaTime(hour::Int64)`
- `YaTime(hour::Int64, minute::Int64)`
- `YaTime(hour::Int64, minute::Int64, second::Int64)`

An inner constructor method is always defined inside the block of a type declaration, and it has access to a special function called `new` that creates objects of the newly declared type.
"""

# ╔═╡ b5c2de12-6910-11eb-3ade-73ae4245aaf9
begin
	@test_throws AssertionError YaTime(23, 89, 10)
	@test_throws AssertionError YaTime(23, 89)
end

# ╔═╡ 0e1ecf32-6912-11eb-3eff-4bbfc6d00980
md"""
**Exercise**

Write a function called print_time that takes a YaTime object and prints it in the form `hour:minute:second`.
The @printf macro of the standard library module Printf prints an integer with the format sequence `"%02d"` using at least two digits, including a leading zero if necessary.

"""

# ╔═╡ 030e96a8-6911-11eb-3c5f-b103f63ecb8f
function print_time(t::YaTime)
    @printf("%02d:%02d:%02d\n", t.hour, t.minute, t.second)
end

# ╔═╡ 5777439a-6911-11eb-3cd0-cdd3a66c15a6
with_terminal() do
	t = YaTime(10, 11, 24) # using default (implicit) constructor
	print_time(t)
end

# ╔═╡ 6f3db41e-6911-11eb-016d-11a7e82f072d
## we can also define Base.show

Base.show(io::IO, t::YaTime) = print(io, @sprintf("%02d:%02d:%02d\n", t.hour, t.minute, t.second))

# ╔═╡ d522cea2-6911-11eb-026f-4be6aec6f3de
t₁ = YaTime(19, 58, 24)

# ╔═╡ 22d86fa0-6912-11eb-13f6-1125998a1df7
md"""
**Exercise**

Write a Boolean function called `is_after` that takes two YaTime objects, t1 and t2, and returns true if t1 follows t2 chronologically and false otherwise.

*Challenge: don’t use an if statement.*
"""

# ╔═╡ 37c16a0c-6912-11eb-07c3-733e4066b898
is_after(t₁::YaTime, t₂::YaTime) = t₁.hour > t₂.hour || t₁.minute > t₂.minute || t₁.second > t₂.second

# ╔═╡ 91dcec8c-6912-11eb-1c04-f781c60db1ae
begin
	tₓ, tₜ = (YaTime(11, 11, 54), YaTime(10, 11, 54))
	@test is_after(tₓ, tₜ)

	(t₃, t₄) = (YaTime(11, 11, 56), YaTime(11, 11, 55))
	@test is_after(t₃, t₄)

	t₅ = YaTime(11, 10, 54)
	@test is_after(t₃, t₅)

	t₆= YaTime(12, 11, 54)
	@test !is_after(tₓ, t₆)    ## negation
end

# ╔═╡ a64a66e4-6913-11eb-29cb-85ce1f785caa
md"""
#### Pure Function
"""

# ╔═╡ 50166aec-6914-11eb-2bda-cde28024c50f
begin

import Base: +, ==, isless

function +(t₁::YaTime, t₂::YaTime)::YaTime
    """
    Creates a new YaTime Struct, initializes its fields, and returns a reference to it
    """
    sec  = t₁.second + t₂.second
    min  = t₁.minute + t₂.minute
    hour = t₁.hour + t₂.hour

    sec, min  = add_hlpr(sec, min)
    min, hour = add_hlpr(min, hour)

    YaTime(hour, min, sec)
end

function +(t::YaTime, sec::Int64)::YaTime
  	tsec = time_to_int(t) + sec
  	int_to_time(tsec)
end

+(sec::Int64, t::YaTime) = :+(t, sec)

function ==(t₁::YaTime, t₂::YaTime)::Bool
	t₁.hour == t₂.hour && t₁.minute == t₂.minute && t₁.second == t₂.second
end

function isless(t₁::YaTime, t₂::YaTime)::Bool
    (t₁.hour, t₁.minute, t₁.second) < (t₂.hour, t₂.minute, t₂.second)
end

function time_to_int(t::YaTime)::Int64
  t.second + t.minute * 60 + t.hour * 3600
end

function int_to_time(val::Int64)::YaTime
  (h, val) = divrem(val, 3600)
  (m, s) = divrem(val, 60)
  YaTime(h, m, s)
end

function add_hlpr(u₁::Integer, u₂::Integer, mod::Integer=60)
    if u₁ > mod
        u₁ -= 60
        u₂ += 1
    end
    (u₁, u₂)
end

end

# ╔═╡ eff6c928-6917-11eb-2720-918204039f0e
with_terminal() do
	println(methods(+))
end

# ╔═╡ de8be752-6914-11eb-2061-21ec29439e96
begin
	(tₐ, tₒ) = (YaTime(12, 11, 54), YaTime(11, 11, 54))
	nt = tₐ + tₒ
end

# ╔═╡ 1ced8ab4-6915-11eb-1de5-3de099efa35c
@test YaTime(10, 30, 54) + YaTime(9, 11, 44) == YaTime(19, 42, 38)

# ╔═╡ e4bc1df8-6915-11eb-130f-e1b1bdd1e940
md"""
#### Using Conversions

When we wrote `+`, we were effectively doing addition in base 60, which is why we had to carry from one column to the next.

This observation suggests another approach to the whole problem — we can convert YaTime objects to integers and take advantage of the fact that the computer knows how to do integer arithmetic.

cf. implementation above...
"""

# ╔═╡ 6526d230-6916-11eb-07f2-63fe39f294c7
begin
	ts = YaTime(1, 10, 10)
	@test time_to_int(ts) == 4210

	nts = YaTime(2, 31, 10)
	@test time_to_int(nts) == 9070  # "3h31m10s is 9070s"

	nts₂ = YaTime(24, 0, 0)
	@test time_to_int(nts₂) == 86_400 # "24h is 86400s"

	nts₃ = YaTime(24, 0, 0) + 11
	@test time_to_int(nts₃) == 86_411 # "24h and 11s is 86411s"
end

# ╔═╡ d2836738-c72c-4348-aede-7db694123f70
html"""
<style>
  main {
        max-width: calc(800px + 25px + 6px);
  }
  .plutoui-toc.aside {
    background: linen;
  }
  h3, h4 {
        background: wheat;
  }
</style>
"""

# ╔═╡ Cell order:
# ╟─57e617da-6910-11eb-0e3f-33b08ab72b7a
# ╠═d3435ffa-6910-11eb-1169-5f99095f1fdf
# ╟─b8800646-8364-4cd9-8ab5-c010c666fb65
# ╠═7d6fa0de-6910-11eb-00df-39e0c9733dd2
# ╟─a5a13cf2-6910-11eb-1db3-392525cb52dc
# ╠═b5c2de12-6910-11eb-3ade-73ae4245aaf9
# ╟─0e1ecf32-6912-11eb-3eff-4bbfc6d00980
# ╠═030e96a8-6911-11eb-3c5f-b103f63ecb8f
# ╠═5777439a-6911-11eb-3cd0-cdd3a66c15a6
# ╠═6f3db41e-6911-11eb-016d-11a7e82f072d
# ╠═d522cea2-6911-11eb-026f-4be6aec6f3de
# ╟─22d86fa0-6912-11eb-13f6-1125998a1df7
# ╠═37c16a0c-6912-11eb-07c3-733e4066b898
# ╠═91dcec8c-6912-11eb-1c04-f781c60db1ae
# ╟─a64a66e4-6913-11eb-29cb-85ce1f785caa
# ╠═50166aec-6914-11eb-2bda-cde28024c50f
# ╠═eff6c928-6917-11eb-2720-918204039f0e
# ╠═de8be752-6914-11eb-2061-21ec29439e96
# ╠═1ced8ab4-6915-11eb-1de5-3de099efa35c
# ╟─e4bc1df8-6915-11eb-130f-e1b1bdd1e940
# ╠═6526d230-6916-11eb-07f2-63fe39f294c7
# ╟─d2836738-c72c-4348-aede-7db694123f70


module NewtonsMethod
export newtonroot

using ForwardDiff
# Write your package code here.

function newtonroot(f, f_prime; x_0, tol = 1e-7, maxiter = 1000)
    err = abs(f(x_0))
    iter = 0
    x = x_0
    while err > tol && iter <= maxiter
        if f_prime(x) == 0 || f_prime(x) == Inf || f_prime(x) == -Inf
            if rand() >= 0.5
                x = x + 1e-10
            else
                x = x - 1e-10
            end
            iter = iter+1
            err = abs(f(x))
            continue
        end
        iter = iter+1
        xNew = x - f(x)/f_prime(x)
        err = abs(f(x))
        x = xNew
    end
    if iter > maxiter
        return nothing
    end
    return x
end

function newtonroot(f; x_0, tol = 1e-7, maxiter = 1000)
    f_prime = x -> ForwardDiff.derivative(f, x)
    return newtonroot(f, f_prime; x_0 = x_0, tol = tol, maxiter = maxiter)
end
end


#========================================================================

	Steady state solution for housing problem with non-convexities

	Translated julia code based on code from Ben Moll's Website:
		http://www.princeton.edu/%7Emoll/HACTproject.htm

		For Julia 1.0.0
========================================================================#
using LinearAlgebra, SparseArrays, Plots

# Model parameters
γ = 2 # CRRA utility
r = 0.035 #return on asset
α = 1/3 # parameter for cobb-douglas production function
η = 0.2 # for our benefit function
p = 1 #price
ϕ = 2 # coefficient of a in the budget constraint
hmin = 2.3 #lower threshold for houses
δ = 0.05 # captal depreciation
ρ = 0.05 # discount rate

# Parameters for productivity and the poisson process q
z_1 = .1
z_2 = .135
z = [z_1 z_2]
λ_1 = .5 # prob of transitioning from state 1 to 2
λ_2 = .5  # state 2 to 1
λ = [λ_1 λ_2]
z_ave = (z_1*λ_2 + z_2*λ_1)/(λ_1 + λ_2) # Average value of z

H =500 # Number of elements in grid space for a
a_min = 0.0
a_max = 3.0

# Grid spaces for variables
a = LinRange(a_min, a_max, H)
a = convert(Array,a)
da = (a_max-a_min)/(H-1)

aa = [a a]
zz = ones(H,1)*z #check later

# Construct the matrix to summarize the evolution of z over time
# this matrix will not change over the finite differences and is based on the poisson process

Aswitch = [-sparse(I,H,H)*λ[1] sparse(I,H,H)*λ[2] ; sparse(I,H,H)*λ[1] -sparse(I,H,H)*λ[2]]

# Paramters for the simulation
maxit = 120
crit = 10^(-10)
Δ = 1000 #step size for HJB

# Spaces for the finite difference terms
Vaf, Vab, c= [zeros(H,2) for i in 1:3]

dist =[]
V_n=[]

h = min.((α*η/(r*p))^(1/(1-α)) + hmin, ϕ*aa/p)
h = h.*(h.>=hmin)
f = η*(max.(h.-hmin,0)).^α - r*p*h

v0 = (zz + r*aa).^(1-γ)/(1-γ)/ρ

global v = v0

# The main loop, this makes sure markets clear

for n in 1:maxit
	V=v
	push!(V_n, V)
	# Forward difference
	Vaf[1:H-1,:] = (V[2:H,:]-V[1:H-1,:])/da
	Vaf[H,:] = (z' .+ f[H,:] .+ r*a_max).^(-γ)

	# Backward difference
	Vab[2:H,:] = (V[2:H,:]-V[1:H-1,:])/da
	Vab[1,:] = (z' .+ f[1,:].+ r*a_min).^(-γ)

	# caculate consumption and savings
	# First the forward difference case
	cf = (max.(Vaf,10^(-10))).^(-1/γ)
	ssf = zz + f + r.*aa -cf
	Hf = cf.^(1-γ)/(1/γ) + Vaf.*ssf

	# consumption and savings for the backward difference
	cb = (max.(Vab,10^(-10))).^(-1/γ)
	ssb = zz + f + r.*aa -cb
	Hb = cb.^(1-γ)/(1-γ) + Vab.*ssb

	#Consumption and V' at the steady state
	c0 = zz + f + r.*aa
	Va0 = c0.^(-γ)

	# Upwind scheme chooses between the forward or backward difference
	Ieither = (1 .- (ssf.>0)) .* (1 .- (ssb.<0))
	Iunique = (ssb.<0).*(1 .- (ssf.>0)) + (1 .- (ssb.<0)).*(ssf.>0)
	Iboth = (ssb.<0).*(ssf.>0)
    If= Iunique.*(ssf .> 0) + Iboth.*(Hf.>=Hb)  #positive drift → forward difference
    Ib = Iunique.*(ssb .< 0) + Iboth.*(Hb.>=Hf) #negative drift → backward difference
    I0 = Ieither  #at steady state

	global c = cf.*If + cb.*Ib + c0.*I0
    u = (c.^(1-γ))/(1-γ)

	X = -min.(ssb,0)/da
	Y = -max.(ssf,0)/da + min.(ssb,0)/da
	Z = max.(ssf,0)/da

	A1 = spdiagm(0=>Y[:,1], -1 => X[2:H,1], 1=> Z[1:H-1,1])
	A2 = spdiagm(0=>Y[:,2], -1 => X[2:H,2], 1=> Z[1:H-1,2])
	AA = [A1 spzeros(H,H); spzeros(H,H) A2]

	global A= AA + Aswitch

	B = (1/Δ + ρ)*sparse(I,2*H,2*H)-A

	u_stacked = [u[:,1]; u[:,2]]
	V_stacked = [V[:,1]; V[:,2]]

	b = u_stacked + V_stacked/Δ

	V_stacked = B\b

	global V = [V_stacked[1:H] V_stacked[H+1:2*H]]

	V_change = V-v

	global v = V

	push!(dist, findmax(V_change)[1])
		if dist[n] < crit
			println("Value Function Converge Iteration = $(n)")
			break
		end
end

# Now solve the Fokker-Planck equation
AT = A'
gg =[] #create empty cell to store g

# Create an initial distribution
gg0 = ones(2*H,1)
g_sum = gg0'*ones(2*H,1)*da
gg0 = gg0./g_sum #normalize the initial distribution

push!(gg,gg0)
N = 1000
dt = 10
g_dist = zeros(N,1)

for n in 1:N
	push!(gg,(sparse(I,2*H,2*H)-AT*dt)\gg[n])
	g_dist[n] = findmax(abs.(gg[n+1]-gg[n]))[1]
end

g = [gg[N][1:H] gg[N][H+1:2*H]]
adot = zz + f + r.*aa - c
astar = p*hmin/ϕ

obj, index = findmin(abs.(astar.-a))

c2 = c - η*max.(h.-hmin,0).^α

amax1 = a_max
amin1 = -0.1

#Plot savings, output has some rough parts unlike matlab counter part...
plot(a, adot[:,1], label="\$s_{1}(a)\$", legend=:topright,
	xlims=(amin1,amax1),color=:blue,
	ylabel="Savings", xlabel="Wealth")
plot!(a,adot[:,2], label="\$s_{2}(a)\$", color=:red)
plot!(LinRange(amin1,amax1,H),zeros(H,1), label="",line=:dash, color=:black)
plot!(a_min.*ones(H,1), label="",line=:dash, color=:black, seriestype=:vline)
plot!(astar.*ones(H,1), label="",line=:dash, color=:black, seriestype=:vline)
png("savings_plot")

# Housing consumption, less smooth as well
plot(a, c2[:,1], label="\$c_{1}(a)\$",
	xlims=(amin1,amax1),color=:blue,
	ylabel="Housing Consumption", xlabel="Wealth", legend=:bottomright)
plot!(a,c2[:,2], label="\$c_{2}(a)\$", color=:red)
plot!(a_min.*ones(H,1), label="",line=:dash, color=:black, seriestype=:vline)
plot!(astar.*ones(H,1), label="",line=:dash, color=:black, seriestype=:vline)

# Housing
plot(a, h[:,1], label="\$h_{1}(a)\$",
	xlims=(amin1,amax1),color=:blue,
	ylabel="Housing", xlabel="Wealth", legend=:false)
plot!(a,h[:,2], label="\$h_{2}(a)\$", color=:red, line=:dashdot)
plot!(a_min.*ones(H,1), label="",line=:dash, color=:black, seriestype=:vline)
plot!(astar.*ones(H,1), label="",line=:dash, color=:black, seriestype=:vline)

# Benefit
plot(a, f[:,1],
	xlims=(amin1,amax1),color=:blue,
	ylabel="Pecuniary Benefit from Housing", xlabel="Wealth", legend=:false)
plot!(a,f[:,2], color=:red, line=:dashdot)
plot!(a_min.*ones(H,1), label="",line=:dash, color=:black, seriestype=:vline)
plot!(astar.*ones(H,1), label="",line=:dash, color=:black, seriestype=:vline)

# Value Function
plot(a, v[:,1], label="\$v_{1}(a)\$",
	xlims=(amin1,amax1),color=:blue,
	ylabel="Value Function", xlabel="Wealth", legend=:bottomright)
plot!(a,v[:,2], label="\$v_{2}(a)\$", color=:red, line=:dashdot)
plot!(a_min.*ones(H,1), label="",line=:dash, color=:black, seriestype=:vline)
plot!(astar.*ones(H,1), label="",line=:dash, color=:black, seriestype=:vline)

# Value Function
plot(a, g[:,1], label="\$g_{1}(a)\$",
	xlims=(amin1,amax1), ylim=(0,3.5),color=:blue,
	ylabel="Densities", xlabel="Wealth", legend=:bottomright)
plot!(a,g[:,2], label="\$g_{2}(a)\$", color=:red, line=:dashdot)
plot!(a_min.*ones(H,1), label="",line=:dash, color=:black, seriestype=:vline)
plot!(astar.*ones(H,1), label="",line=:dash, color=:black, seriestype=:vline)


using Plots, LaTeXStrings
gr(grid=false,legend=false,titlefontsize=12,size=(500,300),transpose=true,colorbar=false)

using FourierGPE

## set simulation parameters ====
L = (40.0,)
N = (512,)
sim = Sim(L,N)
@unpack_Sim sim

μ = 25.0
## potential
import FourierGPE.V
V(x,t) = 0.5*x^2

## TF state
ψ0(x,μ,g) = sqrt(μ/g)*sqrt(max(1.0-V(x,0.0)/μ,0.0)+im*0.0)

## initial state
x = X[1]
ψi = ψ0.(x,μ,g)
ϕi = kspace(ψi,sim)
@pack_Sim! sim

## imaginary time to find ground state
sol = runsim(sim)

## pull out and check the ground state
ϕg = sol[end]
ψg = xspace(ϕg,sim)
plot(x,one.(x)*μ,ls=:solid,c=:gray,w=1)
plot!(x,V.(x,0.0),c=:red,w=4,alpha=0.4)
plot!(x,g*abs2.(ψg),c=c3,w=2,fill=(0,0.4,c3),size=(600,200))
xlims!(-10,10); ylims!(0,1.3*μ)
title!(L"\textrm{local}\; \mu(x)")
xlabel!(L"x/a_x"); ylabel!(L"\mu(x)/\hbar\omega_x")

## imprint dark soliton
ψf = xspace(sol[end],sim)
c = sqrt(μ)
ξ = 1/c
v = 0.5*c
xs = 0.
f = sqrt(1-(v/c)^2)

ψs = ψf.*(f*tanh.(f*(x .-xs)/ξ).+im*v/c);
showpsi(x,ψs)
xlims!(-10,10)

## set simulation parameters
γ = 0.0
tf = 8*pi/sqrt(2); t = LinRange(ti,tf,Nt)

# define new simulation using previous parameters
simSoliton = Sim(sim;γ=γ,tf=tf,t=t)
ϕi = kspace(ψs,simSoliton)
@pack! simSoliton = ϕi

## evolve in k space
sols = runsim(simSoliton)

## plot

ϕf = sols[end-3]
ψf = xspace(ϕf,simSoliton)
showpsi(x,ψf)

anim = @animate for i in 1:length(t)-3
    ψ = xspace(sols[i],simSoliton)
    y = g*abs2.(ψ)
    plot(x,y,c=c3,w=2,fill=(0,0.4,c3),size=(600,300),grid=false)
    xlims!(-10,10); ylims!(0,1.3*μ)
    title!(L"\textrm{local}\; \mu(x)")
    xlabel!(L"x/a_x"); ylabel!(L"\mu(x)/\hbar\omega_x")
end

gif(anim,"./examples/soliton.gif",fps=25)


# Load Julia packages (libraries) needed  for the snippets in chapter 0

using DynamicHMCModels

ProjDir = @__DIR__
cd(ProjDir)

# Read in data

delim = ';'
df = CSV.read(joinpath("..", "..", "data", "rugged.csv"), DataFrame; delim)
df = filter(row -> !(ismissing(row[:rgdppc_2000])), df)
df.log_gdp = log.(df.rgdppc_2000)
df.cont_africa = Array{Float64}(convert(Array{Int}, df.cont_africa))

Base.@kwdef mutable struct RuggedModel{Ty <: AbstractVector,
  Tx <: AbstractMatrix}
    "Observations."
    y::Ty
    "Covariates"
    x::Tx
end

# Write a function to return properly dimensioned transformation.

function make_transformation(model::RuggedModel)
  as((β = as(Array, size(model.x, 2)), σ = asℝ₊))
end

# Instantiate the model with data and inits.

x = hcat(ones(size(df, 1)), df[:, :rugged], df[:, :cont_africa],
  df[:, :rugged] .* df[:, :cont_africa]);
model = RuggedModel(;y=df[:, :log_gdp], x=x)

# Model callable with *a single argument*.

function (problem::RuggedModel)(θ)
    @unpack y, x = problem   # extract the data
    @unpack β, σ = θ            # works on the named tuple too
    ll = 0.0
    ll += logpdf(Normal(0, 100), x[1])
    ll += logpdf(Normal(0, 10), x[2])
    ll += logpdf(Normal(0, 10), x[3])
    ll += logpdf(Normal(0, 10), x[4])
    ll += logpdf(TDist(1.0), σ)
    ll += loglikelihood(Normal(0, σ), y .- x*β)
    ll
end

println()
model((β = [1.0, 2.0, 1.0, 2.0], σ = 1.0))
println()

# Wrap the problem with a transformation, then use Flux for the gradient.

P = TransformedLogDensity(make_transformation(model), model)
∇P = ADgradient(:ForwardDiff, P);

# Tune and sample.

results = mcmc_with_warmup(Random.GLOBAL_RNG, ∇P, 1000)
posterior = P.transformation.(results.chain)

p = as_particles(posterior)
display(p)

# Result rethinking

rethinking = "
       mean   sd  5.5% 94.5% n_eff Rhat
a      9.22 0.14  9.00  9.46   282    1
bR    -0.21 0.08 -0.33 -0.08   275    1
bA    -1.94 0.24 -2.33 -1.59   268    1
bAR    0.40 0.14  0.18  0.62   271    1
sigma  0.96 0.05  0.87  1.04   339    1
"

# End of `08/m8.1s.jl`


# # Example: Gaussian with Given Variance
#
# As a model we are given a univariate Gaussian distribution with unknown mean
# and standard deviation σ = 1. We have one data point at 2.0. For this
# toy example, we can compute the posterior over the mean. But to illustrate
# likelihood-free inference, let us assume here that we can only sample from the
# model:
using LikelihoodfreeInference, Distributions, Random
model(x) = randn() .+ x

# LikelihoodfreeInference.jl passes parameter values as vectors to the model,
# even in the one-dimensional case. In our definition of the model we assume
# that `x[1]` is the mean.
#
# ## Approximate the Posterior
# Our first goal is to find the posterior over the mean given observation and
# a Gaussian prior with mean 0 and standard deviation 5.
data = [2.0]
prior = MultivariateNormal([0.], [5.])

# The true posterior is a Gaussian distribution with mean 25/26*2 and standard
# deviation 26/25
trueposterior = pdf.(Normal.(-1:.01:5, 26/25), 25/26*2.0)
using Plots
gr()
figure = plot(-1:.01:5, trueposterior, label = "posterior")

# Now, we will use an adaptive sequential Monte Carlo method:
smc = AdaptiveSMC(prior = prior, K = 10^4)
result = run!(smc, model, data, verbose = true, maxfevals = 10^6);

# As a Monte Carlo Method the result is a list of particles
particles(smc)
# with corresponding weights
weights(smc)

# The mean of the posterior is given by `weights(smc) .* particles(smc)`, which
# is computed by the `mean` function.
mean(smc)

figure = histogram(vcat(particles(smc)...), weights = weights(smc), normalize = true, label = "AdaptiveSMC")
plot!(figure, -1:.01:5, trueposterior, label = "posterior")

# The `result` above also contains these weights and particles and some
# additional information.
keys(result)

# AdaptiveSMC reduced the epsilon parameter adaptively, as we saw in column
# epsilon of the run above. We can plot this sequence.
scatter(cumsum(result.n_sims)[2:end], result.epsilons,
        yscale = :log10, ylabel = "epsilon", xlabel = "number of model evaluations")

# Alternatively, we may want to use KernelABC.
kabc = KernelABC(prior = prior,
                 kernel = Kernel(),
                 delta = 1e-12,
                 K = 10^4)
result = run!(kabc, model, data, maxfevals = 10^4)
mean(kabc)

figure = histogram(vcat(particles(kabc)...), weights = weights(kabc),
                   xlims = (-1, 5), bins = 100,
                   normalize = true, label = "KernelABC")
plot!(figure, -1:.01:5, trueposterior, label = "posterior")

# ## Point Estimates
# Sometimes we just want a point estimate. We will use BayesianOptimization.jl
# here to minimize the `QDLoss`. We know that the true maximum
# likelihood estimate is at mean = 25/26*2 ≈ 1.923
using BayesianOptimization
p = PointEstimator(optimizer = bo([-10.], [10.]), losstype = QDLoss, prior = prior, K = 100)
result = run!(p, model, data, maxfevals = 5*10^4, verbose = false);
result.x

# KernelRecursiveABC is an alternative method that requires often only few model
# evaluations in low and medium dimensional problems
k = KernelRecursiveABC(prior = prior,
                       kernel = Kernel(),
                       kernelx = Kernel(),
                       delta = 1e-2,
                       K = 100)
result = run!(k, model, data, maxfevals = 2*10^3)
result.x

# ## iid Samples
# Let us suppose here that the data consists of multiple independent and
# identically distributed samples.
data_iid = [[2.0], [1.9], [2.8], [2.1]]

# There are two ways to deal with this data. Either we just assume it is one
# four-dimensional vector
data_onevec = vcat(data_iid...)
# and we define the model as
model_iid_onevec(x) = vcat([model(x) for _ in 1:4]...)
smc = AdaptiveSMC(prior = prior, K = 10^4)
result = run!(smc, model_iid_onevec, data_onevec, verbose = true, maxfevals = 10^6);
histogram(vcat(particles(smc)...), weights = weights(smc), normalize = true, label = "AdaptiveSMC")

# Alternatively, we use another distance function:
model_iid(x) = [model(x) for _ in 1:4]
smc = AdaptiveSMC(prior = prior, K = 10^4, distance = energydistance)
result = run!(smc, model_iid, data_iid, verbose = true, maxfevals = 10^6);
histogram(vcat(particles(smc)...), weights = weights(smc), normalize = true, label = "AdaptiveSMC")


function digitalroot(n::Integer, bs::Integer=10)
    if n < 0 || bs < 2 throw(DomainError()) end
    ds, pers = n, 0
    while bs ≤ ds
        ds = sum(digits(ds, bs))
        pers += 1
    end
    return pers, ds
end

for i in [627615, 39390, 588225, 393900588225, big(2) ^ 100]
    pers, ds = digitalroot(i)
    println(i, " has persistence ", pers, " and digital root ", ds)
end


#=
fft:
- Julia version: 0.6.2
- Author: jerrymei
- Date: 2018-05-05
=#
using FFTW

mutable struct rFFTWrapper{T}
    src::Array{T,2}
    src_hat::Array{Complex{T},2}
    dest_hat::Array{Complex{T},2}
    dest::Array{T,2}
    plan::FFTW.Plan
    inv_plan::FFTW.Plan
    src_hat_calculated::Bool
end

function rFFTWrapper{T}(input::AbstractArray{T,2})
    mx,my=size(input)
    nx,ny=div(mx,2) + 1,my
    src=input
    src_hat=Array{Complex{T}}(nx,ny)
    dest_hat=similar(src_hat)
    dest=similar(input)
    plan=plan_rfft(src)
    inv_plan=plan_irfft(src_hat,mx)
    return rFFTWrapper(src,src_hat,dest_hat,dest,plan,inv_plan,false)
end

function set_src!{T}(fft::rFFTWrapper, input::AbstractArray{T,2})
    if !(input === fft.src)
        assert(size(input) == size(fft.src))
        fft.src=input
        fft.src_hat_calculated=false
    end
end

function set_dest!{T}(fft::rFFTWrapper, output::AbstractArray{T,2})
    assert(size(output) == size(fft.dest))
    fft.dest=output
end

function set_not_calculated(fft::rFFTWrapper)
    fft.src_hat_calculated=false
end

function derivx!(fft::rFFTWrapper)
    mx,my=size(fft.src)
    nx,ny=size(fft.src_hat)
    if !fft.src_hat_calculated
        A_mul_B!(fft.src_hat,fft.plan,fft.src)
        fft.src_hat_calculated=true
    end
    @inbounds for x in 0:nx - 1
        for y in 0:ny - 1
            fft.dest_hat[x + 1,y + 1]=fft.src_hat[x + 1,y + 1] * x * 2pi * im
        end
    end
    A_mul_B!(fft.dest,fft.inv_plan,fft.dest_hat)
end

function derivy!(fft::rFFTWrapper)
    mx,my=size(fft.src)
    nx,ny=size(fft.src_hat)
    if !fft.src_hat_calculated
        A_mul_B!(fft.src_hat,fft.plan,fft.src)
        fft.src_hat_calculated=true
    end
    @inbounds for x in 0:nx - 1
        for y in 0:ny - 1
            fft.dest_hat[x + 1,y + 1]= fft.src_hat[x + 1,y + 1] * (2 * y < my?y:y - my) * 2pi * im
        end
    end
    A_mul_B!(fft.dest,fft.inv_plan,fft.dest_hat)
end

function laplacian!(fft::rFFTWrapper)
    mx,my=size(fft.src)
    nx,ny=size(fft.src_hat)
    if !fft.src_hat_calculated
        A_mul_B!(fft.src_hat,fft.plan,fft.src)
        fft.src_hat_calculated=true
    end
    @inbounds for x in 0:nx - 1
        for y in 0:ny - 1
            fft.dest_hat[x + 1,y + 1]= fft.src_hat[x + 1,y + 1] * (x^2 + (2 * y < my?y:y - my)^2) * (2pi * im)^2
        end
    end
    A_mul_B!(fft.dest,fft.inv_plan,fft.dest_hat)
end

function solve_laplacian!(fft::rFFTWrapper)
    mx,my=size(fft.src)
    nx,ny=size(fft.src_hat)
    if !fft.src_hat_calculated
        A_mul_B!(fft.src_hat,fft.plan,fft.src)
        fft.src_hat_calculated=true
    end
    @inbounds for x in 0:nx - 1
        for y in 1:ny - 1
            fft.dest_hat[x + 1,y + 1]= fft.src_hat[x + 1,y + 1] / (x^2 + (2 * y < my?y:y - my)^2) / (2pi * im)^2
        end
    end
    @inbounds for x in 1:nx - 1
        fft.dest_hat[x + 1,1]=fft.src_hat[x + 1,1] / x^2 / (2pi * im)^2
    end
    A_mul_B!(fft.dest,fft.inv_plan,fft.dest_hat)
end

function derivx!{T}(input::AbstractArray{T,2}, output::AbstractArray{T,2}, fft::rFFTWrapper)
"""Compute y derivative"""
    set_src!(fft,input)
    set_dest!(fft,output)
    derivx!(fft)
end

function derivy!{T}(input::AbstractArray{T,2}, output::AbstractArray{T,2}, fft::rFFTWrapper)
"""Compute y derivative"""
    set_src!(fft,input)
    set_dest!(fft,output)
    derivy!(fft)
end

function laplacian!{T}(input::AbstractArray{T,2}, output::AbstractArray{T,2}, fft::rFFTWrapper)
"""Compute the Laplacian"""
    set_src!(fft,input)
    set_dest!(fft,output)
    laplacian!(fft)
end

function solve_laplacian!{T}(input::AbstractArray{T,2}, output::AbstractArray{T,2}, fft::rFFTWrapper)
"""Solve Poisson equation"""
    set_src!(fft,input)
    set_dest!(fft,output)
    solve_laplacian!(fft)
end


using OrdinaryDiffEq
import OrdinaryDiffEq: ODEProblem
using Petri
using Test
using Catlab
using Catlab.Doctrines
using Catlab.Graphics
using Catlab.WiringDiagrams
using Catlab.Programs
using SemanticModels.CategoryTheory
import SemanticModels.CategoryTheory: undecorate, ⊔
using SemanticModels.PetriModels
using SemanticModels.PetriCospans
import SemanticModels: model

"""    fluxes(m::Petri.Model)

a PetriNet interpreter that computes the mass action kinetics of a petri net.

Usage: pass this to ODEProblem to set up an ODE for a given model.
"""
function fluxes(m::Petri.Model)
    S = m.S
    T = m.Δ
    nS = length(S)
    nT = length(T)
    ϕ = zeros(Float64, nT)
    f(du, u, p, t) = begin
        for (i, t) in enumerate(T)
            ins = t[1]
            # TODO: accomodate multiplicites here
            ϕ[i] = p[i]*prod(u[ins])
        end
        for i in S
            du[i] = 0
        end
        for (i, t) in enumerate(T)
            ins = t[1]
            out = t[2]
            for s in ins
                # TODO: accomodate multiplicites here
                du[s] -= ϕ[i]
            end
            for s in out
                # TODO: accomodate multiplicites here
                du[s] += ϕ[i]
            end
        end
        return du
    end
    return f
end

u₀(m::Petri.Model) = begin
    zeros(Float64, length(m.S))
end

u₀(m::Petri.Model, initialS, initialI=1) = begin
    u0=zeros(Float64, length(m.S))
    u0[1] = initialS
    u0[2] = initialI
    return u0
end

function savedata(sol::ODESolution, colnames, fname::String)
    open(fname, "w") do fp
        println(fp, "time, $colnames")
        map(tuples(sol)) do (u,t)
            print(fp, "$t")
            map(u) do x
                print(fp, ",$x")
            end
            println(fp, "")
        end
    end
end

model(c::PetriCospan) = left(c.f).d[1].model

X = FinSet(1)
# Fseir = compose(exposure(X,X,X),otimes(spontaneous(X,X),id(X)),mmerge(X),mcopy(X),otimes(id(X),spontaneous(X,X)))
# states are [S, I, E, R]
# m_seir = left(Fseir.f).d[1]

Psir = PetriModel(
          Petri.Model(1:3,[
            ([1,2],[2,2]), # exposure
            ([2],[3]),     # recovery
            ], missing, missing))

m = Psir.model

u0 = zeros(Float64, length(m.S))
u0[1]  = 10000
u0[2]  = 1

β = [10/sum(u0), 1/5]

tspan = (0,100.0)
prob = ODEProblem(fluxes(m), u0, tspan, β)
sol = OrdinaryDiffEq.solve(prob, alg=Tsit5())

savedata(sol, "S,I,R", "sirdata.csv")

Psird = PetriModel(
          Petri.Model(1:3,[
            ([1,2],[2,2]), # exposure
            ([2],[3]),     # recovery
            ([2],[4]),     # death
            ], missing, missing))

m = Psir.model

u0 = zeros(Float64, length(m.S))
u0[1]  = 10000
u0[2]  = 1

β = [10/sum(u0), 1/5, 1/10]

tspan = (0,100.0)
prob = ODEProblem(fluxes(m), u0, tspan, β)
sol = OrdinaryDiffEq.solve(prob, alg=Tsit5())

savedata(sol, "S,I,R,D", "sirddata.csv")

Pseir = PetriModel(
          Petri.Model(1:5,[
            ([1,2],[3,2]), # exposure
            ([3],[2]),     # onset
            ([2],[4]),     # recovery
            ], missing, missing))

m = Pseir.model

u0 = zeros(Float64, length(m.S))
u0[1]  = 10000
u0[2]  = 1

β = [10/sum(u0), 1/2, 1/5]

tspan = (0,100.0)
prob = ODEProblem(fluxes(m), u0, tspan, β)
sol = OrdinaryDiffEq.solve(prob, alg=Tsit5())

savedata(sol, "S,I,E,R", "seirdata.csv")
Pseird = PetriModel(
          Petri.Model(1:5,[
            ([1,2],[3,2]), # exposure
            ([3],[2]),     # onset
            ([2],[4]),     # recovery
            ([2],[5]),     # death
            ], missing, missing))

m = Pseird.model

u0 = zeros(Float64, length(m.S))
u0[1]  = 10000
u0[2]  = 1

seirdparams() = begin
    βseird = [10/sum(u0), 1/2, 1/5, 1/16]
end

tspan = (0,100.0)
prob = ODEProblem(fluxes(m), u0, tspan, seirdparams())
sol = OrdinaryDiffEq.solve(prob, alg=Tsit5())

savedata(sol, "S,I,E,R,D", "seirddata.csv")


const LIGHT_SPEED = 299792458 # m/s^2

"""
    lerp(t, a, b)

Linear interpolation on the interval `[a, b]` as `t ∈ [0, 1].

"""
lerp(t, a, b) = (one(t)-t)*a + t*b

"""
    quadratic(A, B, C)

Solve quadratic equation `A x^2 + B x + C = 0` for real coefficients.

If a real solution exists, returns `(true, t0, t1)` where `t0` is the smaller
and `t1` the larger solution. If no real solution exists, returns `(false, Inf, Inf)`.

"""
function quadratic(A::Real, B::Real, C::Real)
    dd = B^2 - 4*A*C
    if dd < 0.0
        return false, Inf, Inf
    end
    q = B<0 ? -0.5 * (B - sqrt(dd)) : -0.5 * (B + sqrt(dd))
    t0 = q/A
    t1 = C/q
    return true, min(t0,t1), max(t0,t1)
end

"""
    round_pow2(n::Integer)

Round an integer up to nearest power of two.
"""
round_pow2(n::Integer) = 2^convert(typeof(n), ceil(log(2, n)))