module effective_interaction

# Optimized functions for calculation of Omega and Gamma

function Omega(a, θ)
    ξ = 2*pi*a
    cosθpow2 = cos(θ)^2
    cosξdivξ = cos(ξ)/ξ
    return 3.0/4*((1.0-3.0*cosθpow2) * (sin(ξ)+cosξdivξ)/ξ^2 - (1.0-cosθpow2)*cosξdivξ)
end

function Gamma(a, θ)
    ξ = 2*pi*a
    cosθpow2 = cos(θ)^2
    if ξ<1e-3
        return 3/2*(2/3 + 1/15*ξ^2 *(-2 + cosθpow2) + ξ^4*(1/140-cosθpow2/210))
    else
        sinξ = sin(ξ)/ξ
        return 3/2*((1-cosθpow2)*sinξ + (1-3*cosθpow2) * (cos(ξ)-sinξ)/ξ^2)
    end
end

function OmegaGamma(a, cosθpow2)
    #cosθpow2 = cos(θ)^2
    cosξ = cospi(2*a)
    sinξ = sinpi(2*a)
    omega = 0.75*((1.0-3.0*cosθpow2) * (sinξ+cosξ/(2*pi*a))/(2*pi*a)^2 - (1.0-cosθpow2)*cosξ/(2*pi*a))
    gamma = 1.5*((1-cosθpow2)*sinξ/(2*pi*a) + (1-3*cosθpow2) * (cosξ-sinξ/(2*pi*a))/(2*pi*a)^2)
    return omega, gamma
end

function OmegaGamma_orthogonal(a)
    cosξ = cospi(2*a)
    sinξ = sinpi(2*a)
    omega = 0.75*((sinξ+cosξ/(2*pi*a))/(2*pi*a)^2 - cosξ/(2*pi*a))
    gamma = 1.5*(sinξ/(2*pi*a) + (cosξ-sinξ/(2*pi*a))/(2*pi*a)^2)
    return omega, gamma
end

function Omega_orthogonal(a)
    ξ = 2*pi*a
    cosξdξ = cos(ξ)/ξ
    return 3.0/4.0*(-cosξdξ + (sin(ξ)+cosξdξ)/ξ^2)
end

function Gamma_orthogonal(a)
    ξ = 2*a
    sincξ = sinc(ξ)
    return 3.0/2.0*(sincξ + (cospi(ξ)-sincξ)/(pi*ξ)^2)
end

function map_effectiveinteraction(a::Vector{Float64}, f::Function)
    omega_list = Float64[]
    gamma_list = Float64[]
    for ai=a
        omega, gamma = f(ai, e_dipole)
        push!(omega_list, omega)
        push!(gamma_list, gamma)
    end
    return omega_list, gamma_list
end

# Finite symmetric systems

"""
    effective_interaction.triangle_orthogonal(a)

Effective Omega and Gamma for a equilateral triangle with edge length `a`.

The polarization axis is orthogonal to the triangle plane.
"""
function triangle_orthogonal(a::Float64)
    omega_eff = 2*Omega_orthogonal(a)
    gamma_eff = 2*Gamma_orthogonal(a)
    return omega_eff, gamma_eff
end

"""
    effective_interaction.square_orthogonal(a)

Effective Omega and Gamma for a square with edge length `a`.

The polarization axis is orthogonal to the square plane.
"""
function square_orthogonal(a::Float64)
    omega_eff = 2*Omega_orthogonal(a) + Omega_orthogonal(sqrt(2.)*a)
    gamma_eff = 2*Gamma_orthogonal(a) + Gamma_orthogonal(sqrt(2.)*a)
    return omega_eff, gamma_eff
end

"""
    effective_interaction.rectangle_orthogonal(a, b)

Effective Omega and Gamma for a rectangle with edge lengths `a` and `b`.

The polarization axis is orthogonal to the rectangle plane.
"""
function rectangle_orthogonal(a::Float64, b::Float64)
    d = sqrt(a^2 + b^2)
    omega_eff, gamma_eff = OmegaGamma_orthogonal(a)
    om, ga = OmegaGamma_orthogonal(b)
    omega_eff += om; gamma_eff+=ga
    om, ga = OmegaGamma_orthogonal(d)
    omega_eff += om; gamma_eff+=ga
    return omega_eff, gamma_eff
end

"""
    effective_interaction.cube_orthogonal(a)

Effective Omega and Gamma for a cube with edge length `a`

The polarization axis is orthogonal to the xy faces.
"""
function cube_orthogonal(a::Float64)
    omega_eff, gamma_eff = square_orthogonal(a)
    sqrt2 = sqrt(2.)
    sqrt3 = sqrt(3.)
    Θdiag = atan(sqrt2, 1.)
    omega_eff += Omega(a, 0.) + 2*Omega(sqrt(2.)*a, pi/4.) + Omega(sqrt3*a, Θdiag)
    gamma_eff += Gamma(a, 0.) + 2*Gamma(sqrt(2.)*a, pi/4.) + Gamma(sqrt3*a, Θdiag)
    return omega_eff, gamma_eff
end

"""
    effective_interaction.box_orthogonal(a, b, c)

Effective Omega and Gamma for a box with edge lengths `a`, `b` and `c`.

The polarization axis is orthogonal to the top face.
"""
function box_orthogonal(a::Float64, b::Float64, c::Float64)
    omega_eff, gamma_eff = rectangle_orthogonal(a, b)

    a2 = a^2
    b2 = b^2
    c2 = c^2

    om, ga = OmegaGamma(c, 1.)
    omega_eff += om; gamma_eff+=ga

    dxz_pow2 = a2+c2
    om, ga = OmegaGamma(sqrt(dxz_pow2), c2/dxz_pow2)
    omega_eff += om; gamma_eff+=ga

    dyz_pow2 = b2+c2
    om, ga = OmegaGamma(sqrt(dyz_pow2), c2/dyz_pow2)
    omega_eff += om; gamma_eff+=ga

    dxyz_pow2 = a2+b2+c2
    om, ga = OmegaGamma(sqrt(dxyz_pow2), c2/dxyz_pow2)
    omega_eff += om; gamma_eff+=ga
    return omega_eff, gamma_eff
end

# Infinite 1D symmetric systems
"""
    effective_interaction.chain(a, Θ, N)

Effective Omega and Gamma for an infinite chain.

The calculation is done by adding N spins left and N spins right of a
central spin.

# Arguments
* `a`: Spin-spin distance.
* `θ`: Angle between polarization axis and spin chain.
* `N`: Number of included spins.
"""
function chain(a::Float64, Θ, N::Int)
    omega_eff::Float64 = 0.
    gamma_eff::Float64 = 0.
    d = 0.
    for j=1:N
        d += a
        omega_eff += Omega(d, Θ)
        gamma_eff += Gamma(d, Θ)
    end
    return 2*omega_eff, 2*gamma_eff
end

"""
    effective_interaction.chain_orthogonal(a, N)

Effective Omega and Gamma for an infinite chain with orthogonal polarization axis.

The calculation is done by adding N spins left and N spins right of a
central spin.

# Arguments
* `a`: Spin-spin distance.
* `N`: Number of included spins.
"""
function chain_orthogonal(a::Float64, N::Int)
    omega_eff::Float64 = 0.
    gamma_eff::Float64 = 0.
    d::Float64 = 0.
    for j=1:N
        d += a
        omega_eff += Omega_orthogonal(d)
        gamma_eff += Gamma_orthogonal(d)
    end
    return 2*omega_eff, 2*gamma_eff
end

# Infinite 2D symmetric systems
"""
    effective_interaction.squarelattice_orthogonal(a, N)

Effective Omega and Gamma for a infinite square lattice.

The polarization axis is orthogonal to the square lattice plane and the
calculation is done by creating a (2N+1)*(2N+1) square lattice
and calculate the combined interaction for the central spin.

# Arguments
* `a`: Spin-spin distance.
* `N`: Number of included spins.
"""
function squarelattice_orthogonal(a::Float64, N::Int)
    omega_eff::Float64 = 0.
    gamma_eff::Float64 = 0.
    for j=1:N
        d = a*j
        omega_eff += 4*Omega_orthogonal(d)
        gamma_eff += 4*Gamma_orthogonal(d)
        dx2 = j^2
        for i=1:j-1
            dy2 = i^2
            d = a*sqrt(dx2+dy2)
            omega_eff += 8*Omega_orthogonal(d)
            gamma_eff += 8*Gamma_orthogonal(d)
        end
        d = a*sqrt(2)*j
        omega_eff += 4*Omega_orthogonal(d)
        gamma_eff += 4*Gamma_orthogonal(d)
    end
    return omega_eff, gamma_eff
end

"""
    effective_interaction.hexagonallattice_orthogonal(a, N)

Effective Omega and Gamma for a infinite hexagonal lattice.

The polarization axis is orthogonal to the square lattice plane and the
calculation is done by creating hexagonal lattice consisting of
N rings and calculate the combined interaction for the central spin.

# Arguments
* `a`: Spin-spin distance.
* `N`: Number of included spins.
"""
function hexagonallattice_orthogonal(a::Float64, N::Int)
    omega_eff::Float64 = 0.
    gamma_eff::Float64 = 0.
    for n=1:2:N
        nv1x = -0.5*n
        nv1y_square = 3.0/4*n^2
        d = a*sqrt(nv1y_square + nv1x^2)
        omega_eff += 6*Omega_orthogonal(d)
        gamma_eff += 6*Gamma_orthogonal(d)
        for i=1:div(n-1,2)
            d = a*sqrt(nv1y_square + (nv1x+1.0*i)^2)
            omega_eff += 12*Omega_orthogonal(d)
            gamma_eff += 12*Gamma_orthogonal(d)
        end
    end
    for n=2:2:N
        nv1x = -0.5*n
        nv1y_square = 3.0/4*n^2
        d = a*sqrt(nv1y_square + nv1x^2)
        omega_eff += 6*Omega_orthogonal(d)
        gamma_eff += 6*Gamma_orthogonal(d)
        for i=1:div(n-2,2)
            d = a*sqrt(nv1y_square + (nv1x+1.0*i)^2)
            omega_eff += 12*Omega_orthogonal(d)
            gamma_eff += 12*Gamma_orthogonal(d)
        end
        i = div(n-2,2)+1
        d = a*sqrt(nv1y_square + (nv1x+1.0*i)^2)
        omega_eff += 6*Omega_orthogonal(d)
        gamma_eff += 6*Gamma_orthogonal(d)
    end
    return omega_eff, gamma_eff
end

# Infinite 3D symmetric systems
"""
    effective_interaction.cubiclattice_orthogonal(a, N)

Effective Omega and Gamma for a infinite cubic lattice.

The polarization axis is orthogonal to the top face of a unit cell and the
calculation is done by creating a (2N+1)*(2N+1)*(2N+1) cubic lattice
and calculate the combined interaction for the central spin.

# Arguments
* `a`: Spin-spin distance.
* `N`: Number of included spins.
"""
function cubiclattice_orthogonal(a::Float64, N::Int)
    omega_eff, gamma_eff = squarelattice_orthogonal(a, N)
    for nx=1:N
        for ny=0:nx
            multiplicity = ((ny==0 || ny==nx) ? 8 : 16)
            nrpow2 = nx^2 + ny^2
            nr = sqrt(nrpow2)
            for nz=1:N
                nzpow2 = nz^2
                d = a*sqrt(nrpow2+nzpow2)
                cosθpow2 = nzpow2/(nrpow2+nzpow2)
                om, ga = OmegaGamma(d, cosθpow2)
                omega_eff += multiplicity*om
                gamma_eff += multiplicity*ga
            end
        end
    end
    d::Float64 = 0.
    for nz=1:N
        d += a
        om, ga = OmegaGamma(d, 1.)
        omega_eff += 2*om
        gamma_eff += 2*ga
    end
    return omega_eff, gamma_eff
end

"""
    effective_interaction.tetragonallattice_orthogonal(a, b, N)

Effective Omega and Gamma for a infinite tetragonal lattice.

The polarization axis is orthogonal to the top face of a unit cell and the
calculation is done by creating a (2N+1)*(2N+1)*(2N+1) tetragonal lattice
and calculate the combined interaction for the central spin.

# Arguments
* `a`: Spin-spin distance for bottom side square.
* `b`: Height of the unit cell.
* `N`: Number of included spins.
"""
function tetragonallattice_orthogonal(a::Float64, b::Float64, N::Int)
    omega_eff, gamma_eff = squarelattice_orthogonal(a, N)
    for nx=1:N
        for ny=0:nx
            multiplicity = ((ny==0 || ny==nx) ? 8 : 16)
            nrpow2 = nx^2 + ny^2
            nr = sqrt(nrpow2)
            for nz=1:N
                nzpow2 = nz^2
                d = sqrt(a^2*nrpow2 + b^2*nzpow2)
                θ = atan(a*nr, b*nz)
                omega_eff += multiplicity*Omega(d, θ)
                gamma_eff += multiplicity*Gamma(d, θ)
            end
        end
    end
    d::Float64 = 0.
    for nz=1:N
        d += b
        omega_eff += 2*Omega(d, 0.)
        gamma_eff += 2*Gamma(d, 0.)
    end
    return omega_eff, gamma_eff
end

"""
    effective_interaction.hexagonallattice3d_orthogonal(a, b, N)

Effective Omega and Gamma for a infinite 3D hexagonal lattice.

The lattice consists of stacked planes of hexagonal lattices where the
the polarization axis is orthogonal to the planes. The
calculation is done by creating hexagonal lattices with N rings, stacking
2N+1 lattices of this kind above each other and calculating the combined
interaction for the central spin.

# Arguments
* `a`: Spin-spin distance for hexagons.
* `b`: Distance between planes of hexagonal lattices
* `N`: Number of included spins.
"""
function hexagonallattice3d_orthogonal(a::Float64, b::Float64, N::Int)
    omega_eff, gamma_eff::Float64 = hexagonallattice_orthogonal(a, N)
    a2 = a^2
    b2 = b^2
    for n=1:2:N
        nv1x = -0.5*n
        nv1y_square = 3.0/4*n^2
        nr_square = a2*(nv1y_square + nv1x^2)
        for iz=1:N
            nz_square = iz^2*b2
            d = sqrt(nr_square+nz_square)
            cosθ_square = nz_square/(nr_square+nz_square)
            om, ga = OmegaGamma(d, cosθ_square)
            omega_eff += 12*om
            gamma_eff += 12*ga
        end
        for i=1:div(n-1,2)
            nr_square = a2*(nv1y_square + (nv1x+1.0*i)^2)
            for iz=1:N
                nz_square = iz^2*b2
                d = sqrt(nr_square+nz_square)
                cosθ_square = nz_square/(nr_square+nz_square)
                om, ga = OmegaGamma(d, cosθ_square)
                omega_eff += 24*om
                gamma_eff += 24*ga
            end
        end
    end
    for n=2:2:N
        nv1x = -0.5*n
        nv1y_square = 3.0/4*n^2
        nr_square = a2*(nv1y_square + nv1x^2)
        for iz=1:N
            nz_square = iz^2*b2
            d = sqrt(nr_square+nz_square)
            cosθ_square = nz_square/(nr_square+nz_square)
            om, ga = OmegaGamma(d, cosθ_square)
            omega_eff += 12*om
            gamma_eff += 12*ga
        end
        for i=1:div(n-2, 2)
            nr_square = a2*(nv1y_square + (nv1x+1.0*i)^2)
            for iz=1:N
                nz_square = iz^2*b2
                d = sqrt(nr_square+nz_square)
                cosθ_square = nz_square/(nr_square+nz_square)
                om, ga = OmegaGamma(d, cosθ_square)
                omega_eff += 24*om
                gamma_eff += 24*ga
            end
        end
        i = div(n-2, 2)+1
        nr_square = a2*(nv1y_square + (nv1x+1.0*i)^2)
        for iz=1:N
            nz_square = iz^2*b2
            d = sqrt(nr_square+nz_square)
            cosθ_square = nz_square/(nr_square+nz_square)
            om, ga = OmegaGamma(d, cosθ_square)
            omega_eff += 12*om
            gamma_eff += 12*ga
        end
    end
    return omega_eff, gamma_eff
