@testset "Single Life" begin
    @testset "issue age 116" begin
        t = MortalityTables.table("2001 VBT Residual Standard Select and Ultimate - Male Nonsmoker, ANB")
        i = Yields.Constant(0.05)
        lc = LifeContingency(SingleLife(mort = t.ultimate, issue_age = 116), i)

        @test l(lc, 0) ≈ 1.0
        @test l(lc, 1) ≈ 0.200120000000000
        @test l(lc, 2) ≈ 0.030764447600000

        @test D(lc, 0) ≈ 1.0
        @test D(lc, 1) ≈ 0.190590476190476
        @test D(lc, 2) ≈ 0.027904260861678

        @test N(lc, 0) ≈ 1.221415195080500
        @test N(lc, 1) ≈ 0.221415195080502
        @test N(lc, 2) ≈ 0.030824718890026

        @test C(lc, 0) ≈ 0.761790476190476000
        @test C(lc, 1) ≈ 0.153610478367347000
        @test C(lc, 2) ≈ 0.023794627623916200

        @test M(lc, 0) ≈ 0.9418373716628330
        @test M(lc, 1) ≈ 0.1800468954723570
        @test M(lc, 2) ≈ 0.0264364171050101

        @test present_value(Insurance(lc)) ≈ 0.9418373716628330
        @test present_value(AnnuityDue(lc)) ≈ 1.2214151950805000

        qs = t.ultimate[116:118]
        @test present_value(Insurance(lc, n=3)) ≈ sum(qs .* [1;cumprod(1 .- qs[1:2])] .* [1.05 ^ -t for t in 1:3])
        @test present_value(AnnuityDue(lc, n=3)) ≈ sum([1;cumprod(1 .- qs[1:2])] .* [1.05 ^ -t for t in 0:2])

        @test LifeContingencies.V(lc,1) == reserve_premium_net(lc,1)
        @test LifeContingencies.v(lc,1) == Yields.discount(lc,1)
        @test LifeContingencies.A(lc) == present_value(Insurance(lc))
        @test LifeContingencies.ä(lc) == present_value(AnnuityDue(lc))
        @test LifeContingencies.a(lc) == present_value(AnnuityImmediate(lc))
        @test LifeContingencies.P(lc) == premium_net(lc)
        @test LifeContingencies.ω(lc) == omega(lc)
    end

    @testset "issue age 30" begin
        t = MortalityTables.table("2001 VBT Residual Standard Select and Ultimate - Male Nonsmoker, ANB")
        i = Yields.Constant(0.05)
        life = SingleLife(mort = t.select[30], issue_age = 30)
        ins = LifeContingency(life, i)

        @test life.issue_age == SingleLife(mort = t.select[30]).issue_age
        @test life.issue_age == SingleLife(t.select[30]).issue_age
        @test life.issue_age == SingleLife(t.select,issue_age= 30).issue_age

        @test l(ins, 0) ≈ 1.0
        @test l(ins, 1) ≈ 0.999670000000000
        @test l(ins, 2) ≈ 0.999210151800000

        @test D(ins, 0) ≈ 1.0
        @test D(ins, 1) ≈ 0.952066666666667
        @test D(ins, 2) ≈ 0.906313062857143

        @test N(ins, 0) ≈ 18.673525637928100
        @test N(ins, 1) ≈ 17.673525637928100
        @test N(ins, 2) ≈ 16.721458971261400

        @test C(ins, 0) ≈ 0.000314285714285764
        @test C(ins, 1) ≈ 0.000417095873015874
        @test C(ins, 2) ≈ 0.000474735413877510

        @test M(ins, 0) ≈ 0.1107844934319970
        @test M(ins, 1) ≈ 0.1104702077177110
        @test M(ins, 2) ≈ 0.1100531118446950

        @test present_value(Insurance(ins)) ≈ 0.1107844934319970
        @test present_value(AnnuityDue(ins)) ≈ 18.6735256379281000
        @test premium_net(ins) ≈ 0.0059327036350854
        @test reserve_premium_net(ins, 1) ≈ 0.0059012862412992
        @test reserve_premium_net(ins, 2) ≈ 0.0119711961204193

        qs = t.select[30][30:55]
        @test present_value(Insurance(ins, n=26)) ≈ sum(qs .* [1;cumprod(1 .- qs[1:25])] .* [1.05 ^ -t for t in 1:26])
        @test present_value(AnnuityDue(ins, n=26)) ≈ sum([1;cumprod(1 .- qs[1:25])] .* [1.05 ^ -t for t in 0:25])

        @test premium_net(ins, 26) ≈ LifeContingencies.A(ins, 26) /  LifeContingencies.ä(ins, 26)

    end
end


@testset "Basic Utilities Tests" begin
  @testset "Utils" begin
    ϵ = log(3); δ = 0.05

    @test gaussianMechConstant(ϵ, δ) ≈ 1.7563398731147597 atol=1e-10
    tmp = gaussianMechConstant2(ϵ, δ)
    @test tmp[1] ≈ 1.7563398731147597 atol=1e-10
    @test tmp[2] ≈ 2.3095249420840473 atol=1e-10
  end

  @testset "Basic mechanisms" begin
    ϵ = log(2); δ = 0.04

    Random.seed!(12345)  # reseed the RNG
    D = Random.rand(100)
    @test gaussianMech(D, Distributions.mean, 2.0/100, ϵ, δ)[1] == 0.4764940748202759
    @test laplaceMech(D, Distributions.mean, 2.0/100, ϵ)[1] == 0.5366973705020675
    @test mean_dp(D, ϵ, δ) == 0.47514367413387
    @test mean_dp(D, ϵ, 0.0) == 0.4920613726730682
    @test mean_dp(D, ϵ) == 0.4886977015049402

    Random.seed!(12345)  # reseed the RNG
    f(x) = [x[1]^2*x[2], sqrt(x[1]), x[2]^3]
    x0 = [3, 4]
    l1sens = 0.5
    r,a = truncatedLaplaceMech(x0, f, l1sens, ϵ, δ)
    @test r[1] == 36.090422106340114
    @test r[2] == 2.50473361660456
    @test r[3] == 63.800895516558114
    @test a == 2.020992083320228
  end
end


@doc raw"""
    Euclidean{T<:Tuple,𝔽} <: Manifold

Euclidean vector space.

# Constructor

    Euclidean(n)

Generate the $n$-dimensional vector space $ℝ^n$.

    Euclidean(n₁,n₂,...,nᵢ; field=ℝ)
    𝔽^(n₁,n₂,...,nᵢ) = Euclidean(n₁,n₂,...,nᵢ; field=𝔽)

Generate the vector space of $k = n_1 \cdot n_2 \cdot … \cdot n_i$ values, i.e. the
manifold $𝔽^{n_1, n_2, …, n_i}$ whose
elements are interpreted as $n_1 × n_2 × … × n_i$ arrays.
For $i=2$ we obtain a matrix space.
The default `field=ℝ` can also be set to `field=ℂ`.
The dimension of this space is $k \dim_ℝ 𝔽$, where $\dim_ℝ 𝔽$ is the
[`real_dimension`](@ref) of the field $𝔽$.
"""
struct Euclidean{N,𝔽} <: Manifold where {N<:Tuple, 𝔽<:AbstractNumbers} end

function Euclidean(n::Vararg{Int,I}; field::AbstractNumbers = ℝ) where {I}
    return Euclidean{Tuple{n...},field}()
end

^(𝔽::AbstractNumbers, n) = Euclidean(n...; field = 𝔽)

"""
    EuclideanMetric <: RiemannianMetric

A general type for any manifold that employs the Euclidean Metric, for example
the [`Euclidean`](@ref) manifold itself, or the [`Sphere`](@ref), where every
tangent space (as a plane in the embedding) uses this metric (in the embedding).

Since the metric is independent of the field type, this metric is also used for
the Hermitian metrics, i.e. metrics that are analogous to the `EuclideanMetric`
but where the field type of the manifold is `ℂ`.

This metric is the default metric for example for the [`Euclidean`](@ref) manifold.
"""
struct EuclideanMetric <: RiemannianMetric end

^(M::Euclidean, n::Int) = ^(M, (n,))
function ^(::Euclidean{T,𝔽}, n::NTuple{N,Int}) where {T,𝔽,N}
    return Euclidean{Tuple{T.parameters...,n...},𝔽}()
end

function allocation_promotion_function(
    M::Euclidean{<:Tuple,ℂ},
    ::Union{typeof(get_vector),typeof(get_coordinates)},
    args::Tuple,
)
    return complex
end

function check_manifold_point(M::Euclidean{N,𝔽}, p) where {N,𝔽}
    if (𝔽 === ℝ) && !(eltype(p) <: Real)
        return DomainError(
            eltype(p),
            "The matrix $(p) is not a real-valued matrix, so it does not lie on $(M).",
        )
    end
    if (𝔽 === ℂ) && !(eltype(p) <: Real) && !(eltype(p) <: Complex)
        return DomainError(
            eltype(p),
            "The matrix $(p) is neither a real- nor complex-valued matrix, so it does not lie on $(M).",
        )
    end
    if size(p) != representation_size(M)
        return DomainError(
            size(p),
            "The matrix $(p) does not lie on $(M), since its dimensions ($(size(p))) are wrong (expected: $(representation_size(M))).",
        )
    end
end

function check_tangent_vector(M::Euclidean{N,𝔽}, p, X; check_base_point = true, kwargs...) where {N,𝔽}
    if check_base_point
        mpe = check_manifold_point(M, p; kwargs...)
        mpe === nothing || return mpe
    end
    if (𝔽 === ℝ) && !(eltype(X) <: Real)
        return DomainError(
            eltype(X),
            "The matrix $(X) is not a real-valued matrix, so it can not be a tangent vector to $(p) on $(M).",
        )
    end
    if (𝔽 === ℂ) && !(eltype(X) <: Real) && !(eltype(X) <: Complex)
        return DomainError(
            eltype(X),
            "The matrix $(X) is neither a real- nor complex-valued matrix, so it can not be a tangent vector to $(p) on $(M).",
        )
    end
    if size(X) != representation_size(M)
        return DomainError(
            size(X),
            "The matrix $(X) does not lie in the tangent space of $(p) on $(M), since its dimensions $(size(X)) are wrong  (expected: $(representation_size(M))).",
        )
    end
end

det_local_metric(M::MetricManifold{<:Manifold,EuclideanMetric}, p) = one(eltype(p))

"""
    distance(M::Euclidean, p, q)

Compute the Euclidean distance between two points on the [`Euclidean`](@ref)
manifold `M`, i.e. for vectors it's just the norm of the difference, for matrices
and higher order arrays, the matrix and ternsor Frobenius norm, respectively.
"""
distance(::Euclidean, p, q) = norm(p .- q)

@doc raw"""
    exp(M::Euclidean, p, X)

Compute the exponential map on the [`Euclidean`](@ref) manifold `M` from `p` in direction
`X`, which in this case is just
````math
\exp_p X = p + X.
````
"""
exp(::Euclidean, ::Any...)

exp!(M::Euclidean, q, p, X) = (q .= p .+ X)

"""
    flat(M::Euclidean, p, X)

Transform a tangent vector `X` into a cotangent. Since they can directly be identified in the
[`Euclidean`](@ref) case, this yields just the identity for a tangent vector `w` in the
tangent space of `p` on `M`.
"""
flat(::Euclidean, ::Any...)

flat!(M::Euclidean, ξ::CoTFVector, p, X::TFVector) = copyto!(ξ, X)

function get_basis(M::Euclidean{<:Tuple,ℝ}, p, B::DefaultOrthonormalBasis)
    vecs = [_euclidean_basis_vector(p, i) for i in eachindex(p)]
    return CachedBasis(B,vecs)
end
function get_basis(M::Euclidean{<:Tuple,ℂ}, p, B::DefaultOrthonormalBasis)
    vecs = [_euclidean_basis_vector(p, i) for i in eachindex(p)]
    return CachedBasis(B,[vecs; im * vecs])
end
function get_basis(M::Euclidean, p, B::DiagonalizingOrthonormalBasis)
    vecs = get_vectors(M, p, get_basis(M, p, DefaultOrthonormalBasis()))
    eigenvalues = zeros(real(eltype(p)), manifold_dimension(M))
    return CachedBasis(B, DiagonalizingBasisData(B.frame_direction, eigenvalues, vecs))
end

function get_coordinates!(M::Euclidean{<:Tuple,ℝ}, Y, p, X, B::DefaultOrDiagonalizingBasis)
    S = representation_size(M)
    PS = prod(S)
    copyto!(Y, reshape(X, PS))
    return Y
end
function get_coordinates!(M::Euclidean{<:Tuple,ℂ}, Y, p, X, B::DefaultOrDiagonalizingBasis)
    S = representation_size(M)
    PS = prod(S)
    Y .= [reshape(real(X), PS)..., reshape(imag(X), PS)...]
    return Y
end

function get_vector!(M::Euclidean{<:Tuple,ℝ}, Y, p, X, B::DefaultOrDiagonalizingBasis)
    S = representation_size(M)
    Y .= reshape(X, S)
    return Y
end
function get_vector!(M::Euclidean{<:Tuple,ℂ}, Y, p, X, B::DefaultOrDiagonalizingBasis)
    S = representation_size(M)
    N = div(length(X), 2)
    Y .= reshape(X[1:N] + im * X[N+1:end], S)
    return Y
end

@doc raw"""
    injectivity_radius(M::Euclidean)

Return the injectivity radius on the [`Euclidean`](@ref) `M`, which is $∞$.
"""
injectivity_radius(::Euclidean) = Inf

@doc raw"""
    inner(M::Euclidean, p, X, Y)

Compute the inner product on the [`Euclidean`](@ref) `M`, which is just
the inner product on the real-valued or complex valued vector space
of arrays (or tensors) of size $n_1 × n_2  ×  …  × n_i$, i.e.

````math
g_p(X,Y) = \sum_{k ∈ I} \overline{X}_{k} Y_{k},
````
where $I$ is the set of vectors $k ∈ ℕ^i$, such that for all
$1 ≤ j ≤ i$ it holds $1 ≤ k_j ≤ n_j$.

For the special case of $i ≤ 2$, i.e. matrices and vectors, this simplifies to
````math
g_p(X,Y) = X^{\mathrm{H}}Y,
````
where $\cdot^{\mathrm{H}}$ denotes the Hermitian, i.e. complex conjugate transposed.
"""
inner(::Euclidean, ::Any...)
@inline inner(::Euclidean, p, X, Y) = dot(X, Y)
@inline inner(::MetricManifold{<:Manifold,EuclideanMetric}, p, X, Y) = dot(X, Y)

inverse_local_metric(M::MetricManifold{<:Manifold,EuclideanMetric}, p) = local_metric(M, p)

default_metric_dispatch(::Euclidean, ::EuclideanMetric) = Val(true)

function local_metric(::MetricManifold{<:Manifold,EuclideanMetric}, p)
    return Diagonal(ones(SVector{size(p, 1),eltype(p)}))
end

@doc raw"""
    log(M::Euclidean, p, q)

Compute the logarithmic map on the [`Euclidean`](@ref) `M` from `p` to `q`,
which in this case is just
````math
\log_p q = q-p.
````
"""
log(::Euclidean, ::Any...)

