using Test
using SparseArrays

using DirectStiffnessMethod

# The textbook [hibbeler] still uses backward imperial units, so we have to do some weird
# convertions from time to time.

@testset "[hibbeler] Example 14.3 (truss)" begin
    nodes = [
        Node(1, 3.0, 4.0, FixedJoint)
        Node(2, 0.0, 0.0, FreeJoint)
        Node(3, 3.0, 0.0, FixedJoint)
    ]

    E = 1.0
    A = 1.0
    I = 1.0

    members = [
        Member(1, 2, 3, PinnedConnection, PinnedConnection, E, A, I)
        Member(2, 2, 1, PinnedConnection, PinnedConnection, E, A, I)
    ]

    s = Structure(nodes, members)

    loads = sparsevec(node_indices(2), [0.0, -2.0, 0.0], 3 * num_nodes(s))
    (d, f) = solve(s, loads)

    # Some values from the textbook are not accurate due to early rounding errors.
    @test round.(d, sigdigits=4) == [0.0, 0.0, 0.0, 4.5, -19.0, 0.0, 0.0, 0.0, 0.0]
    @test round.(f, sigdigits=3) == [1.5, 2.0, 0.0, 0.0, 0.0, 0.0, -1.5, 0.0, 0.0]
end

@testset "[hibbeler] Example 16.1 (frame)" begin
    nodes = [
        Node(1, 0.0, 0.0, SimpleJoint)
        Node(2, 20.0 * 12.0, 0.0, FreeJoint)
        Node(3, 20.0 * 12.0, -20.0 * 12.0, FixedJoint)
    ]

    E = 29.0e3
    A = 10.0
    I = 500.0

    members = [
        Member(1, 1, 2, FixedConnection, FixedConnection, E, A, I)
        Member(2, 2, 3, FixedConnection, FixedConnection, E, A, I)
    ]

    s = Structure(nodes, members)

    loads = sparsevec(node_indices(2), [5.0e3, 0.0, 0.0], 3 * num_nodes(s))
    (d, f) = solve(s, loads)

    # Some values from the textbook are not accurate due to early rounding errors.
    @test round.(d / 1000.0, sigdigits=4) ==
        [0.6958, 0.0, 1.234e-3, 0.6958, -1.551e-3, -2.488e-3, 0.0, 0.0, 0.0]
    @test round.(f / 1000.0, sigdigits=3) ==
        [0.0, -1.87, 0.0, 0.0, 0.0, 0.0, -5.0, 1.87, 750.0]
end

@testset "Truss v.s. Frame" begin
    nodes = [
        Node(1, 0.0, 0.0, FixedJoint)
        Node(2, 0.0, 1.0, FreeJoint)
        Node(3, 1.0, 1.0, FreeJoint)
        Node(4, 1.0, 0.0, FixedJoint)
    ]

    # Steel pipe, length = 1m, external diameter = 2cm, internal diameter = 1cm
    ro = 0.01
    ri = 0.005
    E = 200.0e9
    A = pi * (ro ^ 2 - ri ^ 2)
    I = pi / 4 * (ro ^ 4 - ri ^ 4)

    members = [
        Member(1, 2, 1, PinnedConnection, PinnedConnection, E, A, I)
        Member(2, 2, 3, PinnedConnection, PinnedConnection, E, A, I)
        Member(3, 3, 4, PinnedConnection, PinnedConnection, E, A, I)
        Member(4, 1, 4, PinnedConnection, PinnedConnection, E, A, I)
        Member(5, 2, 4, PinnedConnection, PinnedConnection, E, A, I)
        Member(6, 3, 1, PinnedConnection, PinnedConnection, E, A, I)
    ]

    s = Structure(nodes, members)

    loads = sparsevec(node_indices(2), [10.0, 10.0, 0.0], 3 * num_nodes(s))
    (d, f) = solve(s, loads)

    # Manually calculates the internal axial forces
    SQRT2INV = 1.0 / sqrt(2.0)
    iaf = E * A * [
        d[5], d[7] - d[4], d[8], 0.0,
        (-d[4] * SQRT2INV + d[5] * SQRT2INV) * SQRT2INV,
        (d[7] * SQRT2INV + d[8] * SQRT2INV) * SQRT2INV,
    ]

    @test round.(iaf, sigdigits=3) == [14.4, -5.58, -5.58, 0.0, -6.25, 7.89]

    # Uses fixed connections to construct a frame.
    members_frame = [
        Member(1, 2, 1, FixedConnection, FixedConnection, E, A, I)
        Member(2, 2, 3, FixedConnection, FixedConnection, E, A, I)
        Member(3, 3, 4, FixedConnection, FixedConnection, E, A, I)
        Member(4, 1, 4, FixedConnection, FixedConnection, E, A, I)
        Member(5, 2, 4, FixedConnection, FixedConnection, E, A, I)
        Member(6, 3, 1, FixedConnection, FixedConnection, E, A, I)
    ]

    s_frame = Structure(nodes, members_frame)

    (d_frame, f_frame) = solve(s_frame, loads)

    iaf_frame = begin
        d = d_frame
        E * A * [
            d[5], d[7] - d[4], d[8], 0.0,
            (-d[4] * SQRT2INV + d[5] * SQRT2INV) * SQRT2INV,
            (d[7] * SQRT2INV + d[8] * SQRT2INV) * SQRT2INV,
        ]
    end

    # This shows the relative error when modeling a frame as truss is ignorable in real life.
    rel_err = abs.((iaf - iaf_frame) ./ (iaf_frame .+ eps()))
    @test all(rel_err .< 0.01)
    # (Actually, the maximum relative error is about 0.001 in this setting.)
end

@testset "Portal frame" begin
    offset = 0.5
    # Span = 2m, height = 1m
    nodes = [
        Node(1, 0.0, 0.0, FixedJoint)
        Node(2, 0.0, 1.0, FreeJoint)
        Node(3, 1.0 - offset, 1.0, FreeJoint)
        Node(4, 2.0, 1.0, FreeJoint)
        Node(5, 2.0, 0.0, FixedJoint)
    ]

    # Steel pipe, length = 1m, external diameter = 2cm, internal diameter = 1cm
    ro = 0.01
    ri = 0.005
    E = 200.0e9
    A = pi * (ro ^ 2 - ri ^ 2)  # doesn't matter for frame (matters for truss)
    I = pi / 4 * (ro ^ 4 - ri ^ 4)

    members = [
        Member(1, 1, 2, FixedConnection, FixedConnection, E, A, I)
        Member(2, 2, 3, FixedConnection, FixedConnection, E, A, I)
        Member(3, 3, 4, FixedConnection, FixedConnection, E, A, I)
        Member(4, 4, 5, FixedConnection, FixedConnection, E, A, I)
    ]

    s = Structure(nodes, members)

    loads = sparsevec(node_indices(3), [0.0, -100.0, 0.0], 3 * num_nodes(s))
    (d, f) = solve(s, loads)

    d_x = d[1:3:end]
    # When putting a point load on the left half of the horizontal beam, counterintuitively, the
    # frame leans to the left instead of right.
    #
    # Adjust the relative values of E/I and the frame may have negative-X deformation. And in the
    # case of large deformation/flexible beams, the common intuition is right too.
    @test all(d_x .>= 0.0)

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

    "Offset of the point load" v.s. "The displacement of the horizontal beam":

    The offset ∈ [0.0, 1.0), and `d = (d[4] + d[10]) / 2` is the average horizontal displacement of
    the beam. It can be seen from the table below that the horizontal displacement reaches its
    maximum when the load placed at approximately 0.6m from the axis of symmetry.

        | offset | d * 1e4 |
        |--------|---------|
        |  0.0   |  0.0    |
        |  0.05  |  1.058  |
        |  0.1   |  2.1    |
        |  0.15  |  3.111  |
        |  0.2   |  4.073  |
        |  0.25  |  4.972  |
        |  0.3   |  5.792  |
        |  0.35  |  6.516  |
        |  0.4   |  7.128  |
        |  0.45  |  7.614  |
        |  0.5   |  7.956  |
        |  0.55  |  8.138  |
        |  0.6   |  8.146  |
        |  0.65  |  7.963  |
        |  0.7   |  7.573  |
        |  0.75  |  6.96   |
        |  0.8   |  6.108  |
        |  0.85  |  5.002  |
        |  0.9   |  3.625  |
        |  0.95  |  1.962  |

    ================#
end


module Metrics

using StaticArrays

import ..GPs: sethyperparameters!, gethyperparameters, hyperparametercount

"""
An abstract metric type for those kernels that support the notion of a metric,
e.g., squared exponential kernels. Every child type should implement the
following functions.

    `sqrdistance(metric, x, y)`
    `sethyperparameters!(metric, θ)`
    `gethyperparameters(metric)`
    `hyperparametercount(metric)`
"""
abstract type Metric{T <: Number} end

export Metric

"""
`sqrdistance(metric, x1, x2)` returns the squared distance between `(x1, x2)`
according to the specified metric.
"""
function sqrdistance end

export sqrdistance

"""
`sqrdistance_dθ(metric, x1, x2)` returns the derivative w.r.t. the
hyperparameters of the squared distance between `(x1, x2)` according to the
specified metric.
"""
function sqrdistance_dθ end

export sqrdistance_dθ

"""
`sqrdistance_dx1(metric, x1, x2)` returns the derivative w.r.t. x1 of the
squared distance between `(x1, x2)` according to the specified metric.
"""
function sqrdistance_dx1 end

export sqrdistance_dx1

"""
`sqrdistance_dx2(metric, x1, x2)` returns the derivative w.r.t. x2 of the
squared distance between `(x1, x2)` according to the specified metric.
"""
function sqrdistance_dx2 end

export sqrdistance_dx2

#--- Standard Euclidean metric

"""
Euclidean{N, T}

The standard Euclidean metric in `N` dimensions. (Note that `N=1` corresponds to
an isotropic metric when the input data is multi-dimensional.)
"""
struct Euclidean{N, T} <: Metric{T}
    ℓ::MVector{N, T}
end
Euclidean(ℓ::Union{SVector, Number}) = Euclidean(MVector(ℓ))

export Euclidean

sqrdistance(metric::Euclidean, x1, x2) = sum(((x1 .- x2)./metric.ℓ).^2)

function sqrdistance_dθ(metric::Euclidean{1, T} where {T}, x1, x2)
    return SVector(-2*sum((x1 .- x2).^2)/metric.ℓ[1]^3)
end

function sqrdistance_dθ(metric::Euclidean{N, T} where {N, T}, x1, x2)
    return -2*(x1 .- x2).^2./metric.ℓ.^3
end

function sqrdistance_dx1(metric::Euclidean, x1, x2)
    return 2*(x1 .- x2)./metric.ℓ.^2
end

function sqrdistance_dx2(metric::Euclidean, x1, x2)
    return -2*(x1 .- x2)./metric.ℓ.^2
end

sethyperparameters!(metric::Euclidean, θ) = (metric.ℓ .= θ)
gethyperparameters(metric::Euclidean) = metric.ℓ
hyperparametercount(metric::Euclidean) = length(metric.ℓ)

end


export lsomp
# Based on the algorithm developed by
# Baptista, R., Stolbunov, V., & Nair, P. B. (2019).
# Some greedy algorithms for sparse polynomial chaos expansions.
# Journal of Computational Physics, 387, 303-325.

# Solve ψc = u using the orthogonal matching pursuit with least-square

function lsomp(ψ::AbstractMatrix{T}, u::AbstractVector{T}; invert::Bool=true, verbose::Bool = true, ϵrel::Float64 = 1e-1, maxterms::Int64=typemax(Int64)) where {T}
    m, n = size(ψ)

    residue = copy(u)
    idxset = Int64[]
    dict = collect(1:n)
    ϵu = norm(u)
    ϵhist = [ϵu]
    ϵrel *= ϵu
    ϕ = zeros(T, n)

    for j=1:n
        for i=1:m
        ϕ[j] += ψ[i,j]^2
        end
    end

    ek = zeros(T, n)
    q = zeros(T, m)
    new_idx = 1

    @inbounds for k=1:n
        # Step 6 & 7 update denominators and choose best candidate

        entry = 0.0
        for (j, idx) in enumerate(dict)
            ψj = view(ψ,:,idx)

            ratio = residue'*ψj
            ratio = ratio^2
            ratio /= ϕ[idx]

            if ratio > entry
                entry = copy(ratio)
                new_idx = copy(idx)
            end
        end

        # Step 8 Update set of selected basis
        push!(idxset, new_idx)

        # Step 9 Update candidate dictionary
        filter!(x-> x != new_idx, dict)
        # setdiff!(dict, new_idx)

        # Step 10 Compute residual
        c = view(ψ,:,idxset)\u
        residue = u - view(ψ,:,idxset)*c

        # Step 12 Calculate stopping critera
        ϵ = norm(residue)
        if verbose == true
            push!(ϵhist, copy(ϵ))
        end
        if ϵ < ϵrel || k == maxterms
            break
        end
    end

    # Solve the system with the set of indices
    if verbose == true
        if invert == true
            c = zeros(length(idxset))
            c .= view(ψ,:, idxset)\u
            return idxset, c, ϵhist
        else
            return idxset, ϵhist
        end
    else
        if invert == true
            c = zeros(length(idxset))
            c .= view(ψ,:, idxset)\u
            return idxset, c
        else
            return idxset
        end
    end
end


"""
    add_noise(S,0.05)
    add_noise(-||-, k=0.15)

A function that adds noise to a signal in a consistent way. The function adds both background noise as well as
noise to the signal itself. The level of the background noise and the signal noise is determined by the b
and k variables respectively. The levels are input as a fraction of the mean of the original signal S strength.
"""
function add_noise(s, b; k=0.1)
    sn = max.(s,0.0) .+ k.*(mean(sqrt.(abs.(s)))).*rand.(Normal.(0.0, max.(sqrt.(max.(s,0)), sqrt.(b))))
    err = k.*mean(sqrt.(abs.(s))).*max.(sqrt.(max.(s,0)), sqrt.(b))
    return sn, err
end

"""
    get_orbel_volume(myOrbitGrid, os_equidistant)

Get the orbit element volume of this specific orbit-grid.
If os_equidistant=false, then return a 3D array with all the orbit volumes.
Otherwise, just a scalar.
"""
function get_orbel_volume(og::OrbitGrid, os_equidistant::Bool)

    if os_equidistant
        dRm = abs(og.r[2]-og.r[1]) # The difference between the first and second element should be representative for the whole grid, due to equidistancy
        dE = abs(og.energy[2]-og.energy[1])
        dpm = abs(og.pitch[2]-og.pitch[1])
        dO = dE*dpm*dRm
        return dO
    else
        dO = zeros(length(og.energy), length(og.pitch), length(og.r)) # If not equidistant, create a 3D array with zeros
        for Ei=1:length(og.energy) # For all energies...
            if Ei==length(og.energy) # If at the edge of the orbit grid...
                # Assume edge orbit-element volume to be same as next-to-edge
                dEi = abs(og.energy[end]-og.energy[end-1])
            else
                dEi = abs(og.energy[Ei+1]-og.energy[Ei])
            end
            for pmi=1:length(og.pitch)
                if pmi==length(og.pitch)
                    dpmi = abs(og.pitch[end]-og.pitch[end-1])
                else
                    dpmi = abs(og.pitch[pmi+1]-og.pitch[pmi])
                end
                for Rmi=1:length(og.r)
                    if Rmi==length(og.r)
                        dRmi = abs(og.r[end]-og.r[end-1])
                    else
                        dRmi = abs(og.r[Rmi+1]-og.r[Rmi])
                    end

                    dO[Ei, pmi, Rmi] = dEi*dpmi*dRmi # Replace the zero-element in the 3D array with the resulting orbit-element volume
                end
            end
        end

        return dO
    end
