### A Pluto.jl notebook ###
# v0.12.20

#
# This example assumes below packages are available in your environment
#

using Markdown
using InteractiveUtils

# ╔═╡ dcb4d418-5ec2-11eb-29d8-214f38b4d3ae
using Pkg, DrWatson

# ╔═╡ d20c24f8-5ec2-11eb-3d45-d97fedebee8e
begin
	@quickactivate "StatsModelComparisons"
	using StatsModelComparisons
	using StanSample, StatsFuns, StatsPlots
	using DataFrames, CSV
end

# ╔═╡ e3552750-5e9f-11eb-324b-8df36d671c79
begin
	ProjDir = joinpath(mc_path, "..", "examples", "roaches")
	df = CSV.read(joinpath(ProjDir, "roachdata.csv"), DataFrame)
	df.roach1 = df.roach1 / 100
end;

# ╔═╡ e3691972-5e9f-11eb-20b2-a766ef562598
roaches1_stan = "
data {
  int<lower=0> N;
  int<lower=0> K;
  vector[N] exposure2;
  vector[N] roach1;
  vector[N] senior;
  vector[N] treatment;
  int y[N];
}
transformed data {
  vector[N] log_expo;
  log_expo = log(exposure2);
}
parameters {
  vector[K] beta;
}
transformed parameters {
   vector[N] eta;
   eta = log_expo + beta[1] + beta[2] * roach1 + beta[3] * treatment
                + beta[4] * senior;
}
model {
  y ~ poisson_log(eta);
}
generated quantities {
  vector[N] log_lik;
  for (i in 1:N)
    log_lik[i] = poisson_log_lpmf(y[i] | eta[i]);
}
";

# ╔═╡ e371fde4-5e9f-11eb-39ad-8d25ac3034fc
begin
	data = (N = size(df, 1), K = 4, y = Int.(df.y), roach1=df.roach1,
 		exposure2=df.exposure2, senior=df.senior, treatment=df.treatment)
	n = size(df, 1)
	tmpdir=joinpath(ProjDir, "tmp")
	sm1 = SampleModel("roaches1", roaches1_stan; tmpdir)
	rc1 = stan_sample(sm1; data)
end;

# ╔═╡ e3a796f2-5e9f-11eb-20bf-9de765acb853
if success(rc1)
  read_summary(sm1, true)
  nt1 = read_samples(sm1)

  # Compute LOO and standard error
  log_lik1 = nt1.log_lik'
  loo1, loos1, pk1 = psisloo(log_lik1)
  elpd_loo1 = sum(loos1)
  se_elpd_loo1 = std(loos1) * sqrt(n)
  (elpd_loo = elpd_loo1, se_elpd_loo = se_elpd_loo1)
end

# ╔═╡ abc94b86-5ea9-11eb-1d5d-578926de3257
if success(rc1)
	# Check the shape parameter k of the generalized Pareto distribution
	pk_qualify(pk1) |> display
end

# ╔═╡ e3b82668-5e9f-11eb-1641-a9dfed9eb108
pk_plot(pk1, title="PSIS diagnostic plot for poisson-log model.")

# ╔═╡ e3b8fc58-5e9f-11eb-0607-e5424a04df9c
md" ##### Simple negative binomial regression example using the 2nd parametrization of the negative binomial distribution, see section 40.1-3 in the Stan reference guide."

# ╔═╡ 5824d2ae-5ea2-11eb-3545-dd03c4aa7a69
roaches2_stan = "
data {
  int<lower=0> N;
  int<lower=0> K;
  vector[N] exposure2;
  vector[N] roach1;
  vector[N] senior;
  vector[N] treatment;
  int y[N];
}
transformed data {
  vector[N] log_expo;
  log_expo = log(exposure2);
}
parameters {
  real phi;
  vector[K] beta;
}
transformed parameters {
   vector[N] eta;
   eta = log_expo + beta[1] + beta[2] * roach1 + beta[3] * treatment
                + beta[4] * senior;
}
model {
  phi ~ normal(0, 10);
  beta[1] ~ cauchy(0,10);   //prior for the intercept following Gelman 2008
  for(i in 2:K)
   beta[i] ~ cauchy(0,2.5); //prior for the slopes following Gelman 2008
  y ~ neg_binomial_2_log(eta, phi);
}
generated quantities {
 vector[N] log_lik;
 for(i in 1:N){
  log_lik[i] <- neg_binomial_2_log_lpmf(y[i] | eta[i], phi);
 }
}
";

# ╔═╡ e3cab9ca-5e9f-11eb-3b31-c34fabb9e8d1
begin
	sm2 = SampleModel("roaches2", roaches2_stan; tmpdir)
	rc2 = stan_sample(sm2; data)
end;

# ╔═╡ e3ddad0a-5e9f-11eb-0a7d-150f9f398fe0
if success(rc2)
  read_summary(sm2, true)
  nt2 = read_samples(sm2)

  # Compute LOO and standard error
  log_lik2 = nt2.log_lik'
  loo2, loos2, pk2 = psisloo(log_lik2)
  elpd_loo2 = sum(loos2)
  se_elpd_loo2 = std(loos2) * sqrt(n)
  (elpd_loo = elpd_loo2, se_elpd_loo = se_elpd_loo2)
end

# ╔═╡ ce509a0c-5ea8-11eb-3f2e-01023f29a1e3
if success(rc2)
	  # Check the shape parameter k of the generalized Pareto distribution
	pk_qualify(pk2)
end

# ╔═╡ e3f103f0-5e9f-11eb-159a-452961ed2619
pk_plot(pk2; title="PSIS diagnostic plot for neg-binomial model.", leg=:topright)

# ╔═╡ Cell order:
# ╠═dcb4d418-5ec2-11eb-29d8-214f38b4d3ae
# ╠═d20c24f8-5ec2-11eb-3d45-d97fedebee8e
# ╠═e3552750-5e9f-11eb-324b-8df36d671c79
# ╠═e3691972-5e9f-11eb-20b2-a766ef562598
# ╠═e371fde4-5e9f-11eb-39ad-8d25ac3034fc
# ╠═e3a796f2-5e9f-11eb-20bf-9de765acb853
# ╠═abc94b86-5ea9-11eb-1d5d-578926de3257
# ╠═e3b82668-5e9f-11eb-1641-a9dfed9eb108
# ╟─e3b8fc58-5e9f-11eb-0607-e5424a04df9c
# ╠═5824d2ae-5ea2-11eb-3545-dd03c4aa7a69
# ╠═e3cab9ca-5e9f-11eb-3b31-c34fabb9e8d1
# ╠═e3ddad0a-5e9f-11eb-0a7d-150f9f398fe0
# ╠═ce509a0c-5ea8-11eb-3f2e-01023f29a1e3
# ╠═e3f103f0-5e9f-11eb-159a-452961ed2619


module LssTests

using Revise
using Subsystems
using BlockArrays
using Test
using Infiltrator

@testset "LSS" begin
    @testset "has correct" begin
        N = 5;
        ne = rand(1:5, N)
        me = rand(0:2, N)
        pe = rand(1:5, N)

        As = [rand(ne[i],ne[i]) for i in 1:N]
        Bs = [rand(ne[i],me[i]) for i in 1:N]
        Cs = [rand(pe[i],ne[i]) for i in 1:N]
        K = [[rand(ne[i],ne[j]) for j in 1:N] for i in 1:N]
        Inc = rand(Bool, N, N);

        subs = [LdtiSubsystem(i, As[i], Bs[i], Cs[i]) for i in 1:N]

        Abe = BlockArray{Real}(zeros(sum(ne), sum(ne)), ne, ne)
        Bbe = BlockArray{Real}(zeros(sum(ne), sum(me)), ne, me)
        Cbe = BlockArray{Real}(zeros(sum(pe), sum(ne)), pe, ne)

        for i in 1:N
            Abe[Block(i,i)] = As[i]
            Bbe[Block(i,i)] = Bs[i]
            Cbe[Block(i,i)] = Cs[i]

            for j in vcat(1:i-1, i+1:N)
                if Inc[i,j]
                    addNeighbour(subs[i], subs[j], K[i][j], false)
                    view(Abe, Block(i,j)) .= K[i][j]
                end
            end
        end

        @testset "dimensions" begin
            (ng, mg, pg) = _getdimarrays(subs)
            @test ng == ne
            @test mg == me
            @test pg == pe
        end

        @testset "Block matrices" begin
            ms = tomonolith(subs)
            @test Abe == ms.A
            @test Bbe == ms.B
            @test Cbe == ms.C
        end
    end
end

end


#=
simple integer knapsack problem with convex objective

minₓ f(c ⊙ x) :
b' x ≤ B
x ≥ 1
x ∈ ℤⁿ

where c > 0, b > 0, B > 0, and f : ℝⁿ → ℝ is a convex spectral function
=#

struct Knapsack <: ExampleInstance
    n::Int
    f::VecSpecExt
    relax_int::Bool # relax integrality
end

function build(inst::Knapsack)
    n = inst.n

    # generate parameters
    c = 2 * rand(n)
    b = rand(n)
    B = 3 * sum(b)

    # build model
    model = JuMP.Model()
    JuMP.@variable(model, x[1:n] >= 1)
    JuMP.@constraint(model, dot(b, x) <= B)

    if !inst.relax_int
        JuMP.set_integer.(x)
    end

    JuMP.@variable(model, y)
    JuMP.@objective(model, Min, y)
    add_homog_spectral(inst.f, n, vcat(y, c .* x), model)

    # save for use in tests
    model.ext[:x] = x
    model.ext[:c] = c

    return model
end

function test_extra(inst::Knapsack, model::JuMP.Model)
    stat = JuMP.termination_status(model)
    @test stat in OPT_OR_LIMIT
    (stat == MOI.OPTIMAL) || return

    tol = eps()^0.2
    x = JuMP.value.(model.ext[:x])
    if !inst.relax_int
        @test x ≈ round.(Int, x) atol = tol rtol = tol
    end
    @test all(>(1 - tol), x)
    # check conic feasibility
    c = JuMP.value.(model.ext[:c])
    obj_result = get_val(pos_only(c .* x), inst.f)
    @test JuMP.objective_value(model) ≈ obj_result atol = tol rtol = tol
    return
end


"""
    function modeConvergence(X::AbstractArray, PODfun, stops::AbstractArray{<: AbstractRange}, numModes::Int)
Modal convergence check based on l2-norm of modes. The array `stops` contains
the ranges to investigate where `stops[end]` is used as the reference modes. The `numModes`
largest modes are compared to reduce the computational time. The function used to
POD the data is supplied through `PODfun`
"""
function modeConvergence(X::AbstractArray, PODfun, stops::AbstractArray{<: AbstractRange}, numModes::Int)

    # The full POD which the subsets are compared to
    maxPOD = PODfun(X[:,stops[end]])[1].modes[:,1:numModes]

    numPODs = size(stops,1)
    output = zeros(eltype(X), numModes, numPODs) # Allocate output

    for i = 1:numPODs-1

        podRes = PODfun(X[:,stops[i]])[1].modes[:,1:numModes]

        # Compare the first numModes modes.
        for j = 1:numModes

            # normalize the mode before comparing
            maxPODnorm = norm(maxPOD[:,j])
            podResnorm = norm(podRes[:,j])
            maxComp = maxPOD[:,j]/maxPODnorm
            podComp = podRes[:,j]/podResnorm

            # modes do not have a specific sign, compare the positive and negative
            # to find minimum.
            l2norm1 = norm(maxComp - podComp)
            l2norm2 = norm(maxComp + podComp)
            l2norm = minimum([l2norm1, l2norm2])
            output[j,i] = l2norm
        end

    end

    return output
end

"""
    function modeConvergence!(loadFun, PODfun, stops::AbstractArray{<: AbstractRange}, numModes::Int)
Same as `modeConvergence(X::AbstractArray, PODfun, stops::AbstractArray{<: AbstractRange}, numModes::Int)`
but here the data is reloaded for each comparision so that an inplace POD method can be used to reduce
maximum memory usage.
"""
function modeConvergence!(loadFun, PODfun, stops::AbstractArray{<: AbstractRange}, numModes::Int)

    X = loadFun()

    # The full POD which the subsets are compared to
    maxPOD = PODfun(X[:,stops[end]])[1].modes[:,1:numModes]

    numPODs = size(stops,1)
    output = zeros(eltype(X), numModes, numPODs) # Allocate output

    for i = 1:numPODs-1

        X = loadFun()

        podRes = PODfun(X[:,stops[i]])[1].modes[:,1:numModes]

        # Compare the first numModes modes.
        for j = 1:numModes

            # normalize the mode before comparing
            maxPODnorm = norm(maxPOD[:,j])
            podResnorm = norm(podRes[:,j])
            maxComp = maxPOD[:,j]/maxPODnorm
            podComp = podRes[:,j]/podResnorm

            # modes do not have a specific sign, compare the positive and negative
            # to find minimum.
            l2norm1 = norm(maxComp - podComp)
            l2norm2 = norm(maxComp + podComp)
            l2norm = minimum([l2norm1, l2norm2])
            output[j,i] = l2norm
        end

    end

    return output
end


import Temporal.acf  # used for running autocorrelation function

using Statistics

"""
    diffn(x::Array{Float64}; n::Int64=1)::Array{Float64}

Lagged differencing
"""
function diffn(x::Array{Float64}; n::Int64=1)::Array{Float64}
    @assert n<size(x,1) && n>0 "Argument n out of bounds."
    dx = zeros(size(x))
    dx[1:n] .= NaN
    @inbounds for i=n+1:size(x,1)
        dx[i] = x[i] - x[i-n]
    end
    return dx
end

"""
    diffn(X::Array{Float64,2}; n::Int64=1)::Array{Float64}

Lagged differencing
"""
function diffn(X::Array{Float64,2}; n::Int64=1)::Matrix{Float64}
    @assert n<size(X,1) && n>0 "Argument n out of bounds."
    dx = zeros(size(X))
    @inbounds for j = 1:size(X,2)
        dx[:,j] = diffn(X[:,j], n=n)
    end
    return dx
end

"""
(Adapted from StatsBase: https://raw.githubusercontent.com/JuliaStats/StatsBase.jl/master/src/scalarstats.jl)

`Compute the mode of an arbitrary array::Array{Float64}`
"""
function mode(a::AbstractArray{Float64})
    isempty(a) && error("mode: input array cannot be empty.")
    cnts = Dict{Float64,Int64}()
    # first element
    mc = 1
    mv = a[1]
    cnts[mv] = 1
    # find the mode along with table construction
    @inbounds for i = 2 : length(a)
        x = a[i]
        if haskey(cnts, x)
            c = (cnts[x] += 1)
            if c > mc
                mc = c
                mv = x
            end
        else
            cnts[x] = 1
            # in this case: c = 1, and thus c > mc won't happen
        end
    end
    return mv
end

"""
    runmean(x::Array{Float64}; n::Int64=10, cumulative::Bool=true)::Array{Float64}

Compute a running or rolling arithmetic mean of an array.
"""
function runmean(x::Array{Float64}; n::Int64=10, cumulative::Bool=true)::Array{Float64}
    @assert n<size(x,1) && n>1 "Argument n is out of bounds."
    out = zeros(size(x))
    out[1:n-1] .= NaN
    if cumulative
        fi = 1.0:size(x,1)
        @inbounds for i = n:size(x,1)
            out[i] = sum(x[1:i])/fi[i]
        end
    else
        @inbounds for i = n:size(x,1)
            out[i] = mean(x[i-n+1:i])
        end
    end
    return out