log!(M::Euclidean, X, p, q) = (X .= q .- p)

log_local_metric_density(M::MetricManifold{<:Manifold,EuclideanMetric}, p) = zero(eltype(p))

@generated _product_of_dimensions(::Euclidean{N}) where {N} = prod(N.parameters)

"""
    manifold_dimension(M::Euclidean)

Return the manifold dimension of the [`Euclidean`](@ref) `M`, i.e.
the product of all array dimensions and the [`real_dimension`](@ref) of the
underlying number system.
"""
function manifold_dimension(M::Euclidean{N,𝔽}) where {N,𝔽}
    return _product_of_dimensions(M) * real_dimension(𝔽)
end

mean(::Euclidean{Tuple{1}}, x::AbstractVector{<:Number}; kwargs...) = mean(x)
function mean(
    ::Euclidean{Tuple{1}},
    x::AbstractVector{<:Number},
    w::AbstractWeights;
    kwargs...,
)
    return mean(x, w)
end
mean(::Euclidean, x::AbstractVector; kwargs...) = mean(x)

function mean!(M::Euclidean, p, x::AbstractVector, w::AbstractVector; kwargs...)
    return mean!(M, p, x, w, GeodesicInterpolation(); kwargs...)
end

function mean_and_var(::Euclidean{Tuple{1}}, x::AbstractVector{<:Number}; kwargs...)
    m, v = mean_and_var(x; kwargs...)
    return m, sum(v)
end
function mean_and_var(
    ::Euclidean{Tuple{1}},
    x::AbstractVector{<:Number},
    w::AbstractWeights;
    corrected = false,
    kwargs...,
)
    m, v = mean_and_var(x, w; corrected = corrected, kwargs...)
    return m, sum(v)
end
function mean_and_var(M::Euclidean, x::AbstractVector, w::AbstractWeights; kwargs...)
    return mean_and_var(M, x, w, GeodesicInterpolation(); kwargs...)
end

median(::Euclidean{Tuple{1}}, x::AbstractVector{<:Number}; kwargs...) = median(x)
function median(
    ::Euclidean{Tuple{1}},
    x::AbstractVector{<:Number},
    w::AbstractWeights;
    kwargs...,
)
    return median(x, w)
end

function median!(::Euclidean{Tuple{1}}, p, x::AbstractVector; kwargs...)
    return copyto!(p, [median(vcat(x...))])
end
function median!(::Euclidean{Tuple{1}}, p, x::AbstractVector, w::AbstractWeights; kwargs...)
    return copyto!(p, [median(vcat(x...), w)])
end

@doc raw"""
    norm(M::Euclidean, p, X)

Compute the norm of a tangent vector `X` at `p` on the [`Euclidean`](@ref)
`M`, i.e. since every tangent space can be identified with `M` itself
in this case, just the (Frobenius) norm of `X`.
"""
norm(::Euclidean, p, X) = norm(X)
norm(::MetricManifold{<:Manifold,EuclideanMetric}, p, X) = norm(X)

"""
    normal_tvector_distribution(M::Euclidean, p, σ)

Normal distribution in ambient space with standard deviation `σ`
projected to tangent space at `p`.
"""
function normal_tvector_distribution(M::Euclidean{Tuple{N}}, p, σ) where {N}
    d = Distributions.MvNormal(zero(p), σ)
    return ProjectedFVectorDistribution(TangentBundleFibers(M), p, d, project_vector!, p)
end

@doc raw"""
    project(M::Euclidean, p)

Project an arbitrary point `p` onto the [`Euclidean`](@ref) manifold `M`, which
is of course just the identity map.
"""
project(::Euclidean, ::Any)

project!(M::Euclidean, q, p) = copyto!(q, p)

"""
    project(M::Euclidean, p, X)

Project an arbitrary vector `X` into the tangent space of a point `p` on the
[`Euclidean`](@ref) `M`, which is just the identity, since any tangent
space of `M` can be identified with all of `M`.
"""
project(::Euclidean, ::Any, ::Any)

project!(M::Euclidean, Y, p, X) = copyto!(Y, X)

"""
    projected_distribution(M::Euclidean, d, [p])

Wrap the standard distribution `d` into a manifold-valued distribution. Generated
points will be of similar type to `p`. By default, the type is not changed.
"""
function projected_distribution(M::Euclidean, d, p)
    return ProjectedPointDistribution(M, d, project!, p)
end
function projected_distribution(M::Euclidean, d)
    return ProjectedPointDistribution(M, d, project!, rand(d))
end

"""
    representation_size(M::Euclidean)

Return the array dimensions required to represent an element on the
[`Euclidean`](@ref) `M`, i.e. the vector of all array dimensions.
"""
@generated representation_size(::Euclidean{N}) where {N} = size_to_tuple(N)

"""
    sharp(M::Euclidean, p, ξ)

Transform the cotangent vector `ξ` at `p` on the [`Euclidean`](@ref) `M` to a tangent vector `X`.
Since cotangent and tangent vectors can directly be identified in the [`Euclidean`](@ref)
case, this yields just the identity.
"""
sharp(::Euclidean, ::Any...)

sharp!(M::Euclidean, X::TFVector, p, ξ::CoTFVector) = copyto!(X, ξ)

function show(io::IO, ::Euclidean{N,𝔽}) where {N,𝔽}
    print(io, "Euclidean($(join(N.parameters, ", ")); field = $(𝔽))")
end

"""
    vector_transport_to(M::Euclidean, p, X, q, ::ParallelTransport)

Parallely transport the vector `X` from the tangent space at `p` to the tangent space at `q`
on the [`Euclidean`](@ref) `M`, which simplifies to the identity.
"""
vector_transport_to(::Euclidean, ::Any, ::Any, ::Any, ::ParallelTransport)

vector_transport_to!(M::Euclidean, Y, p, X, q, ::ParallelTransport) = copyto!(Y, X)

var(::Euclidean, x::AbstractVector; kwargs...) = sum(var(x; kwargs...))
function var(::Euclidean, x::AbstractVector{T}, m::T; kwargs...) where {T}
    return sum(var(x; mean = m, kwargs...))
end

"""
    zero_tangent_vector(M::Euclidean, x)

Return the zero vector in the tangent space of `x` on the [`Euclidean`](@ref)
`M`, which here is just a zero filled array the same size as `x`.
"""
zero_tangent_vector(::Euclidean, ::Any...)

zero_tangent_vector!(M::Euclidean, v, x) = fill!(v, 0)


module problem133

description = """
Repunit nonfactors

A number consisting entirely of ones is called a repunit. We shall define R(k) to be a repunit of length k; for example, R(6) = 111111.
Let us consider repunits of the form R(10^n).
Although R(10), R(100), or R(1000) are not divisible by 17, R(10000) is divisible by 17.
Yet there is no value of n for which R(10^n) will divide by 19.
In fact, it is remarkable that 11, 17, 41, and 73 are the only four primes below one-hundred that can be a factor of R(10n).

Find the sum of all the primes below one-hundred thousand that will never be a factor of R(10^n).
"""

# Search the list of primes, for each one either:
# 1) 10^10^n = 1 mod 9p for some n, which means p divides R(10^n)
# 2) A cycle forms in the remainders of 10^10^n mod 9p
# So search the list of primes until a 1 occurs in the remainder or a repeated value is hit.
# IntSet does a pretty good job providing a constant time test for checking if a value has already been seen.
function search(maxprime = 100000)
  total = 0
  for p in primes(maxprime)
    seen = IntSet()
    remainder = (10^10) % (9*p)
    push!(seen, remainder)
    for n = 2:Inf
      remainder = powermod(remainder, 10, 9*p)
      if remainder == 1
        break
      elseif in(remainder, seen)
        total += p
        break
      end
      push!(seen, remainder)
    end
  end
  return total
end

using Base.Test

@test search(100) == sum(primes(100)) - sum([11, 17, 41, 73])

end


module MarketEquilibria

using Random
using JuMP
using Ipopt

export homogeneousfisher, tatonnement, exchange, production

"""
    homogeneousfisher(endowments, A, supplies, form, ρ)

§6.2: Fisher Model with Homogeneous Consumers. `form` specifies which utility
function consumers use, among `:linear`, `:CES`, `:cobb_douglas`, and `:leontief`. `ρ` is the parameter for
`:CES`.
"""
function homogeneousfisher(endowments::Array{Float64,1},
                           A::Array{Float64,2},
                           supplies::Array{Float64,1},
                           form=:linear,
                           ρ::Union{Float64,Nothing}=0.5)::Tuple{Array{Float64,2},
                                                                 Array{Float64,1}}
    @assert form in [:linear, :CES, :cobb_douglas, :leontief] "Unknown form"
    if form == :cobb_douglas
        ρ = 1e-8          # Ideally 0
    end

    (n, m) = size(A)
    @assert size(endowments) == (n, ) "Dim mismatch between A and endowments"
    @assert size(supplies) == (m, ) "Dim mismatch between A and supplies"

    model = Model(Ipopt.Optimizer)
    set_optimizer_attribute(model, "print_level", 0)

    @variable(model, X[1:n, 1:m], lower_bound=0)
    @constraint(model, Supply[j in 1:m], sum(X[:, j]) ≤ supplies[j])

    if form == :linear || ρ > .97
        if form != :linear
            @warn "ρ > .97 tends to fail, so assuming linear"
        end
        # Very bad results if you allow the default scaling 🐸☕
        set_optimizer_attribute(model, "nlp_scaling_method", "none")
        @NLobjective(model, Max, sum(endowments[i] * log(sum(A[i, j] * X[i, j] for j in 1:m)) for i in 1:n))
    elseif form == :leontief
        @variable(model, t)
        @objective(model, Max, t)
        @NLconstraint(model, tLessThan[j in 1:m],
                    t ≤ sum(endowments[i] * log(A[i, j] * X[i, j]) for i in 1:n))
    else
        @NLobjective(model, Max, sum(endowments[i] * log(sum(A[i, j] * X[i, j] ^ ρ for j in 1:m) ^ (1 / ρ)) for i in 1:n))
    end

        # Explicit Cobb–Douglas; very poorly conditioned. I couldn't get it to converge.
        # @NLobjective(model, Max, sum(endowments[i] * log(prod(X[i, j] ^ A[i, j] for j in 1:m)) for i in 1:n))

        # Explicit Leontief is as follows, but minimum() not allowed by JuMP.
        # Got around this using the linearization trick.
        # @NLobjective(model, Max, sum(endowments[i] * log(minimum(A[i, j] * X[i, j] for j in 1:m)) for i in 1:n))

        # Attempt at "softmin" differential substitute; returns errors.
        # @NLobjective(model, Max, sum(endowments[i] * log(-log(sum(exp(-A[i, j] * X[i, j]) for j in 1:m))) for i in 1:n))

    optimize!(model)
    return value.(X), dual.(Supply)
end

"""
    tatonnement(endowments, A, ρ, π0;
                η=1e-4, δ=0.5, α=0.5, verbose=false, maxit=1000, tol=1e-5)

§6.3.1.B: The Discrete Tâtonnement Process. Simulate a bidding process in search
of equilibrium prices for CES exchange.

Price update step is quite different from the book: I *multiply* prices
entrywise by the excess demand: `π .*= exp.(δ * Z * k ^ -α)`. This works pretty well.

`η` is a lower bound on the prices before normalization.
"""
function tatonnement(endowments ::AbstractArray{<:AbstractFloat,2},
                     A          ::AbstractArray{<:AbstractFloat,2},
                     ρ          ::AbstractArray{<:AbstractFloat,1},
                     π0         ::Union{AbstractArray{<:AbstractFloat,1}, Nothing};
                     η          ::AbstractFloat=1e-4,
                     δ          ::AbstractFloat=0.5,
                     α          ::AbstractFloat=0.4,
                     verbose    ::Bool=false,
                     maxit      ::Int=2000,
                     tol        ::AbstractFloat=1e-5)::Tuple{Array{Float64,1},
                                                             Array{Float64,2}}

    (n, m) = size(A)
    @assert (n, m) == size(endowments)
    W = endowments      # sorry

    @assert (n, )  == size(ρ)
    @assert all(ρ .< 1)
    @assert all(0 .< (tol, η, δ, α) .< 1)

    if π0 == nothing
        π0 = rand(m)
    end

    @assert (m, )  == size(π0)

    total_supply = sum(W, dims=1)
    # Express utility as fractions of endowment
    A .*= total_supply
    # So we can normalize W
    W ./= total_supply

    σ = 1 ./ (1 .- ρ)
    π_LB = η / (4 * n)
    π_UB = 1 + π_LB
    π_LB_tight = η / (2 * n)

    π = copy(π0)
    D = zeros(Float64, n, m)
    Z = zeros(Float64, m)

    for k in 1:maxit
        # Oracle that computes player demand under CES (eq. 6.7, p. 149)
        for i in 1:n, j in 1:m
            D[i, j] = A[i, j] ^ σ[i] * sum(π[k] * W[i, k] for k in 1:m) /
                       (π[j] ^ σ[i] * sum((A[i, k] * π[k] ^ -ρ[i]) ^ σ[i] for k in 1:m))
        end

        for j in 1:m
            # Total supply, after normalization, is ones
            Z[j] = sum(D[i, j] for i in 1:n) - 1
        end

        if NaN in Z || all(-tol .< Z .< tol)
            break
        end

        if verbose
            println("Iteration $k")
            println("  Current prices: ", round.(π, digits=5))
            println("  Excess demand:  ", round.(Z, digits=5))
        end

        # This equilibrium step process is completely different from the book.
        # Just something that makes intuitive sense. Book is unclear on how to
        # calculate step size.
        if all(π_LB .≤ π .≤ π_UB)
            verbose ? println("  Step toward equilibrium") : nothing
            π .*= exp.(δ * Z * k ^ -α)
        else                             # Bottom of p. 144
            verbose ? println("  Regularization step") : nothing
            for j in 1:m
                if π[j] > 1
                    π[j] = 1
                elseif π[j] < π_LB_tight
                    π[j] = π_LB_tight
                end
            end
        end
    end

    return π ./ sum(π), D
end