end


const ln2pi = log(2*pi);

"""
    norm2(x::Array{Float64,1})

Computes ||x||^2 more precisely than norm(x)^2.
"""
function norm2(x::Array{Float64,1})
    return sum(x.^2)
end

"""
    norm2(X::Array{Float64,2})

Computes the square of the Frobenius matrix norm ||X||^2_FRO more precisely than vecnorm(A,2)^2.
"""
function norm2(X::Array{Float64,2})
    return sum(X.^2)
end

"""
   delta(new, old)

Computes delta for convergence control for a given variable.
"""
function delta(new, old)
    return norm(old - new)/norm(old)
end

"""
    scaleY(Y::Array{Float64,2})

Scales down a 2 dimensional array so it has approx. standard normal distribution.
"""
function scaleY(Y::Array{Float64,2})
    L, M = size(Y)
    mu = mean(Y,2);
    sigma = var(Y,2);

    # if there are NaN, then sigma is zero for a given row ->
    # the scaled down row is also zero
    # but we treat this more economically by setting the denominator for a given row to one
    # also, deal with numerical zeroes
    den = sigma
    den[abs(den) .<= 1e-15] = 1.0
    den[den .== 0.0] = 1.0
    den = repmat(sqrt(den),1, M)
    nom = Y - repmat(mu, 1, M)
    nom[abs(nom) .<= 1e-8] = 0.0
    Y = nom./den
    return Y
end

"""
    scaleY(Y::Array{Float32,2})

Scales down a 2 dimensional array so it has approx. standard normal distribution.
"""
function scaleY(Y::Array{Float32,2})
    Y64 = convert(Array{Float64,2}, Y)
    Y64 = scaleY(Y64)
    return Y64
    #return convert(Array{Float32,2}, Y64)
end

"""
    preprocess(Y::Array{Float64,2}, lambda::Float64; verb = false)

Scales Y to be normal, deletes unnecessary rows and then multiplies it by lambda
for better algorithm convergence.
"""
function preprocess(Y::Array{Float64,2}, lambda::Float64; verb = false)
    sY = scaleY(Y)
    rowsums = sum(abs(sY), 2)
    L, M = size(Y)
    used_rows = 1:L
    used_rows = used_rows[rowsums .>= 1e-5]
    # there are either zeroes or the variance is so small that the row carries no information

    if verb
        println("Original problem size: $L rows, $(L-size(used_rows)[1]) rows not relevant and are not used.")
    end

    return lambda*sY[used_rows,:]
end

"""
    preprocess(Y::Array{Float32,2}, lambda::Float64; verb = false)

Scales Y to be normal, deletes unnecessary rows and then multiplies it by lambda
for better algorithm convergence.
"""
function preprocess(Y::Array{Float32,2}, lambda::Float64; verb = false)
    Y64 = convert(Array{Float64,2}, Y)
    return preprocess(Y64, lambda, verb = verb)
end

"""
    traceXTY(X, Y)

Computes tr(X^TY) more effectively than trace(X'*Y).
"""
function traceXTY(X::Array{Float64,2}, Y::Array{Float64,2})
    return(sum(X.*Y))
end

"""
    normalEntropy(Sigma)

Computes the entropy of a multivariate normal distribution with covariance Sigma.
"""
function normalEntropy(Sigma::Array{Float64,2})
    m, n = size(Sigma)
    if n != m
        error("Sigma must be square!")
    end
    d = det(Sigma)
    # numerical issues for too small determinants
    if d < eps(0.0) # smallest nonzero float
        d = eps(0.0)
    end

    return m/2 + m/2*ln2pi + 1/2*log(d)
end

"""
    normalEntropy(diagSigma)

Computes the entropy of a multivariate normal distribution with
diagonal of covariance diagSigma.
"""
function normalEntropy(diagSigma::Array{Float64,1})
    n = size(diagSigma)[1]

    return n/2 + n/2*ln2pi + 1/2*sum(log(diagSigma))
end

"""
    gammaEntropy(a, b)

Computes the entropy of a gamma distribution with shape a and rate b.
"""
function gammaEntropy(a::Float64, b::Float64)
    return a + log(b) + lgamma(a) + (1-a)*digamma(a)
end

"""
    gammaELn(a, b)

Computes E[ln(x)] where x follows gamma distribution with shape a and rate b.
"""
function gammaELn(a::Float64, b::Float64)
    return digamma(a) - log(b)
end


using Markdown
using InteractiveUtils

using DrWatson

begin
  @quickactivate "SR2TuringPluto"
  using Turing
  using StatisticalRethinking
end

md"## Clip-05-01-27t.jl"

begin
  df = CSV.read(sr_datadir("WaffleDivorce.csv"), DataFrame)
  df.D = zscore(df.Divorce)
  df.M = zscore(df.Marriage)
  df.A = zscore(df.MedianAgeMarriage)
end

std(df.MedianAgeMarriage)

@model function m5_1_A(A, D)
    a ~ Normal(0, 0.2)
    bA ~ Normal(0, 0.5)
    σ ~ Exponential(1)
    μ = lin(a, A, bA)
    D ~ MvNormal(μ, σ)
end