end

end # module


using JuMP, EAGO

                     m = Model()

                     EAGO.register_eago_operators!(m)

                     @variable(m, -1 <= x[i=1:5] <= 1)

                     @variable(m, -8.966433788908274 <= q <= 11.565431391005125)

                     add_NL_constraint(m, :(log(1 + exp(0.8271637266959466 + -0.29461988642396975*$(x[1]) + 0.9225334019614664*$(x[2]) + -0.12844201057842541*$(x[3]) + -0.9795769523241842*$(x[4]) + 0.7579371331675153*$(x[5]))) + log(1 + exp(0.9202467558270198 + 0.816513326282978*$(x[1]) + 0.1577763582178746*$(x[2]) + 0.37608788225609757*$(x[3]) + 0.8366994089368287*$(x[4]) + 0.28442920912454284*$(x[5]))) + log(1 + exp(-0.23055889086042392 + -0.12029185584830726*$(x[1]) + 0.7385896748366556*$(x[2]) + 0.04575327278060515*$(x[3]) + 0.05400397335314544*$(x[4]) + -0.13287838316667244*$(x[5]))) + log(1 + exp(0.42724629038540085 + 0.286409417560213*$(x[1]) + -0.3599014579136197*$(x[2]) + -0.15981498421265528*$(x[3]) + -0.05370590457282587*$(x[4]) + -0.6772611532220885*$(x[5]))) + log(1 + exp(-0.6445990809995172 + -0.48590839796846375*$(x[1]) + 0.5462761102799756*$(x[2]) + -0.2471758620035076*$(x[3]) + 0.5819097771517305*$(x[4]) + 0.22143679581234954*$(x[5]))) - $q <= 0.0))

                     @objective(m, Min, q)

                     return m


simple = @stochastic_model begin
    @stage 1 begin
        @decision(model, x₁ >= 40)
        @decision(model, x₂ >= 20)
        objective = @expression(model, 100*x₁ + 150*x₂)
        @objective(model, Min, objective)
        @constraint(model, x₁ + x₂ <= 120)
    end
    @stage 2 begin
        @uncertain q₁ q₂ d₁ d₂
        @recourse(model, 0 <= y₁ <= d₁)
        @recourse(model, 0 <= y₂ <= d₂)
        @objective(model, Max, q₁*y₁ + q₂*y₂)
        @constraint(model, 6*y₁ + 10*y₂ <= 60*x₁)
        @constraint(model, 8*y₁ + 5*y₂ <= 80*x₂)
    end
end

ξ₁ = @scenario q₁ = 24.0 q₂ = 28.0 d₁ = 500.0 d₂ = 100.0 probability = 0.4
ξ₂ = @scenario q₁ = 28.0 q₂ = 32.0 d₁ = 300.0 d₂ = 300.0 probability = 0.6

simple_res = SPResult([46.67, 36.25], Dict(1 => [300., 100.], 2 => [300., 100.]), -855.83, -1518.75, 662.92, 286.92, -1445.92, -568.92)
push!(problems, (simple, [ξ₁,ξ₂], simple_res, "Simple"))


module Kernels
in_range(u) = abs(u) ≤ 1

biweight(u)         = in_range(u) ? (1.0 - u ^ 2) ^ 2 : 0.0
cosine(u)           = in_range(u) ? cos(0.5 * π * u) : 0.0
epanechnikov(u)     = in_range(u) ? 1.0 - u ^ 2 : 0.0
triangular(u)       = in_range(u) ? 1.0 - abs(u) : 0.0
tricube(u)          = in_range(u) ? (1.0 - abs(u) ^ 3) ^ 3 : 0.0
triweight(u)        = in_range(u) ? (1.0 - u ^ 2) ^ 3 : 0.0
uniform(u)          = in_range(u) ? 0.5 : 0.0

gaussian(u) = exp(-0.5 * u ^ 2)
logistic(u) = 1.0 / (exp(u) + 2.0 + exp(-u))
end


@testset "Test basis" begin

using jFEMTools
using Tensors
import Tensors: Vec, gradient, norm
jF = jFEMTools

#Test Dubiner 3d basis
@show "Testing Dubiner 3D basis"
interpolation = jF.Dubiner{jF.Tetrahedron,3}()
exact_val = [0.8660254037844386,2.7386127875258306,0.8660254037844386,1.118033988749895,
            7.685213074469699,5.477225575051661,6.48074069840786,1.71846588560844,
            3.7416573867739413,-0.33071891388307384]
for i in 1:10
    @test jF.dubiner_basis(0.5,0.5,0.5,i,2) ≈ exact_val[i]
end

#Test Dubiner 2d base
interpolation = jF.Dubiner{jF.Triangle,4}()
quad_rule = jF.QuadratureRule{jF.Triangle}(jF.Strang(),5)

# Test nodal points for Lagrange base
nodal_points, topology = jF.get_nodal_points(jF.Triangle,1)
@test  nodal_points == [Vec{2, Float64}((0.0, 0.0)),
                        Vec{2, Float64}((1.0, 0.0)),
                        Vec{2, Float64}((0.0, 1.0))]
@test topology == Dict(0=>3,1=>0,2=>0)
nodal_points, topology = jF.get_nodal_points(jF.Triangle,2)
@test nodal_points == [Vec{2, Float64}((0.0, 0.0)),
                        Vec{2, Float64}((1.0, 0.0)),
                        Vec{2, Float64}((0.0, 1.0)),
                        Vec{2, Float64}((0.5, 0.5)),
                        Vec{2, Float64}((0.0, 0.5)),
                        Vec{2, Float64}((0.5, 0.0))]
@test topology == Dict(0=>3,1=>3,2=>0)
nodal_points, topology = jF.get_nodal_points(jF.Triangle, 3)
@test length(nodal_points) == 10
@test topology == Dict(0=>3,1=>6,2=>1)
@test nodal_points[10] ≈ Vec{2, Float64}((1/3, 1/3))
# Test Lagrange base
function ref_value(i::Int, ξ::Vec{2})
    ξ_x = ξ[1]
    ξ_y = ξ[2]
    i == 2 && return ξ_x
    i == 3 && return ξ_y
    i == 1 && return 1. - ξ_x - ξ_y
    throw(ArgumentError("no shape function $i for interpolation $ip"))
end

interpolation = jF.Lagrange(jF.Triangle,1)
ξ_ref = Vec{2}((0.5,0.5))
for j = 1:3
    @test abs(ref_value(j,ξ_ref) - jF.value(interpolation,j,ξ_ref)) < eps(Float64)
    @test norm(gradient(ξ -> ref_value(j, ξ), ξ_ref) - jF.gradient_value(interpolation,j,ξ_ref), Inf) < eps(Float64)
end

function ref_value1(i::Int, ξ::Vec{1})
    ξ_x = ξ[1]
    i == 1 && return 1-ξ_x
    i == 2 && return ξ_x
    throw(ArgumentError("no shape function $i for interpolation $ip"))
end
ξ_ref = Vec{1}((0.5,))
interpolation = jF.Lagrange(jF.Segment,1)
for j = 1:2
    @test abs(ref_value1(j,ξ_ref) - jF.value(interpolation,j,ξ_ref)) < eps(Float64)
    @test norm(gradient(ξ -> ref_value1(j, ξ), ξ_ref) - jF.gradient_value(interpolation,j,ξ_ref), Inf) < eps(Float64)
end

#Test Legendre basis
function ref_value(i::Int, x::Real)
    @assert i <= 3
    if i == 0
        return 1.0
    elseif i == 1
        return sqrt(3)*(2*x - 1)
    elseif i == 2
        return sqrt(5)*(6*x^2 - 6*x + 1)
    else
        return sqrt(7)*(2*x - 1)*(10*x^2 - 10*x + 1)
    end
end
function ref_dvalue(i::Int, x::Real)
    @assert i <= 3
    if i == 0
        return 0
    elseif i == 1
        return sqrt(3)*2
    elseif i == 2
        return sqrt(5)*(12*x - 6)
    else
        return sqrt(7)*12*(5*x^2 - 5*x + 1)
    end
end

# interpolation = jF.Legendre(jF.Segment,3)
# quad_rule = QuadratureRule{1,RefTetrahedron}(GaussLegendre(),3)
# for i = 0:3
#     @test integrate(x->(ref_value(i,x[1]) - value(interpolation, i+1, x))^2,quad_rule) < eps()
#     @test integrate(x->(ref_dvalue(i,x[1]) - gradient_value(interpolation, i+1, x)[1])^2,quad_rule) < eps()
# end

end


using RandomWalks
using BenchmarkTools
using EmpiricalCDFs: EmpiricalCDF, get_data
using Statistics

nsamples = 10^4
nsteps = 10^3
tolerance = 1e-2

ecdf = EmpiricalCDF()
trial_loop = SampleLoopActor(nsamples, ECDFActor(get_x, ecdf))
@time trial!(WalkPlan(ContinuousWalk(WalkB(Point(0.0))), StepLimitActor(nsteps)), trial_loop)

x = get_data(ecdf)
x ./= sqrt(nsteps)

@info "Test standard deviation and mean of an ordinary RW"
@show isapprox(std(ecdf), 1; atol = tolerance)
@show isapprox(mean(ecdf), 0; atol = tolerance)


@doc raw"""
    gaussjacobi(n::Integer) -> Tuple{Vector{Float64},Vector{Float64}}

Return nodes and weights of [Gauss-Jacobi quadrature](https://en.wikipedia.org/wiki/Gauss%E2%80%93Jacobi_quadrature).

```math
\int_{-1}^{1} f(x) (1-x)^\alpha(1+x)^\beta dx \approx \sum_{i=1}^{n} w_i f(x_i)
```

# Examples
```jldoctest
julia> x, w = gaussjacobi(3, 1/3, -1/3);

julia> f(x) = x^4;

julia> I = dot(w, f.(x));

julia> I ≈ 268π/729(√3)
true
```
"""
function gaussjacobi(n::Integer, α::Real, β::Real)
    #GAUSS-JACOBI QUADRATURE NODES AND WEIGHTS
    if n < 0
        throw(DomainError(n, "gaussjacobi($n,$α,$β) not defined: n must be non-negative."))
    elseif α == 0. && β == 0.
        gausslegendre(n)
    elseif α == -0.5 && β == -0.5
        gausschebyshev(n, 1)
    elseif α == 0.5 && β == 0.5
        gausschebyshev(n, 2)
    elseif α == -0.5 && β == 0.5
        gausschebyshev(n, 3)
    elseif α == 0.5 && β == -0.5
        gausschebyshev(n, 4)
    elseif n == 0
        Float64[], Float64[]
    elseif n == 1
        [(β - α) / (α + β + 2)], [2^(α + β + 1) * beta(α + 1, β + 1)]
    elseif min(α,β) ≤ -1.
        throw(DomainError((α,β), "The Jacobi parameters correspond to a nonintegrable weight function"))
    elseif n ≤ 100 && max(α,β) < 5.
        jacobi_rec(n, α, β)
    elseif n > 100 && max(α,β) < 5.
        jacobi_asy(n, α, β)
    elseif n ≤ 4000 && max(α,β) ≥ 5.
        jacobi_gw(n, α, β)
    else
        error("gaussjacobi($n,$α,$β) is not implemented: n must be ≤ 4000 for max(α,β) ≥ 5.")
    end
end

# Convenience function: convert any kind of numbers a and b to a joint floating point type
jacobi_rec(n::Integer, α::Real, β::Real) = jacobi_rec(n, promote(float(α), float(β))...)

function jacobi_rec(n::Integer, α::T, β::T) where {T <: AbstractFloat}
    #Compute nodes and weights using recurrrence relation.
    x11, x12 = HalfRec(n, α, β, 1)
    x21, x22 = HalfRec(n, β, α, 0)

    x = Array{T}(undef,n)
    w = Array{T}(undef,n)
    m1 = length(x11)
    m2 = length(x21)
    sum_w = zero(T)
    @inbounds for i in 1:m2
        idx = m2 + 1 - i
        xi = -x21[i]
        der = x22[i]
        wi = 1 / ((1 - xi^2) * der^2)
        w[idx] = wi
        x[idx] = xi
        sum_w += wi
    end
    @inbounds for i in 1:m1
        idx = m2 + i
        xi = x11[i]
        der = x12[i]
        wi = 1 / ((1 - xi^2) * der^2)
        w[idx] = wi
        x[idx] = xi
        sum_w += wi
    end
    c = (2^(α+β+1)*gamma(2+α)*gamma(2+β)/(gamma(2+α+β)*(α+1)*(β+1)))
    rmul!(w, c / sum_w)
    return x, w
end

function HalfRec(n::Integer, α::T, β::T, flag) where {T <: AbstractFloat}
    # HALFREC  Jacobi polynomial recurrence relation.
    # Asymptotic formula - only valid for positive x.
    r = (flag == 1) ? (ceil(n / 2):-1:1) : (floor(n / 2):-1:1)
    m = length(r)
    c1 = 1 / (2 * n + α + β + 1)
    a1 = one(T)/4 - α^2
    b1 = one(T)/4 - β^2
    c1² = c1^2
    x = Array{T}(undef,m)
    @inbounds for i in 1:m
        C = muladd(2*one(T), r[i], α - one(T)/2) * (T(π) * c1)
        C_2 = C/2
        x[i] = cos(muladd(c1², muladd(-b1, tan(C_2), a1 * cot(C_2)), C))
    end

    P1 = Array{T}(undef,m)
    P2 = Array{T}(undef,m)
    # Loop until convergence:
    for _ in 1:10
        innerjacobi_rec!(n, x, α, β, P1, P2)
        dx2 = 0.0
        @inbounds for i in 1:m
            dx = P1[i] / P2[i]
            _dx2 = abs2(dx)
            dx2 = ifelse(_dx2 > dx2, _dx2, dx2)
            x[i] = x[i] - dx
        end
        dx2 > eps(T) / 1e6 || break
    end
    # Once more for derivatives:
    innerjacobi_rec!(n, x, α, β, P1, P2)
    return x, P2
end