end

"""
    runsum(x::Array{Float64}; n::Int64=10, cumulative::Bool=true)::Array{Float64}

Compute a running or rolling summation of an array.
"""
function runsum(x::Array{Float64}; n::Int64=10, cumulative::Bool=true)::Array{Float64}
    @assert n<size(x,1) && n>1 "Argument n is out of bounds."
    if cumulative
        out = cumsum(x, dims=1)
        out[1:n-1] .= NaN
    else
        out = zeros(size(x))
        out[1:n-1] .= NaN
        @inbounds for i = n:size(x,1)
            out[i] = sum(x[i-n+1:i])
        end
    end
    return out
end

"""
    wilder_sum(x::Array{Float64}; n::Int64=10)::Array{Float64}

Welles Wilder summation of an array
"""
function wilder_sum(x::Array{Float64}; n::Int64=10)::Array{Float64}
    @assert n<size(x,1) && n>0 "Argument n is out of bounds."
    nf = float(n)  # type stability -- all arithmetic done on floats
    out = zeros(size(x))
    out[1] = x[1]
    @inbounds for i = 2:size(x,1)
        out[i] = x[i] + out[i-1]*(nf-1.0)/nf
    end
    return out
end

"""
    runmad(x::Array{Float64}; n::Int64=10, cumulative::Bool=true, fun::Function=median)::Array{Float64}

Compute the running or rolling mean absolute deviation of an array
"""
function runmad(x::Array{Float64}; n::Int64=10, cumulative::Bool=true, fun::Function=median)::Array{Float64}
    @assert n<size(x,1) && n>1 "Argument n is out of bounds."
    out = zeros(size(x))
    out[1:n-1] .= NaN
    center = 0.0
    if cumulative
        fi = collect(1.0:size(x,1))
        @inbounds for i = n:size(x,1)
            center = fun(x[1:i])
            out[i] = sum(abs.(x[1:i].-center)) / fi[i]
        end
    else
        fn = float(n)
        @inbounds for i = n:size(x,1)
            center = fun(x[i-n+1:i])
            out[i] = sum(abs.(x[i-n+1:i].-center)) / fn
        end
    end
    return out
end

"""
    runvar(x::Array{Float64}; n::Int64=10, cumulative=true)::Array{Float64}

Compute the running or rolling variance of an array
"""
function runvar(x::Array{Float64}; n::Int64=10, cumulative=true)::Array{Float64}
    @assert n<size(x,1) && n>1 "Argument n is out of bounds."
    out = zeros(size(x))
    out[1:n-1] .= NaN
    if cumulative
        @inbounds for i = n:size(x,1)
            out[i] = var(x[1:i])
        end
    else
        @inbounds for i = n:size(x,1)
            out[i] = var(x[i-n+1:i])
        end
    end
    return out
end

"""
    runsd(x::Array{Float64}; n::Int64=10, cumulative::Bool=true)::Array{Float64}

Compute the running or rolling standard deviation of an array
"""
function runsd(x::Array{Float64}; n::Int64=10, cumulative::Bool=true)::Array{Float64}
    return sqrt.(runvar(x, n=n, cumulative=cumulative))
end

"""
    runcov(x::Array{Float64}, y::Array{Float64}; n::Int64=10, cumulative::Bool=true)::Array{Float64}

Compute the running or rolling covariance of two arrays
"""
function runcov(x::Array{Float64}, y::Array{Float64}; n::Int64=10, cumulative::Bool=true)::Array{Float64}
    @assert length(x) == length(y) "Dimension mismatch: length of `x` not equal to length of `y`."
    @assert n<size(x,1) && n>1 "Argument n is out of bounds."
    out = zeros(size(x))
    out[1:n-1] .= NaN
    if cumulative
        @inbounds for i = n:length(x)
            out[i] = cov(x[1:i], y[1:i])
        end
    else
        @inbounds for i = n:length(x)
            out[i] = cov(x[i-n+1:i], y[i-n+1:i])
        end
    end
    return out
end

"""
    runcor(x::Array{Float64}, y::Array{Float64}; n::Int64=10, cumulative::Bool=true)::Array{Float64}

Compute the running or rolling correlation of two arrays
"""
function runcor(x::Array{Float64}, y::Array{Float64}; n::Int64=10, cumulative::Bool=true)::Array{Float64}
    @assert length(x) == length(y) "Dimension mismatch: length of `x` not equal to length of `y`."
    @assert n<size(x,1) && n>1 "Argument n is out of bounds."
    out = zeros(size(x))
    out[1:n-1] .= NaN
    if cumulative
        @inbounds for i = n:length(x)
            out[i] = cor(x[1:i], y[1:i])
        end
    else
        @inbounds for i = n:length(x)
            out[i] = cor(x[i-n+1:i], y[i-n+1:i])
        end
    end
    return out
end

"""
    runmax(x::Array{Float64}; n::Int64=10, cumulative::Bool=true, inclusive::Bool=true)::Array{Float64}

Compute the running or rolling maximum of an array
"""
function runmax(x::Array{Float64}; n::Int64=10, cumulative::Bool=true, inclusive::Bool=true)::Array{Float64}
    @assert n<size(x,1) && n>1 "Argument n is out of bounds."
    out = zeros(size(x))
    if inclusive
        if cumulative
            out[n] = maximum(x[1:n])
            @inbounds for i = n+1:size(x,1)
                out[i] = max(out[i-1], x[i])
            end
        else
            @inbounds for i = n:size(x,1)
                out[i] = maximum(x[i-n+1:i])
            end
        end
        out[1:n-1] .= NaN
        return out
    else
        if cumulative
            out[n+1] = maximum(x[1:n])
            @inbounds for i = n+1:size(x,1)-1
                out[i+1] = max(out[i-1], x[i-1])
            end
        else
            @inbounds for i = n:size(x,1)-1
                out[i+1] = maximum(x[i-n+1:i])
            end
        end
        out[1:n] .= NaN
        return out
    end
end

"""
    runmin(x::Array{Float64}; n::Int64=10, cumulative::Bool=true, inclusive::Bool=true)::Array{Float64}

Compute the running or rolling minimum of an array
"""
function runmin(x::Array{Float64}; n::Int64=10, cumulative::Bool=true, inclusive::Bool=true)::Array{Float64}
    @assert n<size(x,1) && n>1 "Argument n is out of bounds."
    out = zeros(size(x))
    if inclusive
        if cumulative
            out[n] = minimum(x[1:n])
            @inbounds for i = n+1:size(x,1)
                out[i] = min(out[i-1], x[i])
            end
        else
            @inbounds for i = n:size(x,1)
                out[i] = minimum(x[i-n+1:i])
            end
        end
        out[1:n-1] .= NaN
        return out
    else
        if cumulative
            out[n+1] = minimum(x[1:n])
            @inbounds for i = n+1:size(x,1)-1
                out[i+1] = min(out[i-1], x[i-1])
            end
        else
            @inbounds for i = n:size(x,1)-1
                out[i+1] = minimum(x[i-n+1:i])
            end
        end
        out[1:n] .= NaN
        return out
    end
end

"""
    runquantile(x::Vector{Float64}; p::Float64=0.05, n::Int=10, cumulative::Bool=true)::Vector{Float64}

Compute the running/rolling quantile of an array
"""
function runquantile(x::Array{Float64}; p::Float64=0.05, n::Int=10, cumulative::Bool=true)::Array{Float64}
    @assert n<size(x,1) && n>1 "Argument n is out of bounds."
    out = zeros(Float64, (size(x,1), size(x,2)))
    if cumulative
        @inbounds for j in 1:size(x,2), i in 2:size(x,1)
            out[i,j] = quantile(x[1:i,j], p)
        end
        out[1,:] .= NaN
    else
        @inbounds for j in 1:size(x,2), i in n:size(x,1)
            out[i,j] = quantile(x[i-n+1:i,j], p)
        end
        out[1:n-1,:] .= NaN
    end
    return out
end

"""
```
function runacf(x::Array{Float64};
                n::Int = 10,
                maxlag::Int = n-3,
                lags::AbstractArray{Int,1} = 0:maxlag,
                cumulative::Bool = true)::Matrix{Float64}
```

Compute the running/rolling autocorrelation of a vector.
"""
function runacf(x::Array{Float64};
                n::Int = 10,
                maxlag::Int = n-3,
                lags::AbstractArray{Int,1} = 0:maxlag,
                cumulative::Bool = true)::Matrix{Float64}
    @assert size(x, 2) == 1 "Autocorrelation input array must be one-dimensional"
    N = size(x, 1)
    @assert n < N && n > 0
    if length(lags) == 1 && lags[1] == 0
        return ones((N, 1))
    end
    out = zeros((N, length(lags))) * NaN
    if cumulative
        @inbounds for i in n:N
            out[i,:] = acf(x[1:i], lags=lags)
        end
    else
        @inbounds for i in n:N
            out[i,:] = acf(x[i-n+1:i], lags=lags)
        end
    end
    return out
end


# Copyright 2017, Iain Dunning, Joey Huchette, Miles Lubin, and contributors    #src
# This Source Code Form is subject to the terms of the Mozilla Public License   #src
# v.2.0. If a copy of the MPL was not distributed with this file, You can       #src
# obtain one at https://mozilla.org/MPL/2.0/.                                   #src

# # The multi-commodity flow problem

# JuMP implementation of the multicommodity transportation model AMPL: A Modeling
# Language for Mathematical Programming, 2nd ed by Robert Fourer, David Gay, and
# Brian W. Kernighan 4-1.
#
# Originally contributed by Louis Luangkesorn, February 26, 2015.

using JuMP
import GLPK
import Test

function example_multi(; verbose = true)
    orig = ["GARY", "CLEV", "PITT"]
    dest = ["FRA", "DET", "LAN", "WIN", "STL", "FRE", "LAF"]
    prod = ["bands", "coils", "plate"]
    numorig = length(orig)
    numdest = length(dest)
    numprod = length(prod)
    ## supply(prod, orig) amounts available at origins
    supply = [
        400    700    800;
        800   1600   1800;
        200    300    300
    ]
    ## demand(prod, dest) amounts required at destinations
    demand = [
        300   300   100    75   650   225   250;
        500   750   400   250   950   850   500;
        100   100     0    50   200   100   250
    ]
    ## limit(orig, dest) of total units from any origin to destination
    limit = [625.0 for j in 1:numorig, i in 1:numdest]
    ## cost(dest, orig, prod) Shipment cost per unit
    cost = reshape([
        [
            [  30,   10,    8,   10,   11,   71,    6];
            [  22,    7,   10,    7,   21,   82,   13];
            [  19,   11,   12,   10,   25,   83,   15]
        ];
        [
            [  39,   14,   11,   14,   16,   82,    8];
            [  27,    9,   12,    9,   26,   95,   17];
            [  24,   14,   17,   13,   28,   99,   20]
        ];
        [
            [  41,   15,   12,   16,   17,   86,    8];
            [  29,    9,   13,    9,   28,   99,   18];
            [  26,   14,   17,   13,   31,  104,   20]
        ]
    ], 7, 3, 3)
    ## DECLARE MODEL
    multi = Model(GLPK.Optimizer)
    ## VARIABLES
    @variable(multi, trans[1:numorig, 1:numdest, 1:numprod] >= 0)
    ## OBJECTIVE
    @objective(
        multi,
        Max,
        sum(
            cost[j, i, p] * trans[i, j, p]
            for i in 1:numorig, j in 1:numdest, p in 1:numprod
        )
    )
    ## CONSTRAINTS
    ## Supply constraint
    @constraint(
        multi,
        supply_con[i in 1:numorig, p in 1:numprod],
        sum(trans[i, j, p] for j in 1:numdest) == supply[p, i]
    )
    ## Demand constraint
    @constraint(
        multi,
        demand_con[j in 1:numdest, p in 1:numprod],
        sum(trans[i, j, p] for i in 1:numorig) == demand[p, j]
    )
    ## Total shipment constraint
    @constraint(
        multi,
        total_con[i in 1:numorig, j in 1:numdest],
        sum(trans[i, j, p] for p in 1:numprod) - limit[i, j] <= 0
    )
    optimize!(multi)
    Test.@test termination_status(multi) == MOI.OPTIMAL
    Test.@test primal_status(multi) == MOI.FEASIBLE_POINT
    Test.@test objective_value(multi) == 225_700.0
    if verbose
        println("RESULTS:")
        for i in 1:length(orig)
            for j in 1:length(dest)
                for p in 1:length(prod)
                    print(" $(prod[p]) $(orig[i]) $(dest[j]) = $(value(trans[i, j, p]))\t")
                end
                println()
            end
        end
    end
    return
end

example_multi()


using Test

include("nfa.jl")

M = NFA(Dict(
             ("q0", '0') => ["q0", "q1"],
             ("q0", '1') => ["q0", "q2"],
             ("q1", '0') => ["q3",],
             ("q2", '0') => ["q2", "q3"],
             ("q2", '1') => ["q3",],
             ("q3", '0') => ["q3",],
             ("q3", '1') => ["q3",]
       ), "q0", Set(["q3"]))

@test !evaluate!(M, "")
@test !evaluate!(M, "01")
@test evaluate!(M, "00")
@test evaluate!(M, "1111110000")

# Taken from here:
# https://www.cs.wcupa.edu/rkline/assets/img/FCS/nfas/n2d2.jpg?1379001651
# regex: (a* + (ab)*) b*
M = NFA(Dict(
             ("1",  λ)  => ["2", "3"],
             ("2",  λ)  => ["5"],
             ("2", 'a') => ["2"],
             ("5", 'b') => ["5"],
             ("3", 'a') => ["4"],
             ("3",  λ)  => ["5"],
             ("4", 'b') => ["3"],
       ), "1", Set(["5"]))

@test evaluate!(M, "")
@test evaluate!(M, "b")
@test evaluate!(M, "a")
@test evaluate!(M, "ab")
@test evaluate!(M, "aa")
@test evaluate!(M, "abab")
@test !evaluate!(M, "aba")


using Splines
using Base.Test

# write your own tests here
include("testBase.jl")
include("testMath.jl")
include("testCalculus.jl")


using Test
using SVGPlots

x = collect(-4:0.1:4)
figure(;width=800, height=600)
subplot(2,3,1)
plot(x, x)
hold(true)
plot(x, x.^2)
subplot(2,3,2)
plot(x, x.^2)
subplot(2,3,3)
plot(x, x.^3)
subplot(2,3,4)
plot(x, x.^4)
subplot(2,3,5)
plot(x, x.^5)
subplot(2,3,6)
plot(x, x.^6)
display(gcf())
figure()
plot(x, x)
hold(true)
plot(x, x.^2, linetype = :dotted)
plot(x, x.^3, linetype = :dashed)
plot(x, x.^4, linetype = :dashdotted)
plot(x, x.^5, linetype = :none)
title("This is a title")
xlabel("x")
ylabel("y")
display(gcf())
figure()
subplot(2,2,1)
plot(x, x)
subplot(2,2,2)
plot(x, x.^2)
subplot(2,2,3)
plot(x, x.^3)
subplot(2,2,4)
plot(x, x.^4)
hold(true)
plot(x, x.^5)
title("test")
display(gcf())