begin
    m5_1_At = m5_1_A(df.A, df.D)
  quap5_1_At = quap(m5_1_At)
  post5_1_At = DataFrame(rand(quap5_1_At.distr, 1000)', quap5_1_At.params)

  A_seq = range(-3, 3.2, length = 30)
  mu5_1_At = lin(post5_1_At.a', A_seq, post5_1_At.bA') |> meanlowerupper

  scatter(df.A, df.D, alpha = 0.4, legend = false)
  plot!(A_seq, mu5_1_At.mean, ribbon =
        (mu5_1_At.mean .- mu5_1_At.lower, mu5_1_At.upper .- mu5_1_At.mean))
end

@model function m5_1_M(M, D)
    a ~ Normal(0, 0.2)
    bM ~ Normal(0, 0.5)
    σ ~ Exponential(1)
    μ = lin(a, M, bM)
    D ~ MvNormal(μ, σ)
end

begin
    m5_1_Mt = m5_1_M(df.M, df.D)
    quap5_1_Mt = quap(m5_1_Mt)
    post5_1_Mt = DataFrame(rand(quap5_1_Mt.distr, 1000)', quap5_1_Mt.params)

    M_seq = range(-3, 3.2, length = 30)
    mu5_1_Mt = lin(post5_1_Mt.a', M_seq, post5_1_Mt.bM') |> meanlowerupper

    scatter(df.M, df.D, alpha = 0.4, legend = false)
    plot!(M_seq, mu5_1_Mt.mean, ribbon =
        (mu5_1_Mt.mean .- mu5_1_Mt.lower, mu5_1_Mt.upper .- mu5_1_Mt.mean))
end

@model function m5_1_A_M(A, M, D)
    a ~ Normal(0, 0.2)
    bM ~ Normal(0, 0.5)
    bA ~ Normal(0, 0.5)
    σ ~ Exponential(1)
    μ = lin(a, A, bA, M, bM)
    D ~ MvNormal(μ, σ)
end

begin
    m5_1_A_Mt = m5_1_A_M(df.A, df.M, df.D)
    quap5_1_A_Mt = quap(m5_1_A_Mt)
    post5_1_A_Mt = DataFrame(rand(quap5_1_A_Mt.distr, 1000)', quap5_1_A_Mt.params)

    A_M_seq = range(-3, 3.2, length = 30)
    mu5_1_A_Mt = lin(post5_1_A_Mt.a', A_M_seq, post5_1_A_Mt.bM') |> meanlowerupper

    scatter(df.M, df.D, alpha = 0.4, legend = false)
    plot!(M_seq, mu5_1_A_Mt.mean, ribbon =
        (mu5_1_A_Mt.mean .- mu5_1_A_Mt.lower, mu5_1_A_Mt.upper .- mu5_1_A_Mt.mean))
end

function plotcoef1(
  models::Vector{T},
  mnames::Vector{String},
  pars::Vector{Symbol};
  fig="", title="", func=nothing,
  sampler=NUTS(0.65), nsamples=2000, nchains=4) where {T <: DynamicPPL.Model}

  # mnames = [nameof(m) for m in models]
  levels = length(models) * (length(pars) + 1)
  colors = [:blue, :red, :green, :darkred, :black]

  s = Vector{NamedTuple}(undef, length(models))
  for (mindx, mdl) in enumerate(models)
    if isnothing(func)
        chns = mapreduce(c -> sample(mdl, sampler, nsamples),
            chainscat, 1:nchains)
        df = DataFrame(Array(chns), names(chns, [:parameters]))
        m, l, u = estimparam(df)
        d = Dict{Symbol, NamedTuple}()
        for (indx, par) in enumerate(names(chns, [:parameters]))
            d[par] = (mean=m[indx], lower=l[indx], upper=u[indx])
        end
        s[mindx] =   (; d...)
    else
        quap_mdl = quap(mdl)
        post = rand(quap_mdl.distr, 10_000)
        df = DataFrame(post', [keys(quap_mdl.coef)...])
        m, l, u = estimparam(df)
        d = Dict{Symbol, NamedTuple}()
        for (indx, par) in enumerate([keys(quap_mdl.coef)...])
            d[par] = (mean=m[indx], lower=l[indx], upper=u[indx])
        end
        s[mindx] =   (; d...)
    end
  end

  xmin = 0; xmax = 0.0
  for i in 1:length(s)
    for par in pars
      syms = Symbol.(keys(s[i]))
      if Symbol(par) in syms
        mp = s[i][par].mean
        xmin = min(xmin, s[i][par].lower)
        xmax = max(xmax, s[i][par].upper)
      end
    end
  end

  ylabs = String[]
  for j in 1:length(models)
    for i in 1:length(pars)
      l = length(String(pars[i]))
      str = repeat(" ", 9-l) * String(pars[i])
      append!(ylabs, [str])
    end
    l = length(mnames[j])
    str = mnames[j] * repeat(" ", 9-l)
    append!(ylabs, [str])
  end

  ys = [string(ylabs[i]) for i = 1:length(ylabs)]
  p = plot(xlims=(xmin, xmax), leg=false, framestyle=:grid)
  title!(title)
  yran = range(1, stop=length(ylabs), length=length(ys))
  yticks!(yran, ys)

  line = 0
  for (mindx, model) in enumerate(models)
    line += 1
    hline!([line] .+ length(pars), color=:darkgrey, line=(2, :dash))
    for (pindx, par) in enumerate(pars)
      line += 1
      syms = Symbol.(keys(s[mindx]))
      if Symbol(par) in syms
        ypos = (line - 1)
        mp = s[mindx][Symbol(par)].mean
        lower = s[mindx][Symbol(par)].lower
        upper = s[mindx][Symbol(par)].upper
        plot!([lower, upper], [ypos, ypos], leg=false, color=colors[pindx])
        scatter!([mp], [ypos], color=colors[pindx])
      end
    end
  end
  if length(fig) > 0
    savefig(p, fig)
  end
  (s, p)
end

s1, p1 = plotcoef(
    [m5_1_At, m5_1_Mt, m5_1_A_Mt],
    [:a, :bM, :bA, :σ]
)

s1 |> display

plot(p1)

gui()

s2, p2 = plotcoef(
    [m5_1_At, m5_1_Mt, m5_1_A_Mt],
    [:a, :bM, :bA, :σ];
    func=quap
)

s2 |> display

plot(p2)

gui()

s3, p3 = plotcoef(
    [m5_1_At],
    [:a, :bM, :bA, :σ]
)

s3 |> display

plot(p3)

gui()

s4, p4 = plotcoef(
    m5_1_At,
    [:a, :bM, :bA, :σ];
    func=quap
)

s4 |> display

plot(p4)

gui()


"""
测试重新切分
"""

using Test
using MARY_fRG.Refine, MARY_fRG.Basics, MARY_fRG.Triangulated
using MARY_fRG.Drawers

@testset "三角形的切分" begin
    rtri = RtTriangle(Point2D(0., 0.), Point2D(1., 1.), Point2D(1., -1.))
    newtris = split_triangle(rtri)
    println(newtris)
    plt = draw_polygon(newtris)
    @savepng plt "newtris"
    eqtri = EqTriangle(Point2D(0., 0.), Point2D(1., sqrt(3)), Point2D(2., 0.))
    newtris = split_triangle(eqtri)
    println(newtris)
    plt = draw_polygon(newtris)
    @savepng plt "newtris2"
end

@testset "正方形的细化" begin
    squ1 = Square(Point2D(0., 0.), 1)
    ltris1, ladjs1 = split_square(squ1, 3)
    plt = draw_polygon(ltris1)
    @savepng plt "refine1"
    ltris2 = refine_triangles(ltris1)
    plt = draw_polygon(ltris2)
    @savepng plt "refine2"
end

@testset "六边型的细化" begin
    hexa1 = EqHexagon(Point2D(0., 0.), 1)
    ltris1, ladjs1 = split_hexagon(hexa1, 4)
    plt = draw_polygon(ltris1)
    @savepng plt "refine3"
    ltris2 = refine_triangles(ltris1)
    plt = draw_polygon(ltris2)
    @savepng plt "refine4"
end

@testset "寻找最近邻" begin
    #正方格子
    squ1 = Square(Point2D(0., 0.), 1.)
    ltris, ladjs = split_square(squ1, 11)
    ladjs2 = find_adjs_by_adjoint(ltris)
    save_triangulated(squ1, ltris, ladjs, "squ1")
    save_triangulated(squ1, ltris, ladjs2, "squ2")
    fl1 = open("squ1.adj", "r")
    fl2 = open("squ2.adj", "r")
    @test read(fl1) == read(fl2)
    #六角格子
    hexa1 = EqHexagon(Point2D(0., 0.), 1.)
    ltris, ladjs = split_hexagon(hexa1, 13)
    ladjs2 = find_adjs_by_adjoint(ltris)
    save_triangulated(squ1, ltris, ladjs, "hexa1")
    save_triangulated(squ1, ltris, ladjs2, "hexa2")
    fl1 = open("hexa1.adj", "r")
    fl2 = open("hexa2.adj", "r")
    @test read(fl1) == read(fl2)
    rm("hexa1.adj")
    rm("hexa2.adj")
    rm("hexa1.ply")
    rm("hexa2.ply")
    rm("hexa1.tri")
    rm("hexa2.tri")
    rm("squ1.adj")
    rm("squ2.adj")
    rm("squ1.ply")
    rm("squ2.ply")
    rm("squ1.tri")
    rm("squ2.tri")
end


struct ADMMOptions
  ρ::Float64
  α::Float64
  maxiter::Int64
  abstol::Float64
  reltol::Float64
end

ADMMOptions(;ρ::Float64=1.,
           α::Float64=1.,
           maxiter::Int64=1000,
           abstol::Float64=1e-4,
           reltol::Float64=1e-2) = ADMMOptions(ρ, α, maxiter, abstol, reltol)

#########################################################
# minimize  trace(S*X) - log det X + lambda*||X||_1
#########################################################

function covsel!(
  X::StridedMatrix{T},
  Z::StridedMatrix{T},
  U::StridedMatrix{T},
  Σ::StridedMatrix{T},
  λ::T;
  options::ADMMOptions = ADMMOptions(),
  penalize_diag::Bool=true
  ) where {T<:AbstractFloat}

  maxiter = options.maxiter
  ρ = options.ρ
  α = options.α
  abstol = options.abstol
  reltol = options.reltol
  γ = one(T) / ρ
  λρ = λ * γ

  p = size(Σ, 1)
  g = ProxGaussLikelihood(Symmetric(Σ))

  tmpStorage = zeros(T, (p, p))
  Zold = copy(Z)

  for iter=1:maxiter
    # x-update
    @. tmpStorage = Z - U
    prox!(g, X, tmpStorage, γ)

    # z-update with relaxation
    copyto!(Zold, Z)
    @. tmpStorage = α*X + (one(T)-α)*Z + U
    @inbounds for c=1:p
      for r=1:c-1
        t = shrink(tmpStorage[r, c], λρ)
        Z[c, r] = t
        Z[r, c] = t
      end
      Z[c, c] = penalize_diag ? shrink(tmpStorage[c, c], λρ) : tmpStorage[c, c]
    end

    # u-update
    @. U = tmpStorage - Z

    # check convergence
    r_norm = norm_diff(X, Z, 2)
    s_norm = norm_diff(Z, Zold, 2) * sqrt(ρ)
    eps_pri = p*abstol + reltol * max( norm(X), norm(Z) )
    eps_dual = p*abstol + reltol * ρ * norm(U)
    if r_norm < eps_pri && s_norm < eps_dual
      break
    end
  end
  Z
end

function covselpath(S::StridedMatrix{T},
                    λarr;
                    options::ADMMOptions = ADMMOptions(),
                    penalize_diag::Bool=true,
                    verbose::Bool=false) where {T<:AbstractFloat}
  p = size(S, 1)

  solutionpath = Array(Array{Float64, 2}, length(λarr))
  X = zeros(p, p)
  Z = zeros(p, p)
  U = zeros(p, p)

  for i=1:length(λarr)
    if verbose
      @printf("lambda = %d/%d\n", i, length(λarr))
    end
    covsel!(X, Z, U, S, λarr[i]; options=options, penalize_diag=penalize_diag)
    solutionpath[i] = copy(Z)
  end
  solutionpath
end


export em_step!, em_run!, mixfit, densfit

"""
    em_step!(est::MixModel, x::Vector{<:Real})

Run a single EM step.

See also: [`em_run!`](@ref), [`mixfit`](@ref)
"""
function em_step!(est::MixModel, x::Vector{<:Real})
    p = Matrix{Float32}(undef, length(est.μ), length(x))
    for i = 1:length(est.μ)
        p[i,:] = est.α[i] * dnorm.(x, est.μ[i], est.σ[i]) ./
            sum(est.α[j] * dnorm.(x, est.μ[j], est.σ[j]) for j=1:length(est.μ))
    end
    for i = 1:length(est.μ)
        est.μ[i] = sum(p[i,:] .* x) / sum(p[i,:])
        est.σ[i] = √(sum(p[i,:] .* (x .- est.μ[i,:]).^2) / sum(p[i,:]))
        est.α[i] = mean(p[i,:])
    end
end

"""
    em_run!(est::MixModel, x::Vector{<:Real}, rtol::AbstractFloat = 0.00001, maxiter::Int = 0)

Run EM steps until the percent increase in log-likelihood is below `rtol` or the number of iterations is greater than `maxiter`.
If `maxiter` is set to zero, EM steps will continue until the increase is below `rtol` regardless of the number of iterations.

See also: [`em_run!`](@ref), [`mixfit`](@ref)
"""
function em_run!(est::MixModel, x::Vector{<:Real}, rtol::AbstractFloat = 0.00001, maxiter::Int = 0)
    ll::Float32 = LL(x, est)
    oll::Float32 = -ll
    i::Int = 1
    while true
        em_step!(est, x)
        oll = ll
        ll = LL(x, est)
        if abs((oll - ll) / oll) < rtol
            break
        end
        if i == maxiter
            @warn "Hit maximum number of iterations"
        end
        i += 1
    end
end

"""
    mixfit(x::Vector{<:Real},
            m::Int;
            rtol::AbstractFloat = 0.00001,
            α::Vector{<:Real} = fill(1/m, m),
            μ::Vector{<:Real} = quantile!(x, (1:m)/m),
            σ::Vector{<:Real} = fill(std(x) / √(m), m),
            maxswap::Int = 5 * m^2,
            maxiter_inner::Int = 0,
            maxiter::Int = 0,
            silent::Bool = false,
            kernel::Function = dnorm)

Get the maximum likelihood estimate of an `m`-component mixture model
using random-swap EM.  Random-swap EM avoids local optimums by randomly
replacing components and using the result with the maximum likelihood [^1].
The component distributions are given by `kernel`.  The maximum number of swaps
is given by `maxswap`.  If results vary across runs, then `maxswap` is too low -
the default is ``5m^2``.  maxiter_inner is the maximum number of iterations for
the estimates that go through the swapping process, and maxiter is the maximum
for the final estimate.  Similarly, `rtol` is for the final estimate, while the
inner estimates use 0.1.  To supress output, simply set `silent` to true.
Starting values can be provided via the α, μ, and σ arguments, but this
shouldn't be necassary due to the use of random swapping.

[^1]: Zhao, Q., Hautamäki, V., Kärkkäinen, I., & Fränti, P. (2012). Random swap EM algorithm for Gaussian mixture models. Pattern Recognition Letters, 33(16), 2120-2126.

See also: [`densfit`](@ref), [`em_run!`](@ref)
"""
function mixfit(x::Vector{<:Real},
                m::Int;
                rtol::AbstractFloat = 0.00001,
                t::Int = 0,
                α::Union{Vector{<:Real}, Nothing} = nothing,
                μ::Union{Vector{<:Real}, Nothing} = nothing,
                σ::Union{Vector{<:Real}, Nothing} = nothing,
                maxswap::Int = 0,
                maxiter_inner::Int = 0,
                maxiter::Int = 0,
                silent::Bool = false,
                kernel::Function = dnorm)
    if maxswap == 0
        maxswap = 5 * m^2
    end

    est = MixModel(fill(1/m, m), quantile!(x, (1:m)/m), fill(std(x) / √(m), m), kernel)

    if μ != nothing
        est.μ = μ
    end
    if σ != nothing
        est.σ = σ
    end
    if α != nothing
        est.α = α
    end

    em_run!(est, x, 0.1, maxiter_inner)
    for i ∈ 1:maxswap
        est_s = copy(est)
        choice = rand(1:m)
        est_s.μ[choice] = rand(x)
        est_s.α = fill(1/m, m)
        est_s.σ[choice] = std(x) / √(m)
        em_run!(est_s, x, 0.1, maxiter_inner)
        if LL(x, est_s) > LL(x, est)
            est = est_s
        end
    end
    em_run!(est, x, rtol, maxiter)

    if silent == false
        describe(est, data = x)
    end
    return est
end

"""
    densfit(x::Vector{<:Real};
            wait::Int = 3,
            rtol_em::AbstractFloat = 0.00001,
            criterion::Function = AIC,
            silent::Bool = false,
            maxiter::Int = 0,
            maxiter_inner::Int = 0,
            maxswap::Int = 0,
            kernel::Function = dnorm)

Estimate the density of `x` by a mixture model.  Successive mixture model estimates
are done via random swap EM with increasing number of clusters until `criterion`
decreases for 3 iterations.  By default, `criterion` uses AIC, which performs
well for density estimation [^2].  However, AIC2 should be used instead if one
wants to actually estimate the number of clusters in a true mixture model.
The relative tolerance for EM convergence is given by `rtol_em`.  To disable output, set `silent` to true.

[^2]: Wang, Y., & Chee, C. S. (2012). Density estimation using non-parametric and semi-parametric mixtures. Statistical Modelling, 12(1), 67-92.

See also: [`densfit`](@ref), [`em_run!`](@ref)
"""
function densfit(x::Vector{<:Real};
                wait::Int = 3,
                rtol_em::AbstractFloat = 0.00001,
                criterion::Function = AIC,
                silent::Bool = false,
                maxiter::Int = 0,
                maxiter_inner::Int = 0,
                maxswap::Int = 0,
                kernel::Function = dnorm)
    new::MixModel = mixfit(x, 1, silent = true)
    current::MixModel = new
    m::Int = 2
    nbad::Int = 0

    while nbad < wait
        new = mixfit(x, m,
                    silent = true,
                    maxiter = maxiter,
                    maxiter_inner = maxiter_inner,
                    maxswap = maxswap,
                    kernel = kernel)

        if (criterion(x, new) >= criterion(x, current))
            nbad += 1
        else
            nbad = 0
            current = new
        end

        m += 1
    end

    if silent == false
        describe(current, data = x)
    end

    return current
end


# Let's try parameter estimation for an ODE, here the famous Lorentz attractor.
# We base our implementation on the code from Paulo Marques available at:
#   https://github.com/pjpmarques/Julia-Modeling-the-World/blob/master/Lorenz%20Attractor.ipynb

function lorentz_ode!(dr::AbstractVector{Float64},
                      params::AbstractVector{Float64},
                      r::AbstractVector{Float64})
    @assert length(dr) == length(r) == 3

    # Get individual params for this ODE
    @inbounds sigma, rho, beta = params

    # Extract the coordinates from the r vector
    @inbounds x, y, z = r

    # The Lorenz equations, put the derivative into dr
    @inbounds dr[1] = sigma*(y - x)       # dx/dt
    @inbounds dr[2] = x*(rho - z) - y     # dy/dt
    @inbounds dr[3] = x*y - beta*z        # dz/dt

    return dr
end

# Define time vector and interval grid.
t = range(0.0, step=0.001, stop=100.0)
# But in order to compare to [Xiang2015] paper we use their choice instead:
t_Xiang2015 = range(0.0, step=0.01, length=301)

# Initial position in space
r0 = [0.1, 0.0, 0.0]

# Constants sigma, rho and beta. In a real ODE problem these would not be known
# and would be estimated.
sigma = 10.0
rho   = 28.0
beta  = 8.0/3.0
real_params = [sigma, rho, beta]

# "Play" an ODE from a starting point and into the future given a sequence of time steps.
# Put the restults into `states`
function calc_states!(states::AbstractMatrix{Float64},
    params::AbstractVector{Float64}, odefunc!::Function,
    startx::AbstractVector{Float64}, times::AbstractVector{Float64})

    numsamples = length(times)
    @assert size(states) == (length(startx), numsamples)

    states[:, 1] = startx
    tprev = times[1]
    deriv = similar(startx)
    for i in 2:numsamples
        tnow = times[i]
        stateprev = view(states, :, i-1)
        odefunc!(deriv, params, stateprev)
        @inbounds states[:, i] .= stateprev .+ deriv .* (tnow - tprev)
        tprev = tnow
    end

    return states
end

calc_states(params::AbstractVector{Float64}, odefunc!::Function,
    startx::AbstractVector{Float64}, times::AbstractVector{Float64}) =
    calc_states!(Matrix{Float64}(undef, length(startx), length(times)),
                 params, odefunc!, startx, times)

# RSS = Residual Sum of Squares, columnwise
rss(actual::AbstractMatrix{Float64}, estimated::AbstractMatrix{Float64}) =
    @inbounds sum(i -> abs2(actual[i] - estimated[i]), eachindex(actual, estimated))

# Calculate the actual/original state vectors that we will use for parameter
# estimation:
origstates = calc_states(real_params, lorentz_ode!, r0, t)
origstates_Xiang2015 = calc_states(real_params, lorentz_ode!, r0, t_Xiang2015)

function subsample(origstates::AbstractMatrix{Float64},
                   times::AbstractVector{Float64};
                   lenratio = 0.25)
    N = length(times)
    @assert size(origstates, 2) == N
    stopidx = round(Int, lenratio*N)
    indexes = 1:stopidx
    return origstates[:, indexes], times[indexes]
end

# The [Xiang2015] paper, https://www.hindawi.com/journals/ddns/2015/740721/,
# used these param bounds:
Xiang2015Bounds = [(9., 11.), (20., 30.), (2., 3.)]

# Now we can optimize using BlackBoxOptim
using BlackBoxOptim

# store temporary states for fitness calculation
const tmpstates_pool = Dict{NTuple{2, Int}, Vector{Matrix{Float64}}}()

# get the RSS between the origstates trajectory and ODE solution for given params
function lorentz_fitness(params::AbstractVector{Float64},
                         origstates::AbstractMatrix{Float64},
                         times::AbstractVector{Float64})
    # get the state matrix from the pool of the proper size
    statesize = size(origstates)
    states_pool = get!(() -> Vector{Matrix{Float64}}(), tmpstates_pool, statesize)
    states = !isempty(states_pool) ? pop!(states_pool) : Matrix{Float64}(undef, statesize...)
    # solve ODE for given params
    calc_states!(states, params, lorentz_ode!, r0, times)
    push!(states_pool, states) # return states matrix back to the pool
    return rss(origstates, states)
end

# But optimizing all states in each optimization step is too much so lets
# sample a small subset and use for first opt iteration.
origstates1, times1 = subsample(origstates, t; lenratio=0.04); # Sample only first 4%
origstates1_Xiang2015, times1_Xiang2015 = subsample(origstates_Xiang2015, t_Xiang2015; lenratio=1.00);

res1 = bboptimize(params -> lorentz_fitness(params, origstates1, times1);
    SearchRange = Xiang2015Bounds, MaxSteps = 8e3)

res2 = bboptimize(params -> lorentz_fitness(params, origstates1_Xiang2015, times1_Xiang2015);
    SearchRange = Xiang2015Bounds, MaxSteps = 11e3) # They allow 12k fitness evals for 3-param estimation

# But lets also try with less tight bounds
LooserBounds = [(0., 22.), (0., 60.), (1., 6.)]
res3 = bboptimize(params -> lorentz_fitness(params, origstates1_Xiang2015, times1_Xiang2015);
    SearchRange = LooserBounds, MaxSteps = 11e3) # They allow 12k fitness evals for 3-param estimation

@info "Results on the long time sequence from Paulo Marques:"
estfitness = lorentz_fitness(best_candidate(res1), origstates, t)
@show estfitness best_candidate(res1) best_fitness(res1)
origfitness = lorentz_fitness(real_params, origstates, t)
@show origfitness real_params

@info "Results on the short time sequence used in [Xiang2015] paper:"
estfitness = lorentz_fitness(best_candidate(res2), origstates_Xiang2015, t_Xiang2015)
@show estfitness best_candidate(res2) best_fitness(res2)
origfitness = lorentz_fitness(real_params, origstates_Xiang2015, t_Xiang2015)
@show origfitness real_params

@info "Results on the short time sequence used in [Xiang2015] paper, but with looser bounds:"
estfitness = lorentz_fitness(best_candidate(res3), origstates_Xiang2015, t_Xiang2015)
@show estfitness best_candidate(res3) best_fitness(res3)
origfitness = lorentz_fitness(real_params, origstates_Xiang2015, t_Xiang2015)
@show origfitness real_params


function getPath(G, idStart, idEnd)
	path = Int[]
	temp = idEnd
	while (temp != idStart)
		path = [temp; path]
		temp = G.V[temp].pred
	end
	return [idStart; path]
end

function DFS!(G, pocetak, kraj)
	v = 1:size(G.AdjMatrix, 1);
	for u = v
		G.V[u].color = 'W'
		G.V[u].pred = -1
	end
	DFS_Visit!(G, pocetak, kraj, pocetak)
end

function DFS_Visit!(G, u, kraj, pocetak)
        if u==kraj
            path = getPath(G, pocetak, kraj)
            println("Put: $path")
        else
   	   G.V[u].color = 'G';
           for v in findall(G.AdjMatrix[u,:] .==1)
		if G.V[v].color == 'W'
                        G.V[v].pred = u
			DFS_Visit!(G, v, kraj, pocetak)
		end
	   end
	   G.V[u].color = 'W';
        end
end


# ---
# title: Max min filters
# cover: assets/max_min.gif
# author: Deeptendu Santra
# date: 2020-10-08
# ---

# In this tutorial we see how can we can effectively use max and min filter to distinguish
# between smooth and texture edges  in grayscale images.

# We will use the [`mapwindow`](@ref) function in `ImageFiltering.jl` which provides a general
# functionality to apply any function to the window around each pixel.

using ImageCore, ImageShow, ImageFiltering
using TestImages

img = Gray.(testimage("house");)      # Original Image

# We can use the `minimum` function to compute the minimum of the grayscale values in the given
# matrix or array. For example:
minimum([Gray(0.7),Gray(0.5),Gray(0.0)]) # Should return Gray(0.0) i.e black.
#
filter_size = (15, 15)
## Using the `mapwindow` function, we create an image of the local minimum.
## `mapwindow` maps the given function over a moving window of given size.
img_min = mapwindow(minimum, img, filter_size)
## Similarly for maximum
img_max = mapwindow(maximum, img, filter_size)
## The max(min) filter
img_max_min = mapwindow(maximum, img_min, filter_size)
## The min(max) filter
img_min_max = mapwindow(minimum, img_max, filter_size)
mosaicview(img_min, img_max, img_max_min, img_min_max; nrow=1)

# Now that we are done with the basic filtered images, we proceed to the next part
# which is edge detection using these filters.

# For edge detection, we need to define thresholds for our image. The threshold is an important tool
# to binarize a grayscale image. The threshold value for a given pixel practically
# decides if the pixel visible or not in the output image. However, appyling a global threshold might not consider the
# variation of colors/brightness within the image. Thus we consider an adaptive type theresholding method
# here using max/min results.

# The max(min) and min(max) filter effectively follows the smooth edges of the image. Therefore,
# their average also follows the smooth parts of the image. If we use this image as a threshold for the original
# image, the smooth parts of the original image will get filtered out, leaving only the texture and/or noise behind.
# So we can use the average of `img_max_min` and `img_min_max` as the texture threshold.

# The average of min and max filters also gives us the smooth edges, but it also includes the noise in the image.
# So using average of `img_min` and `img_max` as a threshold will yeild only texture.

img_texture_noise_threshold = (img_max_min + img_max_min) ./ 2
img_texture_threshold = (img_max + img_min) ./ 2
mosaicview(img_texture_noise_threshold, img_texture_threshold; nrow=1)
#
## The Dynamic Gist is obtained by subtracting the img_texture_threshold from the original image.
## The filtered image gives us the texture of the image.
img_dynamic_gist = img - img_texture_threshold
## The Texture Gist is obtained by subtracting img_texture_noise_threshold from the original image.
## The filtered image gives us the texture along with the noise of the image.
img_texture_gist = img - img_texture_noise_threshold
mosaicview(img_dynamic_gist, img_texture_gist; nrow=1)
#
## We extract out the smooth/ramp parts of the image.
## The darker section of the image consist of the ramp edges. The brighter pixels are mostly noise.
ramp = img_dynamic_gist - img_texture_gist
## Filtered-out edges
edge = img_max - img_min
## Smoothed-out version of edge
edge_smoothed = img_min_max - img_max_min
mosaicview(img, ramp, edge, edge_smoothed; nrow=2)

# # References

# Verbeek, P. W., Vrooman, H. A., & Van Vliet, L. J. (1988). [Low-level image processing by max-min filters](https://core.ac.uk/download/pdf/194053536.pdf). Signal Processing, 15(3), 249-258.

## save covers #src
using ImageMagick #src
mkpath("assets")  #src
ImageMagick.save("assets/max_min.gif", cat(img, edge, edge_smoothed, map(clamp01nan, ramp);dims=3);fps=1) #src


#Linear ARD Covariance Function

"""
    LinArd <: Kernel

ARD linear kernel (covariance)
```math
k(x,x') = xᵀL⁻²x'
```
with length scale ``ℓ = (ℓ₁, ℓ₂, …)`` and ``L = diag(ℓ₁, ℓ₂, …)``.
"""
mutable struct LinArd <: Kernel
    "Length scale"
    ℓ::Vector{Float64}
    "Priors for kernel parameters"
    priors::Array

    """
        LinArd(ll::Vector{Float64})

    Create `LinArd` with length scale `exp.(ll)`.
    """
    LinArd(ll::Vector{Float64}) = new(exp.(ll), [])
end

Statistics.cov(lin::LinArd, x::VecF64, y::VecF64) = dot(x./lin.ℓ, y./lin.ℓ)

struct LinArdData{D} <: KernelData
    XtX_d::D
end

function KernelData(k::LinArd, X::MatF64)
    dim, nobs = size(X)
    XtX_d = Array{Float64}(undef, nobs, nobs, dim)
    @inbounds @simd for d in 1:dim
        for i in 1:nobs
            for j in 1:i
                XtX_d[i, j, d] = XtX_d[j, i, d] = X[d, i] * X[d, j]
            end
        end
    end
    LinArdData(XtX_d)
end
kernel_data_key(k::LinArd, X::MatF64) = "LinArdData"
function Statistics.cov(lin::LinArd, X::MatF64)
    K = (X./lin.ℓ)' * (X./lin.ℓ)
    LinearAlgebra.copytri!(K, 'U')
    return K
end
function cov!(cK::MatF64, lin::LinArd, X::MatF64, data::LinArdData)
    dim, nobs = size(X)
    fill!(cK, 0)
    for d in 1:dim
        LinearAlgebra.axpy!(1/lin.ℓ[d]^2, view(data.XtX_d,1:nobs, 1:nobs ,d), cK)
    end
    return cK
end
function Statistics.cov(lin::LinArd, X::MatF64, data::LinArdData)
    nobs = size(X,2)
    K = zeros(Float64, nobs, nobs)
    cov!(K, lin, X, data)
end
@inline @inbounds function cov_ij(lin::LinArd, X::MatF64, data::LinArdData, i::Int, j::Int, dim::Int)
    ck = 0.0
    for d in 1:dim
        ck += data.XtX_d[i,j,d] * 1/lin.ℓ[d]^2
    end
    return ck
end

get_params(lin::LinArd) = log.(lin.ℓ)
get_param_names(lin::LinArd) = get_param_names(lin.ℓ, :ll)
num_params(lin::LinArd) = length(lin.ℓ)

function set_params!(lin::LinArd, hyp::VecF64)
    length(hyp) == num_params(lin) || throw(ArgumentError("Linear ARD kernel has $(num_params(lin)) parameters"))
    @. lin.ℓ = exp(hyp)
end

@inline dk_dll(lin::LinArd, xy::Float64, d::Int) = -2 * xy / lin.ℓ[d]^2
@inline function dKij_dθp(lin::LinArd, X::MatF64, i::Int, j::Int, p::Int, dim::Int)
    if p<=dim
        return dk_dll(lin, dotijp(X,i,j,p), p)
    else
        return NaN
    end
end
@inline function dKij_dθp(lin::LinArd, X::MatF64, data::LinArdData, i::Int, j::Int, p::Int, dim::Int)
    if p <= dim
        return dk_dll(lin, data.XtX_d[i,j,p],p)
    else
        return NaN
    end
end


runall = false;

@testset "Testing Substitution Models, P and Q matrices" begin

m1 = BinaryTraitSubstitutionModel(1.0, 2.0);
show(m1)
m1 = BinaryTraitSubstitutionModel(1.0,2.0, SVector("carnivory", "non-carnivory"));
@test nStates(m1)==2
show(m1)
@test_throws ErrorException PhyloNetworks.BinaryTraitSubstitutionModel(-1.0,2.0)
m2 = EqualRatesSubstitutionModel(4, 3.0);
@test nStates(m2)==4
m2 = EqualRatesSubstitutionModel(4, 3.0, ["S1","S2","S3","S4"]);
show(m2)
m3 = TwoBinaryTraitSubstitutionModel([2.0,1.2,1.1,2.2,1.0,3.1,2.0,1.1],
["carnivory", "noncarnivory", "wet", "dry"]);
show(m3)

@test Q(m1) == SMatrix{2,2}(-1.0, 2.0, 1.0, -2.0)
@test Q(m2) == SMatrix{4,4}(-9.0, 3, 3, 3, 3, -9, 3, 3, 3, 3, -9, 3, 3, 3, 3, -9)
@test Q(m3) ≈ SMatrix{4,4}(-3.0, 3.1, 1.2, 0.0, 1.0, -4.2, 0.0, 2.2, 2.0, 0.0, -3.2, 1.1, 0.0, 1.1, 2.0, -3.3) atol=1e-4
@test P(m1, 0.5) ≈ SMatrix{2,2}(0.7410433867161432,0.5179132265677134,0.2589566132838567,0.4820867734322865) atol=1e-15
@test P(m2, 0.1) ≈ SMatrix{4,4}(0.4758956589341516,0.17470144702194956,0.17470144702194962,0.17470144702194945,0.17470144702194956,0.47589565893415153,0.17470144702194967,0.17470144702194945,0.17470144702194962,0.17470144702194967,0.4758956589341516,0.17470144702194945,0.17470144702194945,0.17470144702194945,0.17470144702194945,0.4758956589341518) atol=1e-15
@test P(m3, 0.5) ≈ SMatrix{4,4}(0.39839916380463375,0.36847565707393248,0.23055614536582461,0.22576141081414305,0.1545971371424259,0.25768553619230444,0.14816051303688715,0.24300762748855972,0.29194735222005136,0.20198250750421617,0.35349416558860958,0.20267178083856716,0.15505634683288913,0.17185629922954704,0.26778917600867863,0.32855918085873009) atol=1e-15
@test P(m1, [0.02,0.01]) ≈ StaticArrays.SArray{Tuple{2,2},Float64,2,4}[[0.980588 0.0194118; 0.0388236 0.961176], [0.990149 0.00985149; 0.019703 0.980297]] atol=1e-6
@test P(m2, [0.02,0.01]) ≈ Array{Float64,2}[[0.839971 0.053343 0.053343 0.053343; 0.053343 0.839971 0.053343 0.053343; 0.053343 0.053343 0.839971 0.053343; 0.053343 0.053343 0.053343 0.839971], [0.91519 0.0282699 0.0282699 0.0282699; 0.0282699 0.91519 0.0282699 0.0282699; 0.0282699 0.0282699 0.91519 0.0282699; 0.0282699 0.0282699 0.0282699 0.91519]] atol=1e-6

end

@testset "Testing random discrete trait simulation" begin

m1 = BinaryTraitSubstitutionModel(1.0,2.0, SVector("carnivory", "non-carnivory"));
m2 = EqualRatesSubstitutionModel(4, 3.0, ["S1","S2","S3","S4"]);

info("\ton a single branch")
srand(12345);
@test randomTrait(m1, 0.2, [1,2,1,2,2]) == [1,2,1,1,2]
srand(12345);
@test randomTrait(m2, 0.05, [1,3,4,2,1]) == [1,3,4,2,1]

info("\ton a network")
net = readTopology("(A:1.0,(B:1.0,(C:1.0,D:1.0):1.0):1.0);")
srand(12345);
a,b = randomTrait(m1, net)
@test a == [1 2 1 1 1 1 2]
@test b == ["-2", "-3", "-4", "D", "C", "B", "A"]
if runall
    for e in net.edge e.length = 10.0; end
    @time a,b = randomTrait(m1, net; ntraits=100000) # ~ 0.014 seconds
    mean(a[:,1]) # expect 1.5 at root
    mean(a[:,2]) # expect 1.333 at other nodes
    @time a,b = randomTrait(m2, net; ntraits=100000) # ~ 0.02 seconds
    length([x for x in a[:,1] if x==4])/length(a[:,1]) # expect 0.25
    length([x for x in a[:,2] if x==4])/length(a[:,2])
    length([x for x in a[:,3] if x==4])/length(a[:,3])
    length([x for x in a[:,4] if x==4])/length(a[:,4])
    length([x for x in a[:,5] if x==4])/length(a[:,5])
    length([x for x in a[:,6] if x==4])/length(a[:,6])
    length([x for x in a[:,7] if x==4])/length(a[:,7]) # expect 0.25
end

net2 = readTopology("(((A:4.0,(B:1.0)#H1:1.1::0.9):0.5,(C:0.6,#H1:1.0):1.0):3.0,D:5.0);")
srand(12345);
a,b = randomTrait(m1, net2; keepInternal=false)
@test a == [1  1  1  2]
@test b == ["D", "C", "B", "A"]
srand(12345);
a,b = randomTrait(m1, net2; keepInternal=true)
@test a == [1  2  1  1  1  1  1  1  1]
@test b == ["-2", "D", "-3", "-6", "C", "-4", "#H1", "B", "A"]
if runall
    for e in net2.edge
        if e.hybrid
            e.length = 0.0
        end
    end
    a,b = randomTrait(m1, net2; ntraits=100000)
    # plot(net2, showNodeNumber=true) shows: H1 listed 7th, parents listed 4th and 6th
    c = map( != , a[:, 4],a[:, 6] ); # traits when parents have different traits
    n1 = sum(map( ==, a[c,7],a[c,6] )) # 39644 traits: hybrid ≠ major parent
    n2 = sum(map( ==, a[c,7],a[c,4] )) #  4401 traits: hybrid ≠ minor parent
    n1/sum(c) # expected 0.9
    n2/sum(c) # expected 0.1
    for e in net2.edge
        e.length = 0.0
    end
    net2.edge[4].length = 10.0
    a,b = randomTrait(m1, net2; ntraits=100000);
    a[:, 1] == a[:, 2]  # true: root = leaf D, as expected
    a[:, 1] == a[:, 5]  # true: root = leaf C
    mean(a[:, 6]) # expected 1.3333
    a[:, 6] == a[:, 9] # true: major hybrid parent node = leaf A
end

end


function test_volume(npoints = 100, dim = 3)
    x = rand(dim, npoints)
    @time v, p = voronoi_random(x, 10000)
    @time a = adjacency(v)
    @time av = [boundary_area_verts(a, v, p) for a in a]
    @time ae = [boundary_area_edges(a, p) for a in a]
    delta = filter(isfinite, ae - av)
    quota = sum(delta .< 1e-6) / length(delta)
    @assert quota > 2/3
    @assert length(delta) / length(av) > 2/3
end

function test_random(d=6, n=200)
    x = rand(d, n)
    @time v,p = voronoi_random(rand(d,n), 10_000_000; maxstuck=100_000);
    @show length(v)
    a = adjacency(v)
    println("avg. no. of neighbors: ", length(a)/n)
    println("avg. no. of vertices per face: ", mean(length.(values(a))))
    @time c = connectivity_matrix(v, p)
    v,p,c
end

function test_grid(n=5, iter=10000)
    plot(legend=false);
    x = hcat(hexgrid(n)...)
    x .+= randn(2,n*n) .* 0.01
    @time v, P  = voronoi_random(x, iter)

    v = Dict(filter(collect(v)) do (k,v)
        norm(v) < 10
        end)

    #c = extractconn(v)
    @time A, Vs = connectivity_matrix(v, P)

    AA = map(x->x>.0, A)
    plot_connectivity!(AA .* 2, P)
    scatter!(eachrow(hcat(values(v)...))...)
    xlims!(1,6); ylims!(0,5)
end

#=
function benchmark(n=100, d=6, iter=100, particles=10)
    x = rand(d, n)
    @benchmark voronoi($x, $iter, $particles)
end
=#

function hexgrid(n)
    P = []
    for i in 1:n
        for j in 1:n
            x = i + j/2
            y = j * sqrt(3) / 2
            push!(P, [x,y])
        end
    end
    P
end

function tests()
    test_volume()
    test_random()
    test_grid()
end


module _Pi

include("pi_chudnovsky.jl")
include("pi_machin.jl")
include("pi_iterative.jl")

end


"""
    `rotstress(stress,theta)`

Rotates `stress` `theta` degrees. `stress` = [sigma1,sigma2,tau12]
"""
function rotstress!(stress::AbstractArray{<:Real,1},theta::Real)
  c = cosd(theta)
  s = sind(theta)
  c2 = c^2
  s2 = s^2
  cs = c*s

  rot11 = c2
  rot12 = s2
  rot13 = 2*cs
  rot21 = s2
  rot22 = c2
  rot23 = -2*cs
  rot31 = -cs
  rot32 = cs
  rot33 = c2-s2

  tmp1 = stress[1]
  tmp2 = stress[2]
  tmp3 = stress[3]

  stress[1] = rot11*tmp1 + rot12*tmp2 + rot13*tmp3
  stress[2] = rot21*tmp1 + rot22*tmp2 + rot23*tmp3
  stress[3] = rot31*tmp1 + rot32*tmp2 + rot33*tmp3

  return stress
end

function rotstress(stress::AbstractArray{<:Real,1}, theta::Real)
  tmp = copy(stress)
  return rotstress!(tmp, theta)
end

"""
    `rotstrain(strain,theta)`

Rotates `strain` `theta` degrees. `strain` = [eps1,eps2,gamma12]
"""
function rotstrain!(strain::AbstractArray{<:Real,1},theta::Real)
  c = cosd(theta)
  s = sind(theta)
  c2 = c^2
  s2 = s^2
  cs = c*s

  rot11 = c2
  rot12 = s2
  rot13 = cs
  rot21 = s2
  rot22 = c2
  rot23 = -cs
  rot31 = -2*cs
  rot32 = 2*cs
  rot33 = c2-s2

  tmp1 = strain[1]
  tmp2 = strain[2]
  tmp3 = strain[3]

  strain[1] = rot11*tmp1 + rot12*tmp2 + rot13*tmp3
  strain[2] = rot21*tmp1 + rot22*tmp2 + rot23*tmp3
  strain[3] = rot31*tmp1 + rot32*tmp2 + rot33*tmp3

  return strain
end

function rotstrain(strain::AbstractArray{<:Real,1},theta::Real)
  tmp = copy(strain)
  return rotstrain!(tmp, theta)
end

"""
    `getplystrain(nply::AbstractArray{<:Integer,1}, tply::AbstractArray{<:Real,1},
        theta::AbstractArray{<:Real,1}, resultantstrain::AbstractArray{<:Real,1})`

    `getplystrain(lam::Laminate, resultantstrain::AbstractArray{<:Real,1})`

Calculates strains in each ply aligned with principal material direction.
"""
function getplystrain(nply::AbstractArray{<:Integer,1}, tply::AbstractArray{<:Real,1},
    theta::AbstractArray{<:Real,1}, resultantstrain::AbstractArray{<:Real,1})

    # number of laminas
    nlam = length(nply)

    # coordinates of top and bottom of plies
    z = getz(tply, nply)

    # local strains at top and bottom of each lamina
    localstrain = [resultantstrain[1:3]+z[k]*resultantstrain[4:6] for k=1:nlam+1]

    # rotate to align with material principal axes
    lowerplystrain = [rotstrain(localstrain[k], theta[k]) for k = 1:nlam]
    upperplystrain = [rotstrain(localstrain[k+1],theta[k]) for k = 1:nlam]

  return lowerplystrain, upperplystrain
end
getplystrain(lam::Laminate, resultantstrain::AbstractArray{<:Real,1}) =
    getplystrain(lam.nply, lam.tply, lam.theta, resultantstrain)

"""
    `getplystress(plystrain::AbstractArray{<:AbstractArray{<:Real,1},1},
        q::AbstractArray{<:AbstractArray{<:Real,2},1},
        matid::AbstractArray{<:Integer,1})`

    `getplystress(plystrain::AbstractArray{<:AbstractArray{<:Real,1},1},
        q::AbstractArray{<:AbstractArray{<:Real,2},1}, lam::Laminate)`

Calculates ply stresses from ply strains and material stiffness matrices (`q`)
"""
function getplystress(plystrain::AbstractArray{<:AbstractArray{<:Real,1},1},
    q::AbstractArray{<:AbstractArray{<:Real,2},1},
    matid::AbstractArray{<:Integer,1})

    return [q[matid[i]]*plystrain[i] for i in 1:length(matid)]
end
getplystress(plystrain::AbstractArray{<:AbstractArray{<:Real,1},1},
    q::AbstractArray{<:AbstractArray{<:Real,2},1}, lam::Laminate) =
    getplystress(plystrain, q, lam.matid)


using NLsolve

include("functions_withG.jl")

hh = Household();
r = 0.019;
A = 1.72;
T = 0.093;
ϕ = 342;
B = 1.0;

@unpack θ,δ,Amat,Ymat,na,nϵ,τ,γ,η,bc = hh
w = A_guess*(1-θ)*((r+δ)/(A*θ))^(θ/(θ-1));
cnext = @. r*Amat+w*Ymat+T

function getcBinding(a,c,ϵ)
    @. (1-τ)*w*ϵ*getl(c,ϵ,γ,ϕ,η,τ,w)+T+(1+(1-τ)*r)*a-bc-c;
end

cbinding = similar(Amat);

for i = 1:na
    for j = 1:nϵ
        cbinding[i,j] = NewtonRoot(c -> getcBinding(Amat[i,j],c,Ymat[i,j]),1.0);
    end
end

cpol,apol,lpol = egm(hh,T_guess,ϕ_guess,w=w,cnext = cnext,cbinding=cbinding,r=r,rnext=r)

cpol,apol,lpol = iterate_egm(hh,A_guess,T_guess,ϕ_guess;r=r_guess)

plot(hh.Amat[1:10,1], hh.Amat[1:10,1], label = "", color = :black, linestyle = :dash)
plot!(hh.Amat[1:10,1], apol[1:10,1], label = "ϵ = $(round(hh.Ymat[1,1], digits = 3))")
plot!(hh.Amat[1:10,1], apol[1:10,5], label = "ϵ = $(round(hh.Ymat[1,5], digits = 3))")

plot(hh.Amat[:,1], lpol[:,1], label = "ϵ = $(round(hh.Ymat[1,1], digits = 3))")
plot!(hh.Amat[:,1], lpol[:,2], label = "ϵ = $(round(hh.Ymat[1,2], digits = 3))")

r,K,N,w = market_clearing(hh,A_guess,r_guess,T_guess,B_guess,ϕ_guess)
plot_market_clearing(hh,A_guess,T_guess,B_guess,ϕ_guess)

function f!(F,x)
    r = x[1]
    A = x[2]
    T = x[3]
    ϕ = x[4]
    Y = 1.0
    θ = 0.3
    G = 0.2
    B = 1.0
    τ = 0.4
    r_eq,K,N,w = market_clearing(hh,A,r,T,B,ϕ)
    F[1] = r_eq-r
    F[2] = Y-A*(K^θ)*(N^(1-θ))
    F[3] = G+T+r-τ*(w*N+r*(K+B))
    F[4] = N-0.28
end

SS = nlsolve(f!,[0.012,5.0,0.05,30.0])


using ElectrostaticPIC1D, QuadGK, Random, Statistics, Test; Random.seed!(0)

@testset "Particles" begin
  @testset "Basic" begin
    q, m = 1.6, 9.11
    nuclide = Nuclide(q, m)
    pos = 1.0
    vel = 2.0
    particle = Particle(nuclide; position=pos, velocity=vel, weight=0.0)
    dt = 3.0
    pushposition!(particle, dt, PeriodicGridBC(-1e16, 1e16))
    @test position(particle) .== pos + vel * dt
    efield = 4.0
    pushvelocity!(particle, efield, dt)
    @test velocity(particle) .== vel + q / m * efield * dt
  end

  @testset "Periodic BCs" begin
    nuclide = Nuclide(1.0, 1.0)
    pos = 0.5
    vel = 0.3
    particle = Particle(nuclide; position=pos, velocity=vel, weight=0.0)
    dt = 1.0
    pushposition!(particle, dt, PeriodicGridBC(0.0, 1.0))
    @test position(particle) .== mod(pos + vel * dt, 1.0)
    pushposition!(particle, dt, PeriodicGridBC(0.0, 1.0))
    @test position(particle) .== mod(pos + 2vel * dt, 1.0)
  end

  @testset "Particles basis functions are normalised" begin
    nuclide = Nuclide(rand(2)...)
    Δ = rand()
    particleshapes = ((BSpline{0}(Δ), "BSpline0"),
                      (BSpline{1}(Δ), "BSpline1"),
                      (BSpline{2}(Δ), "BSpline2"),
                      (DeltaFunctionShape(), "Delta"),
                      (GaussianShape(Δ), "Gaussian"),)
    for (particleshape, pname) ∈ particleshapes
      part = Particle(nuclide, particleshape; x=rand(), v=rand(), w=rand())
      b = ElectrostaticPIC1D.basis(part)
      area = ElectrostaticPIC1D.integral(b, x->1,
        ElectrostaticPIC1D.PeriodicGridBC(1.0))
      @test area ≈ 1
    end
  end

  @testset "Lots of particles tests" begin
    NG = 128 # number of grid points / basis functions
    NPPC = 3 # number of particle per cell
    NP = NG * NPPC
    L = rand()
    Δ = L / NG
    nuclide = Nuclide(rand(2)...)

    physicalmeandensity = rand()

    particleshapes = ((BSpline{0}(Δ), "BSpline0"),
                      (BSpline{1}(Δ), "BSpline1"),
                      (BSpline{2}(Δ), "BSpline2"),
                      (DeltaFunctionShape(), "Delta"),
                      (GaussianShape(Δ), "Gaussian"),)

    fieldsolvers = (
      (FourierField(NG,L), "Fourier"),
      (LSFEMField(NG,L,BSpline{1}(Δ)), "LSFEM_BSpline1"),
      (LSFEMField(NG,L,BSpline{2}(Δ)), "LSFEM_BSpline2"),
      (LSFEMField(NG,L,GaussianShape(Δ * √2)), "LSFEM_Gaussian"),
      (GalerkinFEMField(NG,L,BSpline{1}(Δ), BSpline{2}(Δ)), "Galerkin_BSpline1_BSpline2"),
      (FiniteDifferenceField(NG,L,order=1), "FiniteDifference1"),
      (FiniteDifferenceField(NG,L,order=2), "FiniteDifference2"),
      (FiniteDifferenceField(NG,L,order=4), "FiniteDifference4"),
      )
    q = charge(nuclide)
    xs = collect(1/NP/2:1/NP:1-1/NP/2) .* L # equi-spaced

    perturbations = (0.0, 1e-3)

    for (particleshape, pname) ∈ particleshapes, pert ∈ perturbations
      pertname = iszero(pert) ? "Uniform" : "Wave"
      ρ0 = q * physicalmeandensity
      chargedensityfun(x) = ρ0 * (1.0 + pert * sin(2π*x/L + π/4))
      weightfun(x) = chargedensityfun(x) / q * L / NP
      E0 = q * physicalmeandensity * pert * L/2π
      expectedelectricfield(x) = - E0 * cos(2π*x/L + π/4)

      species = Species([Particle(nuclide, particleshape; x=x, v=0.0,
        w=weightfun(x)) for x in xs])

      expectedoverallchargedensity = ρ0
      totalweight = sum(ElectrostaticPIC1D.weight.(species))
      @test totalweight * q ≈ expectedoverallchargedensity * L
      plasma = Plasma([species])

      @test expectedoverallchargedensity * L ≈ charge(plasma)

      for (field, fname) ∈  fieldsolvers
        @testset  "$pname-$pertname-$fname" begin
          ElectrostaticPIC1D.zero!(field)
          try
            @inferred deposit!(field, plasma)
            @test true
          catch
            ElectrostaticPIC1D.zero!(field)
            deposit!(field, plasma)
            @test false
          end
          @test chargedensity(field) ≈ expectedoverallchargedensity
          if pert == 0.0
            qs = field.chargedensity.(xs)
            @test sqrt(mean((qs .- mean(qs)).^2)) ./ mean(qs) < 1e-6
          else
            solve!(field)
            # because we're testing against a wave field, but not accounting
            # for particle shape here in the test, then we can only expect
            # approximate equality
            for particle ∈ species
              answer = electricfield(field, particle)
              approxexpected = expectedelectricfield(centre(particle))
              @test answer ≈ approxexpected rtol=0.1 atol=E0/1000
              answer = chargedensity(field, particle)
              approxexpected = chargedensityfun(centre(particle))
              @test answer ≈ approxexpected rtol=0.1 atol=ρ0/1000
            end
          end
        end
      end
    end

  end

end


#=
given data (xᵢ, yᵢ), find a polynomial p to solve
minimize    ∑ᵢℓ(p(xᵢ), yᵢ)
subject to  ρⱼ × dᵏp/dtⱼᵏ ≥ 0 ∀ t ∈ D
where
- dᵏp/dtⱼᵏ is the kᵗʰ derivative of p in direction j,
- ρⱼ determines the desired sign of the derivative,
- D is a domain such as a box or an ellipsoid,
- ℓ is a convex loss function.
see e.g. Chapter 8 of thesis by G. Hall (2018)
=#

import DelimitedFiles
import Distributions

struct ShapeConRegrJuMP{T <: Real} <: ExampleInstanceJuMP{T}
    X::Matrix{T}
    y::Vector{T}
    deg::Int
    use_wsos::Bool # use WSOS cone formulation, else SDP formulation
    use_L1_obj::Bool # in objective function use L1 norm, else L2 norm
    use_monotonicity::Bool # if true add monotonicity constraints, else don't
    use_convexity::Bool # if true add convexity constraints, else don't
    is_fit_exact::Bool
end

function ShapeConRegrJuMP{Float64}(
    data_name::Symbol,
    args...)
    Xy = DelimitedFiles.readdlm(joinpath(@__DIR__, "data", "$data_name.txt"))
    (X, y) = (Xy[:, 1:(end - 1)], Xy[:, end])
    # TODO assert X data is on the domain [-1, 1]^n
    return ShapeConRegrJuMP{Float64}(X, y, args...)
end

function ShapeConRegrJuMP{Float64}(
    n::Int,
    num_points::Int,
    func::Symbol,
    signal_ratio::Real,
    args...)
    X = rand(Distributions.Uniform(-1, 1), num_points, n)
    f = shapeconregr_data[func]
    y = Float64[f(X[p, :]) for p in 1:num_points]
    if !iszero(signal_ratio)
        noise = randn(num_points)
        noise .*= norm(y) / sqrt(signal_ratio) / norm(noise)
        y .+= noise
    end
    return ShapeConRegrJuMP{Float64}(X, y, args...)
end

function build(inst::ShapeConRegrJuMP{T}) where {T <: Float64}
    (X, y, deg) = (inst.X, inst.y, inst.deg)
    n = size(X, 2)
    num_points = size(X, 1)
    mono_dom = PolyUtils.BoxDomain{T}(-ones(size(X, 2)), ones(size(X, 2)))
    conv_dom = mono_dom
    mono_profile = ones(Int, size(X, 2))
    conv_profile = 1

    # setup interpolation (not actually using FreeDomain, just need points here)
    halfdeg = div(deg + 1, 2)
    free_dom = PolyUtils.FreeDomain{T}(n)
    (U, points, Ps, V) = PolyUtils.interpolate(free_dom,
        halfdeg, calc_V = true)
    F = qr!(Array(V'), ColumnNorm()) # TODO reuse QR parts
    V_X = PolyUtils.make_chebyshev_vandermonde(X, 2halfdeg)
    X_points_polys = F \ V_X'

    model = JuMP.Model()
    JuMP.@variable(model, regressor[1:U])
    JuMP.@variable(model, z)
    JuMP.@objective(model, Min, z)

    # objective epigraph
    norm_vec = y - X_points_polys' * regressor
    if inst.use_L1_obj || (num_points <= U)
        obj_K = (inst.use_L1_obj ? MOI.NormOneCone : MOI.SecondOrderCone)(
            1 + num_points)
        JuMP.@constraint(model, vcat(z, norm_vec) in obj_K)
    else
        # using L2 norm objective and number of samples exceeds variables,
        # so use qr trick to reduce dimension
        coef_mat = zeros(num_points, U + 1)
        for (i, expr_i) in enumerate(norm_vec)
            for (c, v) in JuMP.linear_terms(expr_i)
                coef_mat[i, JuMP.index(v).value] = c
            end
            coef_mat[i, end] = JuMP.constant(expr_i)
        end
        coef_R = qr(coef_mat).R
        JuMP.@constraint(model, vcat(z, coef_R * vcat(regressor, 1)) in
            MOI.SecondOrderCone(2 + U))
    end

    # monotonicity
    if inst.use_monotonicity
        gradient_halfdeg = div(deg, 2)
        (mono_U, mono_points, mono_Ps) =
            PolyUtils.interpolate(mono_dom, gradient_halfdeg)
        univ_chebs_der = [PolyUtils.calc_univariate_chebyshev(
            mono_points[:, i], 2halfdeg, calc_gradient = true) for i in 1:n]

        for j in 1:n
            iszero(mono_profile[j]) && continue

            univ_chebs_g = [univ_chebs_der[i][(i == j) ? 2 : 1] for i in 1:n]
            V_g = PolyUtils.make_product_vandermonde(univ_chebs_g,
                PolyUtils.n_deg_exponents(n, 2halfdeg))
            scal = inv(maximum(abs, V_g) / 10)
            if scal < 1e-7
                @warn("model is numerically challenging to set up", maxlog = 1)
            end
            lmul!(scal, V_g)
            g_points_polys = F \ V_g'

            gradient_interp = mono_profile[j] * g_points_polys' * regressor

            if inst.use_wsos
                JuMP.@constraint(model, gradient_interp in
                    Hypatia.WSOSInterpNonnegativeCone{T, T}(mono_U, mono_Ps))
            else
                psd_vars = []
                for (r, Pr) in enumerate(mono_Ps)
                    Lr = size(Pr, 2)
                    psd_r = JuMP.@variable(model, [1:Lr, 1:Lr], Symmetric)
                    push!(psd_vars, psd_r)
                    JuMP.@constraint(model, psd_r in JuMP.PSDCone())
                end
                coeffs_lhs = JuMP.@expression(model, [u in 1:mono_U],
                    sum(sum(Pr[u, k] * Pr[u, l] * psd_r[k, l] * (k == l ? 1 : 2)
                    for k in 1:size(Pr, 2) for l in 1:k)
                    for (Pr, psd_r) in zip(mono_Ps, psd_vars)))
                JuMP.@constraint(model, coeffs_lhs .== gradient_interp)
            end
        end
    end

    # convexity
    if inst.use_convexity && !iszero(conv_profile)
        hessian_halfdeg = div(deg - 1, 2)
        (conv_U, conv_points, conv_Ps) =
            PolyUtils.interpolate(conv_dom, hessian_halfdeg)
        univ_chebs_der = [PolyUtils.calc_univariate_chebyshev(
            conv_points[:, i], 2halfdeg, calc_gradient = true,
            calc_hessian = true) for i in 1:n]

        deriv_num(i, j, k) = (k != i && k != j && return 1;
            k == i && k == j && return 3; return 2)

        V_Hs = Matrix{T}[]
        for i in 1:n, j in 1:i
            univ_chebs_H = [univ_chebs_der[k][deriv_num(i, j, k)] for k in 1:n]
            V_H = PolyUtils.make_product_vandermonde(univ_chebs_H,
                PolyUtils.n_deg_exponents(n, 2halfdeg))
            push!(V_Hs, V_H)
        end
        scal = inv(maximum(maximum(abs, V_H) for V_H in V_Hs))
        if scal < 1e-7
            @warn("model is numerically challenging to set up", maxlog = 1)
        end
        lmul!.(scal, V_Hs)

        H_interp = conv_profile * vcat([(F \ V_H')' *
            regressor for V_H in V_Hs]...)

        if inst.use_wsos
            if n == 1
                conv_K = Hypatia.WSOSInterpNonnegativeCone{T, T}(conv_U, conv_Ps)
            else
                Cones.scale_svec!(H_interp, sqrt(T(2)), incr = conv_U)
                conv_K = Hypatia.WSOSInterpPosSemidefTriCone{T}(n, conv_U, conv_Ps)
            end
            JuMP.@constraint(model, H_interp in conv_K)
        else
            psd_vars = []
            for (r, Pr) in enumerate(conv_Ps)
                Lr = size(Pr, 2)
                psd_r = JuMP.@variable(model, [1:(Lr * n), 1:(Lr * n)], Symmetric)
                push!(psd_vars, psd_r)
                JuMP.@constraint(model, psd_r in JuMP.PSDCone())
            end
            # for readability
            Ls = [size(Pr, 2) for Pr in conv_Ps]
            offset = 0
            for x1 in 1:n, x2 in 1:x1
                offset += 1
                coeffs_lhs = JuMP.@expression(model, [u in 1:conv_U],
                    sum(sum(conv_Ps[r][u, k] * conv_Ps[r][u, l] *
                    psd_vars[r][(x1 - 1) * Ls[r] + k, (x2 - 1) * Ls[r] + l]
                    for k in 1:Ls[r] for l in 1:Ls[r]) for r in eachindex(Ls)))
                JuMP.@constraint(model, coeffs_lhs .==
                    H_interp[conv_U .* (offset - 1) .+ (1:conv_U)])
            end
        end
    end

    return model
end

function test_extra(inst::ShapeConRegrJuMP{T}, model::JuMP.Model) where T
    stat = JuMP.termination_status(model)
    @test stat == MOI.OPTIMAL
    if (stat == MOI.OPTIMAL) && inst.is_fit_exact
        # check objective value is correct
        tol = eps(T)^0.25
        @test JuMP.objective_value(model) ≈ 0 atol=tol rtol=tol
    end
end

shapeconregr_data = Dict(
    :func1 => (x -> sum(x .^ 2)),
    :func2 => (x -> sum(x .^ 3)),
    :func3 => (x -> sum(x .^ 4)),
    :func4 => (x -> exp(norm(x)^2 / length(x)) - 1),
    :func5 => (x -> -inv(1 + exp(-10 * norm(x)))),
    :func6 => (x -> sum((x .+ 1) .^ 4)),
    :func7 => (x -> sum((x / 2 .+ 1) .^ 3)),
    :func8 => (x -> sum((x .+ 1) .^ 5 .- 2)),
    :func9 => (x -> (5x[1] + x[2] / 2 + x[3])^2 + sqrt(x[4]^2 + x[5]^2)),
    :func10 => (x -> sum(exp.(x))),
    )


# Functions for contracting two arrays, both for general StridedArrays as well as
# for SymArrays. we only have a few specialized functions that we have needed up
# to now, but carefully optimize each of them

using TensorOperations
using LinearAlgebra
using Requires

# Array[i]*SymArray[(i,j,k)]
# indices 1, 2, and 3 are exchangeable here
function contract(A::StridedVector{T},S::SymArray{(3,),U},n::Union{Val{1},Val{2},Val{3}}) where {T,U}
    TU = promote_type(T,U)
    @assert size(S,1) == length(A)
    res = SymArray{(2,),TU}(size(S,1),size(S,2))
    contract!(res,A,S,n)
end

# We know that $j\leq k$ (because $R$ is itself exchange symmetric)
# \begin{align}
# R_{jk} &= \sum_{i=1}^N g_i S_{ijk}
# \end{align}

# Matrix elements represented by $S_{ijk}$:
# \begin{equation}
# \begin{cases}
# S_{ijk}, S_{ikj}, S_{jik}, S_{jki}, S_{kij}, S_{kji} & i<j<k\\
# S_{ijk}, S_{jik}, S_{jki} & i<j=k\\
# S_{ijk}, S_{ikj}, S_{kij} & i=j<k\\
# S_{ijk} & i=j=k
# \end{cases}
# \end{equation}
# We only need to take the contributions that show up for each $R_{jk}$

# Array[i]*SymArray[(i,j,k)]
# indices 1, 2, and 3 are exchangeable here
function contract!(res::SymArray{(2,),TU}, A::StridedVector{T}, S::SymArray{(3,),U}, n::Union{Val{1},Val{2},Val{3}}) where {T,U,TU}
    # only loop over S once, and put all the values where they should go
    # R[j,k] = sum_i A[i] B[i,j,k]
    # S[i,j,k] with i<=j<=k represents the 6 (not always distinct) terms: Bijk, Bikj, Bjik, Bjki, Bkij, Bkji
    # since R[j,k] is also exchange symmetric, we only need to calculate j<=k
    # commented out terms below are the ones where j>k
    @assert size(S,1) == length(A)
    @assert size(S,1) == size(res,1)
    res.data .= 0
    @inbounds for (v,inds) in zip(S.data,CartesianIndices(S))
        i,j,k = Tuple(inds)
        res[j,k] += v*A[i]
        if i==j==k
            continue
        elseif i==j # i=j < k
            # exch. terms: Sikj and Skij
            # res[k,j] += v*A[i] i2>i3
            res[i,j] += v*A[k]
        elseif j==k # i < j=k
            # exch. terms Sjik, Sjki
            res[i,k] += v*A[j]
            # res[k,i] += v*A[j] i3>i2
        else # i<j<k
            # exch. terms Sikj, Sjik, Sjki, Skij, Skji
            # res[k,j] += v*A[i] i3>i2
            res[i,k] += v*A[j]
            # res[k,i] += v*A[j] i3>i2
            res[i,j] += v*A[k]
            # res[j,i] += v*A[k] i3>i2
        end
    end
    res
end

# Array[k]*SymArray[(i,j),k]
function contract(A::StridedVector{T},S::SymArray{(2,1),U},n::Val{3}) where {T,U}
    TU = promote_type(T,U)
    sumsize = length(A)
    @assert sumsize == size(S,3)
    res = SymArray{(2,),TU}(size(S,1),size(S,2))
    contract!(res,A,S,n)
end

# Array[k]*SymArray[(i,j),k]
function contract!(res::SymArray{(2,),TU},A::StridedVector{T},S::SymArray{(2,1),U},::Val{3}) where {T,U,TU}
    # use that S[(i,j),k] == S[I,k] (i.e., the two symmetric indices act like a "big" index)
    mul!(res.data,reshape(S.data,:,length(A)),A)
    res
end

# Array[i]*SymArray[(i,j),k)]
# since indices 1 and 2 are exchangeable here, use this
function contract(A::StridedVector{T},S::SymArray{(2,1),U},n::Union{Val{1},Val{2}}) where {T,U}
    TU = promote_type(T,U)
    @assert size(S,1) == length(A)
    # the result is a normal 2D array
    res = Array{TU,2}(undef,size(S,2),size(S,3))
    contract!(res,A,S,n)
end

# Array[i]*SymArray[(i,j),k]
# since indices 1 and 2 are exchangeable here, use this
function contract!(res::StridedArray{TU,2},A::StridedVector{T},S::SymArray{(2,1),U},::Union{Val{1},Val{2}}) where {T,U,TU}
    # only loop over S once, and put all the values where they should go
    @assert size(A,1) == size(S,1)
    @assert size(res,1) == size(S,1)
    @assert size(res,2) == size(S,3)
    res .= zero(TU)
    @inbounds for (v,inds) in zip(S.data,CartesianIndices(S))
        i1,i2,i3 = Tuple(inds)
        res[i2,i3] += v*A[i1]
        # if i1 != i2, we have to add the equal contribution from S[i2,i1,i3]
        if i1 != i2
            res[i1,i3] += v*A[i2]
        end
    end
    res
end

# Array[i]*SymArray[(i,j)]
# this is symmetric in i1 and i2
function contract(A::StridedVector{T},S::SymArray{(2,),U},n::Union{Val{1},Val{2}}) where {T,U}
    TU = promote_type(T,U)
    @assert size(S,1) == length(A)
    # the result is a normal 1D vector
    res = Vector{TU}(undef,size(S,2))
    contract!(res,A,S,n)
end

# Array[i]*SymArray[(i,j)]
# this is symmetric in i1 and i2
function contract!(res::StridedVector{TU},A::StridedVector{T},S::SymArray{(2,),U},::Union{Val{1},Val{2}}) where {T,U,TU}
@assert size(A,1) == size(S,1)
    @assert size(res,1) == size(S,1)
    res .= zero(TU)
    # only loop over S once, and put all the values where they should go
    for (v,inds) in zip(S.data,CartesianIndices(S))
        i1,i2 = Tuple(inds)
        res[i2] += v*A[i1]
        # if i1 != i2, we have to add the contribution from S[i2,i1]
        if i1 != i2
            res[i1] += v*A[i2]
        end
    end
    res
end

# Array[i_n]*Array[i1,i2,i3,...,iN]
function contract(A::StridedVector{T},B::StridedArray{U,N},::Val{n}) where {T,U,N,n}
    TU = promote_type(T,U)
    @assert 1 <= n <= N

    resdims = size(B)[1:N .!= n]
    res = similar(B,TU,resdims)

    A = convert(AbstractArray{TU},A)
    B = convert(AbstractArray{TU},B)
    contract!(res,A,B,Val{n}())
end

mygemv!(args...) = BLAS.gemv!(args...)

function _contract_middle!(res,A,B)
    @inbounds for k=1:size(B,3)
        mul!(@view(res[:,k]), @view(B[:,:,k]), A)
    end
end

function __init__()
    @require CuArrays = "3a865a2d-5b23-5a0f-bc46-62713ec82fae" begin
        mygemv!(tA,alpha,A::CuArrays.CuArray,args...) = CuArrays.CUBLAS.gemv!(tA,alpha,A,args...)
        _contract_middle!(res::CuArrays.CuArray,A,B) = error("please install CuTensorOperations.jl to allow this operation with CuArrays")
    end
    @require CuTensorOperations = "01a401d6-3e09-11e9-13be-6f1bb8674acf" begin
        _contract_middle!(res::CuArrays.CuArray,A,B) = (@tensor res[i,k] = B[i,j,k] * A[j])
    end
end

# Array[i_n]*Array[i1,i2,i3,...,iN]
function contract!(res::StridedArray{TU},A::StridedVector{TU},B::StridedArray{TU,N},::Val{n}) where {TU,N,n}
    nsum = length(A)
    @assert size(B,n) == nsum
    @assert ndims(res)+1 == ndims(B)
    ii = 0
    for jj = 1:ndims(B)
        jj==n && continue
        ii += 1
        @assert size(B,jj) == size(res,ii)
    end

    if n==1      # A[i]*B[i,...]
        mygemv!('T',one(TU),reshape(B,nsum,:),A,zero(TU),vec(res))
    elseif n==N  # B[...,i]*A[i]
        mygemv!('N',one(TU),reshape(B,:,nsum),A,zero(TU),vec(res))
    else
        rightsize = prod(size(B,i) for i=n+1:N)
        Br = reshape(B,:,nsum,rightsize)
        resr = reshape(res,:,rightsize)
        _contract_middle!(resr,A,Br)
    end
    res
end


using ReverseDiff: GradientTape, GradientConfig, gradient, gradient!, compile
using Turing: _hmc_step
using Turing

@model bayes_lr(xs, ys) = begin
    N = length(xs)
    @assert N == length(ys)

    s = 1
    β ~ Normal(0, 1)

    for n = 1:N
        ys[n] ~ Normal(xs[n] * β, sqrt(s))
    end
end

N = 100
xs = collect(range(-10, stop = 10, length = N))
s = 1
β = rand(Normal(0, 1))
ys = xs * β + rand(Normal(0, sqrt(s)), N)

println("s=$s, β=$β")

mf = bayes_lr(xs, ys)
chn = sample(mf, HMC(2000, 0.005, 3))

println("mean of β: $(mean(chn[:β][1000:end]))")

θ_dim = 1
function lj_func(θ)
  N = length(xs)
  _lj = zero(Real)

  s = 1

  β = θ[1]
  _lj += logpdf(Normal(0, 1), β)
  for n = 1:N
    _lj += logpdf(Normal(xs[n] * β, sqrt(s)), ys[n])
  end

  return _lj
end

neg_lj_func(θ) = -lj_func(θ)
const f_tape = GradientTape(neg_lj_func, randn(θ_dim))
const compiled_f_tape = compile(f_tape)

function grad_func(θ)

  inputs = θ
  results = similar(θ)
  all_results = DiffResults.GradientResult(results)

  ReverseDiff.gradient!(all_results, compiled_f_tape, inputs)

  neg_lj = all_results.value
  grad, = all_results.derivs

  return -neg_lj, grad

end

std = ones(θ_dim)
θ = randn(θ_dim)
lj = lj_func(θ)

chn = []
accept_num = 1

total_num = 2000
for iter = 1:total_num
  global θ, chn, lj, lj_func, grad_func, std, accept_num
  push!(chn, θ)
  θ, lj, is_accept, τ_valid, α = _hmc_step(θ, lj, lj_func, grad_func, 3, 0.005, std)
  accept_num += is_accept
#   if (iter % 50 == 0) println(θ) end
end

@show mean(chn[1000:end]), lj
@show accept_num / total_num


function read_input(input_filename)
    lines = readlines(input_filename)
    grid = zeros(Int, length(lines), length(lines[1]))
    for iy in 1:length(lines)
        for ix in 1:length(lines[iy])
            grid[iy, ix] = parse(Int, lines[iy][ix])
        end
    end
    return grid
end

function is_local_minimum(x, y, grid)
    if x == 1
        local_min_x = (grid[x,y] < grid[x+1,y])
    elseif x == size(grid)[1]
        local_min_x = (grid[x,y] < grid[x-1,y])
    else
        local_min_x = (grid[x,y] < grid[x+1,y]) && (grid[x,y] < grid[x-1,y])
    end
    if y == 1
        local_min_y = (grid[x,y] < grid[x,y+1])
    elseif y == size(grid)[2]
        local_min_y = (grid[x,y] < grid[x,y-1])
    else
        local_min_y = (grid[x,y] < grid[x,y+1]) && (grid[x,y] < grid[x,y-1])
    end
    return (local_min_x && local_min_y)
end

function find_basin_coords(x, y, grid, basin_coords)
    # basin coords is a set
    # recursive function walks in all directions and adds coords to list
    # if : value at new grid square is not 9, AND new grid square isn't already in set,
    # AND new grid square is not off the ends of the grid.
    if grid[x,y] != 9
        push!(basin_coords, (x, y))
    end
    if ((x-1) in 1:size(grid)[1]) && (grid[x-1,y] != 9) && (!).(in((x-1,y),basin_coords))

        basin_coords = find_basin_coords(x-1, y, grid, basin_coords)
    end
    if ((x+1) in 1:size(grid)[1]) && (grid[x+1,y] != 9)  && (!).(in((x+1,y),basin_coords))
        basin_coords = find_basin_coords(x+1, y, grid, basin_coords)
    end
    if ((y-1) in 1:size(grid)[2]) && (grid[x,y-1] != 9)  && (!).(in((x,y-1),basin_coords))
        basin_coords = find_basin_coords(x, y-1, grid, basin_coords)
    end
    if ((y+1) in 1:size(grid)[2]) && (grid[x,y+1] != 9) && (!).(in((x,y+1),basin_coords))
        basin_coords = find_basin_coords(x, y+1, grid, basin_coords)
    end
    return basin_coords
end

function find_all_local_minima(grid)
    minima_coords = []
    for ix in 1:size(grid)[1]
        for iy in 1:size(grid)[2]
            if is_local_minimum(ix, iy, grid)
                push!(minima_coords,(ix,iy))
            end
        end
    end
    return minima_coords
end

function find_all_basin_sizes(grid)
    minima = find_all_local_minima(grid)
    basin_sizes = []
    for minimum in minima
        basin_coordinates = Set([])
        push!(basin_sizes, length(find_basin_coords(minimum[1], minimum[2],grid, basin_coordinates)))
    end
    return sort(basin_sizes, rev=true)

end

grid = read_input("input.txt")
basin_sizes = find_all_basin_sizes(grid)
answer = prod(basin_sizes[1:3])
println("answer is ",answer)


read_data(fn::AbstractString) = vec(readdlm(fn, skipstart=1))
read_cpts(fn::AbstractString) = vec([0; readdlm(fn, Int, skipstart=1)])

dir = "Test_Files"

comp_handler(r::Test.Success) = true
comp_handler(r::Test.Failure) = false
comp_handler(r::Test.Error) = rethrow(r)

function test_cpts(r_cpts, cpts)
    Test.with_handler(comp_handler) do
        if @test r_cpts == cpts
            println("\tSuccess!")
        else
            println("Fail:")
            println("\tR Changepoints: $(r_cpts)")
            println("\tJulia Changepoints: $(cpts)")
        end
    end
end

# Normal mean change

println("Test Normal mean change...")
data = read_data("$(dir)/Normal_change_in_mean_data")
r_cpts = read_cpts("$(dir)/Normal_change_in_mean_cpts")
cpts, cost = @PELT data Normal(?, 1.0) 2*log(length(data))
test_cpts(r_cpts, cpts)

# Normal var change
println("Test Normal variance change...")
data = read_data("$(dir)/Normal_change_in_var_data")
r_cpts = read_cpts("$(dir)/Normal_change_in_var_cpts")
cpts, cost = @PELT data Normal(1.0, ?) 2*log(length(data))
test_cpts(r_cpts, cpts)

# Normal meanvar change
println("Test Normal mean and variance chaange...")
data = read_data("$(dir)/Normal_change_in_meanvar_data")
r_cpts = read_cpts("$(dir)/Normal_change_in_meanvar_cpts")
cpts, cost = @PELT data Normal(?, ?) 4*log(length(data))
test_cpts(r_cpts, cpts)

# Poisson change
println("Test Poisson change...")
data = read_data("$(dir)/Poisson_data")
r_cpts = read_cpts("$(dir)/Poisson_cpts")
cpts, cost = @PELT data Poisson(?) 2*log(length(data))
test_cpts(r_cpts, cpts)

# Exponential change
println("Test Exponential  change...")
data = read_data("$(dir)/Exponential_data")
r_cpts = read_cpts("$(dir)/Exponential_change_in_mean_cpts")
cpts, cost = @PELT data Exponential(?) 2*log(length(data))
test_cpts(r_cpts, cpts)

# Gamma rate changes
println("Test Gamma rate change...")
data = read_data("$(dir)/Gamma_data")
r_cpts = read_cpts("$(dir)/Gamma_cpts")
cpts, cost = @PELT data Gamma(1.0,?) 2*log(length(data))
test_cpts(r_cpts, cpts)

println("Test Nonparametric change...")
data = read_data("$(dir)/Nonparametric_data")
r_cpts = read_cpts("$(dir)/Nonparametric_cpts")
cpts, cost = @PELT data Nonparametric(10) 4*log(length(data))
test_cpts(r_cpts, cpts)


using ExtendedKronigPennyMatrix
using LinearAlgebra
using PyPlot

function plot_KP(v0, rho)
   clf()
   cm=get_cmap("tab10")

   pot=KronigPennyPotential(v0, rho)
   pf = get_potential(pot)
   a=1
   xs=-2a:a/100:3a
   plot(xs, pf.(xs), "k")

   nmax=60
   Ka=0
   model=KronigPennyModel(pot, nmax, Ka)

   for Ka in (-18:18)/18*π
      update!(model, Ka=Ka)
      ev = eigvals(model.hnm)
      for i in 1:5
         plot(Ka/ π, ev[i], ".", color=cm(i-1))
      end
   end
   xlim(-1,1)
   ylim(-2,32)
   xlabel(L"$Ka / \pi$")
   ylabel(L"Energy / $E_0$")
   title( L"$\rho =$"*string(rho))
end

function main()
   #
   plot_KP(      0, 0.5)
   ylim(-0.5,25)
   savefig("Pavelich_Fig4a.png")
   #
   plot_KP(     10, 0.5)
   ylim(-0.5,30)
   savefig("Pavelich_Fig4b.png")
   savefig("KP_10_05.png")
   #
   plot_KP(     10, 0.8)
   ylim(-0.5,30)
   savefig("KP_10_08.png")
   #
   plot_KP(20.5607, 0.5)
   ylim(-0.5,40)
   savefig("Pavelich_Fig6.png")
   #
   plot_KP(10.8775, 0.8)
   ylim(-0.5,30)
   savefig("Pavelich_Fig7.png")
end

main()


using DynamicHMCModels, MCMCChains
using Flux

ProjDir = @__DIR__
cd(ProjDir)

include(joinpath(@__DIR__, "LBA_functions.jl"))

Base.@kwdef struct LBAModel{T}
   data::T
   N::Int
   Nc::Int
end

function make_transformation(model::LBAModel)
    as((v=as(Array,asℝ₊,Nc),A=asℝ₊,k=asℝ₊,tau=asℝ₊))
end

N = 20
data=simulateLBA(;Nd=N,v=[1.0,1.5,2.0],A=.8,k=.2,tau=.4)
Nc = 3

#dist = LBA(ν=[1.0,1.5,2.0],A=.8,k=.2,τ=.4)
#data = rand(dist,N)

model = LBAModel(; data=data, N=N, Nc=Nc)

function (model::LBAModel)(θ)
    @unpack data=model
    @unpack v,A,k,tau=θ
    d=LBA(ν=v,A=A,k=k,τ=tau)
    minRT = minimum(x->x[2],data)
    logpdf(d,data)+sum(logpdf.(TruncatedNormal(0,3,0,Inf),v)) +
    logpdf(TruncatedNormal(.8,.4,0,Inf),A)+logpdf(TruncatedNormal(.2,.3,0,Inf),k)+
    logpdf(TruncatedNormal(.4,.1,0,minRT),tau)
end

d = [(c,r) for (c,r) in zip(data.choice,data.rt)]

p = LBAModel(d,N,Nc)
p((v=fill(.5,Nc),A=.8,k=.2,tau=.4))

# Use Flux for the gradient.
P = TransformedLogDensity(make_transformation(p), p)
∇P = ADgradient(:Flux, P)

# Sample from the posterior.
results = mcmc_with_warmup(Random.GLOBAL_RNG, ∇P, 1000)
posterior = P.transformation.(results.chain)


## Look at musical artists
using DelimitedFiles
using SparseArrays
names = readlines("artistsim.names")
data = readdlm("artistsim.smat")
##
A = sparse(Int.(data[2:end,1]).+1,
           Int.(data[2:end,2]).+1,
           Int.(data[2:end,3]),
           Int(data[1,1]),Int(data[1,1]))
##
A = max.(A,A')
##
using PyCall
using Conda
pyimport_conda("igraph","python-igraph","conda-forge")
##
using PyCall
igraph = pyimport("igraph")

function igraph_layout(A::SparseMatrixCSC{T}, layoutname::AbstractString="drl") where T
    ei,ej,ew = findnz(A)
    edgelist = [(ei[i]-1,ej[i]-1) for i = 1:length(ei)]
    nverts = size(A)
    G = igraph.Graph(nverts, edges=edgelist, directed=false)
    xy = G.layout(layoutname)
    xy = [Float64(xy[i][j]) for i in 1:length(xy),  j in 1:length(xy[1])]
end

xy = igraph_layout(A, "drl")
##
using GraphRecipes
using LinearAlgebra
using Plots
ei, ej = findnz(triu(A,1))
graphplot(ei, ej, x =xy[:,1], y=xy[:,2],
  markercolor=:black, markerstrokecolor=:white,
  size=(1200,1200),dpi=300,
  markersize=2, linecolor=1, linealpha=0.01, linewidth=0.5,
  markeralpha=0.2,
  axis_buffer=0.02, background=nothing)
savefig("artistsim-plot.png")


function fit(
    df::DataFrame;
    max_time::Int = 10,
    fps = 0.022,
    loc_error = 0.03,
    p0::AbstractVector = [0.5, 0.5],
)
    @argcheck loc_error >= 0.0

    @. model(x, p) = 4 * p[1] * x^p[2] + 4 * loc_error^2
    result = curve_fit(
        model,
        df.delta_t[1:max_time] .* fps,
        df.msd[1:max_time],
        p0,
        lower = [0.0, 0.0],
        upper = [10.0, 2.0],
    )
    result
end

function estimate_α(df::DataFrame, δ1::Int, δ2::Int)
    @argcheck 1 <= δ2 <= δ1
    log(mean(df[df.Δt.==δ1, :msd]) / mean(df[df.Δt.==δ2, :msd])) / log(δ1 / δ2)
end

function fit_msd(
    df;
    max_time::Int64 = 10,
    loc_error::Float64 = 0.03,
    p0::AbstractVector = [0.5, 0.5],
)
    @argcheck loc_error >= 0.0
    @. model(x, p) = 4 * p[1] * x^p[2] + 4 * loc_error^2
    fit = curve_fit(
        model,
        df.delta_t[1:max_time] .* (1 / 45),
        df.msd[1:max_time],
        p0,
        lower = [0.0, 0.0],
        upper = [10.0, 2.0],
    )
    fit.param
end

function fit_msd(df::DataFrame, δ1::Int, δ2::Int)
    @argcheck 1 <= δ2 <= δ1
    log(mean(df[df.delta_t.==δ1, :msd]) / mean(df[df.delta_t.==δ2, :msd])) / log(δ1 / δ2)
end

fit_msd(df::DataFrame, δ::Int) = fit_msd(df::DataFrame, δ, 1)


"""1INFILTR case from Hydrus-1D, horizontal."""
module Example1INFILTR

using Fronts
using Fronts.PorousModels
using Plots

ϵ = 1e-6

Ks = 25  # cm/h
α = 0.01433  # 1/cm
n = 1.506

θs = 0.3308

model = VanGenuchten(n=n, α=α, Ks=Ks, θs=θs)

prob = DirichletProblem(model, i=0.1003, b=0.3308-ϵ)

θ = solve(prob)

r = range(0, 200, length=500)

plt = plot(r, r -> θ(r,1), label="t=1 h", xguide="r [cm]", yguide="θ")
plot!(r, r -> θ(r,4), label="t=4 h")
display(plt)

plt = plot(r, r -> flux(θ,r,1), label="t=1 h", xguide="r [cm]", yguide="U [m/s]")
plot!(r, r -> flux(θ,r,4), label="t=4 h")
display(plt)

end


function fib(n)
    if n < 0
        throw(ArgumentError("negative arguments not allowed"))
    end
    aux(m) = m < 2 ? one(m) : aux(m-1) + aux(m-2)
    aux(n)
end


#
# Provides Restricted Two-body Problem equations
#

"""
Returns the conic section, as specified by eccentricity `e`.
"""
function conic(e::T) where T<:Real

    !isnan(e) || throw(ArgumentError("Provided eccentricity is not a number (NaN)."))

    if e ≈ zero(T)
        return Circular
    elseif e ≈ one(T)
        return Parabolic
    elseif zero(T) < e && e < one(T)
        return Elliptical
    elseif e > one(T)
        return Hyperbolic
    else
        throw(ArgumentError("Provided eccentricity, $e, is not valid."))
    end

end

"""
Returns a Keplarian representation of a Cartesian orbital state.
Algorithm taught in ENAE601.
"""
function keplerian(rᵢ, vᵢ, μ)

    safe_acos(num) = isapprox(num, one(num)) ?
                            acos(one(num)) :
                                isapprox(num, -one(num)) ?
                                    acos(-one(num)) :
                                        acos(num)

    î = SVector{3}(1, 0, 0)
    ĵ = SVector{3}(0, 1, 0)
    k̂ = SVector{3}(0, 0, 1)

    h̅ = specific_angular_momentum_vector(rᵢ, vᵢ)
    a = semimajor_axis(norm(rᵢ), norm(vᵢ), μ)
    n̅ = k̂ × specific_angular_momentum_vector(rᵢ, vᵢ)
    e̅ = eccentricity_vector(rᵢ, vᵢ, μ)
    e = norm(e̅)

    i = safe_acos((h̅ ⋅ k̂) / norm(h̅))

    Ω = ustrip(n̅ ⋅ ĵ) > 0 ?
            safe_acos((n̅ ⋅ î) / norm(n̅)) :
            2π - safe_acos((n̅ ⋅ î) / norm(n̅))

    ω = ustrip(e̅ ⋅ k̂) > 0 ?
            safe_acos((n̅ ⋅ e̅) / (norm(n̅) * e)) :
            2π - safe_acos((n̅ ⋅ e̅) / (norm(n̅) * e))

    ν = ustrip(rᵢ ⋅  vᵢ) > 0 ?
            safe_acos((e̅ ⋅ rᵢ) / (e * norm(rᵢ))) :
            2π - safe_acos((e̅ ⋅ rᵢ) / (e * norm(rᵢ)))

    return upreferred(e), upreferred(a), upreferred(i),
           upreferred(Ω), upreferred(ω),
           upreferred(ν)

end

"""
Returns a Cartesian representation of a Keplerian two-body orbital state
in an inertial frame, centered at the center of mass of the central body.
Algorithm taught in ENAE601.
"""
function cartesian(e, a, i, Ω, ω, ν, μ)
    rᵢ, vᵢ = spatial(i, Ω, ω, perifocal(a, e, ν, μ)...)
    return  upreferred.(rᵢ),
            upreffered.(vᵢ)
end

"""
Returns a spatial representation of the provied Perifocal state.
"""
function spatial(i, Ω, ω, rₚ, vₚ)

    # Set up Perifocal ⟶ Cartesian conversion
    R_3Ω =  SMatrix{3,3}(
            [cos(Ω)          -sin(Ω)            0.;
             sin(Ω)           cos(Ω)            0.;
             0.               0.                1.])
    R_1i = SMatrix{3,3}(
            [1.               0.                0.;
             0.               cos(i)           -sin(i);
             0.               sin(i)            cos(i)])
    R_3ω = SMatrix{3,3}(
            [cos(ω)          -sin(ω)            0.
             sin(ω)           cos(ω)            0.
             0.               0.                1.])

    ᴵTₚ = R_3Ω * R_1i * R_3ω

    return ᴵTₚ * rₚ, ᴵTₚ * vₚ

end

"""
Returns position and velocity vectors in the Perifocal frame.
"""
function perifocal(a, e, ν, μ)

        p = semi_parameter(a, e)
        r = distance(p, e, ν)

        P̂=SVector{3, Float64}(1, 0, 0)
        Q̂=SVector{3, Float64}(0, 1, 0)
        # Ŵ=SVector{3, Float64}(0, 0, 1)

        rₚ = (r * cos(ν) .* P̂ .+ r * sin(ν) .* Q̂)
        vₚ = √(μ/p) * ((-sin(ν) * P̂) .+ ((e + cos(ν)) .* Q̂))

        return rₚ, vₚ

end

function perifocal(i, Ω, ω, rᵢ, vᵢ)

    # Set up Cartesian ⟶ Perifocal conversion
    R_3Ω =  SMatrix{3,3}(
            [cos(Ω)          -sin(Ω)            0.;
             sin(Ω)           cos(Ω)            0.;
             0.               0.                1.])
    R_1i = SMatrix{3,3}(
            [1.               0.                0.;
             0.               cos(i)           -sin(i);
             0.               sin(i)            cos(i)])
    R_3ω = SMatrix{3,3}(
            [cos(ω)          -sin(ω)            0.
             sin(ω)           cos(ω)            0.
             0.               0.                1.])

    ᵖTᵢ = inv(R_3Ω * R_1i * R_3ω)

    return ᵖTᵢ*rᵢ, ᵖTᵢ*vᵢ

end

"""
Returns semimajor axis parameter, a.
"""
semimajor_axis(r, v, μ) = inv( (2 / r) - (v^2 / μ) )

"""
Returns specific angular momentum vector, h̅.
"""
specific_angular_momentum_vector(rᵢ, vᵢ) = rᵢ × vᵢ

"""
Returns scalar specific angular momentum vector, h.
"""
specific_angular_momentum(rᵢ, vᵢ) = norm(specific_angular_momentum_vector(rᵢ, vᵢ))

"""
Returns specific orbital energy, ϵ.
"""
specific_energy(a, μ) = ( -μ / (2 * a) )
specific_energy(r, v, μ) = (v^2 / 2) - (μ / r)

"""
Returns C3 value.
"""
C3(r, v, μ) = v^2 - 2μ/r

"""
Returns v∞.
"""
v_infinity(r, v, μ) = √C3(r, v, μ)

"""
Returns potential energy for an orbit about a `RestrictedTwoBodySystem`.
"""
specific_potential_energy(r, μ) = (μ/r)
specific_potential_energy(r, μ, R, J₂, ϕ) = (μ/r) * (1 - J₂ * (R/r)^2 * ((3/2) * (sin(ϕ))^2 - (1/2)))

"""
Returns orbital eccentricity vector e̅.
"""
function eccentricity_vector(rᵢ, vᵢ, μ)

    return map(x-> abs(x) < eps(typeof(x)) ? zero(x) : x, (1 / μ) * ((vᵢ × specific_angular_momentum_vector(rᵢ, vᵢ)) - μ * rᵢ / norm(rᵢ)))

end

"""
Returns orbital eccentricity, e.
"""
eccentricity(rᵢ, vᵢ, μ) = norm(eccentricity_vector(rᵢ, vᵢ, μ)) |> upreferred

"""
Returns semilatus parameter, p.
"""
semi_parameter(a, e) = a * (1 - e^2)

"""
Returns distance, r.
"""
distance(p, e, ν) = upreferred(p / (1 + e * cos(ν)))

"""
Returns instantaneous velocity, v, for any orbital representation.
"""
speed(r, a, μ) =  upreferred(√( (2 * μ / r) - (μ / a)))

"""
Returns the instantaneous velocity, `v`, for any orbital representation.
"""
speed(p, e, ν, a, μ) = speed(distance(p,e,ν), a, μ)

"""
Returns periapsis distance, rₚ.
"""
periapsis_radius(a, e) = a * (1 - e)

"""
Returns apoapsis distance, rₐ.
"""
apoapsis_radius(a, e) = a * (1 + e)

"""
Returns the orbital period.
"""
period(a, μ) = 2π * √(upreferred(a^3 / μ))

"""
Returns true anomoly, ν.
"""
function true_anomoly(r, h, e, μ)
    val = (h^2 - μ * r) / (μ * r * e)
    acos(isapprox(val, one(val)) ? one(val) : val)
end

"""
Returns mean motion, n.
"""
mean_motion(a, μ) = √(μ / a^3)

"""
Returns time since periapsis, t.
"""
time_since_periapsis(n, e, E) = (E - e * sin(E)) / (n)

"""
Sphere of influence.
"""
SOI(a, m, M) = a * (m / M)^(2/5)

"""
Sphere of activity.
"""
SOA(a, m, M) = a * (m / 3M)^(1/3)

"""
Computes a Hohmann transfer, and returns the departure and
arrival velocity vectors.
"""
function hohmann(r₁, r₂, μ)

    vₐ = √((2μ/r₁) - (2μ/(r₁+r₂)))
    vₚ = √((2μ/r₂) - (2μ/(r₁+r₂)))

    return vₐ, vₚ

end


# Solves the Poisson equation using the Mixed finite element method
using Revise
using AdFem
using DelimitedFiles
using SparseArrays
using PyPlot

λ = 1.0
μ = 1.0
n = 50
mmesh = Mesh(n, n, 1/n, degree = BDM1)
a = 1/2μ
b = -λ/(2μ*(2μ+2λ))

TestCase = [
    (
        (x,y)->begin;x*y*(x - 1)*(y - 1);end,
        (x,y)->begin;x*y*(x - 1)*(y - 1);end,
        (x,y)->begin;2.0*x^2 + 8.0*x*y - 6.0*x + 6.0*y^2 - 10.0*y + 2.0;end,
        (x,y)->begin;6.0*x^2 + 8.0*x*y - 10.0*x + 2.0*y^2 - 6.0*y + 2.0;end,
    ),
    (
        (x,y)->begin;x^2*y^2*(x - 1)*(y^2 - 1);end,
        (x,y)->begin;x^2*y^2*(x - 1)*(y^2 - 1);end,
        (x,y)->begin;12.0*x^3*y^2 - 2.0*x^3 + 24.0*x^2*y^3 - 12.0*x^2*y^2 - 12.0*x^2*y + 2.0*x^2 + 18.0*x*y^4 - 16.0*x*y^3 - 18.0*x*y^2 + 8.0*x*y - 6.0*y^4 + 6.0*y^2;end,
        (x,y)->begin;36.0*x^3*y^2 - 6.0*x^3 + 24.0*x^2*y^3 - 36.0*x^2*y^2 - 12.0*x^2*y + 6.0*x^2 + 6.0*x*y^4 - 16.0*x*y^3 - 6.0*x*y^2 + 8.0*x*y - 2.0*y^4 + 2.0*y^2;end,
    ),
    (
        (x,y)->begin;x^2*y^2*(x - 1)*(y^2 - 1);end,
        (x,y)->begin;x*y*(x - 1)*(y - 1);end,
        (x,y)->begin;12.0*x^3*y^2 - 2.0*x^3 - 12.0*x^2*y^2 + 2.0*x^2 + 18.0*x*y^4 - 18.0*x*y^2 + 8.0*x*y - 4.0*x - 6.0*y^4 + 6.0*y^2 - 4.0*y + 2.0;end,
        (x,y)->begin;24.0*x^2*y^3 - 12.0*x^2*y + 6.0*x^2 - 16.0*x*y^3 + 8.0*x*y - 6.0*x + 2.0*y^2 - 2.0*y;end,
    )

]

for k = 1:length(TestCase)
    @info "Running TestCase $k..."
    ufunc, vfunc, gfunc, hfunc = TestCase[k]

    A = compute_fem_bdm_mass_matrix(a*ones(get_ngauss(mmesh)), b*ones(get_ngauss(mmesh)), mmesh)
    B = compute_fem_bdm_div_matrix(mmesh)
    C = compute_fem_bdm_skew_matrix(mmesh)

    D = [A B' C'
        B spzeros(2mmesh.nelem, 3mmesh.nelem)
        C spzeros(mmesh.nelem, 3mmesh.nelem)]

    gD = bcedge(mmesh)
    t1 = eval_f_on_gauss_pts(gfunc, mmesh)
    t2 = eval_f_on_gauss_pts(hfunc, mmesh)
    f1 = compute_fvm_source_term(t1, mmesh)
    f2 = compute_fvm_source_term(t2, mmesh)

    rhs = [zeros(2mmesh.ndof); f1; f2; zeros(mmesh.nelem)]

    sol = D\rhs
    u = sol[2mmesh.ndof+1:2mmesh.ndof+2mmesh.nelem]
    close("all")
    figure(figsize=(15, 10))
    subplot(231)
    title("Reference")
    xy = fvm_nodes(mmesh)
    x, y = xy[:,1], xy[:,2]
    uf = ufunc.(x, y)
    visualize_scalar_on_fvm_points(uf, mmesh)
    subplot(232)
    title("Numerical")
    visualize_scalar_on_fvm_points(u[1:mmesh.nelem], mmesh)
    subplot(233)
    title("Absolute Error")
    visualize_scalar_on_fvm_points( abs.(u[1:mmesh.nelem] - uf) , mmesh)

    subplot(234)
    xy = fvm_nodes(mmesh)
    x, y = xy[:,1], xy[:,2]
    uf = vfunc.(x, y)
    visualize_scalar_on_fvm_points(uf, mmesh)
    subplot(235)
    visualize_scalar_on_fvm_points(u[mmesh.nelem+1:end], mmesh)
    subplot(236)
    visualize_scalar_on_fvm_points( abs.(u[mmesh.nelem+1:end] - uf) , mmesh)
    savefig("varying_elasticity$k.png")
end


const MIN_PHRED = Phred(1)
const MAX_PHRED = Phred(Int('~') - 33)

"""Convert error probability to PHRED score"""
function p_to_phred(p::Prob)
    return Phred(min(round(-10.0 * log10(p)), MAX_PHRED))
end

function p_to_phred(x::Vector{LogProb})
    return Phred[p_to_phred(p) for p in x]
end

"""Convert PHRED score to log error probability"""
@generated function phred_to_log_p(x)
    return quote
        return x / (-10.0)
    end
end

"""Convert PHRED score to error probability"""
function phred_to_p(q::Phred)
    return exp10(phred_to_log_p(q))
end

function phred_to_p(x::Vector{Phred})
    return exp10.(phred_to_log_p(x))
end

@generated function normalize(parts)
    return quote
        return parts / sum(parts)
    end
end

function cap_phreds(phreds::Vector{Phred}, max_phred::Phred)
    if max_phred < 1
        error("max phred value must be positive")
    end
    return Phred[min(p, max_phred) for p in phreds]
end