"""
    exchange(endowments, A, ρ=:linear)

§6.4: Convex programs for linear and CES exchange economies. Note that unlike `homogeneousfisher()`,
here `ρ` is a vector of CES parameters used by each player. Or pass `ρ=:linear` for linear model.
"""
function exchange(endowments::Array{Float64,2},
                  A::Array{Float64,2},
                  ρ::Union{Array{Float64,1},Symbol}=:linear)::Tuple{Array{Float64,1},
                                                                    Array{Float64,2}}
    (n, m) = size(A)
    @assert size(endowments) == (n, m)         "Dim mismatch between A and endowments"

    model = Model(Ipopt.Optimizer)
    set_optimizer_attribute(model, "print_level", 0)

    if ρ == :linear
        # set_optimizer_attribute(model, "nlp_scaling_method", "none")

        @variable(model, ψ[1:m])
        @variable(model, X[1:n, 1:m], lower_bound=0)
        @NLconstraint(model, IndividualRationality[i in 1:n, j in 1:m],
                        sum(A[i, k] * X[i, k] for k in 1:m) ≥
                        A[i, j] * sum(endowments[i, k] * exp(ψ[k] - ψ[j]) for k in 1:m))
        @constraint(model, Supply[j in 1:m], sum(X[i, j] for i in 1:n) ==
                                             sum(endowments[i, j] for i in 1:n))

        optimize!(model)

        prices = exp.(value.(ψ))
        demands = value.(X)

    else
        @assert size(ρ) == (n, )                   "Dim mismatch between A and ρ"
        @assert all(-1 .≤ ρ .< 0)                  "Need ρ ∈ [-1, 0)"

        @variable(model, σ[1:m], lower_bound=0)
        @NLconstraint(model, Supply[j in 1:m], sum(A[i, j] ^ (1 / (1 - ρ[i])) *
                                               sum(σ[k] ^ 2 * endowments[i, k] for k in 1:m) /
                                               (
                                                   σ[j] ^ ((ρ[i] - 2) / (1 - ρ[i])) *
                                                   sum(A[i, k] ^ (1 / (1 - ρ[i])) * σ[k] ^ ((-2 * ρ[i]) / (1 - ρ[i]))
                                                       for k in 1:m)
                                               )
                                               for i in 1:n) ≤ σ[j] * sum(endowments[i, j] for i in 1:n))

        optimize!(model)

        prices = value.(σ) .^ 2

        demands = zeros(n, m)
        for i in 1:n, j in 1:m
        demands[i, j] = A[i, j] ^ (1 / (1 - ρ[i])) * sum(prices[k] * endowments[i, k] for k in 1:m) /
                        (
                            prices[j] ^ (1 / 1 - ρ[i]) *
                            sum(A[i, k] ^ (1 / (1 - ρ[i])) * prices[k] ^ (-ρ[i] / (1 - ρ[i])) for k in 1:m)
                        )
        end
    end

    return prices, demands
end

"""
    exchange(endowments, A_consumers, A_producers,
             ρ_consumers, ρ_producers, o_producers)

§6.6: Models with production. `ρ_consumers` describes
each consumer's CES utility function, while `ρ_producers` describes the CES
production function. Convergence is very iffy, but best when `A_consumers` is
tall and entries of `ρ` vectors are not too close to 0 or 1.
"""
function production(endowments::Array{Float64,2},
                    A_consumers::Array{Float64,2},
                    A_producers::Array{Float64,2},
                    ρ_consumers::Array{Float64,1},
                    ρ_producers::Array{Float64,1},
                    o_producers::Array{Int,1})::Tuple{Array{Float64,1},
                                                      Array{Float64,1},
                                                      Array{Float64,2}}

    (n, m) = size(A_consumers)
    (l, ) = size(ρ_producers)

    @assert size(A_producers) == (l, m)   "Dim mismatch between A_consumers and A_producers"
    @assert size(endowments) == (n, m)    "Dim mismatch between A_consumers and endowments"
    @assert size(ρ_consumers) == (n, )    "Dim mismatch between A_consumers and ρ_consumers"
    @assert size(o_producers) == (l, )    "Dim mismatch between o and ρ_producers"
    @assert all(0 .< ρ_consumers .< 1)    "Need ρ ∈ (0, 1)"
    @assert all(0 .< ρ_producers .< 1)    "Need ρ ∈ (0, 1)"

    A = vcat(A_producers, A_consumers)
    w = hcat(endowments, zeros(n, l))
    ρ = vcat(ρ_consumers, ρ_producers)
    σ = 1 ./ (1 .- ρ)
    o = vcat(o_producers, m .+ (1:n))

    # Comparison idx since comparison forbidden in lincon
    O = [o[k] == j for k in 1:l, j in 1:m]

    model = Model(Ipopt.Optimizer)
    set_optimizer_attribute(model, "print_level", 0)
    set_optimizer_attribute(model, "nlp_scaling_method", "none")

    @variable(model, ψ[1:m + n], lower_bound=0)
    @variable(model, x[1:m + n], lower_bound=0)
    @variable(model, Z[1:l + n, 1:m], lower_bound=0)
    @variable(model, q[1:l + n], lower_bound=0)

    @NLconstraint(model, ConsumerRationality[i in 1:n],                   # 6.10
                    exp(ψ[m + i]) * x[m + i] ≥ # or ==
                    sum(exp(ψ[j]) * w[i, j] for j in 1:m))
    @NLconstraint(model, ProducerSolvency[k in 1:l + n],                  # 6.11
                    q[k] ≤
                    sum(A[k, j] * x[j] ^ ρ[k] for j in 1:m) ^ (1 / ρ[k]))
    @NLconstraint(model, ProducerRationality[k in 1:l + n],               # 6.12
                    exp(ψ[o[k]] * (1 - σ[k])) ≥
                    sum(A[k, j] ^ σ[k] * exp(ψ[j] * (1 - σ[k])) for j in 1:m))
    @constraint(model, MarketClearing[j in 1:m],                          # 6.13
                    sum(Z[k, j] for k in 1:l + n) ≤
                    sum(w[i, j] for i in 1:n) + sum(O[k, j] * q[k] for k in 1:l))
                                            # Equiv
                                            # + sum((o[k] == j) * q[k] for k in 1:l + n))
    @constraint(model, Supply[i in 1:n], x[m + i] ≤ q[l + i])             # 6.14

    #=  Final constraint given in final paragraph of 154. This should be redundant, but
        aids computation.                   =#
    @NLconstraint(model, ZeroProfit[k in 1:l + n],
                    sum(exp(ψ[j]) * Z[k, j] for j in 1:m) == exp(ψ[o[k]]) * q[k])

    optimize!(model)
    prices = exp.(value.(ψ)[1:m])

    # Player utility, specifically, demand of each player for her own utility item.
    demands = value.(x)[m + 1:end]
    inputs = value.(Z)[l + 1:end, :]

    return prices, demands, inputs
end

end


# # dempe.mod	QQR2-MN-4-3
# # Original AMPL coding by Sven Leyffer
#
# # An MPEC from S. Dempe, "A necessary and sufficient optimality
# # condition for bilevel programming problems", Optimization 25,
# # pp. 341-354, 1992. From book Math. Progr. with Equil. Constr,
# # by Luo, Pang & Ralph, CUP, 1997, p. 354.
#
# # Number of variables:   2 + 1 multipliers
# # Number of constraints: 2
# # Nonlinear complementarity constraints
#
# var x;
# var z;
# var w >= 0;
#
# minimize f: (x - 3.5)^2 + (z + 4)^2;
#
# subject to
#     con1:  z - 3 + 2*z*w = 0;
#     con2:  0 >= z^2 - x     complements  w >= 0;
#
# data;
#
# # starting point
# let x := 1;
# let z := 1;
# let w := 1;
#
#
# let x :=   0.183193 ;
# let z := 0.428106;
# let w := 3.00379;

using JuMP, Complementarity
using Ipopt

m = Model(solver=IpoptSolver())

@variable(m, x)
@variable(m, z)
@variable(m, w>=0)

@NLobjective(m, Min, (x - 3.5)^2 + (z + 4)^2)

@NLconstraint(m, z - 3 + 2*z*w == 0)

@complements(m, 0 >= z^2 - x,  w >= 0, simple)

# setvalue(x, 50)
# setvalue(z, 50)
# setvalue(w, 1e5)

solve(m)

@show getobjectivevalue(m)
@assert isapprox(getobjectivevalue(m), 28.25, atol=1e-4)


function Coherence(B::Array{Float64,2})
    u=Inf
    for k=1:size(B,2)
        for j=1:(k-1)
            len=(sum(B[:,k].^2).^.5)*(sum(B[:,j].^2).^.5)
            innerProd=abs(sum(B[:,k].*B[:,j]))
            u= minimum([u,innerProd/len])
        end
    end
    return u
end

Coherence(B) = Coherence(convert(Array{Float64,2},B))


"""
Smooth jump function ``\\Phi_{\\omega, T}(t)``

    Returns a planned trajectory.

    smoothjump(t :: Real, T:: Real; ω = 2.0 :: Real)

`t` : actual time

`T` : final time

`ω` : steepness factor of trajectory
"""
function smoothjump(t :: Real, T:: Real; ω = 2.0 :: Real)

    if t <= 0
        return  0.0
    elseif t >= T
        return 1.0
    else

        numer, err1 = quadgk(x-> pulsebell(x, T, ω=ω), 0, t)
        denom, err2 = quadgk(x-> pulsebell(x, T, ω=ω), 0, T)

        return numer/denom
    end
end

"""
Bell-shaped pulse function ``\\Omega_{\\omega, T}(t)``

`t` : actual time

`T` : final time

`ω` : steepness factor of trajectory
"""
function pulsebell(t :: Real, T:: Real; ω = 2.0 :: Real)

    if t < 0 || t > T
        return 0
    else
        exp( -1 / ( (1 - t/T)*(t/T))^ω )
    end

end


function fibo_even_sum(N::Int)
    a = 0
    b = 1
    c = 0
    s = 0
    while(a+b<N)
        c = a + b
        if iseven(c)
            s+=c
        end
        a = b
        b = c
    end
    return s
end

println(fibo_even_sum(4000000))


# Tests for the conditional mutual information function

# Test conditional mutual information formula
@test get_conditional_mutual_information(arr1, arr2, arr3) ≈ get_entropy(arr1, arr3) + get_entropy(arr2, arr3) - get_entropy(arr1, arr2, arr3) - get_entropy(arr3)
println("Conditional mutual information formula passed.")


using JuMP
using Clp

PROPORTION_MAT = [
    [5.0 15.0]
    [20.0 5.0]
    [-15.0 -2.0]
]

LIMITS = [50.0, 40.0, -60.0]

COSTS = [8.0, 4.0]