function innerjacobi_rec!(n, x, α::T, β::T, P, PP) where {T <: AbstractFloat}
    # EVALUATE JACOBI POLYNOMIALS AND ITS DERIVATIVE USING THREE-TERM RECURRENCE.
    N = length(x)
    @inbounds for j = 1:N
        xj = x[j]

        Pj = (α - β + (α + β + 2) * xj)/2
        Pm1 = one(T)
        PPj = (α + β + 2)/2
        PPm1 = zero(T)
        for k in 1:n-1
            k0 = muladd(2*one(T), k, α + β)
            k1 = k0 + 1
            k2 = k0 + 2
            A = 2 * (k + 1) * (k + (α + β + 1)) * k0
            B = k1 * (α^2 - β^2)
            C = k0 * k1 * k2
            D = 2 * (k + α) * (k + β) * k2
            c1 = muladd(C, xj, B)
            Pm1, Pj = Pj, muladd(-D, Pm1, c1 * Pj) / A
            PPm1, PPj = PPj, muladd(c1, PPj, muladd(-D, PPm1, C * Pm1)) / A
        end
        P[j] = Pj
        PP[j] = PPj
    end
    nothing
end

function innerjacobi_rec(n, x, α::T, β::T) where {T <: AbstractFloat}
    # EVALUATE JACOBI POLYNOMIALS AND ITS DERIVATIVE USING THREE-TERM RECURRENCE.
    N = length(x)
    P = Array{T}(undef,N)
    PP = Array{T}(undef,N)
    innerjacobi_rec!(n, x, α, β, P, PP)
    return P, PP
end

function weightsConstant(n, α, β)
    # Compute the constant for weights:
    M = min(20, n - 1)
    C = 1.0
    p = -α * β / n
    for m = 1:M
        C += p
        p *= -(m + α) * (m + β) / (m + 1) / (n - m)
        abs(p / C) < eps(Float64) / 100 && break
    end
    return 2^(α + β + 1) * C
end

function jacobi_asy(n, α, β)
    # ASY  Compute nodes and weights using asymptotic formulae.

    # Determine switch between interior and boundary regions:
    nbdy = 10
    bdyidx1 = n - (nbdy - 1):n
    bdyidx2 = nbdy:-1:1

    # Interior formula:
    x, w = asy1(n, α, β, nbdy)

    # Boundary formula (right):
    xbdy = boundary(n, α, β, nbdy)
    x[bdyidx1], w[bdyidx1] = xbdy

    # Boundary formula (left):
    if α ≠ β
        xbdy = boundary(n, β, α, nbdy)
    end
    x[bdyidx2] = -xbdy[1]
    w[bdyidx2] = xbdy[2]

    rmul!(w, weightsConstant(n, α, β))
    return x, w
end

function asy1(n::Integer, α::Float64, β::Float64, nbdy::Integer)
    # Algorithm for computing nodes and weights in the interior.

    # Approximate roots via asymptotic formula: (Gatteschi and Pittaluga, 1985)
    K = π*(2(n:-1:1).+α.-0.5)/(2n+α+β+1)
    tt = K .+ (1/(2n+α+β+1)^2).*((0.25-α^2).*cot.(K/2).-(0.25-β^2).*tan.(K/2))

    # First half (x > 0):
    t = tt[tt .≤ π/2]
    mint = t[end-nbdy+1]
    idx = 1:max(findfirst(t .< mint)-1, 1)

    # Newton iteration
    for _ in 1:10
        vals, ders = feval_asy1(n, α, β, t, idx)  # Evaluate
        dt = vals./ders
        t += dt  # Next iterate
        if norm(dt[idx],Inf) < sqrt(eps(Float64))/100
            break
        end
    end
    vals, ders = feval_asy1(n, α, β, t, idx)  # Once more for luck
    t += vals./ders

    # Store
    x_right = cos.(t)
    w_right = 1 ./ ders.^2

    # Second half (x < 0):
    α, β = β, α
    t = π .- tt[1:(n-length(x_right))]
    mint = t[nbdy]
    idx = max(findfirst(t .> mint), 1):length(t)

    # Newton iteration
    for _ in 1:10
        vals, ders = feval_asy1(n, α, β, t, idx)  # Evaluate.
        dt = vals./ders  # Newton update.
        t += dt
        if norm(dt[idx],Inf) < sqrt(eps(Float64))/100
            break
        end
    end
    vals, ders = feval_asy1(n, α, β, t, idx)  # Once more for luck.
    t += vals./ders  # Newton update.

    # Store
    x_left = cos.(t)
    w_left = 1 ./ ders.^2

    return vcat(-x_left, x_right), vcat(w_left, w_right)
end