# ---
# title: 1643. Kth Smallest Instructions
# id: problem1643
# author: Tian Jun
# date: 2020-10-31
# difficulty: Hard
# categories: Dynamic Programming
# link: <https://leetcode.com/problems/kth-smallest-instructions/description/>
# hidden: true
# ---
#
# Bob is standing at cell `(0, 0)`, and he wants to reach `destination`: `(row,
# column)`. He can only travel **right** and **down**. You are going to help Bob
# by providing **instructions** for him to reach `destination`.
#
# The **instructions** are represented as a string, where each character is
# either:
#
#   * `'H'`, meaning move horizontally (go **right** ), or
#   * `'V'`, meaning move vertically (go **down** ).
#
# Multiple **instructions** will lead Bob to `destination`. For example, if
# `destination` is `(2, 3)`, both `"HHHVV"` and `"HVHVH"` are valid
# **instructions**.
#
# However, Bob is very picky. Bob has a lucky number `k`, and he wants the `kth`
# **lexicographically smallest instructions** that will lead him to
# `destination`. `k` is **1-indexed**.
#
# Given an integer array `destination` and an integer `k`, return _the_`kth`
# _**lexicographically smallest instructions**  that will take Bob to
# _`destination`.
#
#
#
# **Example 1:**
#
# ![](https://assets.leetcode.com/uploads/2020/10/12/ex1.png)
#
#
#
#     Input: destination = [2,3], k = 1
#     Output: "HHHVV"
#     Explanation: All the instructions that reach (2, 3) in lexicographic order are as follows:
#     ["HHHVV", "HHVHV", "HHVVH", "HVHHV", "HVHVH", "HVVHH", "VHHHV", "VHHVH", "VHVHH", "VVHHH"].
#
#
# **Example 2:**
#
# **![](https://assets.leetcode.com/uploads/2020/10/12/ex2.png)**
#
#
#
#     Input: destination = [2,3], k = 2
#     Output: "HHVHV"
#
#
# **Example 3:**
#
# **![](https://assets.leetcode.com/uploads/2020/10/12/ex3.png)**
#
#
#
#     Input: destination = [2,3], k = 3
#     Output: "HHVVH"
#
#
#
#
# **Constraints:**
#
#   * `destination.length == 2`
#   * `1 <= row, column <= 15`
#   * `1 <= k <= nCr(row + column, row)`, where `nCr(a, b)` denotes `a` choose `b`.
#
#
## @lc code=start
using LeetCode

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


# ---
# title: 327. Count of Range Sum
# id: problem327
# author: Tian Jun
# date: 2020-10-31
# difficulty: Hard
# categories: Binary Search, Divide and Conquer, Sort, Binary Indexed Tree, Segment Tree
# link: <https://leetcode.com/problems/count-of-range-sum/description/>
# hidden: true
# ---
#
# Given an integer array `nums`, return the number of range sums that lie in
# `[lower, upper]` inclusive.
# Range sum `S(i, j)` is defined as the sum of the elements in `nums` between
# indices `i` and `j` (`i` ≤ `j`), inclusive.
#
# **Note:**
# A naive algorithm of _O_ ( _n_ 2) is trivial. You MUST do better than that.
#
# **Example:**
#
#
#
#     Input: _nums_ = [-2,5,-1], _lower_ = -2, _upper_ = 2,
#     Output: 3
#     Explanation: The three ranges are : [0,0], [2,2], [0,2] and their respective sums are: -2, -1, 2.
#
#
#
#
# **Constraints:**
#
#   * `0 <= nums.length <= 10^4`
#
#
## @lc code=start
using LeetCode

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


# Bogumił Kamiński, 2022

# Codes for chapter 9

# Code for section 9.1

using DataFrames
using CSV
using Plots
puzzles = CSV.read("puzzles.csv", DataFrame);

using Statistics
plays_lo = median(puzzles.NbPlays)
puzzles.NbPlays .> plays_lo

puzzles.NbPlays > plays_lo

rating_lo = 1500
rating_hi = quantile(puzzles.Rating, 0.99)
rating_lo .< puzzles.Rating .< rating_hi

row_selector = (puzzles.NbPlays .> plays_lo) .&&
               (rating_lo .< puzzles.Rating .< rating_hi)

sum(row_selector)
count(row_selector)

# Code for listing 9.1

good = puzzles[row_selector, ["Rating", "Popularity"]]

# Code for plotting histograms

plot(histogram(good.Rating; label="Rating"),
     histogram(good.Popularity; label="Popularity"))

# Code for column selectors

puzzles[1, "Rating"]

puzzles[:, "Rating"]

row1 = puzzles[1, ["Rating", "Popularity"]]

row1["Rating"]
row1[:Rating]
row1[1]
row1.Rating
row1."Rating"

good = puzzles[row_selector, ["Rating", "Popularity"]]

good[1, "Rating"]
good[1, :]
good[:, "Rating"]
good[:, :]

names(puzzles, ["Rating", "Popularity"])
names(puzzles, [:Rating, :Popularity])
names(puzzles, [4, 6])
names(puzzles, [false, false, false, true, false, true, false, false, false])
names(puzzles, r"Rating")
names(puzzles, Not([4, 6]))
names(puzzles, Not(r"Rating"))
names(puzzles, Between("Rating", "Popularity"))
names(puzzles, :)
names(puzzles, All())
names(puzzles, Cols(r"Rating", "NbPlays"))
names(puzzles, Cols(startswith("P")))

names(puzzles, startswith("P"))

names(puzzles, Real)

names(puzzles, AbstractString)

puzzles[:, names(puzzles, Real)]

# Code for row subsetting

df_small = DataFrame(id=1:4)
df_small[[1, 3], :]
df_small[[true, false, true, false], :]
df_small[Not([2, 4]), :]
df_small[Not([false, true, false, true]), :]

df1 = puzzles[:, ["Rating", "Popularity"]];
df2 = puzzles[!, ["Rating", "Popularity"]];

df1 == df2
df1 == puzzles
df2 == puzzles

df1.Rating === puzzles.Rating
df1.Popularity === puzzles.Popularity
df2.Rating === puzzles.Rating
df2.Popularity === puzzles.Popularity

using BenchmarkTools
@btime $puzzles[:, ["Rating", "Popularity"]];
@btime $puzzles[!, ["Rating", "Popularity"]];

puzzles[1, 1]
puzzles[[1], 1]
puzzles[1, [1]]
puzzles[[1], [1]]

# Code for making views

@view puzzles[1, 1]

@view puzzles[[1], 1]

@view puzzles[1, [1]]

@view puzzles[[1], [1]]

@btime $puzzles[$row_selector, ["Rating", "Popularity"]];
@btime @view $puzzles[$row_selector, ["Rating", "Popularity"]];

parentindices(@view puzzles[row_selector, ["Rating", "Popularity"]])

# Code for section 9.2

describe(good)

rating_mapping = Dict{Int, Vector{Int}}()
for (i, rating) in enumerate(good.Rating)
    if haskey(rating_mapping, rating)
        push!(rating_mapping[rating], i)
    else
        rating_mapping[rating] = [i]
    end
end
rating_mapping

good[rating_mapping[2108], :]

unique(good[rating_mapping[2108], :].Rating)

using Statistics
mean(good[rating_mapping[2108], "Popularity"])

ratings = unique(good.Rating)

mean_popularities = map(ratings) do rating
    indices = rating_mapping[rating]
    popularities = good[indices, "Popularity"]
    return mean(popularities)
end

scatter(ratings, mean_popularities;
        xlabel="rating", ylabel="mean popularity", legend=false)

using Loess
model = loess(ratings, mean_popularities);
ratings_predict = float.(sort(ratings))
popularity_predict = predict(model, ratings_predict)

plot!(ratings_predict, popularity_predict; width=5, color="black")

combine(groupby(good, :Rating), :Popularity => mean)


# rmsd100

using MIRT: rmsd100

using Test: @test, @test_throws, @inferred

a = [1 2; 4 5; 7 8]
b = [3 2; 4 6; 8 7]
c = [false true; true false; false false] #avg. 0
d = [false false; false false; false true] #avg. 1
e = [true true; true true; true false] # average 1,0,1,0,4 = sqrt(6/5)
tests = cat(c, d, e, dims = 3)
pred = [0, 100, 100*sqrt(6/5)]
for i in 1:size(tests,1)
	@test rmsd100(a, b ; mask=tests[:,:,i]) == pred[i]
end
@test rmsd100([1, 0], [0, 1]) == 100 # no mask test

#=  not important to support or test.  user can pass mask=ig.mask if needed
a = ones(300,400)
b = ones(300,400)
b[100:200,100:200] .= 0
#so there's a blot of zeros from 100,100 to 200,200
b[200:250,300:320] .= 3
#and a blot of threes over here
ig1 = image_geom(nx = 300,ny = 400,dx = 1,dy = 1)
@test my_rmse(a,b,ig1) > 20
=#


# Interconnections of linear systems

# include("StateSpace.jl");
# include("Utilities.jl");

function Negate(S::StateSpace)::StateSpace
    return StateSpace(S.A,S.B,-S.C,-S.D);
end

function Subsystem(S::StateSpace, inputs::Vector{Int64}, outputs::Vector{Int64})::StateSpace
    # Check IO dimensions.
    if (!isempty(filter(x->(x>S.nu),inputs))) error("Input index exceeds input dimension."); end
    if (!isempty(filter(x->(x>S.ny),outputs))) error("Output index exceeds input dimension."); end

    B = S.B[:,inputs];
    C = S.C[outputs,:];
    D = S.D[outputs,inputs];

    return StateSpace(S.A,B,C,D);
end

function Parallel(S1::StateSpace,S2::StateSpace)::StateSpace
    # Check io dimensions.
    if (S1.nu!=S2.nu || S1.ny!=S2.ny)
        error("Input and output dimensions must match.")
    end

    A = [S1.A zeros(S1.ns,S2.ns);zeros(S2.ns,S1.ns) S2.A];
    B = [S1.B;S2.B];
    C = [S1.C S2.C];
    D = (S1.D+S2.D);
    return StateSpace(A,B,C,D);
end

function Series(S1::StateSpace,S2::StateSpace)::StateSpace
    # Check io dimensions.
    if (S1.ny!=S2.nu)
        error("Output of S1 must match input of S2.");
    end

    A = [S1.A zeros(S1.ns,S2.ns);S2.B*S1.C S2.A];
    B = [S1.B;S2.B*S1.D];
    C = [S2.D*S1.C S2.C];
    D = S2.D*S1.D;

    return StateSpace(A,B,C,D);
end

function Feedback(S1::StateSpace,S2::StateSpace)::StateSpace
    # Check io dimensions -- or provide io options for feedback.
    if (S1.ny != S2.nu || S1.nu != S2.ny)
        error("Incompatible I/O dimensions.")
    end

    A1 = S1.A; B1 = S1.B; C1 = S1.C; D1 = S1.D;
    A2 = S2.A; B2 = S2.B; C2 = S2.C; D2 = S2.D;

    I = eye(S2.ny);
    X = I+D2*D1;

    if (det(X) == 0)
        error("Feedback interconnection is illposed.");
    end

    invX = inv(X);
    Cb = [D2*C1 C2];

    AA1 = [A1 zeros(S1.ns,S2.ns)] - B1*inv(X)Cb;
    AA2 = [B2*C1 A2];
    AA = [AA1;AA2];

    BB1 = B1*invX;
    BB2 = B2*D1;
    BB = [BB1;BB2];

    CC = [C1 zeros(S1.ny,S2.ns)] - D1*invX*Cb;
    DD = D1*inv(X);

    return StateSpace(AA,BB,CC,DD);
end

# Need to implement partial interconnections.
# Eventually, these will replace the above code.

function Parallel(S1::StateSpace,S2::StateSpace,
    in1::Vector{Int64}, in2::Vector{Int64},
    out1::Vector{Int64}, out2::Vector{Int64})::StateSpace
    #                               +------+
    #                 v1 ---------->|      |----------> z1
    #                               |  S1  |
    #                        u1 +-->|      |---+ y1
    #                           |   +------+   |
    #                  u ------>+              O------> y
    #                           |   +------+   |
    #                        u2 +-->|      |---+ y2
    #                               |  S2  |
    #                 v2 ---------->|      |----------> z2
    #                               +------+
    # New system maps (v1;u;v2) -> (z1;y;z2)
    #   -- in1 defines channel u1
    #   -- in2 defines channel u2
    #   -- out1 defines channel y1
    #   -- out2 defines channel y2

    # Check input/output indices are within signal dimensions
    if (!isempty(filter(x->(x>S1.nu),in1))) error("S1: Input index exceeds input dimension."); end
    if (!isempty(filter(x->(x>S1.ny),out1))) error("S1: Output index exceeds output dimension."); end
    if (!isempty(filter(x->(x>S2.nu),in2))) error("S2: Input index exceeds input dimension."); end
    if (!isempty(filter(x->(x>S2.ny),out2))) error("S2: Output index exceeds output dimension."); end

    # Check compatibility of interconnection
    if (length(in1)!= length(in2)) error("Size of in1 and in2 must be same."); end
    if (length(out1)!= length(out2)) error("Size of out1 and out2 must be same."); end

    # Get indices for partitions of system matrices
    in1c = setdiff(1:S1.nu,in1); out1c = setdiff(1:S1.ny,out1); # These can be empty
    in2c = setdiff(1:S2.nu,in2); out2c = setdiff(1:S2.ny,out2); # These can be empty

    nv1 = length(in1c);
    nv2 = length(in2c);
    nz1 = length(out1c);
    nz2 = length(out2c);
    ns1 = S1.ns;
    ns2 = S2.ns;

    # Partitions of S1
    B1 = S1.B[:,in1c]; # defines v1 channel
    B2 = S1.B[:,in1]; # defines u1 channel
    C1 = S1.C[out1c,:];
    C2 = S1.C[out1,:];
    D11 = S1.D[out1c,in1c];
    D12 = S1.D[out1c,in1];
    D21 = S1.D[out1,in1c];
    D22 = S1.D[out1,in1];

    # Partitions of S2
    F1 = S2.B[:,in2]; # defines u2 channel
    F2 = S2.B[:,in2c];  # defines v2 channel
    G1 = S2.C[out2,:];
    G2 = S2.C[out2c,:];
    H11 = S2.D[out2,in2];
    H12 = S2.D[out2,in2c];
    H21 = S2.D[out2c,in2];
    H22 = S2.D[out2c,in2c];

    # Construct the new system matrices
    AA = [S1.A zeros(ns1,ns2);
    zeros(ns2,ns1) S2.A];

    # If nz1, nz2, nv1, nv2 are zero, this will result in empty matrices in the partitions.
    # Creating matrices with empty submatrices is a problem.
    # Using hcat and vcat to create the matrices, eliminates this problem.
    BB1 = hcat(B1,B2,zeros(ns1,nv2));
    BB2 = hcat(zeros(ns2,nv1),F1,F2);
    BB  = vcat(BB1,BB2);

    CC1 = hcat(C1,zeros(nz1,ns2));
    CC2 = hcat(C2,G1);
    CC3 = hcat(zeros(nz2,ns1),G2);
    CC  = vcat(CC1,CC2,CC3);

    DD1 = hcat(D11, D12, zeros(nz1,nv2));
    DD2 = hcat(D21, (D22+H11), H12);
    DD3 = hcat(zeros(nz2,nv1), H21, H22);
    DD = vcat(DD1,DD2,DD3);

    return StateSpace(AA,BB,CC,DD);