m = Model(Clp.Optimizer)
@variable(m, x[1:2] >= 0)
@objective(m, Min, COSTS' * x)
@constraint(m, PROPORTION_MAT * x .≥ LIMITS)

print(m)

optimize!(m)

status = termination_status(m)

println("Solution status: ", status)

println("Objective value: ", objective_value(m))

println("Values: ", value.(x))


push!(LOAD_PATH, "./gridmodule/")
push!(LOAD_PATH, "./femmodule/")

using BaseFEM
using BaseGrid
using LinearAlgebra

# 1D ring
#m=grid1D("./mesh1d/mesh2")
#tau = 1e-3
#eps = 0.1

# 3D ball
#m=grid3D("./mesh3d/ball")
#tau = 1e-3
#eps = 0.1

# greenland example
#m=grid2D("./mesh2d/greenland")
#tau = 1e-1
#eps = 3.00

# sphere example
m=grid2D("./mesh2d/sphere_hq05")
tau = 1e-6
eps  = 0.01

top  = m.topology
ndof = ndof_P1(top)
geo  = restrain_geometry(m.geometry,ndof)
A,M=assembly_P1(geo,top)

# initial conditions
u      = 2.0*rand(ndof).-1.0
time = 0.0
S = factorize([M tau*A;-A M])
mprint = 0

for i=1:1000
	global time,u,mprint
	println("iteration $(i) at time $(time)")

	dWdu = (4*u.^3-4*u)/eps

	rhs1 = M*u
	rhs2 = M*dWdu

	sol = S\[rhs1;rhs2]
	u   = sol[1:ndof]
	time = time + tau

	# output to ParaView vtk file
	if mod(i,10) == 0
		mprint = mprint + 1
		outvtk_scalar("output/example_cahnhilliard"*string(mprint)*".vtk",geo,top,"sol",u)
	end
end
# alternative meshes
# m=grid2D("./mesh2d/test")
# m=grid2D("./mesh2d/sphere")


using Soss, DifferentialEquations
using Plots, StatsPlots

function lotka_volterra(du,u,p,t)
  x, y = u
  α, β, γ, δ  = p
  du[1] = (α - β*y)x # dx =
  du[2] = (δ*x - γ)y # dy =
end
p = [1.5, 1.0, 3.0, 1.0]
u0 = [1.0,1.0]
prob1 = ODEProblem(lotka_volterra,u0,(0.0,10.0),p)

sol1 = solve(prob1,Tsit5(),saveat=0.1)
odedata = rand(For(Poisson, Array(sol1)))

plot(sol1, legend = false); scatter!(sol1.t, odedata')

m = @model ode begin
    α ~ Uniform()
    β ~ Uniform()
    γ ~ Uniform()
    δ ~ Uniform()

    # p = [2α+0.5,2β,3γ+1,2δ]
    p = [20α+0.5,20β,30γ+1,20δ]
    prob = remake(ode, p=p)
    predicted = solve(prob,Tsit5(),saveat=0.1)

    data ~ For(Array(predicted)) do λ Poisson(λ) end
end

post = dynamicHMC(m(ode=prob1) | (data=odedata,))

plt = plot(sol1, legend = false)

for x in post
    p = [2 * x.α + 0.5,2 * x.β,3 * x.γ + 1,2 * x.δ]

    thisprob = remake(prob1, p=p)
    thissol = solve(thisprob,Tsit5(),saveat=0.1)
    plot!(plt, thissol, alpha=0.02, linewidth=2, color = :black, legend = false)
end
# display(pl)
plot!(sol1, w=1, legend = false)
scatter!(sol1.t, odedata')


using ProgressMeter, Plots, StatsBase, Statistics, LaTeXStrings, JLD

function TRW_CD(L,P,X0)
    surf = zeros(L+2); surf[X0] = 1.0
    TimePassed = 0
    while surf[1] + surf[end] < 0.99
        Altsurf = copy(surf)
        for i ∈ 2:L+1
            Altsurf[i+1] += surf[i]*(1-P)
            Altsurf[i-1] += surf[i]*P
            Altsurf[i] -= surf[i]
        end
        surf = Altsurf
        TimePassed += 1
    end
    return TimePassed
end

L = 40
PList = hcat(0:0.01:1)
X0List = hcat(1:L+2)

Data = zeros(length(X0List),length(PList))
progress = Progress(length(X0List)*length(PList)*runnum)

for i in 1:length(X0List)
    X0 = X0List[i]
    for j in 1:length(PList)
        P = PList[j]
        Data[i,j] = mean(TRW_CD(L,P,X0))
    end
end

save("../../Data/Q6/Q6-LT.jld", "Data", Data)
Data = load("../../Data/Q6/Q6-LT.jld")["Data"]

P1 = heatmap(0:0.01:1, 1:L+2,Data, xaxis=nothing,legend = nothing,title = L"Life\ time")
P2 = contour(0:0.01:1, 1:L+2,Data, fill = true ,legend = nothing,xlabel = L"P")
P3 = plot(P1,P2,layout = (2,1),ylabel = L"X_{init}",link = :y)
CBar = heatmap((0:1650).*ones(1651,1),xlabel ="",legend=:none, xticks=:none, yticks=(1:165:1651, string.(0:165:1651)))

plot([P3]...,CBar,layout = @layout[grid(1,1) a{0.04w}], size = (600,500))
savefig("../../Figs/Q6/Q6-LT.pdf")

X0List[21]

X0Plot = []
for X0 in [11,21,31]
    PLT = begin
        plot(PList,Data[X0,:,1], c = :black,linewidth = 3)
        plot!(PList,Data[X0,:,1], linewidth = 2, c = :purple, fill = (0, 0.5))
        plot!(title = L"X_{init} = %$(X0List[X0])",
            xlabel = L"P", ylabel = L"Avg\ Life\ Time", legend = :none)
    end
    push!(X0Plot, PLT)
end

X0P = plot(X0Plot...,layout = (1,3))

PPlot = []
for P in [26,51,76]
    PLT = begin
        plot(X0List,Data[:,P], c = :black, linewidth = 3)
        plot!(X0List,Data[:,P], linewidth = 2, c = :steelblue,  fill = (0, 0.5))
        plot!(title = L"P = %$(PList[P])",
            xlabel = L"X0", ylabel = L"Avg\ Life\ Time", legend = :none)
    end
    push!(PPlot, PLT)
end
PPlot[1]
PP = plot(PPlot...,layout = (1,3))

Pmain2 = plot(PP,X0P,layout = (2,1),size = (850,690), plot_title = L"Life\ time\ of\ a\ Random\ Walker\ (%$runnum \ runs)")
savefig(Pmain2,"../../Figs/Q6/Q6-scat.pdf")


# ---
# title: 813. Largest Sum of Averages
# id: problem813
# author: Tian Jun
# date: 2020-10-31
# difficulty: Medium
# categories: Dynamic Programming
# link: <https://leetcode.com/problems/largest-sum-of-averages/description/>
# hidden: true
# ---
#
# We partition a row of numbers `A` into at most `K` adjacent (non-empty)
# groups, then our score is the sum of the average of each group. What is the
# largest score we can achieve?
#
# Note that our partition must use every number in A, and that scores are not
# necessarily integers.
#
#
#
#     **Example:**
#     Input:
#     A = [9,1,2,3,9]
#     K = 3
#     Output: 20
#     Explanation:
#     The best choice is to partition A into [9], [1, 2, 3], [9]. The answer is 9 + (1 + 2 + 3) / 3 + 9 = 20.
#     We could have also partitioned A into [9, 1], [2], [3, 9], for example.
#     That partition would lead to a score of 5 + 2 + 6 = 13, which is worse.
#
#
#
#
# **Note:**
#
#   * `1 <= A.length <= 100`.
#   * `1 <= A[i] <= 10000`.
#   * `1 <= K <= A.length`.
#   * Answers within `10^-6` of the correct answer will be accepted as correct.
#
#
## @lc code=start
using LeetCode

## add your code here:
## @lc code=end


# ---
# title: 307. Range Sum Query - Mutable
# id: problem307
# author: zhwang
# date: 2022-03-12
# difficulty: Medium
# categories: Binary Indexed Tree, Segment Tree
# link: <https://leetcode.com/problems/range-sum-query-mutable/description/>
# hidden: true
# ---
#
# Given an integer array _nums_ , find the sum of the elements between indices
# _i_ and _j_ ( _i_ ≤ _j_ ), inclusive.
#
# The _update(i, val)_ function modifies _nums_ by updating the element at index
# _i_ to _val_.
#
# **Example:**
#
#
#
#     Given nums = [1, 3, 5]
#
#     sumRange(0, 2) -> 9
#     update(1, 2)
#     sumRange(0, 2) -> 8
#
#
#
#
# **Constraints:**
#
#   * The array is only modifiable by the _update_ function.
#   * You may assume the number of calls to _update_ and _sumRange_ function is distributed evenly.
#   * `0 <= i <= j <= nums.length - 1`
#
#
## @lc code=start
using LeetCode

mutable struct SegmentTree{T<:Real}
    n::Int
    tree::Vector{T}

    function SegmentTree(nums::Vector{T}) where T
        n = length(nums)
        tree = append!(zeros(Int, n - 1), nums)
        for i in (n - 1):-1:1
            tree[i] = tree[i << 1] + tree[i << 1 | 1]
        end
        new{T}(n, tree)
    end
end

function update!(ST::SegmentTree, ind::Int, val::Int)::Nothing
    tree = ST.tree
    ind += ST.n - 1
    delta = val - tree[ind]
    while ind > 0
        tree[ind] += delta
        ind >>= 1
    end
end

function sum_range(ST::SegmentTree, left::Int, right::Int)::Int
    left += ST.n - 1
    right += ST.n - 1
    res, tree = 0, ST.tree
    while left <= right
        if isodd(left) ## right child
            res += tree[left] ## record outside value
            left += 1
        end
        left >>= 1
        if iseven(right) ## left child
            res += tree[right] ## record outside value
            right -= 1
        end
        right >>= 1
    end
    res
end
## @lc code=end

## @lc test=start
@testset "307.range-sum-query-mutable.jl" begin
    @testset "307 case 1" begin
        ST = SegmentTree([1, 3, 5])
        @test sum_range(ST, 1, 3) == 9
        update!(ST, 2, 2)
        @test sum_range(ST, 1, 3) == 8
    end
end
## @lc test=end

## @lc info=start
# link: [solution to 307]()
## @lc info=end


# Fig_1_Brusselator_eigenvalues_v3.jl
# ==================================== u.p.  23.8.21

using DrWatson
@quickactivate "NonlinearDynamicsTextbook"
include(srcdir("style.jl"))
using DynamicalSystems, PyPlot, OrdinaryDiffEq
using StaticArrays
using LinearAlgebra

# parameters
# ----------
a = 9
bs = [9.8, 10.2]
dvecs = [
    [4.5, 3.0, 2.0, 1.7],
    reverse([0.1, 1.0, 1.7, 2.0]),
]
fig, axs = subplots(1,2)

for (j, b) in enumerate(bs)

dvec = dvecs[j]
ivec = [1 2 6 4]   # order of colors

qmin = 0.0
qmax = 2.0
nq = 301
qvec = range(qmin,qmax; length=nq)
Re_lambda_vec = zeros(nq)
Im_lambda_vec = zeros(nq)
Re_qvec = zeros(nq)
Im_qvec = zeros(nq)

ax = axs[j]

ax.axhline(0; linewidth=1.0,color="k" )

for id = 1:length(dvec)

    dd = dvec[id]
    jj = kk = 0
    for iq = 1:nq
        cc = qvec[iq] * qvec[iq]
        b11 = b - 1 - cc
        b12 = a
        b21 = -b
        b22 = - a - dd*cc
        traceB = b11+b22
        detB = b11*b22 - b12*b21
        rootarg = traceB*traceB/4 - detB
        if rootarg < 0
            Re_lambda = traceB/2
            jj = jj + 1
            Im_qvec[jj] = qvec[iq]
            Im_lambda_vec[jj] = Re_lambda
        else
            Re_lambda = traceB/2 + sqrt(rootarg)
            if kk == 0
               kk = 1
               Re_qvec[kk] = Im_qvec[jj]
               Re_lambda_vec[kk] = Im_lambda_vec[jj]
            end
            kk = kk + 1
            Re_qvec[kk] = qvec[iq]
            Re_lambda_vec[kk] = Re_lambda
        end
    end
    ax.plot(Re_qvec[1:kk],Re_lambda_vec[1:kk], linewidth=3.0,linestyle="--",color=COLORS[ivec[id]] )
    ax.plot(Im_qvec[1:jj],Im_lambda_vec[1:jj], linewidth=3.0,linestyle="-",color=COLORS[ivec[id]] )
end

ax.set_xlim([qmin,qmax])
if j == 1
    ax.set_ylim([-2.1,2.6])
else
    ax.set_ylim([-2.1,0.4])
end
ax.legend([ "_" ; "_" ; L"d="*string(dvec[1]);
                  "_" ; L"d="*string(dvec[2]);
                  "_" ; L"d="*string(dvec[3]);
                  "_" ; L"d="*string(dvec[4]);
              #    "_" ; L"d="*string(dvec[5]);
                  ],
    fontsize = 26, handlelength = 1, handletextpad = 0.6,
    # loc = "upper left",
)

end

axs[1].set_xlabel(L"q")
axs[2].set_xlabel(L"q")
axs[1].set_ylabel(L"\max (\rm{Re} (\lambda))")

add_identifiers!(fig)
fig.tight_layout(pad=0.3)
wsave(plotsdir("11", "brusselator_eigenvalues"), fig)


export Xcorr
#TODO make more efficient

"""
`Xcorr([domainType=Float64::Type,] dim_in::Tuple, h::AbstractVector)`

`Xcorr(x::AbstractVector, h::AbstractVector)`

Creates a `LinearOperator` which, when multiplied with an array `x::AbstractVector`, returns the cross correlation between `x` and `h`. Uses `xcross`.

"""
struct Xcorr{T,H <:AbstractVector{T}} <: LinearOperator
	dim_in::Tuple{Int}
	h::H
end

# Constructors
function Xcorr(DomainType::Type, DomainDim::NTuple{N,Int},  h::H) where {H<:AbstractVector, N}
	eltype(h) != DomainType && error("eltype(h) is $(eltype(h)), should be $(DomainType)")
	N != 1 && error("Xcorr treats only SISO, check Filt and MIMOFilt for MIMO")
	Xcorr{DomainType,H}(DomainDim,h)
end
Xcorr(x::H, h::H) where {H} = Xcorr(eltype(x), size(x), h)

# Mappings

function mul!(y::H,A::Xcorr{T,H},b::H) where {T,H}
	y .= xcorr(b,A.h)
end

function mul!(y::H,L::AdjointOperator{Xcorr{T,H}},b::H) where {T,H}
    A = L.A
	l =floor(Int64,size(A,1)[1]/2)
	idx = l+1:l+length(y)
	y .= conv(b,A.h)[idx]
end

# Properties

domainType(L::Xcorr{T}) where {T} = T
codomainType(L::Xcorr{T}) where {T} = T

#TODO find out a way to verify this,
is_full_row_rank(L::Xcorr)    = true
is_full_column_rank(L::Xcorr) = true

size(L::Xcorr) = ( 2*max(L.dim_in[1], length(L.h))-1, ), L.dim_in

fun_name(A::Xcorr)  = "◎"


# Compares two methods for simulating a homogeneous Poisson point process,
# where the idea behind one method is faster than the other. In this code,
# the two methods are labelled A and B.
#
# Method A simulates all the Poisson ensembles randomly by first randomly
# generating all the Poisson random variables in one step. It then randomly
# positions all the points (across all ensembles) in one step. All the
# points are then are then separated accordingly into ensembles.
#
# Method B iterates through a for-loop, and for each iteration, it randomly
# generates a Poisson variable and positions the points for each ensemble.
#
# Method A uses more vectorization than Method B so it should be faster
# than Method B in general.
#
# Author: H. Paul Keeler, 2019.
# Website: hpaulkeeler.com
# Repository: github.com/hpaulkeeler/posts
# For more details, see the post:
# hpaulkeeler.com/testing-the-julia-language-with-point-process-simulations/

clearconsole();

using Random;
using Distributions; #for random simulations
using Plots; #for plotting

Random.seed!(1);

###START Parameters START###
numbSim=10^5; #number of simulations

#Point process parameters
lambda=5; #intensity (ie mean density) of Poisson point process

#Simulation window parameters
xMin=0;xMax=1;
yMin=0;yMax=1;
xDelta=xMax-xMin;yDelta=yMax-yMin; #rectangle dimensions
areaTotal=xDelta*yDelta; #area of rectangle
massTotal=areaTotal*lambda;  #total measure/mass of the point process
###END Parameters END###

###START Simulation section START###
###START Method A: Generate *all* ensembles at once START###
@time begin
    xxCellA=Array{Array{Float64}}(undef,numbSim);
    yyCellA=Array{Array{Float64}}(undef,numbSim);

    numbPointsA=rand(Poisson(massTotal),numbSim); #Poisson number of points
    numbPointsCumA=cumsum(numbPointsA);
    numbPointsTotal=numbPointsCumA[end];

    #uniform x/y coordinates of Poisson points
    xxAll=xDelta.*(rand(numbPointsTotal)).+xMin;#x coordinates of Poisson points
    yyAll=yDelta.*(rand(numbPointsTotal)).+yMin;#y coordinates of Poisson points

    #create some indexing for the array reshaping step
    indexFirst=copy(numbPointsCumA);
    indexFirst[2:end]=indexFirst[1:end-1].+1;
    indexFirst[1]=1;
    indexSecond=copy(numbPointsCumA);
    #reshape point arrays into an array of arrays.
    for ii in 1:numbSim
        xxCellA[ii]=(xxAll[indexFirst[ii]:indexSecond[ii]]);
        yyCellA[ii]=(yyAll[indexFirst[ii]:indexSecond[ii]]);
    end

    #yyCellAllA=[(yyAll[indexFirst[ii]:indexSecond[ii]]) for ii in 1:numbSim]
end;
###END Method A: Generate *all* ensembles at once END###

###START Method B: Generate each ensemble separately START###
@time begin
    xxCellB=Array{Array{Float64}}(undef,numbSim);
    yyCellB=Array{Array{Float64}}(undef,numbSim);
    numbPointsB=zeros(numbSim);
    #loop through for all ensembles
    for ss=1:numbSim
        numbPointsTemp=rand(Poisson(massTotal));#Poisson number of points
        xxCellB[ss]=xDelta.*(rand(numbPointsTemp)).+xMin;#x coordinates of Poisson points
        yyCellB[ss]=yDelta.*(rand(numbPointsTemp)).+yMin;#y coordinates of Poisson points
        numbPointsB[ss]=numbPointsTemp;
    end
end;
###END Method B: Generate each ensemble separately END###
###END Simulation section   END###


# indicator of a halfspace

export IndHalfspace

"""
**Indicator of a halfspace**

    IndHalfspace(a, b)

For an array `a` and a scalar `b`, returns the indicator of set
```math
S = \\{x : \\langle a,x \\rangle \\leq b \\}.
```
"""
struct IndHalfspace{R <: Real, T <: AbstractArray{R}} <: ProximableFunction
  a::T
  b::R
  norm_a::R
  function IndHalfspace{R, T}(a::T, b::R) where {R <: Real, T <: AbstractArray{R}}
    norm_a = norm(a)
    if norm_a == 0 && b < 0
        error("function is improper")
    end
    new(a, b, norm_a)
  end
end

IndHalfspace(a::T, b::R) where {R <: Real, T <: AbstractArray{R}} = IndHalfspace{R, T}(a, b)

is_convex(f::IndHalfspace) = true
is_set(f::IndHalfspace) = true
is_cone(f::IndHalfspace) = f.b == 0 || f.b == Inf

function (f::IndHalfspace{R})(x::AbstractArray{R}) where R
  if dot(f.a, x) - f.b <= eps(R)*f.norm_a*(1 + abs(f.b))
    return zero(R)
  end
  return R(Inf)
end

function prox!(y::AbstractArray{R}, f::IndHalfspace{R}, x::AbstractArray{R}, gamma::R=one(R)) where R
  s = dot(f.a, x)
  if s > f.b
    y .= x .- ((s - f.b)/f.norm_a^2) .* f.a
  else
    copyto!(y, x)
  end
  return zero(R)
end

fun_name(f::IndHalfspace) = "indicator of a halfspace"
fun_dom(f::IndHalfspace) = "AbstractArray{Real}"
fun_expr(f::IndHalfspace) = "x ↦ 0 if <a,x> ⩽ b, +∞ otherwise"
fun_params(f::IndHalfspace) =
  string( "a = ", typeof(f.a), " of size ", size(f.a), ", ",
          "b = $(f.b)")

function prox_naive(f::IndHalfspace{R}, x::AbstractArray{R}, gamma::R=one(R)) where R
  s = dot(f.a, x) - f.b
  if s <= 0
    return x, 0.0
  end
  return x - (s/norm(f.a)^2)*f.a, 0.0
end


export buildPOP_1v1c, buildPOP_1v2c, buildPOP_1v2, buildPOP_EllJoszMolc
export buildPOP_WB2, buildPOP_WB5
export lasserre_ex1, lasserre_ex2, lasserre_ex3, lasserre_ex5
export get_WB5cliques, get_case9cliques

############################
### Geometric problems
############################

function buildPOP_1v1c()
    z = Variable("z", Complex)
    problem = Problem()
    add_variable!(problem, z)
    set_objective!(problem, -real(z))
    add_constraint!(problem, "ineq", abs2(z) << 4)
    return problem
end

function buildPOP_1v2c()
    z = Variable("z", Complex)
    problem = Problem()
    add_variable!(problem, z)
    set_objective!(problem, imag(z))
    add_constraint!(problem, "ineq_brn", abs2(z) << 1)
    θ = π/3
    add_constraint!(problem, "ineq_rot", real(z*exp(-im*θ)) >> 0)
    return problem
end

function buildPOP_1v2()
    x1 = Variable("x1", Real)
    x2 = Variable("x2", Real)
    problem = Problem()
    add_variable!(problem, x1); add_variable!(problem, x2)
    set_objective!(problem, -1.0*x1)
    add_constraint!(problem, "ineq", (x1^2+x2^2) << 1)
    θ1 = π/3
    add_constraint!(problem, "eq_rot1", (cos(θ1)*x1+sin(θ1)*x2) == 0)
    # θ2 = -π/3
    # add_constraint!(problem, "ineq_rot2", (cos(θ2)*x1+sin(θ2)*x2) >> 0)
    return problem
end

"""
    problem = buildPOP_2v3c

    Elliptic example problemp from Josz, Molzahn 2018 paper.
"""
function buildPOP_EllJoszMolc()
    z1 = Variable("z1", Complex)
    z2 = Variable("z2", Complex)
    problem = Problem()
    add_variable!(problem, z1); add_variable!(problem, z2);
    set_objective!(problem, 3-abs2(z1)-0.5*im*z1*conj(z2)^2+0.5im*z2^2*conj(z1))
    add_constraint!(problem, "eq1", (abs2(z1)-0.25*z1^2-0.25*conj(z1)^2) == 1)
    add_constraint!(problem, "eq2", (abs2(z1)+abs2(z2)) == 3)
    add_constraint!(problem, "eq3", (im*z2-im*conj(z2)) == 0)
    add_constraint!(problem, "ineq", (z2+conj(z2)) >> 0)
    return problem
end

############################
### OPF problems
############################

function buildPOP_WB2(; v2max = 0.976, rmeqs = false, setnetworkphase=false, addball=false)
    # OPFpbs = load_OPFproblems(MatpowerInput, joinpath("..", "data", "data_Matpower", "matpower", "WB2.m"))
    # problem_c = build_globalpb!(OPFpbs)
    problem_c, pt = import_from_dat(getinstancepath("Matpower", "QCQP", "WB2"))

    ## Converting to real ineq. only problem
    !rmeqs || change_eq_to_ineq!(problem_c)
    problem = pb_cplx2real(problem_c)

    if setnetworkphase
        ## Fixing volt phase of last bus to 0
        lastctr = problem.constraints["BaseCase_2_Volt_VOLTM_Re"]
        rm_constraint!(problem, "BaseCase_2_Volt_VOLTM_Re")

        ## Setting imag part to 0
        # pt = Point(SortedDict(Variable("BaseCase_2_VOLT_Im", Real)=>0.0), isdense=true)
        # infer_problem!(problem, pt)

        add_constraint!(problem, "BaseCase_2_Volt_VOLTM_Re", sqrt(lastctr.lb) << Variable("BaseCase_2_VOLT_Re", Real) << v2max)
        add_constraint!(problem, "BaseCase_2_Volt_VOLTM_Im", Variable("BaseCase_2_VOLT_Im", Real) == 0)
    elseif v2max != 0.976
        problem.constraints["BaseCase_2_Volt_VOLTM_Re"].ub = v2max^2
    end

    ## Adding ball constraint
    if addball
        p = Polynomial()
        for var in problem.variables
            p += Variable(var[1], var[2])^2
        end
        ub = problem.constraints["BaseCase_1_Volt_VOLTM_Re"].ub + v2max^2
        add_constraint!(problem, "Ball_ctr", p << ub)
    end

    return problem
end

function buildPOP_WB5(; q5min = 1.05, rmeqs = false)
    # OPFpbs = load_OPFproblems(MatpowerInput, joinpath("..", "data", "data_Matpower", "matpower", "WB5.m"))
    problem_c, pt = import_from_dat(getinstancepath("Matpower", "QCQP", "WB5"))

    # Sgen = OPFpbs["BaseCase"].ds.bus["BUS_5"]["Gen_1"].power_min
    # OPFpbs["BaseCase"].ds.bus["BUS_5"]["Gen_1"].power_min = real(Sgen) + im*q5min
    # problem_c = build_globalpb!(OPFpbs)

    ## Converting to real ineq. only problem
    !rmeqs || change_eq_to_ineq!(problem_c)
    return pb_cplx2real(problem_c)
end

############################
### Global Optim pbs from Lasserre2001
############################

"""
    problem, relax_ctx = lasserre_ex1()

    From Lasserre2001, global minimum : (3) -0.2428.
"""
function lasserre_ex1()
    x1 = Variable("x1", Real)
    x2 = Variable("x2", Real)
    problem = Problem()
    add_variable!(problem, x1); add_variable!(problem, x2)
    set_objective!(problem, (x1^2+1)^2 + (x2^2+1)^2 + (x1+x2+1)^2)

    relax_ctx = set_relaxation(problem; hierarchykind=:Real,
                                        d = 2)
    return problem, relax_ctx
end

"""
    problem, relax_ctx = lasserre_ex2()

    From Lasserre2001, global minimum : -11.4581.
"""
function lasserre_ex2()
    x1 = Variable("x1", Real)
    x2 = Variable("x2", Real)
    problem = Problem()
    add_variable!(problem, x1); add_variable!(problem, x2)
    set_objective!(problem, (x1^2+1)^2 + (x2^2+1)^2 -2*(x1+x2+1)^2)

    relax_ctx = set_relaxation(problem; hierarchykind=:Real,
                                        d = 2)
    return problem, relax_ctx
end

"""
    problem, relax_ctx = lasserre_ex3()

    From Lasserre2001, global minimum : -1/27, x1*² = x2*² = 1/3.
"""
function lasserre_ex3()
    x1 = Variable("x1", Real)
    x2 = Variable("x2", Real)
    problem = Problem()
    add_variable!(problem, x1); add_variable!(problem, x2)
    set_objective!(problem, x1^2 * x2^2 * (x1^2 + x2^2 - 1))

    relax_ctx = set_relaxation(problem; hierarchykind=:Real,
                                        d = 3)
    return problem, relax_ctx
end

"""
    problem, relax_ctx = lasserre_ex5()

    From Lasserre2001, global minimum : -2, for (1, 2).
    Relaxation : order 1 -> -3; order 2 -> -2.
"""
function lasserre_ex5(;d = 2)
    x1 = Variable("x1", Real)
    x2 = Variable("x2", Real)
    problem = Problem()
    add_variable!(problem, x1); add_variable!(problem, x2)
    set_objective!(problem, -(x1-1)^2 -(x1-x2)^2 -(x2-3)^2)
    add_constraint!(problem, "crt1", (1-(x1-1)^2) >> 0)
    add_constraint!(problem, "crt2", (1-(x1-x2)^2) >> 0)
    add_constraint!(problem, "crt3", (1-(x2-3)^2) >> 0)

    relax_ctx = set_relaxation(problem; hierarchykind=:Real,
                                        d = d)
    return problem, relax_ctx
end

function get_WB5cliques(relax_ctx, problem)
    if !relax_ctx.issparse
        return get_maxcliques(relax_ctx, problem)
    else
        maxcliques = Dict{String, Set{Variable}}()
        maxcliques["clique1"] = Set{Variable}([
            Variable("BaseCase_1_VOLT_Im", Real),
            Variable("BaseCase_1_VOLT_Re", Real),
            Variable("BaseCase_2_VOLT_Im", Real),
            Variable("BaseCase_2_VOLT_Re", Real),
            Variable("BaseCase_3_VOLT_Im", Real),
            Variable("BaseCase_3_VOLT_Re", Real)])
        maxcliques["clique2"] = Set{Variable}([
            Variable("BaseCase_2_VOLT_Im", Real),
            Variable("BaseCase_2_VOLT_Re", Real),
            Variable("BaseCase_3_VOLT_Im", Real),
            Variable("BaseCase_3_VOLT_Re", Real),
            Variable("BaseCase_4_VOLT_Im", Real),
            Variable("BaseCase_4_VOLT_Re", Real),
            Variable("BaseCase_5_VOLT_Im", Real),
            Variable("BaseCase_5_VOLT_Re", Real)])
        return maxcliques
    end
end

function get_case9cliques(relax_ctx, problem)
    if !relax_ctx.issparse
        return get_maxcliques(relax_ctx, problem)
    else
        maxcliques = Dict{String, Set{Variable}}()
        maxcliques["clique1"] = Set{Variable}([
            Variable("BaseCase_1_VOLT_Im", Real),
            Variable("BaseCase_1_VOLT_Re", Real),
            Variable("BaseCase_5_VOLT_Im", Real),
            Variable("BaseCase_5_VOLT_Re", Real),
            Variable("BaseCase_4_VOLT_Im", Real),
            Variable("BaseCase_4_VOLT_Re", Real),
            Variable("BaseCase_9_VOLT_Im", Real),
            Variable("BaseCase_9_VOLT_Re", Real),
            Variable("BaseCase_8_VOLT_Im", Real),
            Variable("BaseCase_8_VOLT_Re", Real)])
        maxcliques["clique2"] = Set{Variable}([
            Variable("BaseCase_2_VOLT_Im", Real),
            Variable("BaseCase_2_VOLT_Re", Real),
            Variable("BaseCase_9_VOLT_Im", Real),
            Variable("BaseCase_9_VOLT_Re", Real),
            Variable("BaseCase_8_VOLT_Im", Real),
            Variable("BaseCase_8_VOLT_Re", Real),
            Variable("BaseCase_7_VOLT_Im", Real),
            Variable("BaseCase_7_VOLT_Re", Real),
            Variable("BaseCase_6_VOLT_Im", Real),
            Variable("BaseCase_6_VOLT_Re", Real)])
        maxcliques["clique3"] = Set{Variable}([
            Variable("BaseCase_3_VOLT_Im", Real),
            Variable("BaseCase_3_VOLT_Re", Real),
            Variable("BaseCase_7_VOLT_Im", Real),
            Variable("BaseCase_7_VOLT_Re", Real),
            Variable("BaseCase_6_VOLT_Im", Real),
            Variable("BaseCase_6_VOLT_Re", Real),
            Variable("BaseCase_5_VOLT_Im", Real),
            Variable("BaseCase_5_VOLT_Re", Real),
            Variable("BaseCase_4_VOLT_Im", Real),
            Variable("BaseCase_4_VOLT_Re", Real)])
        return maxcliques
    end
end


# TODO: I cound that "utils.jl" has been `include`ed in the main `VectorSphericalWaves` module. Is there a more neat way? when I `include("utils.jl")` here, I get lots of warnings.

# exporting functions that calculate for all m,n combinations
export M_N_wave_all_m_n
export B_C_P_mn_of_θ_ϕ_for_all_m_n
export πₘₙ_τₘₙ_all_all_m_n

#############################################################################################
# This version is different from "VectorSphericalHarmonics.jl", as it calculate VSWF for all m,n and all kr, θ, ϕ in one call
# This should be more stable and faster, as it employs recurrence relations in calculating πₘₙ and τₘₙ
# using recurrence relations in calculating πₘₙ and τₘₙ

#############################################################################################
# calculate π(θ) and τ(θ) using recurrence relations
function πₘₙ_τₘₙ_all_all_m_n(n_max::I, θ::R; verbose=false) where {R <: Real, I <: Integer}
    # calculate A with recurrence relation
    A = SortedDict(0 => 1.0)
    for m = 0:n_max - 1
        A[m + 1] = A[m] * sqrt((2m + 1) / (2 * (m + 1)))
    end

    # calculate πₘₙ with recurrence relation
    πₘₙ_all = zeros(get_max_single_index_from_n_max(n_max))
    for m = 1:n_max   # TODO: I think I need to start m from 0, not 1.
        n = m
        πₘₙ_all[single_index_from_m_n(m, n)] = m * A[m] * sin(θ)^(m - 1)
        if verbose; println("m=$m, n=$n, πₘₙ_all=$(πₘₙ_all[single_index_from_m_n(m, n)])"); end

        # calculate π₋ₘₙ from πₘₙ
        if m != 0
            πₘₙ_all[single_index_from_m_n(-m, n)] = (-1)^(m + 1) * πₘₙ_all[single_index_from_m_n(m, n)]
            if verbose; println("m=$(-m), n=$n, πₘₙ_all=$(πₘₙ_all[single_index_from_m_n(-m, n)])"); end
        end

        for n = (m + 1):n_max
            if n == m + 1
                πₘₙ_all[single_index_from_m_n(m, n)] = 1 / sqrt(n^2 - m^2) * ((2n - 1) * cos(θ) * πₘₙ_all[single_index_from_m_n(m, n - 1)])
                if verbose; println("m=$m, n=$n, πₘₙ_all=$(πₘₙ_all[single_index_from_m_n(m, n)])"); end
            else
                πₘₙ_all[single_index_from_m_n(m, n)] = 1 / sqrt(n^2 - m^2) * ((2n - 1) * cos(θ) * πₘₙ_all[single_index_from_m_n(m, n - 1)]) - sqrt((n - 1)^2 - m^2) * πₘₙ_all[single_index_from_m_n(m, n - 2)]
                if verbose; println("m=$m, n=$n, πₘₙ_all=$(πₘₙ_all[single_index_from_m_n(m, n)])"); end
            end

            # calculate π₋ₘₙ from πₘₙ
            if m != 0
                πₘₙ_all[single_index_from_m_n(-m, n)] = (-1)^(m + 1) * πₘₙ_all[single_index_from_m_n(m, n)]
                if verbose; println("m=$(-m), n=$n, πₘₙ_all=$(πₘₙ_all[single_index_from_m_n(-m, n)])"); end
            end
        end
    end
    τₘₙ_all = 0 # TODO: add the code for it
    return πₘₙ_all, τₘₙ_all
end

#############################################################################################
# Legendre and Associated Legendre
function Legendre_polynomials_Pn_array(n_max::I, x::NN) where {I <: Integer, NN <: Number}
    """
    Calculate all Legendre polynomials Pₙ(x) from n = 1 up to n=n_max using recurrence relation
    https://en.wikipedia.org/wiki/Legendre_polynomials

    returns
    dictionary. TODO: find a better way if this cause adjoint calculation trouble
    """
    P = SortedDict(
        0 => 1,
        1 => x,
    )

    for n = 1:(n_max - 1)
        P[n + 1] = ((2n + 1) * x * P[n] - n * P[n - 1]) / (n + 1)
    end

    return P
end

function Legendre_polynomials_Pn(n::I, x::NN) where {I <: Integer, NN <: Number}
    """
    Calculate Legendre polynomials Pₙ(x) at a given n
    """
    return Legendre_polynomials_Pn_array(n, x)[n]
end

function Associated_Legendre_polynomials_Pmn_array(m::I, n_max::I, x) where {I <: Integer}
    """
    Associated Legendre polynomials Pᵐₙ(x) from n = m up to n=n_max using recurrence relation
    https://en.wikipedia.org/wiki/Associated_Legendre_polynomials#Recurrence_formula

    returns
    dictionary. TODO: find a better way if this cause adjoint calculation trouble
    """
    n = m
    P = SortedDict(
        n => (-1)^n * factorial(factorial(2n - 1)) * (1 - x^2)^(n / 2),
    )
    P[n + 1] = x * (2n + 1) * P[n]

    for n = (m + 1):(n_max - 1)
        P[n + 1] = ( (2n + 1) * x * P[n] - (n + m) * P[n - 1] ) / (n - m + 1)
    end

    return P
end

function Associated_Legendre_polynomials_Pmn(m::I, n::I, x) where {I <: Integer}
    """
    Associated Legendre polynomials Pᵐₙ(x) at a given n
    """
    return Associated_Legendre_polynomials_Pmn_array(m, n, x)[n]
end

#############################################################################################
# Wigner-d, using recurrence
# TODO: I think this will not work for large s,m,n values, try to fix it as in `wignerdjmn_ELZOUKA`
function wignerd_and_∂wignerd_for_all_s(s_max::I, m::I, n::I, θ::R; get_derivatives=true, verbose=false) where {R <: Real, I <: Integer}
    """
    Calculate dˢₘₙ(θ) and ∂(dˢₘₙ(θ))/∂θ for all values of s, where s starts from s_min up to s_max. s_min is the maximum of |m| and |n|

    """
    # TODO: special case of m=0, n=0 (eq. B.27)
    s_min = max(abs(m), abs(n))

    # calculate ξ from eq. B.16
    if n >= m
        ξ_mn = 1
    else
        ξ_mn = (-1)^(m - n)
    end

    x = cos(θ)

    # calculate d^s_min__m_n(θ) from eq. B.24
    d_smin_m_n =
        ξ_mn * 2.0^(-s_min) * sqrt(
            factorial(BigInt(2s_min)) / (factorial(BigInt(abs(m - n))) * factorial(BigInt(abs(m + n))))
        ) *
        (1 - x)^(abs(m - n) / 2) *
        (1 + x)^(abs(m + n) / 2)

    d = SortedDict(
        s_min - 1 => 0.0,
        s_min   => d_smin_m_n
    )

    # applying the recurrence relation
    if n == 0
        if m == 0
            if verbose; println("m=$m, n=$n, I will use Legendre_polynomials_Pn_array"); end
            P_s = Legendre_polynomials_Pn_array(s_max + 1, x)
            for s = s_min:s_max + 1
                d[s] = P_s[s]
            end
        else
            if verbose; println("m=$m, n=$n, I will use Associated_Legendre_polynomials_Pmn_array"); end
            P_m_s = Associated_Legendre_polynomials_Pmn_array(m, s_max + 1, x)
            for s = s_min:s_max + 1
                d[s] = sqrt(factorial(s - m) / factorial(s + m)) * P_m_s[s]
            end
        end
    else
        for s = s_min:s_max + 1
            if verbose; println("m=$m, n=$n, I will use the general recurrence"); end
            d[s + 1] = 1 / (s * sqrt((s + 1)^2 - m^2) * sqrt((s + 1)^2 - n^2)) * (
                (2s + 1) * (s * (s + 1) * x - m * n) * d[s]
                - 1 * (s + 1) * sqrt(s^2 - m^2) * sqrt(s^2 - n^2) * d[s - 1]
            ) # eq. B.22
        end
    end

    # calculate the derivative ∂(dˢₘₙ(θ))/∂θ
    if get_derivatives
        ∂d_∂θ = SortedDict(
            s_min - 1 => 0.0,
            s_min   => 0.0
        )

        sin_theta = sin(θ)

        for s = s_min:s_max
            if verbose; println("s=$s"); end
            ∂d_∂θ[s] = 1 / sin_theta * (
                -1 * ((s + 1) * sqrt((s^2 - m^2) * (s^2 - n^2))) / (s * (2s + 1)) * d[s - 1]
                - 1 * (m * n) / (s * (s + 1)) * d[s]
                + 1 * (
                        s *
                        sqrt((s + 1)^2 - m^2) *
                        sqrt((s + 1)^2 - n^2)
                    ) /
                    ((s + 1) * (2s + 1)) * d[s + 1]
            )
        end
        return d, ∂d_∂θ
    else
        return d
    end

end

#############################################################################################
# calculate π(θ) and τ(θ) using Wigner-d that was calculated using recurrence relations
function πₘₙ_τₘₙ_all_all_m_n_using_wigner(n_max::I, θ::R; verbose=false) where {R <: Real, I <: Integer}
    πₘₙ_all = zeros(get_max_single_index_from_n_max(n_max))
    τₘₙ_all = zeros(get_max_single_index_from_n_max(n_max))

    for m = 0:n_max  # TODO: special case of m=0
        d_rec_all_n, ∂d_∂θ_rec_all_n  = wignerd_and_∂wignerd_for_all_s(n_max, 0, m, θ)
        for n = m:n_max
            if n != 0
                if verbose; println("m=$m, n=$n, calculate πₘₙ_all, τₘₙ_all"); end
                πₘₙ_all[single_index_from_m_n(m, n)] = m / sin(θ) * d_rec_all_n[n]
                τₘₙ_all[single_index_from_m_n(m, n)] = ∂d_∂θ_rec_all_n[n]

                if m != 0
                    πₘₙ_all[single_index_from_m_n(-m, n)] = (-1)^(m + 1) * πₘₙ_all[single_index_from_m_n(m, n)]
                    τₘₙ_all[single_index_from_m_n(-m, n)] = (-1)^(m)   * τₘₙ_all[single_index_from_m_n(m, n)]
                end
            end
        end
    end
    return πₘₙ_all, τₘₙ_all
end

function B_C_mn_of_θ_for_all_m_n(n_max::I, θ::R) where {R <: Real, I <: Integer}
    """
    Calculate Bₙₘ(θ), Cₙₘ(θ), Pₙₘ(θ) equations C.19, C.20, C.21
    The order of m,n is according to the function "single_index_from_m_n"
    """
    πₘₙ_all, τₘₙ_all = πₘₙ_τₘₙ_all_all_m_n_using_wigner(n_max, θ)
    B = (_ -> zero(SVector{3,Complex})).(πₘₙ_all)
    C = (_ -> zero(SVector{3,Complex})).(πₘₙ_all)
    for idx in eachindex(πₘₙ_all)
        B[idx] = [0,      τₘₙ_all[idx], im * πₘₙ_all[idx] ]
        C[idx] = [0, im * πₘₙ_all[idx], -1 * τₘₙ_all[idx] ]
    end

    return B, C
end

function P_mn_of_θ_for_all_m_n(n_max::I, θ::R) where {R <: Real, I <: Integer}
    """
    Calculate Pₙₘ(θ) using equations C.19, C.20
    The order of m,n is according to the function "single_index_from_m_n"
    """
    P = fill(zero(SVector{3,Complex}), get_max_single_index_from_n_max(n_max))
    for m = 0:n_max
        d_rec_all_n  = wignerd_and_∂wignerd_for_all_s(n_max, 0, m, θ; get_derivatives=false)
        for n = m:n_max
            if n != 0
                P[single_index_from_m_n(m, n)] = [d_rec_all_n[n], 0, 0]
                P[single_index_from_m_n(-m, n)] = [(-1)^m * d_rec_all_n[n], 0, 0]  # using symmetry relation B.5, we can get d_(0,-m) from d_(0,m)
            end
        end
    end
    return P
end

function B_C_P_mn_of_θ_ϕ_for_all_m_n(n_max::I, θ::R, ϕ::R) where {R <: Real, I <: Integer}
    """
    Calculate Bₙₘ(θ,ϕ), Cₙₘ(θ,ϕ), Pₙₘ(θ,ϕ) for all m and n
    """
    B_of_θ, C_of_θ = B_C_mn_of_θ_for_all_m_n(n_max, θ)
    P_of_θ = P_mn_of_θ_for_all_m_n(n_max, θ)

    for n = 1:n_max
        for m = -n:n
            factor = (-1)^m * convert(R, sqrt_factorial_n_plus_m_over_factorial_n_minus_m(m,n)) * exp(im * m * ϕ)
            B_of_θ[single_index_from_m_n(m, n)] *= factor
            C_of_θ[single_index_from_m_n(m, n)] *= factor
            P_of_θ[single_index_from_m_n(m, n)] *= factor
        end
    end

    return B_of_θ, C_of_θ, P_of_θ
end

function M_N_wave_all_m_n(n_max::I, kr::NN, θ::R, ϕ::R; kind="regular") where {R <: Real, I <: Integer, NN <: Number}
    """
    Parameters
    ==========
    kind: string, either ["regular" or "incoming"] or ["irregular" or "outgoing"]
    """
    radial_function, radial_function_special_derivative  = get_radial_function_and_special_derivative_given_kind(kind)

    B_of_θ_ϕ, C_of_θ_ϕ, P_of_θ_ϕ  = B_C_P_mn_of_θ_ϕ_for_all_m_n(n_max, θ, ϕ)

    M = 0 .* B_of_θ_ϕ
    N = 0 .* B_of_θ_ϕ

    for n = 1:n_max
        for m = -n:n
            M[single_index_from_m_n(m, n)] = convert(R, γ_mn(m, n)) * radial_function(n, kr) * C_of_θ_ϕ[single_index_from_m_n(m, n)]
            N[single_index_from_m_n(m, n)] = convert(R, γ_mn(m, n)) * (
                n * (n + 1) / kr * radial_function(n, kr)    * P_of_θ_ϕ[single_index_from_m_n(m, n)]
                + (radial_function_special_derivative(n, kr) * B_of_θ_ϕ[single_index_from_m_n(m, n)])
            )
        end
    end
    return M, N
end


module pindex

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

using Distributions, Parameters, LinearAlgebra, Compat, Test

function FixedLaspeyres(p,q)
N=size(p,1)
T=size(p,2)
Pl=zeros(N+1,T)
s=zeros(N,T)
   for t in 1:T
        for n in 2:N+1
        s[n-1,t] = (p[n-1,t]*q[n-1,t])/p[:,t]'q[:,t]
        Pl[n,t]=Pl[n-1,t]+p[n-1,t]/p[n-1,1]*s[n-1,1]
        end
    end
return Pl[N+1,:]
end

function FixedPaasche(p,q)
N=size(p,1)
T=size(p,2)
Pp=zeros(N+1,T)
s=zeros(N,T)
   for t in 1:T
        for n in 2:N+1
        s[n-1,t] = p[n-1,t]*q[n-1,t]/p[:,t]'q[:,t]
        Pp[n,t]=Pp[n-1,t]+p[n-1,1]/p[n-1,t]*s[n-1,t]
        end
    end
return Pp[N+1,:].^(-1)
end

function FixedFisher(p,q)
     Pf=(FixedPaasche(p,q).^(1/2)).*(FixedLaspeyres(p,q).^(1/2))
return Pf
end

function FixedTornqvist(p,q)
N=size(p,1)
T=size(p,2)
Pt=zeros(N+1,T)
s=zeros(N,T)
   for t in 1:T
        for n in 2:N+1
        s[n-1,t] = p[n-1,t]*q[n-1,t]/p[:,t]'q[:,t]
        Pt[n,t]=(1/2)*(s[n-1,1]+s[n-1,t])*log(p[n-1,t]/p[n-1,1])
        end
    end
return exp.(Pt[N+1,:])
end

function ChainedLaspeyres(p,q)
N=size(p,1)
T=size(p,2)
Pl=zeros(N+1,T)
Plc=ones(N+1,T)
s=zeros(N,T)
   for t in 2:T
        for n in 2:N+1
        s[n-1,t-1] = p[n-1,t-1]*q[n-1,t-1]/p[:,t-1]'q[:,t-1]
        Pl[n,t]=Pl[n-1,t]+p[n-1,t]/p[n-1,t-1]*s[n-1,t-1]
        Plc[n,t]=Plc[n,t-1]*Pl[n,t]
        end
    end
return Plc[N+1,:]
end

function ChainedPaasche(p,q)
N=size(p,1)
T=size(p,2)
Pp=zeros(N+1,T)
Ppc=ones(N+1,T)
s=zeros(N,T)
   for t in 2:T
        for n in 2:N+1
        s[n-1,t] = (p[n-1,t]*q[n-1,t])/p[:,t]'q[:,t]
        Pp[n,t]=Pp[n-1,t]+p[n-1,t-1]/p[n-1,t]*s[n-1,t]
        Ppc[n,t]=Ppc[n,t-1]*(Pp[n,t])^(-1)
        end
    end
return Ppc[N+1,:]
end

function ChainedFisher(p,q)
     Pfc=(ChainedPaasche(p,q).^(1/2)).*(ChainedLaspeyres(p,q).^(1/2))
return Pfc
end

function ChainedTornqvist(p,q)
N=size(p,1)
T=size(p,2)
Pt=zeros(N+1,T)
Ptc=ones(N+1,T)
s=zeros(N,T)
   for t in 2:T
        for n in 2:N+1
        s[n-1,t-1] = p[n-1,t-1]*q[n-1,t-1]/p[:,t-1]'q[:,t-1]
        s[n-1,t] = (p[n-1,t]*q[n-1,t])/p[:,t]'q[:,t]
        Pt[n,t]=Pt[n-1,t]+(1/2)*(s[n-1,t-1]+s[n-1,t])*log(p[n-1,t]/p[n-1,t-1])
        Ptc[n,t]=Ptc[n,t-1]*exp(Pt[n,t])
        end
    end
return Ptc[N+1,:]
end

function FixLaspeyres(p,q)
N=size(p,1)
T=size(p,2)
Pl=zeros(T)
    for t in 1:T
    Pl[t]=p[:,t]'q[:,1]/p[:,1]'q[:,1]
    end
return Pl
end

function FixPaasche(p,q)
N=size(p,1)
T=size(p,2)
Pl=zeros(T)
    for t in 1:T
    Pl[t]=p[:,t]'q[:,t]/p[:,1]'q[:,t]
    end
return Pl
end

function FixFisher(p,q)
     Pf=(FixPaasche(p,q).^(1/2)).*(FixLaspeyres(p,q).^(1/2))
return Pf
end

function ChainLaspeyres(p,q)
N=size(p,1)
T=size(p,2)
Pl=ones(T)
Plc=ones(T)
    for t in 2:T
    Pl[t]=p[:,t]'q[:,t-1]/p[:,t-1]'q[:,t-1]
    Plc[t]=Plc[t-1]*Pl[t]
    end
return Plc
end

function ChainPaasche(p,q)
N=size(p,1)
T=size(p,2)
Pp=ones(T)
Ppc=ones(T)
    for t in 2:T
    Pp[t]=p[:,t]'q[:,t]/p[:,t-1]'q[:,t]
    Ppc[t]=Ppc[t-1]*Pp[t]
    end
return Ppc
end

function ChainFisher(p,q)
    Pfc=(ChainPaasche(p,q).^(1/2)).*(ChainLaspeyres(p,q).^(1/2))
return Pfc
end

export FixedLaspeyres,  FixedPaasche, FixedFisher, FixedTornqvist, ChainedLaspeyres, ChainedPaasche, ChainedFisher, ChainedTornqvist, FixLaspeyres, FixPaasche, FixFisher, ChainLaspeyres, ChainPaasche, ChainFisher

end # module


using jInv.Mesh
using Test

# setup 3D problem
nc = [5, 7, 8]
x0 = rand(3)
domain = [x0[1], 4, x0[2], 2, x0[3], 6]
h   = (domain[2:2:end]-domain[1:2:end])./nc
h1  = h[1]*ones(nc[1])
h2  = h[2]*ones(nc[2])
h3  = h[3]*ones(nc[3])

Mt = getTensorMesh3D(h1,h2,h3,x0)
Mr = getRegularMesh(domain,nc)
M2D = getRegularMesh(domain[1:4],nc[1:2])
Mt2 = getTensorMesh3D(h1+rand(nc[1]),h2+rand(nc[2]),h3+rand(nc[3]),x0)

Meshes = (Mt,Mr,Mt2,M2D)

for k=1:length(Meshes)
	M = Meshes[k]
	print("\ttesting nodal averaging for $(typeof(M))...")
	Av = getNodalAverageMatrix(M)
	xn = getNodalGrid(M)
	xc = getCellCenteredGrid(M)
	@test norm(Av*xn-xc,Inf)<1e-12

	print("passed!\n")
end


module MyInterpolations

export MyLinInterp

immutable MyLinInterp
    grid::Array
    vals::Array
end

function (f::MyLinInterp)(x)
    i = searchsortedlast(f.grid,x)
    k = searchsortedfirst(f.grid,x)
    if i == 0 || (i == length(f.grid) && k == length(f.grid) + 1)
        return 0
    end

    if i == length(f.grid) && k != length(f.grid) + 1
        return (f.vals[i])
    end

    interpolated_value = (x - f.grid[i])*(f.vals[i+1]-f.vals[i])/(f.grid[i+1]-f.grid[i])+f.vals[i]
    return interpolated_value
end

end


using CSV, DataFrames, StatsPlots, StanSample, Statistics

ProjDir = @__DIR__

old_df = CSV.read(joinpath(ProjDir, "data", "data_final.csv"), delim=' ')

df = DataFrame()
for row in eachrow(old_df)
	vals = split(row[1], "\t")
	yr = Meta.parse(row[1][1:4])
	wk = Meta.parse(row[1][5:6])
	s = Meta.parse(vals[2])
	p = Meta.parse(vals[3])
	q = Meta.parse(vals[4])
	append!(df, DataFrame(year=yr, week=wk, store=s, quantity=q, price=p))
end

df[!, :ppu] = df[:, :price] ./ df[:, :quantity]
df[!, :ppu_l] = log.(df[:, :ppu])
df[!, :quantity_l] = log.(df[:, :quantity])

# Clean up

df = df[df.ppu_l .> 0.04, :]
df = df[df.quantity_l .> 0.0, :]

CSV.write(joinpath(ProjDir, "data", "data.csv"), df)

p1 = scatter(df[:, :ppu], df[:, :quantity],
	xlabel="ppu", ylabel="quantity sold")
p2 = scatter(df[:, :ppu_l], df[:, :quantity_l],
	xlabel="log(ppu)", ylabel="log(quantity sold)")
plot(p1, p2, layout=(2, 1))
savefig(joinpath(ProjDir, "plots", "data.png"))

bglm = "
data {
  int<lower=0> N;
  vector[N] q; // Outcome
  vector[N] p; // Predictor
}
parameters {
  real<lower=0> a;
  real b;
  real<lower=0> sigma;
}
model {
  vector[N] mu;            		// mu is a vector
  a ~ normal(2.0, 10.0);		// intercept positive
  b ~ normal(-2.0, 10.0);		// demand drops as price increases
  sigma ~ uniform(0, 15);
  mu = a + b * p;
  q ~ normal(mu, sigma);
}
";

data = Dict(
  :N => size(df, 1),
  :q => df[:, :quantity_l],
  :p => df[:, :ppu_l]
)

sm = SampleModel("bglm", bglm)

rc = stan_sample(sm, data=data)

if success(rc)
	chns = read_samples(sm)
	plot(chns)
	savefig(joinpath(ProjDir, "plots", "chains.png"))

	dfs = DataFrame(chns)
	show(chns)

	p = 0.6:0.01:2.0
	p1 = scatter(df[:, :ppu_l], df[:, :quantity_l],
		xlabel="ppu_l", ylabel="quantity_l sold")
	a_mean = mean(dfs[:, :a])
	b_mean = mean(dfs[:, :b])
	mu_q = a_mean .+ b_mean .* p
	plot!(p1, p, mu_q)
	savefig(joinpath(ProjDir, "plots", "regression.png"))

	price_points_l = [1.05, 1.35, 1.5, 1,6, 1.78]
	scatter(price_points_l, a_mean .+ b_mean .* price_points_l,
		xlims = (0.6, 2.0), ylims= (0, 7), leg=false,
		xlabel="log(ppu)", ylabel="log(quantity sold)")
	for row in eachrow(dfs)
		plot!(p, row[:a] .+ row[:b] .* p, color=:gray)
	end
	scatter!(price_points_l, a_mean .+ b_mean .* price_points_l)
	savefig(joinpath(ProjDir, "plots", "sampling.png"))
end


@testset "Prim" begin
    g4 = CompleteGraph(4)

    distmx = [
      0  1  5  6
      1  0  4  10
      5  4  0  3
      6  10  3  0
    ]

    vec_mst = Vector{Edge}([Edge(1, 2), Edge(2, 3), Edge(3, 4)])
    for g in testgraphs(g4)
        # Testing Prim's algorithm
        mst = @inferred(prim_mst(g, distmx))
        @test mst == vec_mst
    end

    #second test
    distmx_sec = [
        0       0       0.26    0       0.38    0       0.58    0.16
        0       0       0.36    0.29    0       0.32    0       0.19
        0.26    0.36    0       0.17    0       0       0.4     0.34
        0       0.29    0.17    0       0       0       0.52    0
        0.38    0       0       0       0       0.35    0.93    0.37
        0       0.32    0       0       0.35    0       0       0.28
        0.58    0       0.4     0.52    0.93    0       0       0
        0.16    0.19    0.34    0       0.37    0.28    0       0
    ]

    vec2 = Vector{Edge}([Edge(8, 2), Edge(1, 3), Edge(3, 4), Edge(6, 5), Edge(8, 6), Edge(3, 7), Edge(1, 8)])
    gx = SimpleGraph(distmx_sec)
    for g in testgraphs(gx)
        mst2 = @inferred(prim_mst(g, distmx_sec))
        @test mst2 == vec2
    end

    # Third test: Example from
    #   https://en.wikipedia.org/wiki/Kruskal%27s_algorithm#Example
    # Adapted from networkx
    #   https://github.com/networkx/networkx/blob/master/networkx/algorithms/tree/tests/test_mst.py
    edges_w = [(0, 1, 7), (0, 3, 5), (1, 2, 8), (1, 3, 9), (1, 4, 7),
               (2, 4, 5), (3, 4, 15), (3, 5, 6), (4, 5, 8), (4, 6, 9),
               (5, 6, 11)]
    n = 7

    if false    # verify the example, visually
        for e in edges_w
            println("$('A' + e[1]) -- $('A' + e[2]): $(e[3])")
        end
    end

    gx = LightGraphs.SimpleGraph([Edge(e[1]+1, e[2]+1) for e in edges_w])

    @test nv(gx) == n
    @test ne(gx) == length(edges_w)

    distmx3 = zeros(Float64, n, n)
    for e in edges_w
        distmx3[e[1]+1, e[2]+1] = e[3]
        distmx3[e[2]+1, e[1]+1] = e[3]
    end
    mst3 = @inferred(prim_mst(gx, distmx3))

    expected3 = [Edge(e[1]+1, e[2]+1)
                 for e in [(0, 1, 7),
                           (0, 3, 5),
                           (1, 4, 7),
                           (2, 4, 5),
                           (3, 5, 6),
                           (4, 6, 9)]]

    """
    The cost of the edge selection `es` given the distance matrix `distmx`
    """
    edge_costs(es, distmx) =
        sum(distmx[e.src, e.dst] for e in es)

    @test edge_costs(mst3, distmx3) == edge_costs(expected3, distmx3)

    # maximum spanning tree
    mst3n = @inferred(prim_mst(gx, -distmx3))
    expected3n = [Edge(e[1]+1, e[2]+1)
                  for e in [(0, 1, 7),
                            (1, 2, 8),
                            (1, 3, 9),
                            (3, 4, 15),
                            (4, 6, 9),
                            (5, 6, 11)]]
    @test edge_costs(mst3n, distmx3) == edge_costs(expected3n, distmx3)
end


"""
    interval_norm(A::Matrix{Interval})

computes the infinite norm of interval matrix A
"""
interval_norm(A::AbstractMatrix) = maximum(sum(mag.(A); dims=2)) # TODO: proper expand norm function and add 1-norm
interval_norm(v::AbstractVector) = maximum(mag.(v))
"""
Preconditions the interval system Ax = b by multipling by the (approximate) inverse of Ac,
that is the midpoint of A.
"""
function precondition(A, b)
    Ac = mid.(A)
    return Ac\A, Ac\b
end

"""
    enclose(A::Matrix{Interval}, b::Vector{Interval})

Computes an initial enclosure Σ so that x ⊆ Σ, where x is the solution of the interval
system Ax = b.
"""
function enclose(A::SMatrix{N, N, T, M}, b::SVector{N, T}) where {N, T, M}
    A1 = I - A
    e = interval_norm(b)/(1 - interval_norm(A1))
    x0 = MVector{N, T}(fill(-e..e, N))
    return x0
end
function enclose(A::MMatrix{N, N, T, M}, b::MVector{N, T}) where {N, T, M}
    A1 = I - A
    e = interval_norm(b)/(1 - interval_norm(A1))
    x0 = MVector{N, T}(fill(-e..e, N))
    return x0
end

function enclose(A, b)
    A1 = I - A
    e = interval_norm(b)/(1 - interval_norm(A1))
    x0 = fill(-e..e, length(b)  )
    return x0
end

"""
    comparison_matrix(A::Matrix{Interval})

Computes the comparison matrix ⟨A⟩ of the given matrix A according to the definition
⟨A⟩_ii = mig(A_ii)
⟨A⟩_ij = -mag(A_ij)
"""
function comparison_matrix(A)
    n = size(A, 1)
    compA = -mag.(A)
    @inbounds for (i, idx) in enumerate(diagind(A))
        compA = setindex(compA, mig(A[i, i]), idx)
    end
    return compA
end


module zad2

function zad22(funk,xL,xU;ep=1e-8,golden=0.618)
              fa=funk(xL)
              fb=funk(xU)
              x1=xU-((xU-xL)*(golden))
              x2=xL+((xU-xL)*(golden))
              err=Inf
              while err>ep
                  if funk(x2)<funk(x1)
                      xL=x1
                      x1=x2
                      x2=xU-((xU-xL)*(golden))
                  elseif funk(x2)>funk(x1)
                      xU=x2
                      x2=x1
                  x1=xL+((xU-xL)*(golden))
                  else
                      xL=(x1+x2)/2
                      xU=xL
                  end
                 err=2*abs((xU-xL)/(xU+xL))
                  end
              Xmin=(x1+x2)/2
              print("Minimum: ",round(funk(Xmin),digits=2), " for: ", Xmin, " x1: " ,round(x1,digits=2), " x2: ", round(x2,digits=2))

       return(round(funk(Xmin),digits=2),round(Xmin,digits=2))
       end

export zad22

end # module


# Model m4.3as.jl

using Pkg, DrWatson

begin
    using StanQuap
    using StatisticalRethinking
end

begin
    df = CSV.read(sr_datadir("Howell1.csv"), DataFrame; delim=';')
    df = filter(row -> row[:age] >= 18, df);
    mean_weight = mean(df.weight)
    #df.weight_c = df.weight .- mean_weight
end;

stan4_3 = "
data {
    int<lower=1> N;
    vector[N] weight;
    vector[N] height;
}
parameters {
    real a;
    real<lower=0> b;
    real<lower=0, upper=50> sigma;
}
model {
    // Define mu as a vector.
    vector[N] mu;

    // Priors for mu and sigma
    sigma ~ uniform(0 , 50);
    a ~ normal($(mean_weight), 20);
    b ~ lognormal(0, 1);

    // Observed heights
    for (i in 1:N) {
    mu[i] = a + b * (weight[i] - $(mean_weight));
    }
    height ~ normal(mu, sigma);
}
";

data = (N = length(df.height), height = df.height, weight = df.weight)
init = (a = 180.0, b = 1.0, sigma = 10.0)
q4_3s, m4_3s, o4_3s = stan_quap("m4.2s", stan4_3; data, init);

if q4_3s.converged
    quap4_3s_df = sample(q4_3s)          # DataFrame with samples
    precis(quap4_3s_df)
    post4_3s = read_samples(m4_3s)
    post4_3s |> display
end

# End of m4.3as.jl


# SPDX-License-Identifier: MIT

"""
    struct ProbabilitySemifield{T<:AbstractFloat} <: Semfield
        val::T
    end

Log-semifield is defined as :
  * ``x \\oplus y \\triangleq x + y``
  * ``x \\otimes y \\triangleq x \\cdot y``
  * ``x \\oslash y \\triangleq \\frac{x}{y}``
``\\forall x, y \\in [0, 1]``.
"""
struct ProbabilitySemifield{T<:AbstractFloat} <: Semifield
    val::T
end

Base.:+(x::ProbabilitySemifield{T}, y::ProbabilitySemifield{T}) where T =
    ProbabilitySemifield{T}(x.val + y.val)
Base.:*(x::ProbabilitySemifield, y::ProbabilitySemifield) =
    ProbabilitySemifield(x.val * y.val)
Base.:/(x::ProbabilitySemifield, y::ProbabilitySemifield) =
    ProbabilitySemifield(x.val / y.val)
Base.zero(::Type{ProbabilitySemifield{T}}) where T =
    ProbabilitySemifield{T}(T(0))
Base.one(::Type{ProbabilitySemifield{T}}) where T =
    ProbabilitySemifield{T}(T(1))
Base.isless(x::ProbabilitySemifield, y::ProbabilitySemifield) =
    isless(x.val, y.val)
Base.typemin(x::Type{ProbabilitySemifield{T}}) where T = zero(T)
Base.typemax(x::Type{ProbabilitySemifield{T}}) where T = one(T)


#
# create Lienar FEM Hat functions with ReLU
#
# centred at x=0, width = 1.0 => spacing = 0.5
#
using Flux, NNlib, DiffEqFlux
using Plots
import SEM
#--------------------------------------#
x = SEM.linspace(-2,2,1000)

fm1 = @. 2.0relu(x+0.5)
fm0 = @. 4.0relu(x-0.0)
fp1 = @. 2.0relu(x-0.5)

hat = fm1 - fm0 + fp1

p = plot()
p = plot!(title="FEM Hat Function",xlims=[-1,1],ylims=[0,1.5])
p = plot!(x,hat,width=4,color=:black,style=:solid,legend=false)
savefig(p,"reluHat.png")

# ReLU
p=plot(title="ReLU"); p=plot!(x,relu.(x),width=4,color="black",legend=false,style=:solid); savefig(p,"relu.png")


# Find the slope of two point
find_slope(x1, y1, x2, y2) = (y2 - y1) / (x2 - x1)

# Find the closest whole numbers on both sides of a x
neighbors(x) = round(x), ceil(x)

# Find the predicted value of y given a value x
function find_y(x, known_x)
    # Check if the given value is exactly one in the known list
    if round(x) == x
        return known_x[x]
    end

    left_x, right_x = neighbors(x)

    left_y = known_x[left_x]
    right_y = known_x[right_x]

    slope = find_slope(left_x, left_y, right_x, right_y)

    return left_y + (slope * (x - left_x))
end

println(find_slope(2, 3, 5, 6))
println(neighbors(4.5))
println(find_y(2, [3, 4, 5, 9]))


#=

    deterministic.jl

    Contains code to solve the deterministic ODE model

           β     λ    pβ
        T --→ S --→ E --→ S + F

    Author:     Alexander P. Browning
                ======================
                School of Mathematical Sciences
                Queensland University of Technology
                ======================
                ap.browning@icloud.com
                alexbrowning.me

=#

"""
    solve_ode_model(t,[λ,β,p])

Solve the ODE model using an analytical solution obtained in Mathematica.
Returns [A(t),I(t)].

"""
function solve_ode_model(t,θ::Vector)
    λ,β,p = θ
    A = (exp(-t*β)*(exp(-t*((-1+p)*β+λ))*(-1+p)*p*β^3*λ-(-1+p)*(β-λ)*λ*(p*β+λ)^2+exp(t*β)*((-1+p)*β+λ)*(λ^2*(β+λ)+p^2*β*(-λ^2+β^2*(1+t*λ)+β*λ*(1+t*λ))+p*λ*(-λ^2+β^2*(1+t*λ)+β*λ*(1+t*λ)))))/((β+λ)*((-1+p)*β+λ)*(p*β+λ)^2)
    S = (exp(-t*(p*β+λ))*β*((-1+p)*β*λ+exp(t*(p*β+λ))*p*(β+λ)*((-1+p)*β+λ)-exp(t*((-1+p)*β+λ))*(-1+p)*λ*(p*β+λ)))/((β+λ)*((-1+p)*β+λ)*(p*β+λ))
    I = A - S
    return [A,I]
end
solve_ode_model(t,λ,β,p) = solve_ode_model(t,[λ,β,p])

"""
    solve_ode_model_all_vars(t,[λ,β,p])

Solve the ODE model using matrix exponentiation and return all variables.
Returns [T(t),S(t),E(t),F(t)].

"""
function solve_ode_model_all_vars(t,θ::Vector)
    λ,β,p = θ
    M = [-β 0 0 0; β -λ p*β 0; 0 λ -p*β 0; 0 0 p*β 0]
    x₀ = [λ/(λ+β),β/(λ+β),0,0]
    return exp(M*t) * x₀
end
solve_ode_model_all_vars(t,λ,β,p) = solve_ode_model_all_vars(t,[λ,β,p])

"""
    solve_alt_ode_model(t,[γ,λ,β,p])

Solve the alternate ODE model where antibody binds to receptors at constant rate γ.
"""
function solve_alt_ode_model(t,θ::Vector)
    γ,λ,β,p = θ
    M = [-β 0 0 0 0; β -γ 0 p*β 0; 0 γ -λ 0 0; 0 0 λ -p*β 0; 0 0 0 p*β 0]
    x₀ = [λ/(λ+β),β/(λ+β),0,0,0]
    x = exp(M*t) * x₀
    A = sum(x[3:5])
    I = sum(x[4:5])
    return [A,I]
end
solve_alt_ode_model(t,γ,λ,β,p) = solve_alt_ode_model(t,[γ,λ,β,p])


# ---
# jupyter:
#   jupytext:
#     formats: ipynb,jl:hydrogen
#     text_representation:
#       extension: .jl
#       format_name: hydrogen
#       format_version: '1.3'
#       jupytext_version: 1.11.2
#   kernelspec:
#     display_name: Julia 1.8.0-DEV
#     language: julia
#     name: julia-1.8
# ---

# %%
struct F{p, T<:Integer} <: Integer
    a::T
    F{p, T}(a::Integer) where {p, T<:Integer} = new{p, T}(mod(T(a), p))
end
F{p}(a::Integer) where p = F{p, typeof(a)}(mod(a, p))
F{p, S}(x::F{p, T}) where {S<:Integer, p, T<:Integer} = F{p, S}(x.a)

Base.promote_rule(::Type{F{p, T}}, ::Type{S}) where {p, T<:Integer, S<:Integer} =
    F{p, promote_type(T, S)}

Base.zero(::Type{F{p, T}}) where {p, T<:Integer} = F{p}(mod(zero(T), p))
Base.one(::Type{F{p, T}}) where {p, T<:Integer} = F{p}(mod(one(T), p))
for op in (:-, :+)
    @eval Base.$op(x::F{p}) where p = F{p}(mod($op(x.a), p))
end
for op in (:-, :+, :*)
    @eval Base.$op(x::F{p}, y::F{p}) where p = F{p}(mod($op(x.a, y.a), p))
end
Base.inv(x::F{p}) where p = F{p}(invmod(x.a, p))
Base.:/(x::F{p}, y::F{p}) where p = x * inv(y)
Base.:\(x::F{p}, y::F{p}) where p = inv(x) * y
Base.:(==)(x::F{p}, y::F{p}) where p = x.a == y.a
Base.:<(x::F{p}, y::F{p}) where p = x.a < y.a

Base.show(io::IO, x::F{p}) where p = print(io, "F", p, '(', x.a, ')')

F7 = F{7, BigInt}
x, y = F7(10), F7(-2)
@show(x, y, zero(x), one(x), +x, -x, x + y, x - y, x * y, x / y, x \ y, x^3, x^-5, x == F7(3), x == 3)
println()
A = F7[1 2; 3 4]
@show(A, 4A, A/4)
println()
using LinearAlgebra
L, U = lu(A, NoPivot())
@show(L, U, L * U, det(A), inv(A), inv(A) * A, A * inv(A));

# %%
"""See https://docs.julialang.org/en/v1/manual/interfaces/"""
Base.iterate(Fp::Type{F{p, T}}) where {p, T<:Integer} = (zero(Fp), zero(T))
function Base.iterate(Fp::Type{F{p, T}}, state) where {p, T<:Integer}
    nextstate = state + 1
    nextstate < p ? (Fp(nextstate), nextstate) : nothing
end
Base.IteratorSize(Fp::Type{F{p, T}}) where {p, T<:Integer} = Base.HasLength()
Base.length(Fp::Type{F{p, T}}) where {p, T<:Integer} = p
Base.eltype(Fp::Type{F{p, T}}) where {p, T<:Integer} = Fp

squares(Fp) = Fp[x^2 for x in Fp]
squareroots(k, Fp) = Fp[x for x in Fp if x^2 == k]
@show(collect(F7), squares(F7), squareroots.(0:6, Ref(F7)));

# %%
using GaloisFields, LinearAlgebra
GF7 = @GaloisField 7
B = GF7[1 2; 3 4]
det(B)

# %%
inv(B)

# %%
lu(B)

# %%
using AbstractAlgebra

@show GF7 = GF(7)
x, y = GF7(10), GF7(-2)
@show(x, y, zero(x), one(x), -x, x + y, x - y, x * y, x^3, x^-5, x == GF7(3), x == 3)
println()

@show C = GF7[1 2; 3 4]
@show P, x = PolynomialRing(GF7, "x")
@show(det(C), inv(C), lu(C), charpoly(P, C))
println()

squares(Fp) = [x^2 for x in Fp]
squareroots(k, Fp) = [x for x in Fp if x^2 == k]
@show(collect(GF7), squares(GF7), squareroots.(0:6, Ref(GF7)));

# %%


using DataFrames

# Create test data
INCIDENCE_MAT = [0 1 0 1 1 0 0 0 0 0;
		 0 1 1 0 0 0 1 0 0 0;
		 1 1 0 1 1 0 0 0 1 1;
		 1 0 0 1 0 1 0 0 1 1;
		 0 0 1 1 0 0 0 1 0 0]

C_MAT = INCIDENCE_MAT'INCIDENCE_MAT

@testset "Similarity methods tests" begin

	# Check inputs:
	@test_throws MethodError similarity(DataFrame(INCIDENCE_MAT), method="jaccard", orientation="colwise") # non-matrix input
	@test_throws MethodError get_rca(INCIDENCE_MAT[:, 1], method="jaccard", orientation="colwise") # One-dim array
	@test_throws MethodError similarity(["1" "2"; "3" "4"], method="jaccard", orientation="colwise") # non-numeric matrix
	@test_throws AssertionError
	similarity(INCIDENCE_MAT, method="abc", orientation="colwise") # wrong method
	@test_throws MethodError
	similarity(INCIDENCE_MAT, method="jaccard", orientation="abc") # wrong orientation

	# Check output type:
	@test typeof(similarity(INCIDENCE_MAT, method="jaccard", orientation="colwise")) == Array{Float64, 2}
	@test typeof(similarity(INCIDENCE_MAT, method="jaccard", orientation="rowwise")) == Array{Float64, 2}
	@test typeof(similarity(INCIDENCE_MAT, method="association", orientation="colwise")) == Array{Float64, 2}

	# Check output result:
	# Methods
	@test similarity(INCIDENCE_MAT, method="association", orientation="colwise")[2, 4] == C_MAT[2, 4] / (C_MAT[2, 2] * C_MAT[4, 4])
	@test similarity(INCIDENCE_MAT, method="association", orientation="colwise")[9, 10] == C_MAT[9, 10] / (C_MAT[9, 9] * C_MAT[10, 10])
	@test similarity(INCIDENCE_MAT, method="cosine", orientation="colwise")[2, 4] == C_MAT[2, 4] / sqrt(C_MAT[2, 2] * C_MAT[4, 4])
	@test similarity(INCIDENCE_MAT, method="cosine", orientation="colwise")[9, 10] == C_MAT[9, 10] / sqrt(C_MAT[9, 9] * C_MAT[10, 10])
	@test similarity(INCIDENCE_MAT, method="inclusion", orientation="colwise")[2, 4] == C_MAT[2, 4] / min(C_MAT[2, 2], C_MAT[4, 4])
	@test similarity(INCIDENCE_MAT, method="inclusion", orientation="colwise")[9, 10] == C_MAT[9, 10] / min(C_MAT[9, 9], C_MAT[10, 10])
	@test similarity(INCIDENCE_MAT, method="jaccard", orientation="colwise")[2, 4] == C_MAT[2, 4] / (C_MAT[2, 2] + C_MAT[4, 4] - C_MAT[2, 4])
	@test similarity(INCIDENCE_MAT, method="jaccard", orientation="colwise")[9, 10] == C_MAT[9, 10] / (C_MAT[9, 9] + C_MAT[10, 10] - C_MAT[9, 10])
	@test similarity(INCIDENCE_MAT, method="proximity", orientation="colwise")[2, 4] == C_MAT[2, 4] / max(C_MAT[2, 2], C_MAT[4, 4])
	@test similarity(INCIDENCE_MAT, method="proximity", orientation="colwise")[9, 10] == C_MAT[9, 10] / max(C_MAT[9, 9], C_MAT[10, 10])

	# Orientation
	@test similarity(INCIDENCE_MAT, method="association", orientation="rowise") == similarity(INCIDENCE_MAT', method="association", orientation="colwise")
end

@testset "Density tests" begin

	# Define
	JACCARD_SIM = similarity(INCIDENCE_MAT, orientation="colwise", method = "jaccard")
	# Check inputs:
	@test_throws MethodError density(M = DataFrame(INCIDENCE_MAT), SIM = JACCARD_SIM) # INCIDENCE_MAT should be a matrix
	@test_throws MethodError density(M = INCIDENCE_MAT, SIM = DataFrame(JACCARD_SIM)) # SIM should be a matrix
	@test_throws AssertionError density(M = INCIDENCE_MAT, SIM = vcat(JACCARD_SIM, JACCARD_SIM)) # SIM should be square
	@test_throws AssertionError density(M = INCIDENCE_MAT, SIM = hcat(JACCARD_SIM, JACCARD_SIM[:, 1])) # number of cols in M should be == number of cols and rows in SIM

	# Check output:
	# Check that output is a matrix
	@test typeof(density(M = INCIDENCE_MAT, SIM = JACCARD_SIM)) == Array{Float64, 2}
	# Check that output is symmetrical
	@test size(density(M = INCIDENCE_MAT, SIM = JACCARD_SIM), 1) == size(density(M = INCIDENCE_MAT, SIM = JACCARD_SIM), 2)

	# Check output result:
	@test density(
		      M = INCIDENCE_MAT,
		      SIM = JACCARD_SIM
		      )[2, 5] == sum(INCIDENCE_MAT[2, :] .* JACCARD_SIM[5, :]) / sum(JACCARD_SIM[5,:])

	@test density(
		      M = INCIDENCE_MAT,
		      SIM = JACCARD_SIM
		      )[5, 9] == sum(INCIDENCE_MAT[5, :] .* JACCARD_SIM[9, :]) / sum(JACCARD_SIM[9, :])

end


using Random

# CVIT - Create itr and itst indeces for k-fold-cv
#
#    Description
#     [ITR,ITST]=CVITR(N,K) returns 1xK cell arrays ITR and ITST holding
#      cross-validation indeces for train and test sets respectively.
#      K-fold division is balanced with all sets having floor(N/K) or
#      ceil(N/K) elements.
#
#     [ITR,ITST]=CVITR(N,K,RS) with integer RS=true also makes random
#      permutation, using substream RS. This way different permutations
#      can be produced with different RS values, but same permutation is
#      obtained when called again with same RS. Function restores the
#      previous random stream before exiting.
#

# Copyright (c) 2010 Aki Vehtari

# This software is distributed under the GNU General Public
# License (version 2 or later); please refer to the file
# License.txt, included with the software, for details.

function cvit(n, k=10, rsubstream=false)

    a = k-rem(n,k)
    b = floor(Int, n/k);

    itst = Any[]
    itr = Any[]

    for cvi in 1:a
        push!(itst, collect(1:b) .+ (cvi-1) * b)
        push!(itr, setdiff(1:n,itst[cvi]))
    end
    for cvi in (a+1):k
        push!(itst, (a * b) + collect(1:(b + 1)) + (cvi - a - 1) * (b + 1))
        push!(itr, setdiff(1:n,itst[cvi]))
    end

    if rsubstream
        rng = MersenneTwister()
        rii = randperm(rng, n)
        for cvi in 1:k
            itst[cvi] = rii[itst[cvi]]
            itr[cvi] = rii[itr[cvi]]
        end
    end
    itr, itst
end