"""
Evaluate the interior asymptotic formula at x = cos(t).
Assumption:
* `length(t) == n ÷ 2`
"""
function feval_asy1(n::Integer, α::Float64, β::Float64, t::AbstractVector, idx)
    # Number of terms in the expansion:
    M = 20

    # Number of elements in t:
    N = length(t)

    # Some often used vectors/matrices:
    onesM = ones(M)

    # The sine and cosine terms:
    A = repeat((2n+α+β).+(1:M),1,N).*repeat(t',M)/2 .- (α+1/2)*π/2  # M × N matrix
    cosA = cos.(A)
    sinA = sin.(A)

    sinT = repeat(sin.(t)',M)
    cosT = repeat(cos.(t)',M)
    cosA2 = cosA.*cosT .+ sinA.*sinT
    sinA2 = sinA.*cosT .- cosA.*sinT

    sinT = hcat(ones(N), cumprod(repeat((csc.(t/2)/2),1,M-1), dims=2))  # M × N matrix
    secT = sec.(t/2)/2

    _vec = [(α+j-1/2)*(-α+j-1/2)/(2n+α+β+j+1)/j for j in 1:M-1]
    P1 = [1;cumprod(_vec)]
    P1[3:4:end] = -P1[3:4:end]
    P1[4:4:end] = -P1[4:4:end]
    P2 = Matrix(1.0I, M, M)
    for l in 1:M
        _vec = [(β+j-1/2)*(-β+j-1/2)/(2n+α+β+j+l)/j for j in 1:M-l-2]
        P2[l,l+1:M-2] = cumprod(_vec)
    end
    PHI = repeat(P1,1,M).*P2

    _vec = [(α+j-1/2)*(-α+j-1/2)/(2n+α+β+j-1)/j for j in 1:M-1]
    P1 = [1;cumprod(_vec)]
    P1[3:4:end] = -P1[3:4:end]
    P1[4:4:end] = -P1[4:4:end]
    P2 = Matrix(1.0I, M, M)
    for l in 1:M
        _vec = [(β+j-1/2)*(-β+j-1/2)/(2n+α+β+j+l-2)/j for j in 1:M-l-2]
        P2[l,l+1:M-2] = cumprod(_vec)
    end
    PHI2 = repeat(P1,1,M).*P2

    S = zeros(N)
    S2 = zeros(N)
    for m in 1:M
        l = 1:2:m
        phi = PHI[l, m]
        dS1 = (sinT[:, l]*phi) .* cosA[m, :]
        phi2 = PHI2[l, m]
        dS12 = (sinT[:, l]*phi2) .* cosA2[m, :]
        l = 2:2:m
        phi = PHI[l, m]
        dS2 = (sinT[:, l]*phi) .* sinA[m, :]
        phi2 = PHI2[l, m]
        dS22 = (sinT[:, l]*phi2) .* sinA2[m, :]
        if m - 1 > 10 && norm(dS1[idx] + dS2[idx], Inf) < eps(Float64) / 100
            break
        end
        S .+= dS1
        S .+= dS2
        S2 .+= dS12
        S2 .+= dS22
        sinT[:,1:m] .*= secT
    end

    # Constant out the front:
    dsa = α^2/2n
    dsb = β^2/2n
    dsab = (α+β)^2/4n
    ds = dsa + dsb - dsab
    s = ds
    i = 1
    dsold = ds # to fix α = -β bug.
    while abs(ds/s)+dsold > eps(Float64)/10
        dsold = abs(ds/s)
        i += 1
        tmp = -(i-1)/(i+1)/n
        dsa = tmp*dsa*α
        dsb = tmp*dsb*β
        dsab = tmp*dsab*(α+β)/2
        ds = dsa + dsb - dsab
        s = s + ds
    end
    p2 = exp(s)*sqrt(2π*(n+α)*(n+β)/(2n+α+β))/(2n+α+β+1)
    # g is a vector of coefficients in
    # ``\Gamma(z) = \frac{z^{z-1/2}}{\exp(z)}\sqrt{2\pi} \left(\sum_{i} B_i z^{-i}\right)``, where B_{i-1} = g[i].
    # https://math.stackexchange.com/questions/1714423/what-is-the-pattern-of-the-stirling-series
    g = [1, 1/12, 1/288, -139/51840, -571/2488320, 163879/209018880,
         5246819/75246796800, -534703531/902961561600,
         -4483131259/86684309913600, 432261921612371/514904800886784000]
    f(g,z) = dot(g, [1;cumprod(ones(9)./z)])

    # Float constant C, C2
    C = 2*p2*(f(g,n+α)*f(g,n+β)/f(g,2n+α+β))/π
    C2 = C*(α+β+2n)*(α+β+1+2n)/(4*(α+n)*(β+n))

    vals = C*S

    # Use relation for derivative:
    ders = (n*((α-β).-(2n+α+β)*cos.(t)).*vals .+ (2*(n+α)*(n+β)*C2).*S2)./(2n+α+β)./sin.(t)
    denom = 1 ./ (sin.(abs.(t)/2).^(α+0.5).*cos.(t/2).^(β+0.5))
    vals .*= denom
    ders .*= denom

    return vals, ders
end

function boundary(n::Integer, α::Float64, β::Float64, npts::Integer)
# Algorithm for computing nodes and weights near the boundary.

    # Use Newton iterations to find the first few Bessel roots:
    smallK = min(30, npts)
    jk = approx_besselroots(α, smallK)

    # Approximate roots via asymptotic formula: (see Olver 1974)
    phik = jk/(n + .5*(α + β + 1))
    x = cos.( phik .+ ((α^2-0.25).*(1 .-phik.*cot.(phik))./(8*phik) .- 0.25.*(α^2-β^2).*tan.(0.5.*phik))./(n + 0.5*(α + β + 1))^2 )

    # Newton iteration:
    for _ in 1:10
        vals, ders = innerjacobi_rec(n, x, α, β)  # Evaluate via asymptotic formula.
        dx = -vals./ders  # Newton update.
        x += dx  # Next iterate.
        if norm(dx,Inf) < sqrt(eps(Float64))/200
            break
        end
    end
    vals, ders = innerjacobi_rec(n, x, α, β)  # Evaluate via asymptotic formula.
    dx = -vals./ders
    x += dx

    # flip:
    x = reverse(x)
    ders = reverse(ders)

    # Revert to x-space:
    w = 1 ./ ((1 .- x.^2) .* ders.^2)
    return x, w
end

function jacobi_jacobimatrix(n, α, β)
    s = α + β
    ii = 2:n-1
    si = 2*ii .+ s
    aa = [(β - α)/(2 + s);
          (β^2 - α^2) ./ ((si .- 2).*si);
          (β^2 - α^2) ./ ((2n - 2+s).*(2n+s))]
    bb = [2*sqrt( (1 + α)*(1 + β)/(s + 3))/(s + 2) ;
          2 .*sqrt.(ii.*(ii .+ α).*(ii .+ β).*(ii .+ s)./(si.^2 .- 1))./si]
    return SymTridiagonal(aa, bb)
end

function jacobimoment(α, β)
    s = α + β
    T = float(typeof(s))
    # Same as 2^(α+β+1) * beta(α+1,β+1)
    return exp((s+1)*log(convert(T,2)) + loggamma(α+1)+loggamma(β+1)-loggamma(2+s))
end

function jacobi_gw(n::Integer, α, β)
    # Golub-Welsh for Gauss--Jacobi quadrature. This is used when max(α,β)>5.
    x, V = eigen(jacobi_jacobimatrix(n, α, β))  # Eigenvalue decomposition.
    # Quadrature weights:
    w = V[1,:].^2 .* jacobimoment(α, β)
    return x, w
end


# Apply a log-transform on a Gamma-distributed variable to sample from it using a normal proposal

using Klara

plogtarget(p::Float64, v::Vector) = (v[1]-1)*p-exp(p)/v[2]+p

p = BasicContUnvParameter(:p, logtarget=plogtarget, nkeys=3)

model = likelihood_model([Constant(:k), Constant(:θ), p], isindexed=false)

mcsampler = MH()

mcrange = BasicMCRange(nsteps=100000, burnin=10000)

v0 = Dict(:k=>2., :θ=>1., :p=>log(10.))

job = BasicMCJob(model, mcsampler, mcrange, v0)

run(job)

chain = output(job)

mean(exp.(chain.value))

acceptance(chain, diagnostics=false)


using FinancialToolbox, Test, DualNumbers, HyperDualNumbers, FinancialFFT, FinancialMonteCarlo

A = 600.0;
N = 18;

method_cm = CarrMadanMethod(A, N);
S0 = 100.0;
K = 100.0;
r = 0.02;
T = 1.0;
d = 0.01
sigma = dual(0.2, 1.0);
lam = 15.0;
mu1 = 0.03;
sigma1 = 0.02;
zero_rate = ZeroRate(r);

Model = BlackScholesProcess(sigma, Underlying(S0, d));

EUData = EuropeanOption(T, K)
Nsim = 100000;
Nstep = 30;
mc = MonteCarloConfiguration(Nsim, Nstep, FinancialMonteCarlo.AntitheticMC());
analytic_result = blsprice(S0, K, r, T, sigma, d);
function test_dual(x, y, toll)
    @test abs(x.value - y.value) < toll
    @test abs(x.epsilon - y.epsilon) < toll
    nothing
end

result_cm = pricer(Model, zero_rate, method_cm, EUData);
method_lewis = LewisMethod(700.0, 200000);
result_lewis = pricer(Model, zero_rate, method_lewis, EUData);
method_cm_lewis = CarrMadanLewisMethod(A, N);
result_cm_lewis = pricer(Model, zero_rate, method_cm_lewis, EUData);
toll = 1e-2;
@testset "dual test" begin
    test_dual(analytic_result, result_cm, toll)
    test_dual(analytic_result, result_lewis, toll)
    test_dual(analytic_result, result_cm_lewis, toll)
end


n = rand(1:10)
p = -2*Float64[1:n;] + (n + 1.5);
C = matrixdepot("cauchy", p)
@test_approx_eq matrixdepot("dingdong", n) C
println("'dingdong' passed test...")


# # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # #
#
# Desription
# ==============================================================================
#
#   Tests related to conversion from quaternions to direction cosine matrices.
#
# # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # #

# File: ./src/conversions/quaternion_to_dcm.jl
# ============================================

# Functions: quat_to_dcm
# ----------------------

@testset "Quaternion => DCM (Float64)" begin
    T = Float64

    # Create a random quaternion.
    q = Quaternion(randn(T), randn(T), randn(T), randn(T))
    q = q / norm(q)

    # Convert to quaternion and test the conversion.
    D = quat_to_dcm(q)
    @test eltype(D) === T

    q0 = q.q0
    q1 = q.q1
    q2 = q.q2
    q3 = q.q3

    @test D[1, 1] ≈ q0^2 + q1^2 - q2^2 - q3^2
    @test D[1, 2] ≈ 2(q1 * q2 + q0 * q3)
    @test D[1, 3] ≈ 2(q1 * q3 - q0 * q2)
    @test D[2, 1] ≈ 2(q1 * q2 - q0 * q3)
    @test D[2, 2] ≈ q0^2 - q1^2 + q2^2 - q3^2
    @test D[2, 3] ≈ 2(q2 * q3 + q0 * q1)
    @test D[3, 1] ≈ 2(q1 * q3 + q0 * q2)
    @test D[3, 2] ≈ 2(q2 * q3 - q0 * q1)
    @test D[3, 3] ≈ q0^2 - q1^2 - q2^2 + q3^2
end

@testset "Quaternion => DCM (Float32)" begin
    T = Float32

    # Create a random quaternion.
    q = Quaternion(randn(T), randn(T), randn(T), randn(T))
    q = q / norm(q)

    # Convert to quaternion and test the conversion.
    D = quat_to_dcm(q)
    @test eltype(D) === T

    q0 = q.q0
    q1 = q.q1
    q2 = q.q2
    q3 = q.q3

    @test D[1, 1] ≈ q0^2 + q1^2 - q2^2 - q3^2
    @test D[1, 2] ≈ 2(q1 * q2 + q0 * q3)
    @test D[1, 3] ≈ 2(q1 * q3 - q0 * q2)
    @test D[2, 1] ≈ 2(q1 * q2 - q0 * q3)
    @test D[2, 2] ≈ q0^2 - q1^2 + q2^2 - q3^2
    @test D[2, 3] ≈ 2(q2 * q3 + q0 * q1)
    @test D[3, 1] ≈ 2(q1 * q3 + q0 * q2)
    @test D[3, 2] ≈ 2(q2 * q3 - q0 * q1)
    @test D[3, 3] ≈ q0^2 - q1^2 - q2^2 + q3^2
end


function nlsolve!(f, ∇f!, scheme, x0; inverse=true, c=0.001, g_tol=1e-8, maxiter=10^4, show_trace=false)
    n = length(x0)
    x_curr = copy(x0)
    ∇f_next, ∇f_curr = similar(x0), similar(x0)
    ∇f!(∇f_curr, x_curr)

    B = initial_approximation(inverse, scheme, n)

    d = .-solve(B, ∇f_curr)
    α, f_α, ls_success = backtracking(f, x_curr, d, ∇f_curr; α_0=1.0, c=c)
    Δx = α .* d
    x_next = x_curr .+ Δx

    ∇f!(∇f_next, x_next)
    y = ∇f_next .- ∇f_curr
    update!(scheme, B, Δx, y)

    for i = 1:maxiter
        copyto!(x_curr, x_next)
        ∇f!(∇f_curr, x_curr)

        d = solve!(d, B, ∇f_curr)

        α, f_α, ls_success = backtracking(f, x_curr, d, ∇f_curr; α_0=1.0, c=c, verbose=show_trace)
        Δx .= α .* d
        x_next .= x_curr .+ Δx
        ∇f!(∇f_next, x_next)
        y .= ∇f_next .- ∇f_curr

        if norm(∇f_next) < g_tol
            return x_next, ∇f_next
        end

        update!(scheme, B, Δx, y)
    end
    return x_next, ∇f_next
end

function nlsolve(f, ∇f, scheme, x0; inverse=true, c=0.001, g_tol=1e-8, maxiter=10^4, show_trace=false)
    n = length(x0)
    x_curr = copy(x0)
    ∇f_curr = ∇f(x_curr)

    B = initial_approximation(inverse, scheme, n)

    d = -solve(B, ∇f_curr)
    α, f_α, ls_success = backtracking(f, x_curr, d, ∇f_curr; α_0=1.0, c=c)
    Δx = α .* d
    x_next = x_curr + Δx
    ∇f_next = ∇f(x_next)
    y = ∇f_next - ∇f_curr
    update!(scheme, B, Δx, y)

    for i = 1:maxiter
        x_curr = copy(x_next)
        ∇f_curr = ∇f(x_curr)
        d = .- solve(B, ∇f_curr)
        α, f_α, ls_success = backtracking(f, x_curr, d, ∇f_curr; α_0=1.0, c=c)
        Δx = α .* d
        x_next = x_curr + Δx
        ∇f_next = ∇f(x_next)
        y = ∇f_next - ∇f_curr

        if norm(∇f_next) < g_tol
            return x_next, ∇f_next
        end
        if show_trace
            println("Objective value (curr): ", f(x_next))
            println("Objective value (prev): ", f(x_curr))
            println("Step size: ", α)
        end
        update!(scheme, B, Δx, y)
    end
    return x_next, ∇f_next
end


function  grids(nr,nx,rbottom)

#   constructs an equi-spaced grid in radius and x=cos(theta)
#   This particular version uses the shell thickness as a length
#   scale. This means that the radius at the top boundary is not r = 1;
#
#   nr       -  number of grid points in radius
#   nx       -  number of grid point in cos(theta)
#   rbottom   -  radius of lower layer  (dimensionless)
#
#   output
#   r        - radial grid (excludes boundaries)
#   x        - x grid (excludes poles and equator)
#   theta    - colatitude (radian)

   rcmb = 1.0;
   dr = (rcmb-rbottom) / nr;
   r = (rbottom+dr/2:dr:rcmb-dr/2);     # exclude boundaries

#  avoid grid point on equator
   if (rem(nx,2) ==1)
      dx = 2/(nx-1);
   else
       dx = 2/nx;
   end

   x = (-1 + dx/2 : dx : 1 - dx/2);
   theta = acos.(x);   # colatitude (in radian)

   return r, x, theta

end


# # Testing different approximations for the rate-of-change law

# Ultimately we want to evolve, in time, an equation of the form $u_t = F(u)$. For the 2D NSE, the most computationally demanding term is the bilinear one: $B(\omega) = \left(\partial_x^2 - \partial_y^2\right)(uv) + \partial_{xy}\left(v^2 - u^2\right)$, where $\mathbf{u} = (u(x, y), v(x, y))$ is the velocity field, which, in the two-dimensional periodic case, can easily be obtained from the vorticity $\omega = \boldsymbol{\nabla} \times \mathbf{u}$.

# Here, we test different ways of approximating $B(u)$

# Here are the packages we are gonna need.

using FFTW
using Plots
using LinearAlgebra: mul!
using Test
using BenchmarkTools
using Random

@info "Threads: $(FFTW.nthreads())"

# ## Operators

# We first define a method to build the various operators acting in spectral space:

function get_operators(N, κ₀)
    ## Differentiation in spectral space
    Dx_hat = im * κ₀ * [
        ifelse(k1 ≤ div(N, 2) + 1, k1 - 1, k1 - 1 - N) for k2 in 1:div(N, 2)+1, k1 in 1:N
    ]
    Dy_hat = im * κ₀ * [k2 - 1 for k2 in 1:div(N, 2)+1, k1 in 1:N]

    Delta_hat = - κ₀^2 * [
        ifelse(k1 ≤ div(N, 2) + 1, (k1 - 1)^2 + (k2 - 1)^2, (k1 - 1 - N)^2 + (k2 - 1)^2)
        for k2 in 1:div(N, 2)+1, k1 in 1:N
    ]

    ## Recovering of the velocity field from the vorticity
    Hu_hat = - Dy_hat ./ Delta_hat
    Hu_hat[1, 1] = 0.0
    Hv_hat = Dx_hat ./ Delta_hat
    Hv_hat[1, 1] = 0.0

    ## For the Basdevant formulation
    DxsqDysq_hat = Dx_hat.^2 .- Dy_hat.^2
    Dxy_hat = Dx_hat .* Dy_hat

    return Dx_hat, Dy_hat, Delta_hat, Hu_hat, Hv_hat, DxsqDysq_hat, Dxy_hat
end

# Methods to generate a scalar field from a certain list of excitable wavenumbers.

function field_from_spectrum(L, N, modes::Matrix{<:Integer}, amps::Matrix{<:Real})
    κ₀ = 2π/L
    x = y = (L/N):(L/N):L
    field = sum(
        [
            2κ₀^2 * (k[1]^2 + k[2]^2) * (
                a[1] * cos.(κ₀ * (k[1] * one.(y) * x' + k[2] * y * one.(x)'))
                - a[2] * sin.(κ₀ * (k[1] * one.(y) * x' + k[2] * y * one.(x)'))
            )
            for (k, a) in zip(eachrow(modes), eachrow(amps))
        ]
    )
    return field
end

function field_from_spectrum(rng::AbstractRNG, L, N, num_modes::Int)

    modes = rand(rng, 1:div(N,10), num_modes, 2)
    amps = rand(rng, num_modes, 2)

    field = field_from_spectrum(L, N, modes, amps)

    return field
end

field_from_spectrum(L, N, num_modes::Int) = field_from_spectrum(Xoshiro(), L, N, num_modes)

# Evolution steps

function step_naive!(vort_hat, dt, params)
    operators, vars = params
    Dx_hat, Dy_hat, Delta_hat, Hu_hat, Hv_hat,
        DxsqDysq_hat, Dxy_hat, Exp_nu_dt_Delta_hat = operators
    g_hat, N, Nsub = vars
    u_hat = Hu_hat .* vort_hat
    v_hat = Hv_hat .* vort_hat
    u = irfft(u_hat, N)
    v = irfft(v_hat, N)
    vort = irfft(vort_hat, N)
    wu_hat = rfft(vort .* u)
    wv_hat = rfft(vort .* v)
    vort_hat .= Exp_nu_dt_Delta_hat .* (
        vort_hat .+ dt * (
            g_hat .- Dx_hat .* wu_hat .- Dy_hat .* wv_hat
        )
    )
    ## dealiasing
    vort_hat[div(Nsub,2) + 1:end, :] .= 0.0im
    vort_hat[:, div(Nsub,2) + 1:div(N,2) + div(Nsub,2)] .= 0.0im
    return vort_hat
end

function step_Basdevant!(vort_hat, dt, params)
    operators, vars = params
    Dx_hat, Dy_hat, Delta_hat, Hu_hat, Hv_hat,
        DxsqDysq_hat, Dxy_hat, Exp_nu_dt_Delta_hat = operators
    g_hat, N, Nsub = vars
    u_hat = Hu_hat .* vort_hat
    v_hat = Hv_hat .* vort_hat
    u = irfft(u_hat, N)
    v = irfft(v_hat, N)
    uv_hat = rfft(u .* v)
    v2u2_hat = rfft(v.^2 - u.^2)
    vort_hat .= Exp_nu_dt_Delta_hat .* (
        vort_hat .+ dt * (
            g_hat .- DxsqDysq_hat .* uv_hat .- Dxy_hat .* v2u2_hat
        )
    )
    ## dealiasing
    vort_hat[div(Nsub,2) + 1:end, :] .= 0.0im
    vort_hat[:, div(Nsub,2) + 1:div(N,2) + div(Nsub,2)] .= 0.0im
    return vort_hat
end

function step_Basdevant_plan!(vort_hat, dt, params)
    operators, vars, auxs, plans = params

    Dx_hat, Dy_hat, Delta_hat, Hu_hat, Hv_hat,
        DxsqDysq_hat, Dxy_hat, Exp_nu_dt_Delta_hat = operators
    g_hat, N, Nsub = vars
    u_hat, v_hat, u, v, uv, v2u2, uv_hat, v2u2_hat = auxs
    plan, plan_inv = plans

    u_hat .= Hu_hat .* vort_hat
    v_hat .= Hv_hat .* vort_hat

    mul!(u, plan_inv, u_hat)
    mul!(v, plan_inv, v_hat)

    uv .= u .* v
    v2u2 .= v.^2 .- u.^2
    mul!(uv_hat, plan, uv)
    mul!(v2u2_hat, plan, v2u2)

    @. vort_hat = Exp_nu_dt_Delta_hat * (
        vort_hat + dt * (
            g_hat - DxsqDysq_hat * uv_hat - Dxy_hat * v2u2_hat
        )
    )

#=     vort_hat .= Exp_nu_dt_Delta_hat .* (
        vort_hat .+ dt .* (
            g_hat .- DxsqDysq_hat .* uv_hat .- Dxy_hat .* v2u2_hat
        )
    ) =#
    ## dealiasing
    vort_hat[div(Nsub,2) + 1:end, :] .= 0.0im
    vort_hat[:, div(Nsub,2) + 1:div(N,2) + div(Nsub,2)] .= 0.0im
    return vort_hat
end

function step_Basdevant_plan_loop!(vort_hat, dt, params)
    operators, vars, auxs, plans = params

    Dx_hat, Dy_hat, Delta_hat, Hu_hat, Hv_hat,
        DxsqDysq_hat, Dxy_hat, Exp_nu_dt_Delta_hat = operators
    g_hat, N, Nsub = vars
    u_hat, v_hat, u, v, uv, v2u2, uv_hat, v2u2_hat = auxs
    plan, plan_inv = plans

    @inbounds for i in eachindex(u_hat)
        u_hat[i] = Hu_hat[i] * vort_hat[i]
        v_hat[i] = Hv_hat[i] .* vort_hat[i]
    end

    mul!(u, plan_inv, u_hat)
    mul!(v, plan_inv, v_hat)

    @inbounds for i in eachindex(u)
        uv[i] = u[i] * v[i]
        v2u2[i] = v[i]^2 - u[i]^2
    end

    mul!(uv_hat, plan, uv)
    mul!(v2u2_hat, plan, v2u2)

    for i in eachindex(vort_hat)
        vort_hat[i] = Exp_nu_dt_Delta_hat[i] * (
            vort_hat[i] + dt * (
                g_hat[i] - DxsqDysq_hat[i] * uv_hat[i] - Dxy_hat[i] * v2u2_hat[i]
            )
        )
    end
    ## dealiasing
    @inbounds for j in 1:N, i in div(Nsub,2) + 1:div(N, 2) + 1
        vort_hat[i, j] = 0.0im
    end
    @inbounds for j in div(Nsub,2) + 1:div(N,2) + div(Nsub,2), i in 1:div(N, 2) + 1
        vort_hat[i, j] = 0.0im
    end
    return vort_hat
end

function step_Basdevant_plan_doubleloop!(vort_hat, dt, params)
    operators, vars, auxs, plans = params

    Dx_hat, Dy_hat, Delta_hat, Hu_hat, Hv_hat,
        DxsqDysq_hat, Dxy_hat, Exp_nu_dt_Delta_hat = operators
    g_hat, N, Nsub = vars
    u_hat, v_hat, u, v, uv, v2u2, uv_hat, v2u2_hat = auxs
    plan, plan_inv = plans

    M, N = size(vort_hat)

    @inbounds for j in 1:N, i in 1:M
        u_hat[i, j] = Hu_hat[i, j] * vort_hat[i, j]
        v_hat[i, j] = Hv_hat[i, j] .* vort_hat[i, j]
    end

    mul!(u, plan_inv, u_hat)
    mul!(v, plan_inv, v_hat)

    @inbounds for j in 1:N, i in 1:N
        uv[i, j] = u[i, j] * v[i, j]
        v2u2[i, j] = v[i, j]^2 - u[i, j]^2
    end

    mul!(uv_hat, plan, uv)
    mul!(v2u2_hat, plan, v2u2)

    for j in 1:N, i in 1:M
        vort_hat[i, j] = Exp_nu_dt_Delta_hat[i, j] * (
            vort_hat[i, j] + dt * (
                g_hat[i, j] - DxsqDysq_hat[i, j] * uv_hat[i, j] - Dxy_hat[i, j] * v2u2_hat[i, j]
            )
        )
    end
    ## dealiasing
    @inbounds for j in 1:N, i in div(Nsub,2) + 1:M
        vort_hat[i, j] = 0.0im
    end
    @inbounds for j in div(Nsub,2) + 1:div(N,2) + div(Nsub,2), i in 1:M
        vort_hat[i, j] = 0.0im
    end
    return vort_hat
end

# ## The spatial domain

L = 2π
κ₀ = 2π/L
N = 128
Nsub = 84
x = y = (L/N):(L/N):L

ν = 1.0

dt = 1.0e-5 # 1.0e-3 a 1.0e-5
t_final = 1.0

num_steps = Int(round(t_final / dt))

Dx_hat, Dy_hat, Delta_hat, Hu_hat, Hv_hat, DxsqDysq_hat, Dxy_hat = get_operators(N, κ₀)
Exp_nu_dt_Delta_hat = exp.(ν * dt * Delta_hat)

vort_init = field_from_spectrum(L, N, 4)
vort_init_hat = rfft(vort_init)

g_steady = field_from_spectrum(L, N, 1)
g_steady_hat = rfft(g_steady)

vort_steady_hat = - g_steady_hat ./ Delta_hat
vort_steady_hat[1, 1] = 0.0im
vort_steady = irfft(vort_steady_hat, N)

vort_hat = copy(vort_init_hat)
vort = irfft(vort_hat, N)

u_hat = zero(vort_hat)
v_hat = zero(vort_hat)
u = zeros(N, N)
v = zeros(N, N)
uv = zeros(N, N)
v2u2 = zeros(N, N)
uv_hat = zero(vort_hat)
v2u2_hat = zero(vort_hat)
plan = plan_rfft(vort, flags=FFTW.MEASURE)
plan_inv = plan_irfft(vort_hat, N, flags=FFTW.MEASURE)

operators =
    Dx_hat, Dy_hat, Delta_hat, Hu_hat, Hv_hat, DxsqDysq_hat, Dxy_hat, Exp_nu_dt_Delta_hat
vars = g_steady_hat, N, Nsub
auxs = u_hat, v_hat, u, v, uv, v2u2, uv_hat, v2u2_hat
plans = plan, plan_inv

params = (
    operators,
    vars,
    auxs,
    plans
)

vort_hat = copy(vort_init_hat)
vort = irfft(vort_hat, N)

println("Enstrophy convergence:")
println(sum(abs2, vort_hat - vort_steady_hat) * (1/N)^4)
for n in 1:num_steps
    ## step_naive!(vort_hat, dt, (operators, vars))
    ## step_Basdevant!(vort_hat, dt, (operators, vars))
    ## step_Basdevant_plan!(vort_hat, dt, (operators, vars, auxs, plans))
    step_Basdevant_plan_doubleloop!(vort_hat, dt, (operators, vars, auxs, plans))
    if rem(n, 1000) == 0
        println(sum(abs2, vort_hat - vort_steady_hat) * (1/N)^4)
    end
end

vort = irfft(vort_hat, N)
display(heatmap(x, y, vort, xlabel="x", ylabel="y", title="vorticity", titlefont=12))

display(surface(x, y, vort - vort_steady, xlabel="x", ylabel="y", zlabel="error", title="difference `vort_final .- vort_steady`", titlefont=12))

@info "step_naive!"
@btime step_naive!(vh, τ, p) setup = (vh = copy(vort_hat); τ = $dt; p = $(operators, vars));
@info "step_Basdevant!"
@btime step_Basdevant!(vh, τ, p) setup = (vh = copy(vort_hat); τ = $dt; p = $(operators, vars));
@info "step_Basdevant_plan!"
@btime step_Basdevant_plan!(vh, τ, p) setup = (vh = copy(vort_hat); τ = $dt; p = $(operators, vars, auxs, plans));
@info "step_Basdevant_plan_loop!"
@btime step_Basdevant_plan_loop!(vh, τ, p) setup = (vh = copy(vort_hat); τ = $dt; p = $(operators, vars, auxs, plans));
@info "step_Basdevant_plan_doubleloop!"
@btime step_Basdevant_plan_doubleloop!(vh, τ, p) setup = (vh = copy(vort_hat); τ = $dt; p = $(operators, vars, auxs, plans));

#=
[ Info: step_naive!
  261.500 μs (211 allocations: 2.16 MiB)
[ Info: step_Basdevant!
  222.791 μs (174 allocations: 2.03 MiB)
[ Info: step_Basdevant_plan!
  136.834 μs (0 allocations: 0 bytes)
[ Info: step_Basdevant_plan_loop!
  131.042 μs (0 allocations: 0 bytes)
[ Info: step_Basdevant_plan_doubleloop!
  137.000 μs (0 allocations: 0 bytes)
=#


#!/usr/bin/env julia

@printf "Make data/spikes.csv...\n"
include("julia/pprocess_gamma.jl")
x_gamma, rs_gamma = pprocess_gamma(15.0, 3.5);
writecsv("data/spikes.csv", x_gamma)

@printf "Make data/firing_rate_[true|estimated].csv...\n"
# データとモデルで両方作る．
# データはx_gamma使う．モデルはssgamma使う
include("julia/intensity_gamma.jl")
rs_mle_gamma = intensity_gamma(x_gamma);
writecsv("data/firing_rate_true.csv", rs_gamma);
writecsv("data/firing_rate_estimated.csv", rs_mle_gamma);


### A Pluto.jl notebook ###
# v0.14.0

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

# ╔═╡ 0e592c68-9167-11eb-30fd-efdb809f47f7
#Set up the working environment
begin
	using Pkg
	Pkg.activate("..")
	using Plots, PlutoUI,LaTeXStrings, DifferentialEquations
end

# ╔═╡ e81c951c-9166-11eb-1dd3-51fe9fcb57db
md"""
# Modelling Predator-Prey Dynamics - Coupled Differential Equations

So far, we've looked at a couple of simple models that described the growth or decline of a single population. Now, let's turn towards a more complicated scenario, modelling of **predator-prey** dynamics...

A famous and very simple model of pedator-prey dynamics is the so-called **Lotka-Volterra** model:

$$\begin{eqnarray}
	\frac{dN}{dt} = N (a-bP)\\
	\frac{dP}{dt} = P (cN - d)
\end{eqnarray}$$

Here, $N$ represents the prey population and $P$ represents the population of predators.

Let us take a closer look at the assumptions baked into this model:
1. In the abesence of predators, the prey grows exponentially
2. Predation has the effect that it reduces the growth rate of the prey population in a manner that depends on $N$ and $P$.

3. If no prey is present, then the predator population declines in an exponential manner.

4. The growth rate of the predator depends on prey in a directly proportional manner.
"""

# ╔═╡ 4611ebd2-3a1e-4c62-80c7-e7dfc56483e6
md"""
> **Exercise - Equilibrium/Stability:** The given parameter set gives rise oscillating (regularly fluctuating) population of predators and prey.
> 1. Try to find one parameter set that gives rise to stable populations (equilibrium) of predators and prey by systematically varying the parameters.
> 2. What would equilibrium mean in terms of the differential equations that underlie the model?
> 3. Try to find the expressions for $N$ and $P$ that correspond to their equilibrium populations.
> 4. Equilibria can be stable or unstable, depending on how they respond to tiny changes in the parameters. What happens if any of the two initial condition deviates a little bit from the equilibrium condition?
> 5. In reality, population sizes, as well as the factors that are represented by the parameters are subject to stochastic (random) variation. This randomness is not represented in the model. Discuss the potential effects of randomness on real populations that would otherwise be well described by the model.
"""

# ╔═╡ 0c6eb46c-a03e-4950-abe6-368b85a4b55e
md"""
> **Exercise - Interspecies Relationships**: Unter certain parameters, the model gives results that appear strange at first sight: Peaks in predator populations vastly outnumber those of the prey population.
>
>1. Using the concept of energy flow in ecosystems, suggest why a predator population that vastly outnumbers the prey population is generally not what we would expect to see in an ecosystem
>2. Find a parameter set that gives you a predator population that is substantially larger than the prey population.
>3. Which other type of interspecies relationship could be adequately described by this model? Would the relative population sizes make more sense given this type of interspecies relationship?
"""

# ╔═╡ f1a68086-9169-11eb-2721-652c6732d36f
md""" ## Simulation
*Change the sliders to adjust the corresponding values*:

| Parameter               | Value                  |
|-------------------------|------------------------|
| prey growth constant $a$ | $(@bind a Slider(0:0.001:0.25, default=0.2, show_value=true)) |
| prey death constant $b$ | $(@bind b Slider(0:0.001:0.25, default=0.02, show_value=true)) |
| predator growth constant $c$ | $(@bind c Slider(0:0.001:0.25, default=0.005, show_value=true)) |
| predator death constant $d$ | $(@bind d Slider(0:0.001:0.25, default=0.1, show_value=true)) |
| initial prey population size $N_0$ | $(@bind N0 Slider(1:1:100, default=50, show_value=true)) |
| initial predator population size $P_0$ | $(@bind P0 Slider(1:1:100, default=10, show_value=true)) |
| simulation time $t_{max}$ | $(@bind tmax Slider(1:5:5000, default=100, show_value=true)) |
"""

# ╔═╡ 0f573286-916e-11eb-157e-1d0139bd363c
function lotkavolterra!(du,u,p,t)
	du[1] = u[1]*(p[1] - p[2]*u[2])
	du[2] = u[2]*(p[3]*u[1] - p[4])
end

# ╔═╡ 2400191e-916e-11eb-2d01-3b3a48757ff4
p = [a, b, c, d]

# ╔═╡ 2c79a5ea-916e-11eb-1abb-7933f2a1eae1
u0 = [N0, P0]

# ╔═╡ 2c7a590c-916e-11eb-227b-97d8098a22b3
tspan = (0, tmax)

# ╔═╡ 33519b3e-916e-11eb-1cb5-7fbc17fba1d8
prob = ODEProblem(lotkavolterra!,u0,tspan,p)

# ╔═╡ 6c09dc9c-916a-11eb-1bff-fb5495955ad7
begin
	# Defining the system of ODEs in paramtetrised manner and using in-place notation 		(! after function name) to improve performance
	# u[1] = N, u[2] = P

	sol = solve(prob)
	simplot = plot(sol, labels = ["N(t)" "P(t)"], title = "Timecourse")
	phaseplot = plot(sol, vars = (1,2), xlabel = "N(t)", ylabel = "P(t)", labels= "", title = "Phaseplane")
	plot(simplot, phaseplot, layout = (1, 2),fontfamily="Computer Modern")
end

# ╔═╡ 6acb8812-fd98-492b-ae20-38efd9cf8cba
md"""
#### Notebook options
Show solutions: $(@bind show_solutions CheckBox(default=false))
"""

# ╔═╡ 3bb1ca46-0993-47bf-9ba7-ad65185bac96
begin
	if show_solutions == true
md"""
!!! tip "Solution"

	**1.** A simple equilibrium can be found by setting any of the two growth rate constants to zero. This would correspond to the extinction of the two populations.

	**2.** In terms of the equations, equilibrium corresponds to a situation where neither of the two populations grows or declines, hence both expressions ($\frac{dN}{dt}$ and $\frac{dP}{dt}$) are zero:

	$$\begin{eqnarray}
		\frac{dN}{dt} = 0 = N (a-bP)\\
		\frac{dP}{dt} = 0 = P (cN - d)
	\end{eqnarray}$$

	**3.** Using the insight from above and solving for $N$ and $P$, we find:

	$$\{N_{eq,1} = 0, P_{eq,1} = 0\}$$

	and

	$$\{N_{eq,2} = \frac{d}{c}, P_{eq,2} = \frac{a}{b}\}$$

	**4.** Tiny deviations instantly give rise to oscillations around the equilibrium point. The populations never reach equilibrium.

	**5.** Due to the tendency of the populations to oscillate around the non-zero equilibrium, we would not expect to see non-oscillation populations.
"""
	end
end

# ╔═╡ f6a109fd-232e-42b6-a2c2-88c7eb9befce
begin
	if show_solutions == true
md"""
!!! tip "Solution"
	**1.** Since only a fraction of the energy is transfered from lower to higher trophic levels, the situation described by the model does not reflect expected population sizes for typical predator-prey relationships.
	**2.** Many combinations work, e.g. $a$ = 0.2, $b$ = 0.02, $c$ = 0.2, $d$ = 0.2, $N_0$ = 10, $P_0$ = 40
	**3.** Another antagonistic (-/-) relationship that could be described by this model is parasitism. In this case the relative population sizes would make more sense.
"""
	end
end

# ╔═╡ Cell order:
# ╟─0e592c68-9167-11eb-30fd-efdb809f47f7
# ╟─e81c951c-9166-11eb-1dd3-51fe9fcb57db
# ╟─4611ebd2-3a1e-4c62-80c7-e7dfc56483e6
# ╟─3bb1ca46-0993-47bf-9ba7-ad65185bac96
# ╟─0c6eb46c-a03e-4950-abe6-368b85a4b55e
# ╟─f6a109fd-232e-42b6-a2c2-88c7eb9befce
# ╟─f1a68086-9169-11eb-2721-652c6732d36f
# ╟─6c09dc9c-916a-11eb-1bff-fb5495955ad7
# ╟─0f573286-916e-11eb-157e-1d0139bd363c
# ╟─2400191e-916e-11eb-2d01-3b3a48757ff4
# ╟─2c79a5ea-916e-11eb-1abb-7933f2a1eae1
# ╟─2c7a590c-916e-11eb-227b-97d8098a22b3
# ╟─33519b3e-916e-11eb-1cb5-7fbc17fba1d8
# ╟─6acb8812-fd98-492b-ae20-38efd9cf8cba


# Example that performes a NonLinear 2-D Wavelet Approximation
# via thresholding

#push path
push!(LOAD_PATH,pwd()*"/src") # for julia 0.4 you need to put your path
#import libraries
using FWT
using PyPlot
using Images

n = 1024;
name = "data_samples/ernst_reuter_haus.bmp";
f = load_image(name, n);
f = rescale(sum(f,3));
f = f[:,:,1];

Jmax = round(Int64, n/eta);
Jmin = 1;

# Display the threshold coefficients

clf;
subplot(1,2,1);
plot_wavelet(fW,Jmin);
title("Original coefficients");
subplot(1,2,2);
plot_wavelet(fWT,Jmin);
title("Thresholded coefficients");

# number of kept coefficients
m = round(n^2/16);
# compute the threshold T
Jmin = 1;
fW = FWT.perform_wavortho_transf(f,Jmin,+1, h);
# select threshold
v = sort(abs(fW[:]));
if v[1]<v[n^2]
    v = FWT.reverse(v);
end
# inverse transform
T = v[round(Int64,m)];
fWT = fW .* (abs(fW).>T);
# inverse
fnlin = FWT.perform_wavortho_transf(fWT,Jmin,-1, h);
# display
clf;
u1 = @sprintf("Linear, SNR=%.3fdB", snr(f,fLin));
u2 = @sprintf("Non-linear, SNR=%.3fdB", snr(f,fnlin));
imageplot(FWT.clamp(fLin),u1, 1,2,1 );
imageplot(clamp(fnlin),u2, 1,2,2 );


"""
    StableNumber(base::Number, exponent)

Stores results from munipulations of `StableArray`.
"""
struct StableNumber{E, A <: Number}
    base::A
    exponent::E
end

base(x::StableNumber) = x.base
exponent(x::StableNumber) = x.exponent

function _plus_ordered(x1::StableNumber, x2::StableNumber)
    xbase = base(x1) * exp(exponent(x1) - exponent(x2)) + base(x2)
    StableNumber(xbase, exponent(x2))
end

function Base.:(+)(x1::StableNumber, x2::StableNumber)
    indx = argmax((x1.exponent, x2.exponent))
    if indx == 2
        _plus_ordered(x1, x2)
    else
        _plus_ordered(x2, x1)
    end
end

function Base.:(+)(x1::StableNumber, x2::Number)
    x1 + StableNumber(x2, 1)
end

function Base.:(+)(x1::Number, x2::StableNumber)
    StableNumber(x1, 1) + x2
end

function Base.:(-)(x::StableNumber)
    StableNumber(-base(x), exponent(x))
end

function Base.:(-)(x1::StableNumber, x2::StableNumber)
    x1 + (-x2)
end

function Base.:(-)(x1::StableNumber, x2)
    x1 + (-x2)
end

function Base.:(-)(x1, x2::StableNumber)
    x1 + (-x2)
end

function Base.:(*)(x1::StableNumber, x2::StableNumber)
    StableNumber(base(x1) * base(x2), exponent(x1) + exponent(x2))
end

function Base.:(/)(x1::StableNumber, x2::StableNumber)
    StableNumber(base(x1) / base(x2), exponent(x1) - exponent(x2))
end

function Base.:(/)(x1::StableNumber, x2::Number)
    StableNumber(base(x1) * sign(x2), exponent(x1) - log(abs(x2)))
end

function Base.:(<)(x1::StableNumber, x2::StableNumber)
    xd = x1 - x2
    base(xd) < 0
end

function Base.:(<)(x1, x2::StableNumber)
    xd = x1 - x2
    base(xd) < 0
end

function Base.:(<)(x1::StableNumber, x2)
    xd = x1 - x2
    base(xd) < 0
end

function Base.log(x::StableNumber)
    log(base(x)) + exponent(x)
end

function Base.abs(x::StableNumber)
    StableNumber(abs(base(x)), exponent(x))
end

function Base.isnan(x::StableNumber)
    isnan(base(x)) || isnan(exponent(x))
end

function Base.show(io::IO, ::MIME"text/plain", x::StableNumber)
    print(io, "Stabled($(base(x))×exp($(exponent(x))))")
end

function Base.sign(x::StableNumber)
    sign(base(x))
end

function logview(x::StableNumber)
    StableNumber(sign(x), exponent(x)+log(abs(base(x))))
end


@testset "LP Atoms: $solver" for solver in solvers
    @testset "abs atom" begin
        x = Variable()
        p = minimize(abs(x), x<=-1)
        @test vexity(p) == ConvexVexity()
        solve!(p, solver)
        @test p.optval ≈ 1 atol=TOL
        @test evaluate(abs(x)) ≈ 1 atol=TOL

        x = Variable(2,2)
        p = minimize(sum(abs(x)), x[2,2]>=1, x[1,1]>=1, x>=0)
        @test vexity(p) == ConvexVexity()
        solve!(p, solver)
        @test p.optval ≈ 2 atol=TOL
        @test evaluate(sum(abs(x))) ≈ 2 atol=TOL
    end

    @testset "maximum atom" begin
        x = Variable(10)
        a = shuffle(collect(0.1:0.1:1.0))
        p = minimize(maximum(x), x >= a)
        @test vexity(p) == ConvexVexity()
        solve!(p, solver)
        @test p.optval ≈ maximum(a) atol=TOL
        @test evaluate(maximum(x)) ≈ maximum(a) atol=TOL
    end

    @testset "minimum atom" begin
        x = Variable(1)
        a = reshape(shuffle(collect(0.01:0.01:1.0)), (10, 10))
        p = maximize(minimum(x), x <= a)
        @test vexity(p) == ConvexVexity()
        solve!(p, solver)
        @test p.optval ≈ minimum(a) atol=TOL
        @test evaluate(minimum(x)) ≈ minimum(a) atol=TOL

        x = Variable(4, 4)
        y = Variable(4, 6)
        z = Variable(1)
        c = ones(4, 1)
        d = fill(2.0, (6, 1))
        constraints = [[x y] <= 2, z <= 0, z <= x, 2z >= -1]
        objective = sum(x + z) + minimum(y) + c' * y * d
        p = maximize(objective, constraints)
        @test vexity(p) == ConvexVexity()
        solve!(p, solver)
        @test p.optval ≈ 130 atol=TOL
        @test (evaluate(objective))[1] ≈ 130 atol=TOL
    end

    @testset "max atom" begin
        x = Variable(10, 10)
        y = Variable(10, 10)
        a = reshape(shuffle(collect(0.01:0.01:1.0)), (10, 10))
        b = reshape(shuffle(collect(0.01:0.01:1.0)), (10, 10))
        p = minimize(maximum(max(x, y)), [x >= a, y >= b])
        @test vexity(p) == ConvexVexity()
        solve!(p, solver)
        max_a = maximum(a)
        max_b = maximum(b)
        @test p.optval ≈ max(max_a, max_b) atol=10TOL
        @test evaluate(maximum(max(x, y))) ≈ max(max_a, max_b) atol=10TOL
    end

    @testset "min atom" begin
        x = Variable(10, 10)
        y = Variable(10, 10)
        a = reshape(shuffle(collect(0.01:0.01:1.0)), (10, 10))
        b = reshape(shuffle(collect(0.01:0.01:1.0)), (10, 10))
        p = maximize(minimum(min(x, y)), [x <= a, y <= b])
        @test vexity(p) == ConvexVexity()
        solve!(p, solver)
        min_a = minimum(a)
        min_b = minimum(b)
        @test p.optval ≈ min(min_a, min_b) atol=10TOL
        @test evaluate(minimum(min(x, y))) ≈ min(min_a, min_b) atol=10TOL
    end

    @testset "pos atom" begin
        x = Variable(3)
        a = [-2; 1; 2]
        p = minimize(sum(pos(x)), [x >= a, x <= 2])
        @test vexity(p) == ConvexVexity()
        solve!(p, solver)
        @test p.optval ≈ 3 atol=TOL
        @test evaluate(sum(pos(x))) ≈ 3 atol=TOL
    end

    @testset "neg atom" begin
        x = Variable(3)
        p = minimize(1, [x >= -2, x <= -2, neg(x) <= 3])
        @test vexity(p) == ConvexVexity()
        solve!(p, solver)
        @test p.optval ≈ 1 atol=TOL
        @test evaluate(sum(neg(x))) ≈ 6 atol=TOL
    end

    @testset "sumlargest atom" begin
        x = Variable(2)
        p = minimize(sumlargest(x, 2), x >= [1; 1])
        @test vexity(p) == ConvexVexity()
        solve!(p, solver)
        @test p.optval ≈ 2 atol=TOL
        @test evaluate(sumlargest(x, 2)) ≈ 2 atol=TOL

        x = Variable(4, 4)
        p = minimize(sumlargest(x, 3), x >= eye(4), x[1, 1] >= 1.5, x[2, 3] >= 2.1)
        @test vexity(p) == ConvexVexity()
        solve!(p, solver)
        @test p.optval ≈ 4.6 atol=TOL
        @test evaluate(sumlargest(x, 2)) ≈ 3.6 atol=TOL
    end

    @testset "sumsmallest atom" begin
        x = Variable(4, 4)
        p = minimize(sumlargest(x, 2), sumsmallest(x, 4) >= 1)
        @test vexity(p) == ConvexVexity()
        solve!(p, solver)
        @test p.optval ≈ 0.5 atol=TOL
        @test evaluate(sumsmallest(x, 4)) ≈ 1 atol=TOL

        x = Variable(3, 2)
        p = maximize(sumsmallest(x, 3), x >= 2, x <= 5, sumlargest(x, 3) <= 12)
        @test vexity(p) == ConvexVexity()
        solve!(p, solver)
        @test p.optval ≈ 12 atol=TOL
        @test evaluate(sumsmallest(x, 3)) ≈ 12 atol=TOL
    end

    @testset "dotsort atom" begin
        x = Variable(4, 1)
        p = minimize(dotsort(x, [1, 2, 3, 4]), sum(x) >= 7, x >= 0, x <= 2, x[4] <= 1)
        @test vexity(p) == ConvexVexity()
        solve!(p, solver)
        @test p.optval ≈ 19 atol=TOL
        @test vec(x.value) ≈ [2; 2; 2; 1] atol=TOL
        @test evaluate(dotsort(x, [1, 2, 3, 4])) ≈ 19 atol=TOL

        x = Variable(2, 2)
        p = minimize(dotsort(x, [1 2; 3 4]), sum(x) >= 7, x >= 0, x <= 2, x[2, 2] <= 1)
        @test vexity(p) == ConvexVexity()
        solve!(p, solver)
        @test p.optval ≈ 19 atol=TOL
        @test evaluate(dotsort(x, [1, 2, 3, 4])) ≈ 19 atol=TOL
    end

    @testset "hinge loss atom" begin
        # TODO: @davidlizeng. We should finish this someday.
    end

    @testset "norm inf atom" begin
        x = Variable(3)
        p = minimize(norm_inf(x), [-2 <= x, x <= 1])
        @test vexity(p) == ConvexVexity()
        solve!(p, solver)
        @test p.optval ≈ 0 atol=TOL
        @test evaluate(norm_inf(x)) ≈ 0 atol=TOL
    end

    @testset "norm 1 atom" begin
        x = Variable(3)
        p = minimize(norm_1(x), [-2 <= x, x <= 1])
        @test vexity(p) == ConvexVexity()
        solve!(p, solver)
        @test p.optval ≈ 0 atol=TOL
        @test evaluate(norm_1(x)) ≈ 0 atol=TOL
    end
end


function total_derivative(params_baseline, ϵ = 0.000001, settings = settings_defaults())
    @unpack ρ, d, γ = params_baseline
    sol_baseline = stationary_algebraic(params_baseline, settings)
    if γ == 1
        U_1 = 1/(ρ*sol_baseline.c)
        U_2 = (1/ρ^2)
    else
        U_1 = sol_baseline.c^(-γ) / (ρ + (γ-1)* sol_baseline.g)
        U_2 = sol_baseline.c^(1-γ) * (ρ + (γ-1)* sol_baseline.g)^(-2)
    end

    d_counterfactual = d + ϵ
    params_counterfactual = merge(params_baseline, (d = d_counterfactual,))
    sol_counterfactual = stationary_algebraic(params_counterfactual, settings)

	ACR = (-1/params_baseline.θ)*log(sol_counterfactual.λ_ii/sol_baseline.λ_ii)

    sol_fc_d = steady_state_from_g(sol_baseline.g, sol_baseline.z_hat, sol_baseline.Ω, params_counterfactual, settings)
    sol_fc_g = steady_state_from_g(sol_baseline.g + ϵ, sol_baseline.z_hat, sol_baseline.Ω, params_baseline, settings)
    sol_fc_zhat = steady_state_from_g(sol_baseline.g, sol_baseline.z_hat + ϵ, sol_baseline.Ω, params_baseline, settings)
    sol_fc_Omega = steady_state_from_g(sol_baseline.g, sol_baseline.z_hat, sol_baseline.Ω + ϵ, params_baseline, settings)

    partial_fc_d = (sol_fc_d.c - sol_baseline.c)/(d_counterfactual - params_baseline.d)
    partial_fc_g = (sol_fc_g.c - sol_baseline.c)/(sol_fc_g.g - sol_baseline.g)
    partial_fc_zhat = (sol_fc_zhat.c - sol_baseline.c)/(sol_fc_zhat.z_hat - sol_baseline.z_hat)
    partial_fc_Omega = (sol_fc_Omega.c - sol_baseline.c)/(sol_fc_Omega.Ω - sol_baseline.Ω)

    partial_Omega_d = (sol_counterfactual.Ω - sol_baseline.Ω)/(d_counterfactual - params_baseline.d)
    partial_zhat_d = (sol_counterfactual.z_hat - sol_baseline.z_hat)/(d_counterfactual - params_baseline.d)
    partial_g_d = (sol_counterfactual.g - sol_baseline.g)/(d_counterfactual - params_baseline.d)

    d_U_d_total = (sol_counterfactual.U_bar - sol_baseline.U_bar)/(d_counterfactual - params_baseline.d)
    total_decomp = U_1*(partial_fc_Omega*partial_Omega_d + partial_fc_zhat*partial_zhat_d + partial_fc_g*partial_g_d + partial_fc_d) + U_2*partial_g_d
    check = total_decomp/d_U_d_total
	total_decomp_CE = ρ*total_decomp

    decomp_fc_d = partial_fc_d
    decomp_fc_Omega_Omega_d = partial_fc_Omega*partial_Omega_d
    decomp_fc_zhat_zhat_d = partial_fc_zhat*partial_zhat_d
    decomp_fc_g_g_d = partial_fc_g*partial_g_d
    decomp_g_d = partial_g_d

    U1_partial_fc_d_frac = U_1*partial_fc_d/total_decomp
    U1_decomp_fc_Omega_Omega_d_frac = U_1*decomp_fc_Omega_Omega_d/total_decomp
	U1_decomp_fc_zhat_zhat_d_frac = U_1*decomp_fc_zhat_zhat_d/total_decomp
	U1_decomp_fc_g_g_d_frac = U_1*decomp_fc_g_g_d/total_decomp
	U2_decomp_g_d_frac = U_2*decomp_g_d/total_decomp

    planner_0_g_CE = ρ*(U_1*partial_fc_g +U_2)
    planner_0_Omega_CE = ρ*U_1*partial_fc_Omega
    planner_0_zhat_CE = ρ*U_1*partial_fc_zhat
	MU_g_CE = (1/ρ)
	semi_elasticity_g_d = params_baseline.d*partial_g_d
	U1_partial_fc_d_ACRfrac = U1_partial_fc_d_ACRfrac = ρ*U_1*partial_fc_d*ϵ/ACR
	d_T = 1 + 0.9*(params_baseline.d - 1)
	d_d = d_T - params_baseline.d
	CE_extrap10p = total_decomp_CE * d_d
	g_extrap10p = partial_g_d * d_d

    return (check = check, d_U_d_total = d_U_d_total, U_1 = U_1, U_2 = U_2,
										ACR = ACR,
										∂_fc_d = partial_fc_d,
										∂_fc_g = partial_fc_g,
										∂_fc_zhat = partial_fc_zhat,
										∂_fc_Omega = partial_fc_Omega,
										∂_g_d = partial_g_d,
										∂_zhat_d = partial_zhat_d,
										∂_Omega_d = partial_Omega_d,
										total_decomp = total_decomp,
										decomp_fc_d = decomp_fc_d,
										decomp_fc_Omega_Omega_d = decomp_fc_Omega_Omega_d ,
										decomp_fc_zhat_zhat_d = decomp_fc_zhat_zhat_d,
										decomp_fc_g_g_d = decomp_fc_g_g_d,
										decomp_g_d = decomp_g_d,
										U1_partial_fc_d_frac = U1_partial_fc_d_frac,
										U1_decomp_fc_Omega_Omega_d_frac = U1_decomp_fc_Omega_Omega_d_frac,
										U1_decomp_fc_zhat_zhat_d_frac = U1_decomp_fc_zhat_zhat_d_frac,
										U1_decomp_fc_g_g_d_frac = U1_decomp_fc_g_g_d_frac,
										U2_decomp_g_d_frac = U2_decomp_g_d_frac,
										planner_0_g_CE = planner_0_g_CE,
										planner_0_Omega_CE = planner_0_Omega_CE,
										planner_0_zhat_CE = planner_0_zhat_CE,
										MU_g_CE = MU_g_CE,
										semi_elasticity_g_d = semi_elasticity_g_d,
										U1_partial_fc_d_ACRfrac = U1_partial_fc_d_ACRfrac,
										d_T = d_T,
										d_d = d_d,
										CE_extrap10p = CE_extrap10p,
										g_extrap10p = g_extrap10p)
end

function steady_state_from_g(g, z_hat, Ω, parameters, settings)
    @assert parameters.υ > 0 && parameters.κ > 0
    @assert z_hat > 1 && Ω > 0 && g > 0 # input validation still required
    staticvalues = staticvals([g, z_hat, Ω], parameters)
    return merge(staticvalues, merge((g = g, z_hat = z_hat, Ω = Ω,), welfare([g, z_hat, Ω], parameters)))
end

function stationary_algebraic(parameters, settings)
    x0 = settings.stationary_x0(parameters, settings)
    @assert parameters.υ > 0 && parameters.κ > 0
    sol = nlsolve(x -> stationary_algebraic_aux(x, parameters), x0, inplace = false) # uses distorted parameters
    converged(sol) || error("Solver didn't converge.")
    g, z_hat, Ω = sol.zero
    @assert z_hat > 1 && Ω > 0 && g > 0
    staticvalues = staticvals([g, z_hat, Ω], parameters) # uses undistorted parameters (except x)

    # calculate value function using differential objects based on the grid
    @unpack z_ex = settings
    z = z_ex[2:end-1]
    @unpack θ, σ = parameters
    bc = (Mixed(ξ = σ-1), Mixed(ξ = σ-1)) # boundary conditions for differential operators
    L_1_minus = L₁₋bc(z_ex, bc) # use backward difference as the drift is negative
    L_2 = L₂bc(z_ex, bc)
    ω = ω_weights(z_ex, θ, σ-1) # get quadrature weights for the distribution on the rescaled grid.
    Ξ₁ = 1/(1 - (σ-1)*(z[1] - z_ex[1])) # (24), with ξ = (σ-1)
    p = (L_1_minus = L_1_minus, L_2 = L_2, ω = ω, Ξ₁ = Ξ₁, z = z)
    v_tilde = stationary_numerical_given_vals([g, z_hat, Ω], p, parameters, settings).v_tilde

    return merge(staticvalues, merge((g = g, z_hat = z_hat, Ω = Ω, v_tilde = v_tilde), welfare([g, z_hat, Ω], parameters))) # calculate quantities and return
end

function welfare(vals, parameters)
    g, z_hat, Ω = vals
    @unpack ρ, σ, N, θ, γ, d, κ, ζ, η, Theta, χ, υ, μ, δ = parameters
    @unpack F, r, ν, a, b, S, L_tilde, z_bar, w, π_min = staticvals(vals, parameters)
    c_T = c(L_tilde, Ω, z_bar)
    return (y = c_T, # y = c by (C.78)
            c = c_T,
            λ_ii = λ_ii(z_hat, parameters),
            U_bar = γ == 1 ? (ρ*log(c_T) + g) / ρ^2 : 1/(1-γ) * (c_T^(1-γ))/(ρ + (γ-1)*g)) # (C.16)
end

# Kernel for the algebraic stationary.
function stationary_algebraic_aux(vals, parameters)
    @unpack ρ, σ, N, θ, γ, d, κ, η, Theta, υ, μ, δ, χ, ζ = parameters
    @unpack F, r, ν, a, b, S, L_tilde, z_bar, w, x, π_min = staticvals(vals, parameters)
    g, z_hat, Ω = vals

    big_denom = ν*(θ + ν)*(θ - σ + 1) # (C.19)
    denom_1 = a*(g - r) # (C.19)
    num_1 = ν*(N-1)*(θ - σ + 1)*(d^(1 - σ)*(θ + ν)*z_hat^(-θ + σ - 1)-b*θ*z_hat^(-θ-ν)) # (C.19)
    num_2 = θ*(ν*(N-1)*d^(1-σ)*(θ+ν)*z_hat^(-θ + σ -1) + (ν + σ - 1)*(θ + ν - σ + 1)) # (C.19)
    return [x/π_min - a*(χ/(1-χ))*(σ + ν - 1)/ν, # (C.18)
            1 + (σ-1)/ν - (num_1/denom_1 + num_2)/big_denom + (χ/(1-χ))*(σ + ν - 1)/(ν), # (C.19)
            π_min - (1- L_tilde)/((σ -1)*z_bar^(σ-1))] # (C.20)
end

function staticvals(vals, parameters)
    g, z_hat, Ω = vals
    @assert parameters.η ≈ 0  # we think that the L_tilde formulas are wrong for η != 0
    @unpack ρ, σ, N, θ, γ, d, κ, ζ, η, Theta, χ, υ, μ, δ = parameters

    F(z) = 1 - z^(-θ) # (C.1)
    r = ρ + γ*g + δ # (C.6)
    ν = (μ-g)/υ^2 + sqrt(((g-μ)/υ^2)^2 + (r-g)/(υ^2/2)) # (C.3)
    a = (r - g - (σ - 1)*(μ - g + (σ - 1)*υ^2/2))^(-1) # (C.4)
    b = (1 - a*(r-g))*d^(1-σ)*z_hat^(ν + σ - 1) # (C.5)
    S_T = S(g, parameters) # calls like this rely on utils/static.jl
    L_tilde_T = L_tilde(g, z_hat, Ω, δ, S_T, parameters)
    L_tilde_x_T = L_tilde_x(z_hat, Ω, parameters)
    L_tilde_E_T = L_tilde_E(Ω, δ, parameters)
    L_tilde_a_T = L_tilde_a(Ω, S_T, parameters)
    z_bar_T = z_bar(z_hat, Ω, parameters)
    w_T = w(z_bar_T, parameters)
    x_T = x(w_T, parameters)
    π_min_T = (d^(σ-1) * κ)/(z_hat^(σ-1)) # (C.12, inverted to express π_min as a function of parameters and z_hat) This is where the export threshold is used.  pi_min(L_tilde, z_bar)  is already used in the system of equations
    π_rat_T = π_rat(z_hat, parameters)

    return (F = F, r = r, ν = ν, a = a, b = b, S = S_T, L_tilde = L_tilde_T, L_tilde_x = L_tilde_x_T, L_tilde_E = L_tilde_E_T, L_tilde_a = L_tilde_a_T,
            z_bar = z_bar_T, w = w_T, x = x_T, π_min = π_min_T, π_rat = π_rat_T)
end

# Kernel of the stationary numerical.
function stationary_numerical_given_vals(vals, p, parameters, settings)
    g, z_hat, Ω = vals
    @unpack F, r, ν, a, b, S, L_tilde, z_bar, w, π_min = staticvals([g, z_hat, Ω], parameters)
    @unpack ρ, σ, N, θ, γ, d, κ, η, Theta, ζ, χ, υ, μ, δ = parameters
    @unpack Ξ₁, L_1_minus, L_2, ω, z = p
    r_tilde = r - g - 0 # (C.59, and g_w = 0 at steady state)
    ρ_tilde = r_tilde - (σ - 1)*(μ - g + (σ-1)*(υ^2/2)) # (C.41)
    A_T = ρ_tilde * I - (μ - g + (σ-1)*υ^2)*L_1_minus - υ^2/2 * L_2 # (52)
    i(z) = z >= log(z_hat) ? 1 : 0 # indicator function for next equation.
    π(z) = π_min * (1 + (N-1)*d^(1-σ)*i(z)) - (N-1)*κ*exp(-(σ-1)*z)*i(z) # (39)
    v_tilde = A_T \ π.(z) # discretized system of ODE for v, where v'(T) = 0 (53)

    # System of equations, given the numerically solved ODE
    return (residuals = [Ξ₁*v_tilde[1] - dot(v_tilde, ω) + ζ, # (54)
            Ξ₁*v_tilde[1] - ζ*(1-χ)/χ, # (56)
            π_min - (1 - L_tilde)/((σ-1)*z_bar^(σ-1))], # (38)
            v_tilde = v_tilde)
end

# Solve the ODE by root-finding. Should use the analytic most of the time, but this is a useful check.
function stationary_numerical(parameters, settings)
    @unpack z_ex = settings
    x0 = settings.stationary_x0(parameters, settings)
    @unpack ρ, σ, N, θ, γ, d, κ, η, Theta, υ, μ, δ = parameters
    @assert υ > 0 && κ > 0 # Parameter validation
    z = z_ex[2:end-1] # form a uniform extended grid

    # Differential objects we need for the residuals
    bc = (Mixed(ξ = σ-1), Mixed(ξ = σ-1)) # boundary conditions for differential operators
    L_1_minus = L₁₋bc(z_ex, bc) # use backward difference as the drift is negative
    L_2 = L₂bc(z_ex, bc)
    ω = ω_weights(z_ex, θ, σ-1) # Get quadrature weights for the distribution on the rescaled grid.
    Ξ₁ = 1/(1 - (σ-1)*(z[1] - z_ex[1])) # (24), with ξ = (σ-1)

    p = (Ξ₁ = Ξ₁, L_1_minus = L_1_minus, L_2 = L_2, ω = ω, z = z)
    sol = nlsolve(vals -> stationary_numerical_given_vals(vals, p, parameters, settings).residuals, x0, inplace = false)
    converged(sol) || error("Solver didn't converge.")
    g_T, z_hat_T, Ω_T = sol.zero

    staticvalues = staticvals([g_T, z_hat_T, Ω_T], parameters)
    v_tilde = stationary_numerical_given_vals([g_T, z_hat_T, Ω_T], p, parameters, settings).v_tilde
    return merge(staticvalues, merge((g = g_T, z_hat = z_hat_T, Ω = Ω_T, v_tilde = v_tilde), welfare([g_T; z_hat_T; Ω_T], parameters)))
end

# Utilities for comparison, display, etc.
function compare_steady_states(parameters, settings; verbose = false, algebraic = false)

    if algebraic
        stationary = stationary_algebraic
    else
        stationary = stationary_numerical
    end

    parameters_0 = merge(parameters, (d = parameters.d_0,))
    parameters_T = merge(parameters, (d = parameters.d_T,))
    stationary_0 = stationary(parameters_0, settings)
    stationary_T = stationary(parameters_T, settings)
    change_welfare = 100*consumption_equivalent(stationary_T.U_bar, stationary_0.U_bar, parameters)
    change_trade = (1-stationary_T.λ_ii) - (1-stationary_0.λ_ii)

    if verbose
        println("SS to SS welfare gain: $(change_welfare)")
        println("Change in Trade: $(change_trade)")
        println("Growth Rates across SS: $(100*stationary_T.g) (time T) vs $(100*stationary_0.g) (time 0)")
    end

    return (stationary_0 = stationary_0,
            stationary_T = stationary_T,
            change_welfare = change_welfare,
            change_trade = change_trade)
end


using DiscreteChoice
using ForwardDiff
using Test

# Rosenbrock.
f(x) = (1.0 - x[1])^2 + 100.0 * (x[2] - x[1]^2)^2;
g(x) = ForwardDiff.gradient(f, x);
H(x) = ForwardDiff.hessian(f, x);

# Gradient descent.
@test optimize(f, g, zeros(2), 100000)[1] ≈ ones(2)

# Newton.
@test optimize(f, g, H, zeros(2))[1] ≈ ones(2)

# Quasi-Newton using inverse BFGS.
#@test optimize(f, g, inv_BFGS, zeros(2), true)[1] ≈ ones(2)

# Trust region using Steihaug-Toint's method and exact Hessian.
@test optimize(f, g, H, tcg, zeros(2))[1] ≈ ones(2)

# Trust region using Steihaug-Toint's method and approximative Hessian using BFGS.
@test optimize(f, g, BFGS!, tcg, zeros(2), true)[1] ≈ ones(2)

# Trust region using Steihaug-Toint's method and approximative Hessian using SR1.
@test optimize(f, g, SR1!, tcg, zeros(2), true)[1] ≈ ones(2)


using DynamicPolynomials

function bench()
    o = one(Rational{BigInt})
    @polyvar x y z
    a = (o * x + o * y^2) * (o * z^3 + o * y^2 + o * x)
    @show a
    b = (o * x + o * y + o * z) * (o * x^2 + o * y)
    @show b
    @time gcd(a, b)
end


Base.angle(d1, d2) = acos(min(dot(normalize(d1), normalize(d2)), 1))
Base.angle(a::T, b::T) where {T<:Union{CalibratedHit, PMT, MCTrack}} = Base.angle(a.dir, b.dir)
Base.angle(a, b::Union{CalibratedHit, PMT, MCTrack}) = Base.angle(a, b.dir)
Base.angle(a::Union{CalibratedHit, PMT, MCTrack}, b) = Base.angle(a.dir, b)

"""
$(SIGNATURES)

Calculate the azimuth angle for a given direction.
"""
azimuth(d) = atan(d[2], d[1])

"""
$(SIGNATURES)

Calculate the zenith angle for a given direction.
"""
zenith(d) = acos(-d[3]/norm(d))

"""
$(SIGNATURES)

Calculate the cartesian coordinates for given `ϕ`, `θ` and radius `r`.
"""
function cartesian(ϕ, θ; r=1.0)
    return Position(r*sin(θ) * cos(ϕ), r*sin(θ) * sin(ϕ), cos(θ))
end

"""
$(SIGNATURES)

Calculate the distance between a point (p1) and a line (given by p2 and d2).
"""
function pld3(p1, p2, d2)
    norm(cross(d2, (p2 - p1))) / norm(d2)
end

"""
$(SIGNATURES)

Calculate the distance between two lines.
"""
function lld3(P, u, Q, v)
    R = Q - P
    n = cross(u, v)
    return norm(R⋅n) / norm(n)
end

"""
$(SIGNATURES)

Calculate the distance between two tracks.
"""
function lld3(t₁::Track, t₂::Track)
    return lld3(t₁.pos, t₁.dir, t₂.pos, t₂.dir)
end

"""
$(SIGNATURES)

Projects a point to a track.
"""
function project(P, t::Track)
    A = t.pos
    B = A + t.dir
    project(P, A, B)
end

"""
$(SIGNATURES)

Project P onto a line spanned by A and B.
"""
function project(P, A, B)
    Position(A + ((P-A)⋅(B-A))/((B-A)⋅(B-A)) * (B-A))
end

"""
$(SIGNATURES)

The function to fit time residual plot distributions.
"""
function langauss(x, LA, Lμ, Lσ, GA, Gμ, Gσ, offset)
    LA * pdf(Landau(Lμ, Lσ), x) + GA * pdf(Normal(Gμ, Gσ), x) + offset
end


using Plots, LaTeXStrings, Measures; pyplot()

a, c, m = 69069, 1, 2^32
next(z) = (a*z + c) % m

N = 10^6
data = Array{Float64,1}(undef, N)

x = 808
for i in 1:N
    data[i] = x/m
    global x = next(x)
end

p1 = scatter(1:1000, data[1:1000],
    c=:blue, m=4, msw=0, xlabel=L"n", ylabel=L"x_n")
p2 = histogram(data, bins=50, normed=:true,
    ylims=(0,1.1), xlabel="Support", ylabel="Density")
plot(p1, p2, size=(800, 400), legend=:none, margin = 5mm)


using Test

push!(LOAD_PATH, expanduser(".")) # Assumed to be ran from the repl folder.

#include("RadioPropagation.jl")
using RadioPropagation

# Unit test methods.
@testset "RadioPropagation.jl" begin
    @test RadioPropagation.two_ray_propagation( 2e3, 10, 12, 20e9, -1 ) == 1.9993946082759417 + 0.03479104696580754im;
    @test RadioPropagation.rain_attenuation_db_per_km( 'v', 30, 20 ) == 3.3400000000000003;
    @test RadioPropagation.rain_attenuation_db_per_km_circular_pol( 30, 20 ) == 3.6755027981960815; # dB/km
    @test RadioPropagation.fog_attenuation_db_per_km( 10, 0.8, 23 ) == 4.68976; # dB/km
end


using BosonSampling
using Test
using LinearAlgebra

@testset "BosonSampling.jl" begin

    @testset "scattering matrix" begin

        input_state = [2,0]
	    output_state = [1,1]
	    U = matrix_test(2)
	    @test scattering_matrix(U, input_state, output_state) == [1 2; 1 2]

	    input_state = [1,1]
	    output_state = [0,2]
	    U = matrix_test(2)
	    @test scattering_matrix(U, input_state, output_state) == [2 2; 4 4]

	end

	@testset "process_probability: HOM" begin

		@test process_probability(fourier_matrix(2), [1,1], [0,2]) ≈ 0.5 atol = 1e-8
	    @test process_probability(fourier_matrix(2), [1,1], [2,0]) ≈ 0.5 atol = 1e-8
	    @test process_probability(fourier_matrix(2), [1,1], [1,1]) ≈ 0. atol = 1e-8

		S_bosonic = ones(2,2)
		S_dist = [1 0; 0 1]

		@test process_probability_partial(fourier_matrix(2), S_bosonic, [1,1], [0,2]) ≈ 0.5 atol = 1e-8
		@test process_probability_partial(fourier_matrix(2), S_bosonic, [1,1], [2,0]) ≈ 0.5 atol = 1e-8
		@test process_probability_partial(fourier_matrix(2), S_bosonic, [1,1], [1,1]) ≈ 0. atol = 1e-8
		@test process_probability_partial(fourier_matrix(2), S_dist, [1,1], [0,2]) ≈ 0.25 atol = 1e-8
		@test process_probability_partial(fourier_matrix(2), S_dist, [1,1], [2,0]) ≈ 0.25 atol = 1e-8
		@test process_probability_partial(fourier_matrix(2), S_dist, [1,1], [1,1]) ≈ 0.5 atol = 1e-8

	end

	@testset "probability distribution of photons in modes : HOM" begin

		U = fourier_matrix(2)

		occupancy_vector = [1, 0]
		@test proba_partition_bosonic(U = U, occupancy_vector = occupancy_vector) ≈ [0.5,0,0.5] atol = 1e-5

		occupancy_vector = [0, 1]
		@test proba_partition_bosonic(U = U, occupancy_vector = occupancy_vector) ≈ [0.5,0,0.5] atol = 1e-5

		occupancy_vector = [1, 1]
		@test proba_partition_bosonic(U = U, occupancy_vector = occupancy_vector) ≈ [0,0,1.] atol = 1e-5

	end

	@testset "probability distribution of distinguishable photons in modes : HOM" begin

	    U = fourier_matrix(2)
	    part = [1,0]
	    @test proba_partition_distinguishable(occupancy_vector = part, U = U) ≈ [0.25,0.5,0.25] atol = 1e-8

	end

	@testset "partial distinguishability partitions" begin

		m = 10
		n = 4
		input_state = zeros(Int, m)
		occupancy_vector = zeros(Int, m)
		for i=1:n
			input_state[i] = 1
		end
		occupancy_vector[1] = 1

		U = fourier_matrix(10)

		@test proba_partition_partial(U = U, S = ones(n, n), occupancy_vector = occupancy_vector, input_state = input_state) == proba_partition_bosonic(U = U, occupancy_vector = occupancy_vector, input_state = input_state)
	    @test proba_partition_partial(U = U,  S = Matrix{Float64}(I, n, n), occupancy_vector = occupancy_vector, input_state = input_state) ≈ proba_partition_distinguishable(occupancy_vector = occupancy_vector, U = U, input_state = input_state)

	end

	@testset "theoretical_distribution" begin

		n = 3
		occupation = ModeOccupation(random_occupancy(n,n))
		i = Input{Bosonic}(occupation)
		interf = Fourier(i.r.m)
		p_theo = theoretical_distribution(input=i, interf=interf)

		@test sum(p_theo) ≈ 1 atol=1e-9
		@test all(p -> p>=0, p_theo)

		output_events = output_mode_occupation(n,n)
	    for i = 1:length(output_events)
	        if check_suppression_law(output_events[i])
	            @test p_theo[i] ≈ 0 atol=1e-6
	        end
	    end

	end

	@testset "noisy statistics" begin

		n = 3
		occupation = ModeOccupation(random_occupancy(n,n))
    	i = Input{OneParameterInterpolation}(occupation, 0.5)
	    interf = Fourier(i.r.m)
	    res = noisy_distribution(input=i, reflectivity=0.5, interf=interf)

		p_exact = res[1]
		p_approx = res[2]
		p_samp = res[3]

		@test sum(p_exact) ≈ 1 atol=1e-9
		@test sum(p_approx) ≈ 1 atol=1e-9
		@test sum(p_samp) ≈ 1 atol=1e-9
		@test all(p->p>=0, p_exact)
		@test all(p->p>=0, p_approx)
		@test all(p->p>=0, p_samp)

		i = Input{Bosonic}(ModeOccupation(random_occupancy(n,n)))
		res = noisy_distribution(input=i, reflectivity=0.999, interf=interf, approx=false, samp=false)
		p_exact = res[1]
		output_events = output_mode_occupation(n,n)
		for i = 1:length(output_events)
            if check_suppression_law(output_events[i])
		        @test p_exact[i] ≈ 0 atol=1e-6
		    end
		end

	end

	@testset "suppression law boson samplers" begin

		n = 3
	    interf = Fourier(n)

		input_clifford_sampler = Input{Bosonic}(first_modes(n,n))
	   	input_noisy_sampler = Input{OneParameterInterpolation}(first_modes(n,n), 1.0)

	    out_clifford_sampler = cliffords_sampler(input=input_clifford_sampler, interf=interf)
		out_noisy_sampler = noisy_sampler(input=input_noisy_sampler, reflectivity=1.0, interf=interf)

	    @test !check_suppression_law(out_clifford_sampler)
		@test !check_suppression_law(out_noisy_sampler)

	end

	@testset "examples usage" begin
		@test_nowarn include("example_usage.jl")

	end

end


using Test

using CompScienceMeshes
using BEAST
using StaticArrays
using LinearAlgebra

c = 3e8
μ = 4*π*1e-7
ε = 1/(μ*c^2)
f = 5e7
λ = c/f
k = 2*π/λ
ω = k*c
η = sqrt(μ/ε)

a = 1
Γ_orig = CompScienceMeshes.meshcuboid(a,a,a,0.1)
Γ = translate(Γ_orig,SVector(-a/2,-a/2,-a/2))

Φ, Θ = [0.0], range(0,stop=π,length=100)
pts = [point(cos(ϕ)*sin(θ), sin(ϕ)*sin(θ), cos(θ)) for ϕ in Φ for θ in Θ]

# This is an electric dipole
# The pre-factor (1/ε) is used to resemble
# (9.18) in Jackson's Classical Electrodynamics
E = (1/ε) * dipolemw3d(location=SVector(0.4,0.2,0),
                       orientation=1e-9.*SVector(0.5,0.5,0),
                       wavenumber=k)

n = BEAST.NormalVector()

𝒆 = (n × E) × n
H = (-1/(im*μ*ω))*curl(E)
𝒉 = (n × H) × n

𝓣 = Maxwell3D.singlelayer(wavenumber=k)

X = raviartthomas(Γ)

T = Matrix(assemble(𝓣,X,X))
e = Vector(assemble(𝒆,X))
j_EFIE = T\e

nf_E_EFIE = potential(MWSingleLayerField3D(wavenumber=k), pts, j_EFIE, X)
nf_H_EFIE = potential(BEAST.MWDoubleLayerField3D(wavenumber=k), pts, j_EFIE, X) ./ η
ff_E_EFIE = potential(MWFarField3D(wavenumber=k), pts, j_EFIE, X)

accuracy_default_quadrature_nf_e = norm(nf_E_EFIE - E.(pts))/norm(E.(pts))
accuracy_default_quadrature_nf_h = norm(nf_H_EFIE - H.(pts))/norm(H.(pts))
accuracy_default_quadrature_ff_e =
    norm(ff_E_EFIE - E.(pts, isfarfield=true))/norm(E.(pts, isfarfield=true))

function fastquaddata(op::BEAST.MaxwellOperator3D,
    test_local_space::BEAST.RefSpace, trial_local_space::BEAST.RefSpace,
    test_charts, trial_charts)

    a, b = 0.0, 1.0
    # CommonVertex, CommonEdge, CommonFace rules
    println("Fast quadrule is called")
    tqd = quadpoints(test_local_space, test_charts, (1,2))
    bqd = quadpoints(trial_local_space, trial_charts, (1,2))
    leg = (BEAST._legendre(3,a,b), BEAST._legendre(4,a,b), BEAST._legendre(5,a,b),)

    return (tpoints=tqd, bpoints=bqd, gausslegendre=leg)
end

function fastquaddata(fn::BEAST.Functional, refs, cells)
    println("Fast RHS quadrature")
    return quadpoints(refs, cells, [1])
end

fastT = Matrix(assemble(𝓣,X,X, quaddata=fastquaddata))
faste = Vector(assemble(𝒆,X, quaddata=fastquaddata))
fastj_EFIE = fastT\faste

fastnf_E_EFIE = potential(MWSingleLayerField3D(wavenumber=k), pts, fastj_EFIE, X)
fastnf_H_EFIE = potential(BEAST.MWDoubleLayerField3D(wavenumber=k), pts, fastj_EFIE, X) ./ η
fastff_E_EFIE = potential(MWFarField3D(wavenumber=k), pts, fastj_EFIE, X)

accuracy_fast_quadrature_nf_e = norm(fastnf_E_EFIE - E.(pts))/norm(E.(pts))
accuracy_fast_quadrature_nf_h = norm(fastnf_H_EFIE - H.(pts))/norm(H.(pts))
accuracy_fast_quadrature_ff_e =
    norm(fastff_E_EFIE - E.(pts, isfarfield=true))/norm(E.(pts, isfarfield=true))

@test accuracy_fast_quadrature_nf_e > accuracy_default_quadrature_nf_e
@test accuracy_fast_quadrature_nf_h > accuracy_default_quadrature_nf_h
@test accuracy_fast_quadrature_ff_e > accuracy_default_quadrature_ff_e


#Playing around with the mean polynomial function

using Gadfly, GaussianProcesses

#Training data
x=[-4.0,-3.0,-1.0,0.0,2.0];

# Create quadratic function
y = 0.5*x.^2 + 2.0*x + 0.5*rand(5)

#Test data
xpred = collect(-5.0:0.1:5.0);

#Specify mean and covariance function
beta = [0.5 2.0]
mPoly = MeanPoly(beta)
kern = SE(0.0,0.0)
gp = GP(x,y,mPoly,kern)  #Fit the GP

plot(gp)   #Plot the GP


#=  After Azevedo and Leshno (2016, §IV.A), an example of how in some situations
	schools may be disincentivized to improve their quality.   =#

using DeferredAcceptance

function argsort(vec::AbstractArray{<:Real, 1})::AbstractArray{<:Real, 1}
    return invperm(sortperm(vec))
end

# Might add these functions to the module later.
function quality(assn, capacities, scores)
	(m, ) = size(capacities)
	res = zeros(m)
	for (s, c) in enumerate(assn)
		if c ≤ m
			res[c] += scores[s, c]
		end
	end
	res ./= capacities
	return res
end

function cutoffs(assn, capacities, scores)
	(m, ) = size(capacities)
	res = zeros(m)
	for c in 1:m
		res[c] = minimum(scores[:, c][assn .== c])
	end
	return res
end

students = [2 1 1 1 1;
			1 2 2 2 2]

# School 2 improves, winning over one applicant
students_delta = [2 2 1 1 1;
				  1 1 2 2 2]

# For illustration purposes
scores = [ 0.95  0.96  0.89  0.83  0.61;
 		   0.97  0.84  0.82  0.9   0.68 ]

schools = mapslices(argsort, 1 .- scores, dims=2)

@assert schools == [2 1 3 4 5;
		            1 3 4 2 5]

capacities = [2, 3]

out_normal = DA(students, schools', capacities)
out_delta = DA(students_delta, schools', capacities)

println("Normal assn:  ", out_normal[1])
println("    Quality:  ", quality(out_normal[1], capacities, scores'))
println("    Cutoffs:  ", cutoffs(out_normal[1], capacities, scores'))
println("Delta:        ", out_delta[1])
println("    Quality:  ", quality(out_delta[1], capacities, scores'))
println("    Cutoffs:  ", cutoffs(out_delta[1], capacities, scores'))
# Normal assn:  [2, 1, 1, 2, 2]
#     Quality:  [0.925, 0.8500000000000001]
#     Cutoffs:  [0.89, 0.68]
# Delta:        [2, 2, 1, 1, 2]
#     Quality:  [0.86, 0.8300000000000001]
#     Cutoffs:  [0.83, 0.68]

#=  School 2 ends up losing a great student (4), because persuading student 2
	away from school 1 meant that school 1 lowered its cutoff and 4 got in. =#

# Shows that the stable match is unique.
println("Normal-rev:   ",DA(students, schools', capacities, rev=true)[1])
println("Delta-rev:    ",DA(students_delta, schools', capacities, rev=true)[1])
# Normal-rev:   [2, 1, 1, 2, 2]
# Delta-rev:    [2, 2, 1, 1, 2]