end

function Series(S1::StateSpace,S2::StateSpace, out1::Vector{Int64}, in2::Vector{Int64})::StateSpace
    #                                       +------+
    #                                v2 --->|      |
    #                       +------+        |  S2  |-----> y2
    #                       |      |------->|      |
    #              u1 ----->|      |y1   u2 +------+
    #                       |  S1  |
    #                       |      |---> z1
    #                       +------+
    # New system maps (v2;u1) -> (y2;z1)
    # -- out1 defines y1
    # -- in2 defines u2

    # Check input/output indices are within signal dimensions
    if (!isempty(filter(x->(x>S1.ny),out1))) error("S1: Output index exceeds output dimension."); end
    if (!isempty(filter(x->(x>S2.nu),in2)))  error("S2: Input index exceeds input dimension."); end

    # Check compatibility of interconnection
    if (length(out1)!= length(in2)) error("Incompatible I/O dimension for series connection."); end

    # Get indices for partitions of system matrices
    out1c = setdiff(1:S1.ny,out1); # Defines z, These can be empty
    in2c = setdiff(1:S2.nu,in2);   # These can be empty

    nv2 = length(in2c);
    nz1 = length(out1c);

    ns1 = S1.ns;
    ns2 = S2.ns;

    # Partitions of S1 -- follows the diagram
    B = S1.B;
    C1 = S1.C[out1,:]; # x1->y1
    C2 = S1.C[out1c,:]; # x1->z1
    D1 = S1.D[out1,:]; # u->y1
    D2 = S1.D[out1c,:]; # u->z1

    # Partitions of S2
    F1 = S2.B[:,in2c]; # defines v2 channel
    F2 = S2.B[:,in2];  # defines u2 channel
    G = S2.C; # y2
    H1 = S2.D[:,in2c]; # v2->y2
    H2 = S2.D[:,in2];  # u2->y2

    # Construct the new system matrices
    AA = [S1.A zeros(ns1,ns2);
          F2*C1 S2.A];

    # If nz1, nv2 are zero, this will result in empty matrices in the partitions.
    # Creating matrices with empty submatrices is a problem.
    # Using hcat and vcat to create the matrices, eliminates this problem.
    BB1 = hcat(zeros(ns1,nv2),B);
    BB2 = hcat(F1,F2*D1);
    BB  = vcat(BB1,BB2);

    CC1 = hcat(H2*C1,G);
    CC2 = hcat(C2,zeros(nz1,ns2));
    CC  = vcat(CC1,CC2);

    DD1 = hcat(H1,H2*D1);
    DD2 = hcat(zeros(nz1,nv2),D2);
    DD = vcat(DD1,DD2);

    return StateSpace(AA,BB,CC,DD);
end

function Feedback(S1::StateSpace,S2::StateSpace, feedin::Vector{Int64}, feedout::Vector{Int64},FeedbackSign::Int64=-1)::StateSpace

    #                            +------+
    #                v --------->|      |--------> z
    #                            |  S1  |
    #                u --->O---->|      |----+---> y
    #                      |  u1 +------+    |
    #                      |                 |
    #                      +-----[  S2  ]<---+
    # Maps  (v;u) -> (z;y)
    #   -- feedin defines u1
    #   -- feedout defines y

    # Check input/output indices are within signal dimensions
    if (!isempty(filter(x->(x>S1.nu),feedin))) error("S1: Feedback input index exceeds input dimension."); end
    if (!isempty(filter(x->(x>S1.ny),feedout))) error("S1: Feedback output index exceeds output dimension."); end

    # Check compatibility of interconnection
    if (length(feedin)!= S2.ny || length(feedout)!=S2.nu)
        error("Incompatible I/O dimensions for feedback.");
    end

    # Get indices for partitions of system matrices
    in1c = setdiff(1:S1.nu,feedin); # defines v, can be empty
    out1c = setdiff(1:S1.ny,feedout); # defines z, can be empty

    nv = length(in1c);
    nz = length(out1c);
    ns1 = S1.ns;
    ns2 = S2.ns;
    ny = length(feedout);
    nu = length(feedin);

    # Partitions of S1
    A = S1.A;
    B1 = S1.B[:,in1c]; # defines v channel
    B2 = S1.B[:,feedin]; # defines u1 channel
    C1 = S1.C[out1c,:]; # defines z channel
    C2 = S1.C[feedout,:]; # defines y channel
    D11 = S1.D[out1c,in1c]; # v->z
    D12 = S1.D[out1c,feedin]; # u1->z
    D21 = S1.D[feedout,in1c]; # v->y
    D22 = S1.D[feedout,feedin]; # u1->y

    # System matrices of S2
    E = S2.A;
    F = S2.B;
    G = S2.C;
    H = S2.D;

    I = eye(nu);
    K = (I-FeedbackSign*H*D22);
    if (det(K) == 0)
        error("Feedback interconnection is illposed.");
    end
    invK = inv(K);

    X = FeedbackSign*invK*hcat(H*C2,G);
    Y = invK*hcat(FeedbackSign*H*D21,I);

    # Construct the new system matrices
    AA1 = hcat(A,zeros(ns1,ns2)) + B2*X;
    AA2 = hcat(F*C2,E) + F*D22*X;
    AA = vcat(AA1,AA2);

    BB1 = hcat(B1,zeros(ns1,nu)) + B2*Y;
    BB2 = hcat(F*D21,zeros(ns2,nu)) + F*D22*Y;
    BB = vcat(BB1,BB2);

    CC1 = hcat(C1,zeros(nz,ns2)) + D12*X;
    CC2 = hcat(C2,zeros(ny,ns2)) + D22*X;
    CC = vcat(CC1,CC2);

    DD1 = hcat(D11,zeros(nz,nu)) + D12*Y;
    DD2 = hcat(D21,zeros(ny,nu)) + D22*Y;
    DD = vcat(DD1,DD2);

    return StateSpace(AA,BB,CC,DD);
end

# TO DO
# 1. Linear fractional transformation, and LFT algebra.
# LinearFractionalTransformation(...) -- uppper and lower.


using Diagonalizations, LinearAlgebra, PosDefManifold, Test

const err=1e-6

# method (1) real
t, n, k=50, 10, 10
A=randn(n, n) # mixing matrix in model x=As
Xset = [genDataMatrix(t, n) for i = 1:k]
Xfixed=randn(t, n)./1
for i=1:length(Xset) Xset[i]+=Xfixed end
Cset = ℍVector([ℍ((Xset[s]'*Xset[s])/t) for s=1:k])
aC=ajd(Cset; algorithm=:OJoB, simple=true)
aC2=ajd(Cset; algorithm=:NoJoB, simple=true)
aC3=ajd(Cset; algorithm=:LogLike, simple=true)
aC4=ajd(Cset; algorithm=:LogLikeR, simple=true)
aC5=ajd(Cset; algorithm=:JADE, simple=true)
aC6=ajd(Cset; algorithm=:JADEmax, simple=true)
aC7=ajd(Cset; algorithm=:GAJD, simple=true)
aC8=ajd(Cset; algorithm=:QNLogLike, simple=true)

# a=ajd(Cset; algorithm=:GAJD2, simple=true, verbose=true)

# method (2) real
aX=ajd(Xset; algorithm=:OJoB, simple=true)
aX2=ajd(Xset; algorithm=:NoJoB, simple=true)
aX3=ajd(Xset; algorithm=:LogLike, simple=true)
aX4=ajd(Xset; algorithm=:LogLikeR, simple=true)
aX5=ajd(Xset; algorithm=:JADE, simple=true)
aX6=ajd(Xset; algorithm=:JADEmax, simple=true)
aX7=ajd(Xset; algorithm=:GAJD, simple=true)
aX8=ajd(Xset; algorithm=:QNLogLike, simple=true)

@test aX≈aC
@test aX2≈aC2
@test aX3≈aC3
@test aX4≈aC4
@test aX5≈aC5
@test aX6≈aC6
@test aX7≈aC7
@test aX8≈aC8

# method (1) complex
t, n, k=50, 10, 10
Ac=randn(ComplexF64, n, n) # mixing matrix in model x=As
Xcset = [genDataMatrix(ComplexF64, t, n) for i = 1:k]
Xcfixed=randn(ComplexF64, t, n)./1
for i=1:length(Xcset) Xcset[i]+=Xcfixed end
Ccset = ℍVector([ℍ((Xcset[s]'*Xcset[s])/t) for s=1:k])
aCc=ajd(Ccset; algorithm=:OJoB, simple=true)
aCc2=ajd(Ccset; algorithm=:NoJoB, simple=true)
aCc3=ajd(Ccset; algorithm=:LogLike, simple=true)
aCc4=ajd(Ccset; algorithm=:JADE, simple=true)
aCc5=ajd(Ccset; algorithm=:JADEmax, simple=true)

# method (2) complex
aXc=ajd(Xcset; algorithm=:OJoB, simple=true)
aXc2=ajd(Xcset; algorithm=:NoJoB, simple=true)
aXc3=ajd(Xcset; algorithm=:LogLike, simple=true)
aXc4=ajd(Xcset; algorithm=:JADE, simple=true)
aXc5=ajd(Xcset; algorithm=:JADEmax, simple=true)

@test aXc≈aCc
@test aXc2≈aCc2
@test aXc3≈aCc3
@test aXc4≈aCc4
@test aXc5≈aCc5

# create 20 REAL random commuting matrices
# they all have the same eigenvectors
Cset2=PosDefManifold.randP(3, 20; eigvalsSNR=Inf, commuting=true)
# estimate the approximate joint diagonalizer (AJD)
a=ajd(Cset2; algorithm=:OJoB)
# the orthogonal AJD must be equivalent to the eigenvector matrix
# of any of the matrices in Cset
@test norm([spForm(a.F'*eigvecs(C)) for C ∈ Cset2])/20 < err
# do the same for JADE algorithm
a=ajd(Cset2; algorithm=:JADE)
@test norm([spForm(a.F'*eigvecs(C)) for C ∈ Cset2])/20 < err
# do the same for JADEmax algorithm
a=ajd(Cset2; algorithm=:JADEmax)
@test norm([spForm(a.F'*eigvecs(C)) for C ∈ Cset2])/20 < err

# generate positive definite matrices with model A*D_κ*D, where
# A is the mixing matrix and D_κ, for all κ=1:k, are diagonal matrices.
# The estimated AJD matrix must be the inverse of A
# and all transformed matrices bust be diagonal
n, k=3, 10
Dest=PosDefManifold.randΛ(eigvalsSNR=10, n, k)
# make the problem identifiable
for i=1:k Dest[k][1, 1]*=i/(k/2) end
for i=1:k Dest[k][3, 3]/=i/(k/2) end
A=randn(n, n) # non-singular mixing matrix
Cset3=Vector{Hermitian}([Hermitian(A*D*A') for D ∈ Dest])
a=ajd(Cset3; algorithm=:NoJoB, eVarC=n)
@test spForm(a.F'*A)<√err
@test mean(nonD(a.F'*Cset3[i]*a.F) for i=1:k)<err
a=ajd(Cset3; algorithm=:LogLike, eVarC=n)
@test spForm(a.F'*A)<√err
@test mean(nonD(a.F'*Cset3[i]*a.F) for i=1:k)<err
a=ajd(Cset3; algorithm=:LogLikeR, eVarC=n)
@test spForm(a.F'*A)<√err
@test mean(nonD(a.F'*Cset3[i]*a.F) for i=1:k)<err
a=ajd(Cset3; algorithm=:GAJD, eVarC=n)
@test spForm(a.F'*A)<√err
@test mean(nonD(a.F'*Cset3[i]*a.F) for i=1:k)<err
a=ajd(Cset3; algorithm=:QNLogLike, eVarC=n)
@test spForm(a.F'*A)<√err
@test mean(nonD(a.F'*Cset3[i]*a.F) for i=1:k)<err

# Do the same thing for orthogonal diagonalizers:
# now A will be orthogonal
O=randU(n) # orthogonal mixing matrix
Cset4=Vector{Hermitian}([Hermitian(O*D*O') for D ∈ Dest])
a=ajd(Cset4; algorithm=:OJoB, eVarC=n)
@test spForm(a.F'*O)<√err
@test mean(nonD(a.F'*Cset4[i]*a.F) for i=1:k)<err
a=ajd(Cset4; algorithm=:JADE, eVarC=n)
@test spForm(a.F'*O)<√err
@test mean(nonD(a.F'*Cset4[i]*a.F) for i=1:k)<√err
a=ajd(Cset4; algorithm=:JADEmax, eVarC=n)
@test spForm(a.F'*O)<√err
@test mean(nonD(a.F'*Cset4[i]*a.F) for i=1:k)<√err

# repeat the test adding noise; now the model is no more exactly identifiable
for k=1:length(Cset3) Cset3[k]+=randP(n)/1000 end
a=ajd(Cset3; algorithm=:NoJoB, eVarC=n)
@test spForm(a.F'*A)<err^(1/6)
@test mean(nonD(a.F'*Cset3[i]*a.F) for i=1:k)<√err
a=ajd(Cset3; algorithm=:LogLike, eVarC=n)
@test spForm(a.F'*A)<err^(1/6)
@test mean(nonD(a.F'*Cset3[i]*a.F) for i=1:k)<√err
a=ajd(Cset3; algorithm=:LogLikeR, eVarC=n)
@test spForm(a.F'*A)<err^(1/6)
@test mean(nonD(a.F'*Cset3[i]*a.F) for i=1:k)<√err
a=ajd(Cset3; algorithm=:GAJD, eVarC=n)
@test spForm(a.F'*A)<err^(1/6)
@test mean(nonD(a.F'*Cset3[i]*a.F) for i=1:k)<√err
a=ajd(Cset3; algorithm=:QNLogLike, eVarC=n)
@test spForm(a.F'*A)<err^(1/6)
@test mean(nonD(a.F'*Cset3[i]*a.F) for i=1:k)<√err

# the same thing for orthogonal diagonalizers
for k=1:length(Cset4) Cset4[k]+=randP(n)/1000 end
a=ajd(Cset4; algorithm=:OJoB, eVarC=n)
@test spForm(a.F'*O)<err^(1/6)
@test mean(nonD(a.F'*Cset4[i]*a.F) for i=1:k)<√err
a=ajd(Cset4; algorithm=:JADE, eVarC=n)
@test spForm(a.F'*O)<err^(1/6)
@test mean(nonD(a.F'*Cset4[i]*a.F) for i=1:k)<√err
a=ajd(Cset4; algorithm=:JADEmax, eVarC=n)
@test spForm(a.F'*O)<err^(1/6)
@test mean(nonD(a.F'*Cset4[i]*a.F) for i=1:k)<√err

# create 20 COMPLEX random commuting matrices
# they all have the same eigenvectors
Ccset2=PosDefManifold.randP(ComplexF64, 3, 20; eigvalsSNR=Inf, commuting=true)
# estimate the approximate joint diagonalizer (AJD)
ac=ajd(Ccset2; algorithm=:OJoB)
# he AJD must be equivalent to the eigenvector matrix of any of the matrices in Cset
# just a sanity check as rounding errors appears for complex data
@test norm([spForm(ac.F'*eigvecs(C)) for C ∈ Ccset2])/20<√err
# do the same for JADE algorithm
ac=ajd(Ccset2; algorithm=:JADE)
@test norm([spForm(ac.F'*eigvecs(C)) for C ∈ Ccset2])/20<√err
# do the same for JADEmax algorithm
ac=ajd(Ccset2; algorithm=:JADEmax)
@test norm([spForm(ac.F'*eigvecs(C)) for C ∈ Ccset2])/3<√err

# the same thing using the NoJoB and LogLike algorithms. Require less precision
# as the NoJoB solution is not constrained in the orthogonal group
ac=ajd(Ccset2; algorithm=:NoJoB)
@test norm([spForm(ac.F'*eigvecs(C)) for C ∈ Ccset2])/20<√err
ac=ajd(Ccset2; algorithm=:LogLike)
@test norm([spForm(ac.F'*eigvecs(C)) for C ∈ Ccset2])/20<√err

# REAL data:
# normalize the trace of input matrices,
# give them weights according to the `nonDiagonality` function
# apply pre-whitening and limit the explained variance both
# at the pre-whitening level and at the level of final vector selection
Cset=PosDefManifold.randP(20, 80; eigvalsSNR=10, SNR=10, commuting=false)

a=ajd(Cset; algorithm=:OJoB, trace1=true, w=nonD, preWhite=true, eVarC=4, eVar=0.99)
a=ajd(Cset; algorithm=:NoJoB, trace1=true, w=nonD, preWhite=true, eVarC=4, eVar=0.99)
a=ajd(Cset; algorithm=:LogLike, w=nonD, preWhite=true, eVarC=4, eVar=0.99)
a=ajd(Cset; algorithm=:LogLikeR, w=nonD, preWhite=true, eVarC=4, eVar=0.99)
a=ajd(Cset; algorithm=:JADE, w=nonD, preWhite=true, eVarC=4, eVar=0.99)
a=ajd(Cset; algorithm=:JADEmax, w=nonD, preWhite=true, eVarC=4, eVar=0.99)
a=ajd(Cset; algorithm=:GAJD, w=nonD, preWhite=true, eVarC=4, eVar=0.99)
a=ajd(Cset; algorithm=:QNLogLike, w=nonD, preWhite=true, eVarC=4, eVar=0.99)

# AJD for plots below
a=ajd(Cset; algorithm=:QNLogLike, verbose=true, preWhite=true)

using Plots
# plot the original covariance matrices
# and their transformed counterpart
CMax=maximum(maximum(abs.(C)) for C ∈ Cset);
 h1 = heatmap(Cset[1], clim=(-CMax, CMax), title="C1", yflip=true, c=:bluesreds);
 h2 = heatmap(Cset[2], clim=(-CMax, CMax), title="C2", yflip=true, c=:bluesreds);
 h3 = heatmap(Cset[3], clim=(-CMax, CMax), title="C3", yflip=true, c=:bluesreds);
 h4 = heatmap(Cset[4], clim=(-CMax, CMax), title="C4", yflip=true, c=:bluesreds);
 📈=plot(h1, h2, h3, h4, size=(700,400))
# savefig(📈, homedir()*"\\Documents\\Code\\julia\\Diagonalizations\\docs\\src\\assets\\FigAJD1.png")

Dset=[a.F'*C*a.F for C ∈ Cset];
 DMax=maximum(maximum(abs.(D)) for D ∈ Dset);
 h5 = heatmap(Dset[1], clim=(-DMax, DMax), title="F'*C1*F", yflip=true, c=:bluesreds);
 h6 = heatmap(Dset[2], clim=(-DMax, DMax), title="F'*C2*F", yflip=true, c=:bluesreds);
 h7 = heatmap(Dset[3], clim=(-DMax, DMax), title="F'*C3*F", yflip=true, c=:bluesreds);
 h8 = heatmap(Dset[4], clim=(-DMax, DMax), title="F'*C4*F", yflip=true, c=:bluesreds);
 📉=plot(h5, h6, h7, h8, size=(700,400))
# savefig(📉, homedir()*"\\Documents\\Code\\julia\\Diagonalizations\\docs\\src\\assets\\FigAJD2.png")

# COMPLEX data:
# normalize the trace of input matrices,
# give them weights according to the `nonDiagonality` function
# apply pre-whitening and limit the explained variance both
# at the pre-whitening level and at the level of final vector selection
Ccset=PosDefManifold.randP(3, 20; eigvalsSNR=10, SNR=2, commuting=false)

ac=ajd(Ccset; trace1=true, w=nonD, preWhite=true,
       algorithm=:OJoB, eVarC=8, eVar=0.99)
ac=ajd(Ccset; eVarC=8, eVar=0.99)
ac=ajd(Ccset; algorithm=:LogLike, eVarC=8, eVar=0.99)
ac=ajd(Ccset; algorithm=:JADE, eVarC=8, eVar=0.99)
ac=ajd(Ccset; algorithm=:JADEmax, eVarC=8, eVar=0.99)


using Compat

const LIBMBHFMM2D = string(Base.Filesystem.dirname(Base.source_path()), "/../bin/libmbhfmm2d")

function modbh_qn(x::Array{Float64,1},
                   lambda::Float64,rscale::Float64,n::Int)

    ifders = zero(Int32)
    ders = zeros(Float64,1)
    n = convert(Int32,n)

    n2 = length(x)

    qtemp = zeros(Float64,n+2)
    qn = zeros(Float64,n+1,n2)

    kvec = zeros(Float64,n+5)

    for i = 1:n2
        ccall( (:diffslogbk_fast_,LIBMBHFMM2D), Cvoid,
               (Ref{Float64},Ref{Float64},Ref{Float64},
                Ref{Float64},Ref{Int32},Ref{Float64},
                Ref{Float64},Ref{Int32}),
               x[i],lambda,rscale,qtemp,ifders,ders,
               kvec,n)
        qn[:,i] = qtemp[1:n+1]
    end

    return qn
end

function besselq(n::Int,x::Float64,lambda::Float64;
                 rscale::Float64=1.0)
    xvec = zeros(Float64,1)
    xvec[1] = x
    qn = modbh_qn(xvec,lambda,rscale,n)

    return qn[n+1,1]
end

function modbh_qn_test()
    xtest1 = [0.1]
    lamtest1 = 1.0
    q0test1 = 0.1244839317079709285005145657

    rscale = 1.0
    n = 1

    qn = modbh_qn(xtest1,lamtest1,rscale,n)

    err = abs(q0test1-qn[1,1])
    relerr = err/abs(q0test1)

    return err, relerr
end

function modbh_pn(x::Array{Float64,1},
                   lambda::Float64,rscale::Float64,n::Int)

    ifders = zero(Int32)
    ders = zeros(Float64,1)
    n = convert(Int32,n)

    n2 = length(x)

    ptemp = zeros(Float64,n+2)
    pn = zeros(Float64,n+1,n2)

    ivec = zeros(Float64,n+5)

    for i = 1:n2
        ccall( (:diffszkik_fast_,LIBMBHFMM2D), Cvoid,
               (Ref{Float64},Ref{Float64},Ref{Float64},
                Ref{Float64},Ref{Int32},Ref{Float64},
                Ref{Float64},Ref{Int32}),
               x[i],lambda,rscale,ptemp,ifders,ders,
               ivec,n)
        pn[:,i] = ptemp[1:n+1]
    end

    return pn
end

function modbh_pn_test()
    xtest1 = [0.1]
    lamtest1 = 1.0
    p0test1 = 0.00250156293409560140021055

    rscale = 1.0
    n = 1

    pn = modbh_pn(xtest1,lamtest1,rscale,n)

    err = abs(p0test1-pn[1,1])
    relerr = err/abs(p0test1)

    return err, relerr
end

function besselp(n::Int,x::Float64,lambda::Float64;
                 rscale::Float64=1.0)
    xvec = zeros(Float64,1)
    xvec[1] = x
    pn = modbh_pn(xvec,lambda,rscale,n)

    return pn[n+1,1]
end

function besselj0m1(x::Float64)
    # evaluate J_0(x) - 1 stably

    xh2 = (0.5*x)^2
    if (xh2 <= 1)
        tot = 0.0
        xterm = 1.0
        for i = 1:12
            xterm = -xterm*xh2/(i*i)
            tot = tot + xterm
        end
    else
        tot = besselj(0,x)-1
    end

    return tot
end

function besselj0m1_test()
    testin = 0.1
    testval = -0.002498437933959967718713101525
    funval = besselj0m1(testin)
    err = abs(testval-funval)
    relerr = err/abs(testval)

    return err, relerr
end


function optimalportfolio_greedy_nopermute(mkt::VariedCostsMarket)::Vector{Int}
    priority_order = [j for j in 1:mkt.m if mkt.g[j] ≤ mkt.H]
    sort!(
        priority_order,
        by=function (j)
            mkt.ft[j] / mkt.g[j]
        end,
        rev=true
    )

    X = Int[]
    H = mkt.H
    for j in priority_order
        if mkt.g[j] ≤ H
            push!(X, j)
            H -= mkt.g[j]
        else
            break
        end
    end

    return X
end

"""
    optimalportfolio_greedy(mkt::VariedCostsMarket) -> X, v

Use a greedy heuristic that adds schools in decreasing order of `mkt.ft ./ mkt.g`
to compute a heuristically optimal portfolio for the `VariedCostsMarket`
defined by `mkt`.
"""
function optimalportfolio_greedy(mkt::VariedCostsMarket)::Tuple{Vector{Int},Float64}
    X = optimalportfolio_greedy_nopermute(mkt)
    sort!(X)
    return mkt.perm[X], valuation_nopermute_sorted(X, mkt)
end

"""
    neighbor(X, mkt) -> X_neighbor

Generate a random neighbor of `X` for the market `mkt`. If `X` is feasible,
preserve feasibility; if not, move to a feasible solution. `X` is assumed
to refer to the sorted indices of the market.
"""
function neighbor(X::AbstractVector{<:Integer}, mkt::VariedCostsMarket)::Vector{Int}
    X_neighbor = copy(X)
    Y = setdiff(1:mkt.m, X)
    shuffle!(X_neighbor)
    shuffle!(Y)

    H = sum(mkt.g[X])

    while H ≤ mkt.H && !isempty(Y)
        j = pop!(Y)
        push!(X_neighbor, j)
        H += mkt.g[j]
    end

    while H > mkt.H
        j = popfirst!(X_neighbor)
        H -= mkt.g[j]
    end

    return X_neighbor
end

"""
    optimalportfolio_simulatedannealing(
        mkt::VariedCostsMarket;
        temp::Float64=0.25,
        red::Float64=0.0625,
        nit::Integer=500,
        verbose::Bool=false
    ) -> X, v

Use a simulated annealing procedure to compute a heuristically optimal portfolio
for the `VariedCostsMarket` defined by `mkt`. `temp` is the initial temperature,
`red` is the reduction factor, and `nit` is the total number of iterations.
"""
function optimalportfolio_simulatedannealing(
    mkt::VariedCostsMarket;
    temp::Float64=0.25,
    red::Float64=0.0625,
    nit::Integer=500,
    verbose::Bool=false
)::Tuple{Vector{Int},Float64}
    X = optimalportfolio_greedy_nopermute(mkt)
    sort!(X)
    v = valuation_nopermute_sorted(X, mkt)

    X_best, v_best = X, v

    for i in 1:nit
        X_neighbor = neighbor(X, mkt)
        sort!(X_neighbor)
        v_neighbor = valuation_nopermute_sorted(X_neighbor, mkt)

        if verbose
            println("Iteration $i, v_best = $v_best, v_neighbor = $v_neighbor")
        end

        if (delta = v_neighbor - v) ≥ 0
            X, v = X_neighbor, v_neighbor
            if v > v_best
                X_best, v_best = X, v
            end
        elseif rand() < exp(delta / temp)
            verbose && @show exp(delta / temp)
            X, v = X_neighbor, v_neighbor
        end

        temp *= red
    end

    return mkt.perm[X_best], v_best
end


# # Simple example of a homogeneous ensemble using learning networks

# In this simple example, no bagging is used, so every atomic model
# gets the same learned parameters, unless the atomic model training
# algorithm has randomness, eg, DecisionTree with random subsampling
# of features at nodes.

# Note that MLJ has a built in model wrapper called `EnsembleModel`
# for creating bagged ensembles with a few lines of code.

# ## Definition of composite model type

using MLJ, PyPlot
import Statistics

# learning network (composite model spec):

Xs = source()
ys = source(kind=:target)

atom = @load DecisionTreeRegressor
atom.n_subfeatures = 4 # to ensure diversity among trained atomic models

machines = (machine(atom, Xs, ys) for i in 1:100)

# overload `mean` for nodes:
Statistics.mean(v...) = mean(v)
Statistics.mean(v::AbstractVector{<:AbstractNode}) = node(mean, v...)

yhat = mean([predict(m, Xs) for  m in machines]);

# new composite model type and instance:

one_hundred_models = @from_network OneHundredModels(atom=atom) <= yhat

# ## Application to data

X, y = @load_boston;

# tune regularization parameter for a *single* tree:

r = range(atom,
          :min_samples_split,
          lower=2,
          upper=100, scale=:log)

mach = machine(atom, X, y)

curve = learning_curve!(mach,
                        range=r,
                        measure=mav,
                        resampling=CV(nfolds=9),
                        verbosity=0)

plot(curve.parameter_values, curve.measurements)
xlabel(curve.parameter_name)

savefig(joinpath(@OUTPUT, "e1.svg")) # hide

# \fig{e1.svg}

# tune regularization parameter for all trees in ensemble simultaneously:

r = range(one_hundred_models,
          :(atom.min_samples_split),
          lower=2,
          upper=100, scale=:log)

mach = machine(one_hundred_models, X, y)

curve = learning_curve!(mach,
                        range=r,
                        measure=mav,
                        resampling=CV(nfolds=9),
                        verbosity=0)

plot(curve.parameter_values, curve.measurements)
xlabel(curve.parameter_name)

savefig(joinpath(@OUTPUT, "e2.svg")) # hide

# \fig{e2}

#-

PyPlot.close_figs() # hide


###############################################################################
#
# Reflectance related
#
###############################################################################
"""
    BLUE(can_rad::CanopyRads{FT},
         wls::WaveLengths{FT}
    ) where {FT<:AbstractFloat}

Return the BLUE @ 469 nm, given
- `can_rad` [`CanopyRads`](@ref) type struct
- `wls` [`WaveLengths`](@ref) type struct
"""
function BLUE(
            can_rad::CanopyRads{FT},
            wls::WaveLengths{FT}
) where {FT<:AbstractFloat}
    return REF_WL(can_rad, wls, FT(469))
end

"""
    EVI(can_rad::CanopyRads{FT},
        wls::WaveLengths{FT}
    ) where {FT<:AbstractFloat}

Return the EVI, given
- `can_rad` [`CanopyRads`](@ref) type struct
- `wls` [`WaveLengths`](@ref) type struct
"""
function EVI(
            can_rad::CanopyRads{FT},
            wls::WaveLengths{FT}
) where {FT<:AbstractFloat}
    _BLUE = BLUE(can_rad, wls);
    _NIR  = NIR(can_rad, wls);
    _RED  = RED(can_rad, wls);

    return FT(2.5) * (_NIR - _RED) / (_NIR + 6 * _RED - FT(7.5) * _BLUE + 1)
end

"""
    EVI2(can_rad::CanopyRads{FT},
        wls::WaveLengths{FT}
    ) where {FT<:AbstractFloat}

Return the EVI2, given
- `can_rad` [`CanopyRads`](@ref) type struct
- `wls` [`WaveLengths`](@ref) type struct
"""
function EVI2(
            can_rad::CanopyRads{FT},
            wls::WaveLengths{FT}
) where {FT<:AbstractFloat}
    _BLUE = BLUE(can_rad, wls);
    _NIR  = NIR(can_rad, wls);
    _RED  = RED(can_rad, wls);

    return FT(2.5) * (_NIR - _RED) / (_NIR + FT(2.4) * _RED + 1)
end

"""
    LSWI(can_rad::CanopyRads{FT},
         wls::WaveLengths{FT}
    ) where {FT<:AbstractFloat}

Return the LSWI, given
- `can_rad` [`CanopyRads`](@ref) type struct
- `wls` [`WaveLengths`](@ref) type struct
"""
function LSWI(
            can_rad::CanopyRads{FT},
            wls::WaveLengths{FT}
) where {FT<:AbstractFloat}
    _BLUE = BLUE(can_rad, wls);
    _NIR  = NIR(can_rad, wls);
    _RED  = RED(can_rad, wls);
    _SWIR = SWIR(can_rad, wls);

    return (_NIR - _SWIR) / (_NIR + _SWIR)
end

"""
    NDVI(can_rad::CanopyRads{FT},
         wls::WaveLengths{FT}
    ) where {FT<:AbstractFloat}

Return the NDVI, given
- `can_rad` [`CanopyRads`](@ref) type struct
- `wls` [`WaveLengths`](@ref) type struct
"""
function NDVI(
            can_rad::CanopyRads{FT},
            wls::WaveLengths{FT}
) where {FT<:AbstractFloat}
    _NIR = NIR(can_rad, wls);
    _RED = RED(can_rad, wls);

    return (_NIR - _RED) / (_NIR + _RED)
end

"""
    NIR(can_rad::CanopyRads{FT},
        wls::WaveLengths{FT}
    ) where {FT<:AbstractFloat}

Return the NIR @ 858.5 nm, given
- `can_rad` [`CanopyRads`](@ref) type struct
- `wls` [`WaveLengths`](@ref) type struct
"""
function NIR(
            can_rad::CanopyRads{FT},
            wls::WaveLengths{FT}
) where {FT<:AbstractFloat}
    return REF_WL(can_rad, wls, FT(858.5))
end

"""
    NIRv(can_rad::CanopyRads{FT},
         wls::WaveLengths{FT}
    ) where {FT<:AbstractFloat}

Return the NIRv, given
- `can_rad` [`CanopyRads`](@ref) type struct
- `wls` [`WaveLengths`](@ref) type struct
"""
function NIRv(
            can_rad::CanopyRads{FT},
            wls::WaveLengths{FT}
) where {FT<:AbstractFloat}
    _NIR = NIR(can_rad, wls);
    _RED = RED(can_rad, wls);

    return (_NIR - _RED) / (_NIR + _RED) * _NIR
end

"""
    RED(can_rad::CanopyRads{FT},
        wls::WaveLengths{FT}
    ) where {FT<:AbstractFloat}

Return the RED @ 645 nm, given
- `can_rad` [`CanopyRads`](@ref) type struct
- `wls` [`WaveLengths`](@ref) type struct
"""
function RED(
            can_rad::CanopyRads{FT},
            wls::WaveLengths{FT}
) where {FT<:AbstractFloat}
    return REF_WL(can_rad, wls, FT(645))
end

"""
    REF_WL(wls::WaveLengths{FT},
           can_rad::CanopyRads{FT}
           wls::WaveLengths{FT},
           twl::FT
    ) where {FT<:AbstractFloat}

Return the Reflectance, given
- `can_rad` [`CanopyRads`](@ref) type struct
- `wls` [`WaveLengths`](@ref) type struct
- `twl` Target wave length in nm
"""
function REF_WL(
            can_rad::CanopyRads{FT},
            wls::WaveLengths{FT},
            twl::FT
) where {FT<:AbstractFloat}
    @unpack WL,nWL = wls;
    @unpack alb_obs = can_rad;

    # find the index where twl nm is
    ind = 0;
    for i in 1:(nWL-1)
        if WL[i] <= twl <= WL[i+1]
            ind = i;
            break;
        end
    end

    # warning if ind == 0
    if ind==0
        @warn "target wave length out of bounds, please check the set up!";
        return FT(NaN)
    else
        return (WL[ind+1] - twl) / (WL[ind+1] - WL[ind]) * alb_obs[ind] +
               (twl - WL[ind]) / (WL[ind+1] - WL[ind]) * alb_obs[ind+1]
    end
end

"""
    SWIR(can_rad::CanopyRads{FT},
        wls::WaveLengths{FT}
    ) where {FT<:AbstractFloat}

Return the SWIR @ 2130 nm, given
- `can_rad` [`CanopyRads`](@ref) type struct
- `wls` [`WaveLengths`](@ref) type struct
"""
function SWIR(
            can_rad::CanopyRads{FT},
            wls::WaveLengths{FT}
) where {FT<:AbstractFloat}
    return REF_WL(can_rad, wls, FT(2130))
end

###############################################################################
#
# SIF related
#
###############################################################################
"""
    SIF_WL(wls::WaveLengths{FT},
           can_rad::CanopyRads{FT}
           wls::WaveLengths{FT},
           twl::FT
    ) where {FT<:AbstractFloat}

Return the SIF, given
- `can_rad` [`CanopyRads`](@ref) type struct
- `wls` [`WaveLengths`](@ref) type struct
- `twl` Target SIF wave length in nm
"""
function SIF_WL(
            can_rad::CanopyRads{FT},
            wls::WaveLengths{FT},
            twl::FT
) where {FT<:AbstractFloat}
    @unpack WLF,nWLF = wls;
    @unpack SIF_obs = can_rad;

    # find the index where twl nm is
    ind = 0;
    for i in 1:(nWLF-1)
        if WLF[i] <= twl <= WLF[i+1]
            ind = i;
            break;
        end
    end

    # warning if ind == 0
    if ind==0
        @warn "target wave length out of bounds, please check the set up!";
        return FT(NaN)
    else
        return (WLF[ind+1] - twl) / (WLF[ind+1] - WLF[ind]) * SIF_obs[ind] +
               (twl - WLF[ind]) / (WLF[ind+1] - WLF[ind]) * SIF_obs[ind+1]
    end
end

"""
    SIF_740(can_rad::CanopyRads{FT},
            wls::WaveLengths{FT}
    ) where {FT<:AbstractFloat}

Return the SIF @ 740 nm, given
- `can_rad` [`CanopyRads`](@ref) type struct
- `wls` [`WaveLengths`](@ref) type struct
"""
function SIF_740(
            can_rad::CanopyRads{FT},
            wls::WaveLengths{FT}
) where {FT<:AbstractFloat}
    return SIF_WL(can_rad, wls, FT(740))
end

"""
    SIF_757(can_rad::CanopyRads{FT},
            wls::WaveLengths{FT}
    ) where {FT<:AbstractFloat}

Return the SIF @ 757 nm, given
- `can_rad` [`CanopyRads`](@ref) type struct
- `wls` [`WaveLengths`](@ref) type struct
"""
function SIF_757(
            can_rad::CanopyRads{FT},
            wls::WaveLengths{FT}
) where {FT<:AbstractFloat}
    return SIF_WL(can_rad, wls, FT(757))
end

"""
    SIF_771(can_rad::CanopyRads{FT},
            wls::WaveLengths{FT}
    ) where {FT<:AbstractFloat}

Return the SIF @ 771 nm, given
- `can_rad` [`CanopyRads`](@ref) type struct
- `wls` [`WaveLengths`](@ref) type struct
"""
function SIF_771(
            can_rad::CanopyRads{FT},
            wls::WaveLengths{FT}
) where {FT<:AbstractFloat}
    return SIF_WL(can_rad, wls, FT(771))
end


#### naive sampling

immutable CategoricalDirectSampler <: Sampleable{Univariate,Discrete}
    prob::Vector{Float64}

    function CategoricalDirectSampler(p::Vector{Float64})
        isempty(p) && error("p is empty.")
        new(p)
    end
end
ncategories(s::CategoricalDirectSampler) = length(s.prob)

function rand(s::CategoricalDirectSampler)
    p = s.prob
    n = length(p)
    i = 1
    c = p[1]
    u = rand()
    while c < u && i < n
        c += p[i += 1]
    end
    return i
end

##### Alias Table #####

immutable AliasTable <: Sampleable{Univariate,Discrete}
    accept::Vector{Float64}
    alias::Vector{Int}
    isampler::RandIntSampler
end
ncategories(s::AliasTable) = length(s.accept)

function AliasTable{T<:Real}(probs::AbstractVector{T})
    n = length(probs)
    n > 0 || error("The input probability vector is empty.")
    accp = Array(Float64, n)
    alias = Array(Int, n)
    StatsBase.make_alias_table!(probs, 1.0, accp, alias)
    AliasTable(accp, alias, RandIntSampler(n))
end

rand(s::AliasTable) =
    (i = rand(s.isampler); u = rand(); @inbounds r = u < s.accept[i] ? i : s.alias[i]; r)

show(io::IO, s::AliasTable) = @printf(io, "AliasTable with %d entries", ncategories(s))


"GARCH LLH benchmark, Julia version."

import JSON
using Test
using Statistics

data = JSON.parsefile("data.json")
out = Dict()

function garchSim(ɛ²::Vector, ω, α, β)
    h = similar(ɛ²)
    h[1] = mean(ɛ²)
    for i = 2:length(ɛ²)
        @inbounds h[i] = ω + α*ɛ²[i-1] + β*h[i-1]
    end
    h
end

function garchLLH(y::Vector, params::Vector)
    ɛ² = y.^2
    T = length(y)
    h = garchSim(ɛ², params...)
    -(T-1)*log(2π)/2 - sum(log.(h) + (y./sqrt.(h)).^2)/2
end

function benchmark(n)
    #price = Array{Float32}(data["price"])
    ret = Array{Float32}(data["ret"])
    #ret = diff(log.(price))
    #ret = ret .- mean(ret)
    x0 = Array{Float32}(data["x0"])
    tmp = garchLLH(ret, x0)  # warm-up
    @test tmp ≈ data["llh"] atol=1e-3
    t = @elapsed for i = 1:n; tmp = garchLLH(ret, x0); end
end

n_ix = findfirst(ARGS.=="-n")
if n_ix < length(ARGS)
    n = tryparse(Int64, ARGS[n_ix+1])
    n==nothing && error("Wrong number of iterations!")
else
    error("Number of iterations is not set!")
end

out["julia"] = benchmark(n)
println(JSON.json(out))


# f(x)=mean(abs(x))

#TODO sum(abs)

function FmeanAbs(x...)
    tmp=0.0
    for i in 1:length(x)
        tmp+=mean(abs(x[i]))
    end
    return ([tmp/length(x)],nothing)
end

function FmeanAbs_inplace(value::Array,auxvalue,x...) # inplace
    tmp=0.0
    for i in 1:length(x)
        tmp+=mean(abs(x[i]))
    end
    fill!(value,tmp/length(x))
end

function DmeanAbs(derivativeIDX,f_c,faux_c,grad_c,grad_n,x...)
    axpy!(grad_c[1]/(length(x)*length(x[derivativeIDX])),sign(x[derivativeIDX]),grad_n)

end

if PROC=="GPU"

    function meanAbs(A::CudaArray)
        return CudaArray(CUBLAS.asum(A)/length(A)*ones(1,1))
    end
    export meanAbs

    function meanAbs!(A::CudaArray,Out::CudaArray)
        copy!(Out,meanAbs(A))
    end
    export meanAbs!

    function FmeanAbs(x::CudaArray...)
        tmp=CudaArray(zeros(1,1))
        for i in 1:length(x)
            axpy!(1.0/length(x),meanAbs(x[i]),tmp)
        end
        return (tmp,nothing)
    end

    function FmeanAbs_inplace(value::CudaArray,auxvalue,x::CudaArray...) # inplace
        fill!(value,0.0)
        for i in 1:length(x)
            axpy!(1.0/length(x),meanAbs(x[i]),value)
        end
    end

    function DmeanAbs(derivativeIDX,f_c,faux_c,grad_c,grad_n,x::CudaArray...)
        tmp=CudaArray(Float64,size(grad_n))
        vsign!(x[derivativeIDX],tmp)
        alphaaxpy!(1.0/(length(x)*length(x[derivativeIDX])),grad_c,tmp,grad_n)
        free(tmp)
    end

end

Derivative[FmeanAbs]=DmeanAbs # Define dictionary lookup
Inplace[FmeanAbs]=FmeanAbs_inplace

meanAbs(n::ADnode)=ADnode(FmeanAbs,n)

function meanAbs(n::ArrayADnode)
    return ADnode(FmeanAbs,n)
end

#mean(A::ADtrans)=ADnode(Fmean, ftranspose(node[A.parent]))
meanAbs(A::ADtrans)=ADnode(FmeanAbs, node[A.parent]) # mean(A')=mean(A)

export meanAbs


using VortexDistributions, JLD2, Interpolations;
using NearestNeighbors, LinearAlgebra, Distances;
using GLMakie, Colors

include("../util.jl");
# @load "Vortices3D/data/box_vorts.jld2" psi_chaos1 psi_chaos2 psi_knots1 psi_knots2 psi_knots3 psi_knots4 psi_tubes1 psi_tubes2 psi_ringtube X
@load "data/box_rings.jld2" psi_ring1 psi_ring2 psi_ring3 X # make sure julia repl is in base directory

psi = psi_ring1;
plot_iso(psi,X)

## Find vortex points
N = 100
@time vorts_3d = findvortices3D_itp(psi, X, N); # Find vortex points with interpolation depth N
v_matrix = vcat(vorts_3d'...)[:,1:3]' # Convert to matrix for kdtree

scatterVortsOnIso(vorts_3d)
## Find filaments
x = X[1]; y = X[2]; z = X[3];
Δx = x[2]-x[1]; Δy = y[2]-y[1]; Δz = z[2]-z[1];

α = .2
ϵ = (1+α)*sqrt(Δx^2+Δy^2+Δz^2)/N
@time sets3 = setMethod3(v_matrix, ϵ)
@time setsPeriodic = setMethodPeriodic(v_matrix, X, ϵ, true)

# @time nearest = uniqueVortices(ϵ, psi, X, N)
# @time sets1 = setMethod(psi, X, N, ϵ)
# @time sets2 = setMethod2(v_matrix, ϵ);

plot_iso(psi, X)
using Colors
colors = distinguishable_colors(length(sets3),[RGB(0,0.35,0.25)],dropseed=true)
for i in 1:length(setsPeriodic)
    vi = v_matrix[:, collect(setsPeriodic[i])]
    scatter!(vi[1,:],vi[2,:],vi[3,:],markersize=200,color=colors[i])
end

sortslices(v_matrix, dims=2, by=x->x[1])

A = [v_matrix[:, i] for i in 1:length(v_matrix[1,:]) if v_matrix[1, i] == -7.75 || v_matrix[1,i] == 8.0]

for i in 1:length(A)
    println(A[i])
end

## x[end] points
scatter!([8.0], [-0.875], [1.2947], markersize=400, color="red") #A[end-9]

scatter!([8.0], [-3.34913], [0.741825], markersize=400, color="red") #A[end-11]

## x[1] points

scatter!([-7.75], [-1.03938], [0.9677467], markersize=400, color="blue") #A[9]

scatter!([-7.75], [-3.707936], [0.317925], markersize=400, color="blue") #A[6]

euclidean(A[end-9][2:3], A[9][2:3])
euclidean(A[end-11][2:3], A[6][2:3])

sqrt(3)Δx

@time psi_itp1 = interpolate(psi, BSpline(Cubic(Periodic(OnCell()))));
@time psi_itp2 = interpolate(psi, BSpline(Quadratic(Periodic(OnCell()))));

psi_exp = extrapolate(psi_itp, Periodic(OnCell()))

psi_exp[0.1,0:65,0:65]

psi_itp[1.1,1.1,1.1]
psi_exp[65.1,65.1,65.1]

x_itp = interpolate(X[1], BSpline(Linear()))
x_etp = extrapolate(x_itp, Line())
x_etp[64]

x_etp[1:65]

r1 = LinRange(1,10,10)
r2 = LinRange(10,11,10)
vcat(r1, r2)

X[1]

x = collect(LinRange(1,2,64))
X[1]


module Triangulate

type triangulateio64
  pointlist::Array{Float64,1}
  numberofpoints::Int32
  trianglelist::Array{Int32,1}
  numberoftriangles::Int32
  edgelist::Array{Int32,1}
  numberofedges::Int32
end

const TrianglelibPath = abspath(joinpath(splitdir(Base.source_path())[1],"..","lib","libtriangle.so"))

function triangulate(size::Float64,numberofpoints::Int64,pointlist::Array{Float64, 1})
  if numberofpoints * 2 != length(pointlist)
    error("Triangulate: Inconsistent Mesh Input.")
  end
  ret = ccall((:triangleio, TrianglelibPath), Any, (Ptr{Void}, Ptr{UInt8}, Ptr{UInt8}, Int64, Ptr{Float64}),(pointer_from_objref(Triangulate)),"QpcAevn",string("Qprq34.0","a",size,"BPe"), Int32(numberofpoints), pointlist)
  return ret
end

function UnitSquareMesh(tol::Float64)
  ret = triangulate(tol, 4, [0., 0., 1., 0., 1., 1., 0., 1.])
  return ret
end

function UnitCircle(tol::Float64, size::Int64)
  Θ = 2 * π/size
  pointlist = zeros(2 * size)
  for idx in 1:size
    pointlist[2 * idx - 1] = cos(idx * Θ)
    pointlist[2 * idx    ] = sin(idx * Θ)
  end
  ret = triangulate(tol, size, pointlist)
end

export triangulate, UnitSquareMesh, UnitCircle

end


export FixedHomotopy

"""
    FixedHomotopy(H, t) <: AbstractSystem

Fix a homotopy `H(x,t)` at `t`
"""
mutable struct FixedHomotopy{Hom<:AbstractHomotopy,T} <: AbstractSystem
    H::Hom
    t::T
end

struct FixedHomotopyCache{HC<:AbstractHomotopyCache} <: AbstractSystemCache
    cache::HC
end

cache(FH::FixedHomotopy, x) = FixedHomotopyCache(cache(FH.H, x, FH.t))

Base.size(F::FixedHomotopy) = size(F.H)

evaluate!(u, F::FixedHomotopy, x, c::FixedHomotopyCache) =
    evaluate!(u, F.H, x, F.t, c.cache)
evaluate(F::FixedHomotopy, x, c::FixedHomotopyCache) = evaluate(F.H, x, F.t, c.cache)

jacobian!(U, F::FixedHomotopy, x, c::FixedHomotopyCache) =
    jacobian!(U, F.H, x, F.t, c.cache)
jacobian(F::FixedHomotopy, x, c::FixedHomotopyCache) = jacobian(F.H, x, F.t, c.cache)

evaluate_and_jacobian!(u, U, F::FixedHomotopy, x, c::FixedHomotopyCache) =
    evaluate_and_jacobian!(u, U, F.H, x, F.t, c.cache)
evaluate_and_jacobian(F::FixedHomotopy, x, c::FixedHomotopyCache) =
    evaluate_and_jacobian(F.H, x, F.t, c.cache)


@showbanner "Basic tria membrane test"
st=Structure()
lcset=LoadCaseSet()

add_uniaxial_metal!(st,"steel",2e11,0.3,7849.0474)
add_static_case!(lcset,"DL",0)

add_node!(st,1,0,0,0)
add_node!(st,2,6,0,0)
add_node!(st,3,12,0,0)
add_node!(st,4,18,0,0)
add_node!(st,5,0,6,0)
add_node!(st,6,6,6,0)
add_node!(st,7,12,6,0)
add_node!(st,8,18,6,0)

add_tria!(st,1,1,2,5,"steel",1e-3)
add_tria!(st,2,2,6,5,"steel",1e-3)
add_tria!(st,3,2,3,6,"steel",1e-3)
add_tria!(st,4,3,7,6,"steel",1e-3)
add_tria!(st,5,3,4,7,"steel",1e-3)
add_tria!(st,6,4,8,7,"steel",1e-3)

add_nodal_force!(lcset,"DL",8,0,-1e5,0,0,0,0)

set_nodal_restraint!(st,1,true,true,true,true,true,true)
set_nodal_restraint!(st,5,true,true,true,true,true,true)
for i in [2,3,4,6,7,8]
    set_nodal_restraint!(st,i,false,false,true,true,true,true)
end

assembly=assemble!(st,lcset,path=PATH)

solve(assembly)

r=result_nodal_displacement(assembly,"DL",8)

@test r≈[0.00345715, -0.0164571, 0.0, 0.0, 0.0, 0.0] atol=1e-4


function xxz_longrange(lx::Int, delta::Float64=1.0, r::Float64=0.5)
    (lx < 2) && error("The system size is too small : ", lx)

    sz, sp, sm = sparse.(spinoperators(1/2))

    heis_term = 0.5 * (sp ⊗ sm + sm ⊗ sp) + delta * sz ⊗ sz

    Hmat = heis_term
    # recursive generation of Hamiltonian
    if lx > 2
        for n=3:lx
            Hmat = (Hmat ⊗ I(2)) + (I(2^(n-2)) ⊗ heis_term)
            for l in n-2:-1:1
                d = n-l-1
                Hmat += (r^d) * I(2^(l-1)) ⊗ (
                    delta * (sz ⊗ I(2^d) ⊗ sz) +
                    0.5 * (sp ⊗ I(2^d) ⊗ sm + sm ⊗I(2^d) ⊗ sp)
                )
            end
        end
    end
    Hmat
end


using Distributions
using Turing
using Test

# Define model
@model ad_test2(xs) = begin
  s ~ InverseGamma(2,3)
  m ~ Normal(0,sqrt.(s))
  xs[1] ~ Normal(m, sqrt.(s))
  xs[2] ~ Normal(m, sqrt.(s))
  s, m
end

# Run HMC with chunk_size=1
chain = sample(ad_test2([1.5 2.0]), HMC(300, 0.1, 1))


#--------------------------------------------------------------
# MultilevelTT package
# example code for https://arxiv.org/abs/1802.09062
#--------------------------------------------------------------
# Code for 1D example with highly oscillating coefficient,
# using AMEn solver via Python interface (see TTPy.jl)
module MultiscaleAMEn

using TT, TTFEM, TTAlg, TTPy

export amenOscillatory1D

function amenOscillatory1D(L::Integer, ε::AbstractFloat, K::Number)
    Dinv = ttdiagm(add(2., ttones(L), 1., ttcos(L, K*π, .5)))
    Drhs = ttones(L)
    d, _ = stsolve(L, 1, x->add(1., matvec(Dinv,x), -1., Drhs), 3., 1., .5*ɛ, .75)
    D = ttdiagm(d)
    C = ttBPXdnC(L, 1)
    F = ttBPXdnQLaplace(L,1,1)
    p = size(F[L],1)
    q = size(F[L],4)
    T = reshape(F[L],(p*4,q))
    S = reshape(F[L+1],(q,1))
    F[L] = reshape(T*S,(p,2,2,1))
    pop!(F)
    Ft = transpose(F)
    f = ttrhsdn(L)
    g = matvec(C, f)
    svdtrunc!(g, svd!(g), .5*ɛ)
    B = matmat(Ft, matmat(D, F))
    u = amen(B, g, g, ε)
    Cu = matvec(C, u)
    Fu = matvec(F, u)
    return u, Cu, Fu
end

end


include("C:\\Users\\daniel.coutinho\\github\\DSGE-HMC\\src\\priors.jl")

using Distributions
using Plots

x = randn(100)

y = 0.5*x + randn(100)

prior = Normal(0,1)

par = zeros(10000)

likelihood(y,x,par) = sum(logpdf.(Normal(0,1),y - x*par))

unif = Uniform(0,1)

for i in 2:10000
    kernel = Normal(par[i-1],0.5)
    novo_par = rand(kernel)
    kernel_novo = Normal(novo_par,0.5)

    num = likelihood(y,x,novo_par) + logpdf(prior,novo_par) #+ logpdf(kernel_novo,par[i-1])

    dem = likelihood(y,x,par[i-1]) + logpdf(prior,par[i-1]) #+ logpdf(kernel,novo_par)
    alpha = exp(min(0,num - dem))
    p = rand(unif)
    if alpha < p
        par[i] = par[i-1]
    else
        par[i] = novo_par
    end
end

histogram(par)


function [ varp,varm ] = diffusion(varb,varp,varm,cmixh,cmixv,dx,dz,dt,nx,nz)
%UNTITLED9 Summary of this function goes here
%   Detailed explanation goes here

% Horizontal Diffusion

kmixh=cmixh*dx*dx/dt;

for k = 2 : nz-1
  for i = 2 : nx-1
varp[i,k] = varp[i,k] + 2.0*dt*kmixh*(varm[i+1,k]-2.0*varm[i,k]+varm[i-1,k])/(dx*dx);

  end
end

% Vertical Diffusion
kmixv=cmixv*dz*dz/dt;

for k = 3 : nz-1
  for i = 2 : nx-1

varp(i,k) = varp[i,k] + 2*dt*(kmixv*((varm[i,k+1]-varb[k+1])     ...
                 -2.0*(varm[i,k]-varb[k])                        ...
                 +(varm[i,k-1]-varb[k-1]) )/(dz*dz));

  end
end

end


using DrWatson
quickactivate(@__DIR__)
@show projectdir()
##
using CSV
using DataFrames
# using DataFramesMeta
using Dates
using StatsBase
using Plots
using CovidSurvey

using Stheno, KernelFunctions
using AbstractGPs
using TemporalGPs
using StaticArrays

# Load standard packages from the Julia ecosystem
using Optim # Standard optimisation algorithms.
using ParameterHandling # Helper functionality for dealing with model parameters.
using Zygote # Algorithmic Differentiation

using ParameterHandling: value, flatten
plotlyjs()
#----------------------------------------------------------------------------
# survey data
survey = National.readsurveys()
Rt = load(projectdir("out", "Rt.csv"))|>DataFrame
plot(t, Rt.mq, ribbon=(Rt.mq - Rt.llq, Rt.uuq - Rt.mq), linewidth=0, label="95% quantiles")
plot!(t, Rt.mq, ribbon=(Rt.mq - Rt.lq, Rt.uq - Rt.mq), linewidth=2, label="50% quantiles", hover=mq)
plot!(survey.date, survey.colleagues, ticks=:native, label="colleagues")
plot!(survey.date, survey.family, ticks=:native, label="family")
plot!(survey.date, survey.strangers, ticks=:native, label="strangers")
plot!(survey.date, survey.friends, ticks=:native, label="friends")
##
d = survey.friends
y = zscore(d)
x = RegularSpacing(0.0, 1., length(survey.date))
# x = GPPPInput(:f3, collect(1:size(survey,1)))
# Declare model parameters using `ParameterHandling.jl` types.
params = (
    s1 = positive(1.),
    l1 = positive(50.),
    s2 = positive(1.),
    l2 = positive(10.),
    s3 = positive(1.),
    l3 = positive(2.),
    v  = positive(0.1),
)

# function build_gp(ps)
#     f = @gppp let
#         f1 = ps.s1 * stretch(GP(Matern52Kernel()), 1 / ps.l1)
#         f2 = ps.s2 * stretch(GP(SEKernel()), 1 / ps.l2)
#         f3 = f1 + f2
#     end
# end
function build_kernels(ps)
    k1 = ps.s1 * KernelFunctions.transform(Matern52Kernel(), 1/ps.l1)
    k2 = ps.s2 * KernelFunctions.transform(Matern52Kernel(), 1/ps.l2)
    k3 = ps.s3 * KernelFunctions.transform(Matern32Kernel(), 1/ps.l3)
    return k1, k2, k3
end
function build_gp(ps)
    f_naive = GP( sum( build_kernels(ps) ) )
    return to_sde(f_naive, SArrayStorage(Float64))
end

# Construct mapping between structured and Vector representation of parameters.
flat_initial_params, unflatten = flatten(params)

# Specify an objective function for Optim to minimise in terms of x and y.
# We choose the usual negative log marginal likelihood (NLML).
function objective(flat_params)
    params = value(unflatten(flat_params))
    f = build_gp(params)
    return -logpdf(f(x, params.v), y)
end

# Check that the objective function works:
objective(flat_initial_params)

# Optimise using Optim. This optimiser often works fairly well in practice,
# but it's not going to be the best choice in all situations. Consult
# Optim.jl for more info on available optimisers and their properties.
training_results = Optim.optimize(
    objective,
    θ -> only(Zygote.gradient(objective, θ)),
    flat_initial_params, # Add some noise to make learning non-trivial
    BFGS(
        alphaguess = Optim.LineSearches.InitialStatic(scaled=true),
        linesearch = Optim.LineSearches.BackTracking(),
    ),
    Optim.Options(show_trace = true);
    inplace=false,
)

# Extracting the final values of the parameters.
# Should be close to truth.
flat_final_params = training_results.minimizer
final_params = value(unflatten(flat_final_params))
# f = Stheno.wrap(GP( build_kernel(final_params) ), GPC())
f = build_gp(final_params)

objective(flat_final_params)
fp = posterior(f(x, final_params.v),y)

# plot!(fx(x, final_params.var_noise))
ms = marginals(fp(x));
p = plot(x,y)
plot!(
    x, mean.(ms);
    ribbon=3std.(ms), label="", color=:green, fillalpha=0.2, linewidth=2,
)

function build_marginals(ps)
    build_posterior(f) = posterior(f(x, ps.v),y)
    fs = [GP(k) for k in build_kernels(ps)]
    fps = [build_posterior(f) for f in fs]
    [marginals(fp(x)) for fp in fps]
end

ms = build_marginals(final_params)
p = plot(x,y)
for m in ms
    plot!(p, x, mean.(m))
end
display(p)


export Id, monoSymp, beamSplitter, circulator
export randomSymmetric, randomSymp, randomGenericSymp
export randomPassiveLocalSymp, randomLocalSymp, randomizeSymp

# From Number of Modes to Identity Symplectic Matrix
Id(n::Int)::Symp = Matrix(1I, 2*n, 2*n)

monoPassiveSymp(θ::Real)::Symp =  exp(θ*Ω(1))
monoSymp(θ1::Real, l::Real, θ2::Real)::Symp =  monoPassiveSymp(θ1) * Symp([l 0 ; 0  1/l ]) * monoPassiveSymp(θ2)

# Generte a Two-Mode BeamSplitter
beamSplitter(angle::Real)::Symp = monoPassiveSymp(angle) ⊗ Id(1)
# Generate a BeamSplitter between two modes in a multi-mode System
function beamSplitter(angle::Real, mode1, mode2, nModes)::Symp
    mode1, mode2 = sort([mode1, mode2])
    lst = [ [[2 * (i + 2) - 1, 2 * (i + 2)] for i in 1:mode1-1]... ; [1, 2] ; [[2 * (i + mode1 + 1) - 1, 2 * (i + mode1 + 1)] for i in 1: mode2 - mode1 - 1]...; [3, 4] ; [[2 * (i + mode2 ) - 1, 2 * (i + mode2 )]  for i in 1:nModes-mode2]...  ];
    return dsum( beamSplitter(angle), Id(nModes - 2) )[lst, lst]
end

# Generate a Circulator-Like Symplectic Matrix according to the given permutation
circulator(perm::Vector)::Symp = hvcat(length(perm), [i == perm[j] ? Id(1).S : zeros(2, 2)   for i in 1:length(perm), j in 1:length(perm) ]...)

# From Number of Modes to Random Symmetric Matrix
function randomSymmetric(n::Int, range::Real = 3)::Matrix{BigFloat}
    if range <= 0
        return error("range must be positive")
    end
    return Symmetric( range*(rand(2*n, 2*n) - I ) )
end

# From Number of Modes to Random Symplectic Matrix
function randomSymp(n::Int, range::Real = 3)::Symp
    S1 = Id(n) - (Ω(n) * randomSymmetric(n, range) + Id(n) / 2  )^-1
    S2 = Id(n) - (Ω(n) * randomSymmetric(n, range) + Id(n) / 2  )^-1
    return  S1 * S2
end
function randomGenericSymp(n::Int, range::Real = 3)::Symp
        m = randomSymp(n, range)
        return isGeneric(m) ? m : randomGenericSymp(n, range)
end
randomPassiveLocalSymp(n::NofModes)::Symp = n == 1 ? monoPassiveSymp(2π*rand()) : ⊕([monoPassiveSymp(2π*rand()) for i in 1:n]...)

randomLocalSymp(n::NofModes, range::Real = 0.2)::Symp = n == 1 ? monoSymp(2π*rand(), range*( rand() -1/2) + 1.0 , 2π*rand()) : ⊕([ monoSymp(2π*rand(), range*( rand() -1/2) + 1.0 , 2π*rand()) for i in 1:n]...)

function randomizeSympPassive(S::Symp, l::Int)::Symp
    n = nModes(S)
    if l == 1
        return randomPassiveLocalSymp(n) * S  * randomPassiveLocalSymp(n)
        # return randomLocalSymp(n) * S  * randomLocalSymp(n)
    else
        return *([randomizeSympPassive(S, 1) for i in 1:l]...)
    end
end

function randomizeSymp(S::Symp, l::Int)::Symp
    n = nModes(S)
    if l == 1
        return randomLocalSymp(n) * S  * randomLocalSymp(n)
        # return randomLocalSymp(n) * S  * randomLocalSymp(n)
    else
        return *([randomizeSymp(S, 1) for i in 1:l]...)
    end
end


struct Ellipse{Dim,T} <: Primitive{Dim,T}
    center::Point{Dim,T}
    weights::Point{Dim,T}
    radius::T
end

Ellipse(center::NTuple{Dim,T}, weights::NTuple{Dim,T}, radius::T) where {Dim,T} = Ellipse(Point(center), Point(weights), radius)

struct HalfSpace{Dim,T}
    p::Vec{Dim,T}
    offset::T
end

norm2(x) = sum(map(*, x, x))

sdf(h::HalfSpace) = TrueSDF(x -> sum(coordinates(x) .* h.p) - h.offset)
sdf(e::Ellipse) = TrueSDF(x -> hypot(((e.center - x) ./ coordinates(e.weights))...) - e.radius)
sdf(s::Sphere) = TrueSDF(x -> hypot((s.center - x)...) - s.radius)

function sdf(b::Box)
    function _sdf(x)
        c = center(b)
        radii = b.max - c
        p = x - c
        z = zero(typeof(p))
        q = abs.(p) - radii
        hypot(maximum.(zip(q, z))...) + min(maximum(q), zero(eltype(p)))
    end
    TrueSDF(_sdf)
end


using JuMP, Gurobi, Distributions, StatsFuns, StatsBase

include("resource_nonlinear_20_30.jl")

const gurobi_env = Gurobi.Env()

const numVariables = numResources

# generate random samples of the demands
function generateRandomSamples(numSamples::Int64)

	lambda = rand(MvNormal(lambda_mean[1,:],lambda_covariance),numSamples)
	rho = Array{Float64}(numResources,numSamples)
	for i = 1:numResources
		rho[i,:] = rand(Normal(rho_mean[i],rho_stdev),numSamples)
		for samp = 1:numSamples
			rho[i,samp] = min(rho[i,samp],rho_max)
		end
	end

	return lambda, rho
end

srand(1234)

# generate samples to estimate actual risk level of solution
const numAnalyticalSamples = 100000
lambda_analytical, rho_analytical = generateRandomSamples(numAnalyticalSamples)

srand()

# compute conservative bound on the probability that
# constraints aren't satisfied
function boundRiskLevel(x::Array{Float64},reliabilityLevel::Float64=1E-06)

	const numSamples = size(lambda_analytical,2)
	simpleEst::Bool = false

	numViolated::Int64 = checkScenarioConstraints(x,lambda_analytical,rho_analytical)

	riskLevel::Float64 = Inf

	if(simpleEst)
		riskLevel = numViolated*1.0/numSamples
		return riskLevel
	else
		if(numViolated == 0)
			riskLevel = -log(reliabilityLevel)/numSamples
			return riskLevel
		end

		 # termination criterion for bisection
		gamma_tolerance::Float64 = 0.1/numSamples

		# use tail bounds on binomial distribution to estimate lower and upper bounds
		gamma_low::Float64 = (numViolated - sqrt(-numSamples*log(reliabilityLevel)/2.0))/numSamples
		if(gamma_low < 0.0)
			gamma_low = 0.0
		end
		gamma_up::Float64 = (numViolated + sqrt(-numSamples*log(reliabilityLevel)/2.0))/numSamples
		if(gamma_up > 1.0)
			gamma_up = 1.0
		end

		gamma::Float64 = 0.0
		func_value::Float64 = 0.0

		# use bisection to solve the nonlinear equation
		while (gamma_up - gamma_low > gamma_tolerance)

			gamma = (gamma_low + gamma_up)/2.0
			func_value = -reliabilityLevel
			for i = 1:numViolated
				tmp_func = i*log(gamma) + (numSamples-i)*log(1-gamma) + lfact(numSamples) - lfact(numSamples-i) - lfact(i)
				func_value += exp(tmp_func)
				if(func_value > 0)
					break
				end
			end
			if(func_value > 0)
				gamma_low = (gamma_low + gamma_up)/2.0
			else
				gamma_up = (gamma_low + gamma_up)/2.0
			end

		end

		riskLevel = gamma_up
		return riskLevel
	end

end

# determines if all constraints are satisfied for a given single scenario
# returns one if scenario constraints are all satisfied, zero otherwise
function checkScenarioConstraints(x::Array{Float64},lambda::Array{Float64},rho::Array{Float64})

	const numSamples = size(lambda,2)
	numViolated::Int64 = 0

	mod = Model(solver=GurobiSolver(gurobi_env,Presolve=0,OutputFlag=0))

	@variable(mod, y[1:numResources,1:numCustomers] >= 0)

	@objective(mod, Min, 0)
	@constraint(mod, con1[i=1:numResources], sum(y[i,j] for j = 1:numCustomers) <= 0)
	@constraint(mod, con2[j=1:numCustomers], sum(nu[i,j]*y[i,j] for i = 1:numResources) >= 0)

	for iter = 1:numSamples
		for i = 1:numResources
			JuMP.setRHS(con1[i],rho[i,iter]*(x[i])^2)
		end

		for j = 1:numCustomers
			JuMP.setRHS(con2[j],lambda[j,iter])
		end

		status = solve(mod,suppress_warnings=true)

		if(status != Symbol("Optimal"))
			numViolated += 1
		end
	end

	mod = 0

	return numViolated
end

numTrials = 10

# solution file name
const baseDirName = "C:/Users/rkannan/Desktop/SA for CCP/Stochastic approximation/experiments/resource_stochapprox/"

for trial = 1:numTrials

	const dirName = baseDirName * "rep" * string(trial) * "/"

	const fileName = dirName * "solution.txt"

	const outputFile = dirName * "trueRiskLevels.txt"
	open(outputFile, "w") do f
	end

	const timeFile = dirName * "trueRiskTime.txt"
	open(timeFile, "w") do f
	end

	soln = readdlm(fileName)

	const numPoints = size(soln,1)

	for iter = 1:numPoints
		tic()
		trueRiskLevel::Float64 = boundRiskLevel(soln[iter,:])
		riskTime = toq()
		println("Risk level #",iter," : ",trueRiskLevel,"  time: ",riskTime)
		open(outputFile, "a") do f
			write(f,"$trueRiskLevel \n")
		end
		open(timeFile, "a") do f
			write(f,"$riskTime \n")
		end
		gc()
	end
end


export SGD,DebugOptAlg,ADAM

# ------------------------------------------------------------------------------
# SGD
"""
  SGD(;η=t -> 1/(1+t), λ=2)

Stochastic Gradient Descent algorithm (default)

# Fields:
- `η`: Learning rate, as a function of the current epoch [def: t -> 1/(1+t)]
- `λ`: Multiplicative constant to the learning rate [def: 2]
"""
struct SGD <: OptimisationAlgorithm
    η::Function
    λ::Float64
    function SGD(;η=t -> 1/(1+t), λ=2)
        return new(η,λ)
    end
end

function singleUpdate!(θ,▽,optAlg::SGD;nEpoch,nBatch,nBatches,xbatch,ybatch)
    η    = optAlg.η(nEpoch)*optAlg.λ
    #newθ = gradSub.(θ,gradMul.(▽,η))
    θ =  θ - ▽ .* η
    #newθ = gradientDescentSingleUpdate(θ,▽,η)
    return (θ=θ,stop=false)
end

#gradientDescentSingleUpdate(θ::Number,▽::Number,η) = θ .- (η .* ▽)
#gradientDescentSingleUpdate(θ::AbstractArray,▽::AbstractArray,η) = gradientDescentSingleUpdate.(θ,▽,Ref(η))
#gradientDescentSingleUpdate(θ::Tuple,▽::Tuple,η) = gradientDescentSingleUpdate.(θ,▽,Ref(η))

#maxEpochs=1000, η=t -> 1/(1+t), λ=1, rShuffle=true, nMsgs=10, tol=0

# ------------------------------------------------------------------------------
# ADAM
#

"""
  ADAM(;η, λ, β₁, β₂, ϵ)

The [ADAM](https://arxiv.org/pdf/1412.6980.pdf) algorithm, an adaptive moment estimation optimiser.

# Fields:
- `η`:  Learning rate (stepsize, α in the paper), as a function of the current epoch [def: t -> 0.001 (i.e. fixed)]
- `λ`:  Multiplicative constant to the learning rate [def: 1]
- `β₁`: Exponential decay rate for the first moment estimate [range: ∈ [0,1], def: 0.9]
- `β₂`: Exponential decay rate for the second moment estimate [range: ∈ [0,1], def: 0.999]
- `ϵ`:  Epsilon value to avoid division by zero [def: 10^-8]
"""
mutable struct ADAM <: OptimisationAlgorithm
    η::Function
    λ::Float64
    β₁::Float64
    β₂::Float64
    ϵ::Float64
    m::Vector{Learnable}
    v::Vector{Learnable}
    function ADAM(;η=t -> 0.001, λ=1.0, β₁=0.9, β₂=0.999, ϵ=1e-8)
        return new(η,λ,β₁,β₂,ϵ,[],[])
    end
end

"""
   initOptAlg!(optAlg::ADAM;θ,batchSize,x,y,rng)

Initialize the ADAM algorithm with the parameters m and v as zeros and check parameter bounds
"""
function initOptAlg!(optAlg::ADAM;θ,batchSize,x,y,rng = Random.GLOBAL_RNG)
    optAlg.m = θ .- θ # setting to zeros
    optAlg.v = θ .- θ # setting to zeros
    if optAlg.β₁ <= 0 || optAlg.β₁ >= 1 @error "The parameter β₁ must be ∈ [0,1]" end
    if optAlg.β₂ <= 0 || optAlg.β₂ >= 1 @error "The parameter β₂ must be ∈ [0,1]" end
end

function singleUpdate!(θ,▽,optAlg::ADAM;nEpoch,nBatch,nBatches,xbatch,ybatch)
    β₁,β₂,ϵ  = optAlg.β₁, optAlg.β₂, optAlg.ϵ
    η        = optAlg.η(nEpoch)*optAlg.λ
    t        = (nEpoch-1)*nBatches+nBatch
    optAlg.m = @. β₁ * optAlg.m + (1-β₁) * ▽
    optAlg.v = @. β₂ * optAlg.v + (1-β₂) * (▽*▽)
    #optAlg.v = [β₂ .* optAlg.v.data[i] .+ (1-β₂) .* (▽.data[i] .* ▽.data[i]) for i in 1:size(optAlg.v.data)]
    m̂        = @. optAlg.m /(1-β₁^t)
    v̂        = @. optAlg.v /(1-β₂^t)
    θ        = @. θ - (η * m̂) /(sqrt(v̂)+ϵ)
    return     (θ=θ,stop=false)
end

# ------------------------------------------------------------------------------
# DebugOptAlg

struct DebugOptAlg <: OptimisationAlgorithm
    dString::String
    function DebugOptAlg(;dString="Hello World, I am a Debugging Algorithm. I done nothing to your Net.")
        return new(dString)
    end
end

function singleUpdate!(θ,▽,optAlg::DebugOptAlg;nEpoch,nBatch,batchSize,ϵ_epoch,ϵ_epoch_l)
    println(optAlg.dString)
    return (θ=θ,stop=false)
end