#=
Some functions from *Universal resilience patterns in complex networks*
by Gao et al. (2016)
=#

using Statistics: mean, std

"""
    s(N::AbstractUnipartiteNetwork; dims::Union{Nothing,Integer}=nothing)

Computes the average weighted degree. This is proportional to the (weighted)
density of interactions.

If dims is provided, the incoming (`dims=1`) or outgoing (`dims=2`) is computed.
"""
function s(N::AbstractUnipartiteNetwork; dims::Union{Nothing,Integer}=nothing)
    dims == 1 && return sum(N.A, dims=2)
    dims == 2 && return sum(N.A, dims=1)'
    if dims == nothing
        return sum(N.A) / size(N)[1]
    end
end
#=
"""
    s(N::AbstractUnipartiteNetwork; dims)

Computes the vector of incoming weighted degrees of an unipartite network.
"""
s_in(N::AbstractUnipartiteNetwork) = sum(N.A, dims=2)

"""
    s_out(N::AbstractUnipartiteNetwork)

Computes the vector of outgoing weighted degrees of an unipartite network.
"""
s_out(N::AbstractUnipartiteNetwork) = sum(N.A, dims=1)'

=#

"""
    σ_in(N::AbstractUnipartiteNetwork)

Computes the standard deviation of the ingoing weighted degree of an unipartite
network.
"""
σ_in(N::AbstractUnipartiteNetwork) = std(s(N, dims=1), corrected=false)

"""
    σ_out(N::AbstractUnipartiteNetwork)

Computes the standard deviation of the outgoing weighted degree of an unipartite
network.
"""
σ_out(N::AbstractUnipartiteNetwork) = std(s(N, dims=2), corrected=false)

"""
    symmetry(N::AbstractUnipartiteNetwork)

Computes the symmetry between s^in and s^out (the in- and outgoing weighted
degree of an unipartite network). This is computed as the Pearson correlation
between the s^in and s^out. It is hence a value between -1 and 1, where high
positive values indicate that species with many outgoing degrees tend to have
many ingoing degrees and negative values mean the opposite. An undirected network
is perfectly symmetric but, for example, a food web where predators are less likely
to be prey would have a negative symmetry.

> Goa, J., Barzael, B. and Barabási 2016. Universal resilience patterns in complex networks.
> Nature 530(7590), 307-312. doi:10.1038/nature16948

"""
symmetry(N::AbstractUnipartiteNetwork) = (mean(s(N, dims=1) .* s(N, dims=2)) - mean(s(N, dims=1)) * mean(s(N, dims=2))) / (σ_in(N) * σ_out(N))

"""
    heterogeneity(N::AbstractUnipartiteNetwork)

Computes the heterogeneity for an unipartite network, a topological characteristic
which quantifies the difference in in- and outgoing degrees between species. It
is computed as σ_in * σ_out / s_mean. A value of 0 indicates that all species
have the same (weighted) in- and outdegrees.

> Goa, J., Barzael, B. and Barabási 2016. Universal resilience patterns in complex networks.
> Nature 530(7590), 307-312. doi:10.1038/nature16948

"""
heterogeneity(N::AbstractUnipartiteNetwork) = (σ_in(N) * σ_out(N)) / s(N)

"""
    resilience(N::AbstractUnipartiteNetwork)

A resilience parameters described by Gao et al. (2016). It is a global parameters
describing the dynamics of an unipartite network as an effective 1D equation of
the form

f(xeff) = F(xeff) + βeff G(xeff, xeff)

i.e. describing a second-order term representing the effect of the network on the
dynamics of the 'effective state' xeff of the system.

> Goa, J., Barzael, B. and Barabási 2016. Universal resilience patterns in complex networks.
> Nature 530(7590), 307-312. doi:10.1038/nature16948

"""
resilience(N::AbstractUnipartiteNetwork) = dot(s(N, dims=1), s(N, dims=2)) / sum(N)


using Test, SDCO

@testset "SDCO instance, feasibility, objectives" begin

    @testset "testcase1a sparse - symmetric = $issym" for issym in [true, false]
        pb, z = testcase1a(symmetric = issym);

        @test evaluate(pb.A, z.x) == [1., 1., 1.]

        @test get_primobj(pb, z) == 3.
        @test get_dualobj(pb, z) == 0.

        @test get_primslacks(pb, z) == [0., 0., 0.]
        @test norm(get_dualslacks(pb, z.y, z.s)) == 0

        @test get_primfeaserr(pb, z) == 0.
        @test get_dualfeaserr(pb, z) == 0.
    end
end

@testset "Nesterov Todd step" begin

    @testset "testcase1a sparse - symmetric = $issym" for issym in [true, false]
        pb, z = testcase1a(symmetric = issym)

        @testset "mu = $mu" for mu in [0.2 0.4 0.6]
            dz = NesterovToddstep(pb, z, mu)

            @test norm(dz.x - PointE([ (mu-1)*[0. -1/3 -1/3 ; -1/3 2/3 0.; -1/3 0. 2/3] ], Float64[])) < 1e-15
            @test norm(dz.y - (1-mu) .* [1., 1/3, 1/3]) < 1e-15
            @test norm(dz.s - PointE([ (mu-1)*[1. 1/3 1/3 ; 1/3 1/3 0.; 1/3 0. 1/3] ], Float64[])) < 1e-15
        end
    end

    @testset "testcase1b" begin
        pb, z = testcase1b()

        @testset "mu = $mu" for mu in [1. 0.2 0.4 0.6]
            dz = NesterovToddstep(pb, z, mu)

            @test norm(dz.y - [0.6 * (1-3*mu)]) < 1e-15

            @test norm(dz.s - PointE(SDCO.Dense{Float64}[], Float64[ 0.6*(3*mu-1), 1.2*(3*mu-1) ])) < 1e-15

            @test norm(dz.x - PointE(SDCO.Dense{Float64}[], Float64[ 0.4*mu - 2/15, -0.2*mu + 1/15 ])) < 1e-15
        end
    end

end

@testset "Central path following" begin

    @testset "testcase1a sparse - symmetric = $issym" for issym in [true, false]
        pb, z = testcase1a(symmetric = issym)

        @assert delta(pb, z) < 1/3

        zfinal = NTpredcorr(pb, z; outlev=0)

        epsilon = 1e-10
        @test mu(zfinal) < 1e-10

        @test get_primfeaserr(pb, zfinal) < epsilon
        @test get_dualfeaserr(pb, zfinal) < epsilon
        @test dot(zfinal.x, zfinal.s) < epsilon
    end

    @testset "testcase1b" begin
        pb, z = testcase1b()

        @assert delta(pb, z) < 1/3

        zfinal = NTpredcorr(pb, z; outlev=0)

        epsilon = 1e-10
        @test mu(zfinal) < 1e-10

        @test get_primfeaserr(pb, zfinal) < epsilon
        @test get_dualfeaserr(pb, zfinal) < epsilon
        @test dot(zfinal.x, zfinal.s) < epsilon
    end
end

@testset "Getting central point" begin
    @testset "testcase1c sparse - symmetric = $issym" for issym in [true, false]
        pb, z = testcase1c(symmetric = issym)

        @assert delta(pb, z) > 1/3

        zcentralpath = NTgetcentralpathpoint(pb, z; outlev=0)

        @test delta(pb, zcentralpath) < 1/3
    end
end

@testset "Minimum Matrix Norm, feasible starting point" begin
    @testset "mmn-$p-$q-$r" for (p, q, r) in [ (2, 2, 2), (4, 2, 2) ]
        ε = 1e-10

        pb, z = testcase2(p, q, r)
        pb.options[:opt_ε] = ε
        pb.options[:opt_outlev] = 0

        zfinal = solve(pb, z)

        @test mu(zfinal) < ε

        # @test get_primfeaserr(pb, zfinal) < epsilon
        # @test get_dualfeaserr(pb, zfinal) < epsilon
        # @test dot(zfinal.x, zfinal.s) < epsilon
    end
end


struct 𝕊{n} <: Manifold end

dim(::Type{𝕊{n}}) where n = n
embdim(::Type{𝕊{n}}) where n = n+1

# Creates aliases S1, TS1, etc.
for n = 1:3
    @eval const $(Symbol("S$n")) = 𝕊{$n}
    @eval const $(Symbol("TS$n")) = T{𝕊{$n}}
    @eval export $(Symbol("S$n"))
    @eval export $(Symbol("TS$n"))
end
TS{n} = T{𝕊{n}}

# Sterographic Projection
abstract type SterProj <: ChartID end
isglobal(::Chart{<:SterProj,𝕊{n}}) where n = false
struct SterProjSouth <: SterProj end  # South pole
struct SterProjNorth <: SterProj end  # North pole
emb_to_chart(pe, ::Chart{SterProjSouth,𝕊{n}}) where n = SVector{n}(pe[1:n])./(1+pe[n+1])
emb_to_chart(pe, ::Chart{SterProjNorth,𝕊{n}}) where n = SVector{n}(pe[1:n])./(1-pe[n+1])
function chart_to_emb(pn, ::Chart{SterProjSouth,S2})
    den = 1 + pn[1]^2 + pn[2]^2
    SA[2*pn[1]/den, 2*pn[2]/den, (2-den)/den]
end
function chart_to_emb(ps, ::Chart{SterProjNorth,S2})
    den = 1 + ps[1]^2 + ps[2]^2
    SA[2*ps[1]/den, 2*ps[2]/den, (den-2)/den]
end
chart_choice_coord_rep(::Chart{<:SterProj,S2}) = EmbRep()
choose_chart_emb(::EmbRep, pe, ::Chart{<:SterProj,S2}) =
    (pe[3] < 0)[1] ? Chart{SterProjSouth,S2}() : Chart{SterProjNorth,S2}()

# Same both ways
chart_transition(ps, ::Chart{SterProjSouth,𝕊{n}}, ::Chart{SterProjNorth,𝕊{n}}) where n =
    ps/norm(ps)^2
chart_transition(pn, ::Chart{SterProjNorth,𝕊{n}}, ::Chart{SterProjSouth,𝕊{n}}) where n =
    pn/norm(pn)^2

# Angle representation for S1
abstract type AngleChart <: ChartID end
isglobal(::Chart{<:AngleChart,S1}) = false
struct Angleπ <: AngleChart end    # Angles in (-π, π)
struct Angle2π <: AngleChart end   # Angles in (0, 2π)
default_chart(::Type{S1}) = Chart{Angleπ,S1}()
emb_to_chart(pe, ::Chart{Angleπ,S1}) = SA[atan(pe[2],pe[1])]
function emb_to_chart(pe, C2π::Chart{Angle2π,S1})
    Cπ = Chart{Angleπ,S1}()
    pπ = emb_to_chart(pe, Cπ) # Between -π and π
    chart_transition(pπ, Cπ, C2π)
end

chart_to_emb(p, ::Chart{<:AngleChart,S1}) = [cos.(p); sin.(p)]

chart_transition(p2π, ::Chart{Angle2π,S1}, ::Chart{Angleπ,S1}) = (p2π .< π)[1] ? p2π : p2π .- 2π
chart_transition(pπ, ::Chart{Angleπ,S1}, ::Chart{Angle2π,S1}) = (pπ .> 0.)[1] ? pπ : pπ .+ 2π

chart_choice_coord_rep(::Chart{<:AngleChart,S1}) = EmbRep()
choose_chart_emb(::EmbRep, pe, ::Chart{<:AngleChart,S1}) =
    (pe[1] .< 0.)[1] ? Chart{Angle2π,S1}() : Chart{Angleπ,S1}()


using ChainRulesCore

function foo(x)
    a = sin(x)
    b = 0.2 + a
    c = asin(b)
    return c
end

# Define rules (alternatively get them for free via `using ChainRules`)
@scalar_rule(sin(x), cos(x))
@scalar_rule(+(x, y), (1.0, 1.0))
@scalar_rule(asin(x), inv(sqrt(1 - x^2)))

#### Find dfoo/dx via rrules
#### First the forward pass, gathering up the pullbacks
x = 3;
a, a_pullback = rrule(sin, x); # a_pullback(x) = (NoTangent(), cos(x) * x)
b, b_pullback = rrule(+, 0.2, a);
c, c_pullback = rrule(asin, b)

#### Then the backward pass calculating gradients
c̄ = 1;                    # ∂c/∂c
_, b̄ = c_pullback(c̄);     # ∂c/∂b = ∂c/∂b ⋅ ∂c/∂c
println(b̄)
println(inv(sqrt(1 - b^2)))
_, ā2, ā = b_pullback(b̄);  # ∂c/∂a = ∂c/∂b ⋅ ∂b/∂a
println((ā2, ā))
println(b̄ .* (1.0, 1.0))
_, x̄ = a_pullback(ā);     # ∂c/∂x = ∂c/∂a ⋅ ∂a/∂x
println(x̄)                # ∂c/∂x = ∂foo/∂x
# ∂c/∂x = ∂c/∂b ⋅ ∂b/∂a ⋅ ∂a/∂x
#### Find dfoo/dx via frules
x = 3;
ẋ = 1;              # ∂x/∂x
nofields = ZeroTangent();  # ∂self/∂self

a, ȧ = frule((nofields, ẋ), sin, x);                    # ∂a/∂x = ∂a/∂x ⋅ ∂x/∂x
println(ȧ)
println(cos(x))
b, ḃ = frule((nofields, ZeroTangent(), ȧ), +, 0.2, a);  # ∂b/∂x = ∂b/∂a ⋅ ∂a/∂x

c, ċ = frule((nofields, ḃ), asin, b);                   # ∂c/∂x = ∂c/∂b ⋅ ∂b/∂x
ċ                                                       # ∂c/∂x = ∂foo/∂x
# ∂c/∂x = ∂c/∂b ⋅ ∂b/∂a ⋅ ∂a/∂x
#### Find dfoo/dx via FiniteDifferences.jl
using FiniteDifferences
central_fdm(5, 1)(foo, x)
# output
-1.0531613736418257

#### Find dfoo/dx via ForwardDiff.jl
using ForwardDiff
ForwardDiff.derivative(foo, x)
# output
-1.0531613736418153

#### Find dfoo/dx via Zygote.jl
using Zygote
Zygote.gradient(foo, x)


using MXNet

mlp = @mx.chain mx.Variable(:data)             =>
  mx.FullyConnected(name=:fc1, num_hidden=128) =>
  mx.Activation(name=:relu1, act_type=:relu)   =>
  mx.FullyConnected(name=:fc2, num_hidden=64)  =>
  mx.Activation(name=:relu2, act_type=:relu)   =>
  mx.FullyConnected(name=:fc3, num_hidden=10)  =>
  mx.SoftmaxOutput(name=:softmax)

# data provider
batch_size = 100
include(Pkg.dir("MXNet", "examples", "mnist", "mnist-data.jl"))
train_provider, eval_provider = get_mnist_providers(batch_size)

# setup model
model = mx.FeedForward(mlp, context=mx.gpu())

# optimization algorithm
optimizer = mx.SGD(lr=0.1, momentum=0.9)

# fit parameters
mx.fit(model, optimizer, train_provider, n_epoch=20, eval_data=eval_provider)

You can also predict using the model in the following way:

probs = mx.predict(model, eval_provider)

# collect all labels from eval data
labels = Array[]
for batch in eval_provider
    push!(labels, copy(mx.get(eval_provider, batch, :softmax_label)))
end
labels = cat(1, labels...)

# Now we use compute the accuracy
correct = 0
for i = 1:length(labels)
    # labels are 0...9
    if indmax(probs[:,i]) == labels[i]+1
        correct += 1
    end
end
accuracy = 100correct/length(labels)
println(mx.format("Accuracy on eval set: {1:.2f}%", accuracy))


# ### m14.7t.jl

using Pkg, DrWatson

using MCMCChains
using StanSample
using StatisticalRethinking

Dmat = CSV.read(sr_datadir("islandDistMatrix.csv"), DataFrame);
df = CSV.read(sr_datadir("Kline2.csv"), DataFrame);
df.society = 1:10

stan14_7 = "
functions{
    matrix cov_GPL2(matrix x, real sq_alpha, real sq_rho, real delta) {
        int N = dims(x)[1];
        matrix[N, N] K;
        for (i in 1:(N-1)) {
          K[i, i] = sq_alpha + delta;
          for (j in (i + 1):N) {
            K[i, j] = sq_alpha * exp(-sq_rho * square(x[i,j]) );
            K[j, i] = K[i, j];
          }
        }
        K[N, N] = sq_alpha + delta;
        return K;
    }
}
data{
    int T[10];
    int society[10];
    int P[10];
    matrix[10,10] Dmat;
}
parameters{
    vector[10] k;
    real<lower=0> g;
    real<lower=0> b;
    real<lower=0> a;
    real<lower=0> etasq;
    real<lower=0> rhosq;
}
model{
    vector[10] lambda;
    matrix[10,10] SIGMA;
    rhosq ~ exponential( 0.5 );
    etasq ~ exponential( 2 );
    a ~ exponential( 1 );
    b ~ exponential( 1 );
    g ~ exponential( 1 );
    SIGMA = cov_GPL2(Dmat, etasq, rhosq, 0.01);
    k ~ multi_normal( rep_vector(0,10) , SIGMA );
    for ( i in 1:10 ) {
        lambda[i] = (a * P[i]^b/g) * exp(k[society[i]]);
    }
    T ~ poisson( lambda );
}
";

m14_7s_data = Dict(
    :T => df.total_tools,
    :P => df.population,
    :society => df.society,
    :Dmat => Array(Dmat[:, 2:end])
)

m14_7s = SampleModel("m14.7s", stan14_7)
rc14_7s = stan_sample(m14_7s; data = m14_7s_data)
if success(rc14_7s)
  chns14_7s = read_samples(m14_7s, :mcmcchains)
  chns14_7s |> display

  println()

  read_summary(m14_7s) |> display
end

m14_7s_results = "
       mean   sd  5.5% 94.5% n_eff Rhat4
k[1]  -0.16 0.34 -0.71  0.33   615  1.01
k[2]  -0.02 0.33 -0.55  0.47   550  1.01
k[3]  -0.07 0.32 -0.58  0.40   566  1.01
k[4]   0.35 0.29 -0.05  0.81   557  1.01
k[5]   0.08 0.29 -0.36  0.51   509  1.01
k[6]  -0.38 0.30 -0.88  0.03   585  1.00
k[7]   0.14 0.28 -0.29  0.56   567  1.00
k[8]  -0.21 0.29 -0.66  0.20   526  1.01
k[9]   0.26 0.27 -0.15  0.65   534  1.01
k[10] -0.17 0.36 -0.73  0.36   768  1.00
g      0.61 0.60  0.07  1.70  1321  1.00
b      0.28 0.09  0.14  0.42  1011  1.00
a      1.38 1.05  0.25  3.30  1800  1.00
etasq  0.21 0.22  0.03  0.61   982  1.00
rhosq  1.28 1.56  0.08  4.32  2020  1.00
";

# End of m14_7s.jl


using AugmentedGaussianProcesses;
using Distributions, LinearAlgebra
using Random: seed!;
seed!(123);
using Plots;
gr();

N = 100;
noise = 1e-1;
N_test = 1000
X = sort(rand(N))
X_test = collect(range(0; length=N_test, stop=1))
k = SqExponentialKernel() ∘ ScaleTransform(10.0)
K = kernelmatrix(k, X) + noise * I
f = rand(MvNormal(K))
λ = 10
y = rand.(Poisson.(λ * AGP.logistic.(f)))

poissonmodel = VGP(X, y, k, PoissonLikelihood(Float64(λ)), AnalyticVI(); optimiser=false);
train!(poissonmodel, 300)
poissonpred = predict_y(poissonmodel, X_test)

negbinmodel = VGP(X, y, k, NegBinomialLikelihood(λ), AnalyticVI(); optimiser=false);
train!(negbinmodel, 50)
negbinpred = predict_y(negbinmodel, X_test)
##
default(; legendfontsize=14.0, xtickfontsize=10.0, ytickfontsize=10.0)
p = plot(X, y; t=:scatter, lab="Training Points")
plot!(X_test, poissonpred; lab="Poisson Expectation", lw=7.0)
plot!(X_test, negbinpred; lab="Negative Binomial Expectation", lw=7.0, legend=:topright)

display(p)
savefig(p, joinpath(@__DIR__, "Events.png"))
# a,b = logitmodel.fstar(X_test)
# using Expectations, QuadGK, BenchmarkTools

# @btime begin; for i in 1:N_test; d = Normal(a[i],sqrt(max(b[i],1e-8))); expectation(logit,d); end; end;
# @btime begin; for i in 1:N_test; d = Normal(a[i],sqrt(max(b[i],1e-8))); quadgk(x->pdf(d,x)*logit(x),-Inf,Inf); end; end;


#Figure 2-i
include("fig2b.jl")

#calculating evenness
function pielou(n)
    np = filter(x -> x > eps(), n)
    p = np./sum(np)
    ev = length(p) == 1 ? 0.0 : -sum(p.*log.(p))*(1/log(length(p)))
    return ev
end;
ev = mapslices(pielou, Np[:,2:end]; dims=2);
plot(ev, c=:black, title = "Evenness", xlabel = "Time", ylabel = "Relative abundance (log)", leg = false, ylims = (0,1))

# png("Figure 2/graph_2j.png")


gausslegendre(n::Int) = gausslegendre(Float64, n)
gausslegendre(::Type{Float64}, n) = FastGaussQuadrature.gausslegendre(n)
gausslegendre(::Type{T}, n) where {T} = GaussQuadrature.legendre(T, n)

gausslaguerre(n::Int) = gausslaguerre(Float64, n)
gausslaguerre(n::Int, α::T) where {T<:AbstractFloat} = gausslaguerre(T, n, α)
gausslaguerre(n::Int, α) = gausslaguerre(n, float(α))
gausslaguerre(::Type{T}, n::Int) where {T} = gausslaguerre(T, n, zero(T))
gausslaguerre(::Type{Float64}, n, α) = FastGaussQuadrature.gausslaguerre(n, α)

gaussjacobi(n::Int) = gaussjacobi(Float64, n)
gaussjacobi(n::Int, α::Number, β::Number) = gaussjacobi(n, promote(α, β)...)
gaussjacobi(n::Int, α::N, β::N) where {N<:Number} = gaussjacobi(n, float(α), float(β))
gaussjacobi(n::Int, α::T, β::T) where {T<:AbstractFloat} =
    gaussjacobi(T, n, α, β)
gaussjacobi(::Type{T}, n::Int) where {T} = gaussjacobi(T, n, zero(T), zero(T))
gaussjacobi(::Type{Float64}, n, α, β) = FastGaussQuadrature.gaussjacobi(n, α, β)

gausshermite(n::Int) = gausshermite(Float64, n)
gausshermite(::Type{Float64}, n) = FastGaussQuadrature.gausshermite(n)


#!/usr/local/bin/julia
# last tested Julia version: 1.6.1
# Hindmarsh-Rose model on a 2D lattice
# FvW 03/2018

using NPZ
using PyCall
using PyPlot
using Statistics
using VideoIO
@pyimport matplotlib.animation as anim

function animate_pyplot(fname, data)
    """
    Animate 3D array as .mp4 using PyPlot, save as `fname`
    array dimensions:
        1: time
        2, 3: space
    """
    nt, nx, ny = size(data)
    vmin = minimum(data)
    vmax = maximum(data)
    # setup animation image
    fig = figure(figsize=(6,6))
    axis("off")
    t = imshow(data[1,:,:], origin="lower", cmap=ColorMap("gray"),
			   vmin=vmin, vmax=vmax)
    tight_layout()
    # frame generator
    println("[+] animate")
    function animate(i)
        (i%100 == 0) && print("    t = ", i, "/", nt, "\r")
        t.set_data(data[i+1,:,:])
    end
    # create animation
    ani = anim.FuncAnimation(fig, animate, frames=nt, interval=10)
    println("\n")
    # save animation
    ani[:save](fname, bitrate=-1,
               extra_args=["-vcodec", "libx264", "-pix_fmt", "yuv420p"])
    show()
end

function animate_video(fname, data)
    """
    Animate 3D array as .mp4 using VideoIO, save as `fname`
    array dimensions:
        1: time
        2, 3: space
    """
    # BW
    y = UInt8.(round.(255*(data .- minimum(data))/(maximum(data)-minimum(data))))
    # BW inverted
    #y = UInt8.(round.(255 .- 255*(data .- minimum(data))/(maximum(data)-minimum(data))))
    encoder_options = (color_range=2, crf=0, preset="medium")
    framerate=90
    T = size(data,1)
    open_video_out(fname, y[1,end:-1:1,:], framerate=framerate,
                   encoder_options=encoder_options) do writer
        for i in range(2,stop=T,step=1)
            write(writer, y[i,end:-1:1,:])
        end
    end
end

function L(x)
    # Laplace operator
    # periodic boundary conditions
    xU = circshift(x, [-1 0])
    xD = circshift(x, [1 0])
    xL = circshift(x, [0 -1])
    xR = circshift(x, [0 1])
    Lx = xU + xD + xL + xR - 4x
    # non-periodic boundary conditions
    Lx[1,:] .= 0.0
    Lx[end,:] .= 0.0
    Lx[:,1] .= 0.0
    Lx[:,end] .= 0.0
    return Lx
end

function hr2d(N, T, t0, dt, sd, D, a, b, c, d, s, r, x_1, I0, stim, blocks)
	# initialize Hindmarsh-Rose system
	x, y, z = zeros(Float64,N,N), zeros(Float64,N,N), zeros(Float64,N,N)
	dxdt, dydt, dzdt = zeros(Float64,N,N), zeros(Float64,N,N), zeros(Float64,N,N)
	sd_sqrt_dt = sd*sqrt(dt)
	X = zeros(Float64,T,N,N)
	I = zeros(Float64,t0+T,N,N)
	# stimulation protocol
	for st in stim
		t_on, t_off = st[1]
		x0, x1 = st[2]
		y0, y1 = st[3]
		I[t0+t_on:t0+t_off, x0:x1, y0:y1] .= I0
		I[1:t0, x0:x1, y0:y1] .= I0
	end
	# iterate
	for t in range(1, stop=t0+T, step=1)
		(t%100 == 0) && print("    t = ", t, "\r")
		# Hindmarsh-Rose equations
		dx = -a*x.*x.*x + b*x.*x + y - z + I[t,:,:] + D*L(x)
		dy = c .- d*x.*x - y
		dz = r*(s*(x .- x_1) - z)
		# Ito stochastic integration
		x += (dx*dt + sd_sqrt_dt*randn(N,N))
		y += (dy*dt)
		z += (dz*dt)
		# dead block(s):
		for bl in blocks
			x[bl[1][1]:bl[1][2], bl[2][1]:bl[2][2]] .= 0.0
			y[bl[1][1]:bl[1][2], bl[2][1]:bl[2][2]] .= 0.0
			z[bl[1][1]:bl[1][2], bl[2][1]:bl[2][2]] .= 0.0
		end
		(t > t0) && (X[t-t0,:,:] = x)
	end
	println("\n")
	return X
end

function L(x)
	# Laplace operator
    xU = circshift(x, [-1 0])
    xD = circshift(x, [1 0])
    xL = circshift(x, [0 -1])
    xR = circshift(x, [0 1])
    Lx = xU + xD + xL + xR - 4x
	nx = size(x)[1]
	ny = size(x)[2]
    Lx[1,:] .= 0.0
	Lx[nx,:] .= 0.0
	Lx[:,1] .= 0.0
	Lx[:,ny] .= 0.0
    return Lx
end

function main()
    println("[+] Hindmarsh-Rose 2D lattice:")
	N = 64
    T = 15000
	t0 = 2500
	dt = 0.05
	sd = 0.05
	D = 2.0
	a = 1.0
	b = 3.0
	c = 1.0
	d = 5.0
	s = 2.0 # s=1: spiking, s=4: strong accomodation
	r = 0.001
    x_1 = -1.6
	I0 = 3.50 # 1.3
	#I1 = 0.05 # 1.3
    println("[+] Lattice size N: ", N)
	println("[+] Time steps T: ", N)
    println("[+] Warm-up steps t0: ", t0)
	println("[+] Integration time step dt: ", dt)
	println("[+] Noise intensity sd: ", sd)
	println("[+] Diffusion constant D: ", D)
    println("[+] HR parameter a: ", a)
    println("[+] HR parameter b: ", b)
    println("[+] HR parameter c: ", c)
    println("[+] HR parameter d: ", d)
    println("[+] HR parameter s: ", s)
    println("[+] HR parameter r: ", r)
    println("[+] HR parameter x_1: ", x_1)
    println("[+] Stimulation current I0: ", I0)
    #println("[+] Stimulation current I1: ", I1)

	# stim protocol, array of elements [[t0,t1], [x0,x1], [y0,y1]]
	stim = [ [[1,15000], [5,N-5], [5,N-5]] ]
    #stim = []# dead blocks, array of elementy [[x0,x1], [y0,y1]]

	# dead blocks, array of elementy [[x0,x1], [y0,y1]]
	#blocks = [ [[2*n_4,3*n_4], [15,20]], [[2*n_4+10,3*n_4+10], [40,45]] ]
	blocks = []

	# run simulation
	data = hr2d(N, T, t0, dt, sd, D, a, b, c, d, s, r, x_1, I0, stim, blocks)
    println("[+] Data dimensions: ", size(data))

	# plot mean voltage
    m = mean(reshape(data, (T,N*N)), dims=2)
    plot(m, "-k"); show()

	# save data
	I_str = rpad(I0, 4, '0') # stim. current amplitude as 4-char string
	s_str = rpad(s, 4, '0') # noise as 4-char string
	sd_str = rpad(sd, 4, '0') # noisy intensity as 4-char string
    D_str = rpad(D, 4, '0') # diffusion coefficient as 4-char string
	fname1 = string("hr2d_I_", I_str, "_s_", s_str, "_sd_", sd_str,
					"_D_", D_str, ".npy")
    #npzwrite(data_filename, data)
	#println("[+] Data saved as: ", data_filename)

	# video
	fname2 = string("hr2d_I_", I_str, "_s_", s_str, "_sd_", sd_str,
					"_D_", D_str, ".mp4")
	#animate_pyplot(fname2, data) # slow
    animate_video(fname2, data) # fast
	println("[+] Data saved as: ", fname2)
end

main()


struct UnionFind
    par::Vector{Int}
    size::Vector{Int}
    UnionFind(n) = new(collect(1:n), ones(Int, n))
end

root!(uf, x) = uf.par[x] == x ? x : uf.par[x] = root!(uf, uf.par[x])

same(uf, x, y) = root!(uf, x) == root!(uf, y)

function unite!(uf, x, y)
    x = root!(uf, x)
    y = root!(uf, y)
    x == y && return false

    if uf.size[x] > uf.size[y]
        uf.par[y] = x
        uf.size[x] += uf.size[y]
    else
        uf.par[x] = y
        uf.size[y] += uf.size[x]
    end
    return true
end
size!(uf, x) = uf.size[root!(uf, x)]

function main()
    N, M = parse.(Int, split(readline()))
    AB = [parse.(Int, split(readline())) for _ in 1:M]

    anss = zeros(Int, M)
    uf = UnionFind(N)

    for i in M:-1:1
        A, B = AB[i][1], AB[i][2]
        if !same(uf, A, B)
            anss[i] = size!(uf, A) * size!(uf, B)
            unite!(uf, A, B)
        end
    end

    for ans in cumsum(anss)
        println(ans)
    end

end

main()


"""
Based on
https://hyperopt.github.io/hyperopt/
"""
module TestQuantNormal

using Test
using TreeParzen

# Simple quantnormal
# Called HP.QuantNormal multiple number of times to check if the distribution of values
# matched the expected given input parameters and also a test that the values appear at
# correct intervals
d = []
for i in 1:10000
    a = HP.QuantNormal(:x, 0.0, 1.0, 2.0)
    b = TreeParzen.Resolve.node(a, TreeParzen.Trials.ValsDict())
    push!(d, b)
end
@test all(in.(d, [[-4.0, -2.0, 0, 2.0, 4.0]]))

end # module TestQuantNormal
true


using Test
using SumOfSquares
using DynamicPolynomials
using SemialgebraicSets

# Adapted from:
# SOSDEMO4 --- Matrix Copositivity
# Section 3.4 of SOSTOOLS User's Manual
#
# See also (3.41) of [BPT12].
#
# [BPT12] Blekherman, G.; Parrilo, P. A. & Thomas, R. R.
# *Semidefinite Optimization and Convex Algebraic Geometry*.
# Society for Industrial and Applied Mathematics, **2012**.

function horn_test(optimizer,
                   config::MOIT.TestConfig,
                   cone::SumOfSquares.PolyJuMP.PolynomialSet)
    atol = config.atol
    rtol = config.rtol

    # Horn matrix
    H = [1 -1  1  1 -1;
        -1  1 -1  1  1;
         1 -1  1 -1  1;
         1  1 -1  1 -1;
        -1  1  1 -1  1]

    @polyvar x[1:5]

    model = _model(optimizer)
    cref = @constraint(model, x' * H * x in CopositiveInner(cone))
    optimize!(model)

    @test termination_status(model) == MOI.INFEASIBLE

    model = _model(optimizer)
    orthant = @set x[1] ≥ 0 && x[2] ≥ 0 && x[3] ≥ 0 && x[4] ≥ 0 && x[5] ≥ 0
    cref = @constraint(model, sum(x) * x' * H * x in cone, domain = orthant)
    optimize!(model)

    @test termination_status(model) == MOI.INFEASIBLE

    model = _model(optimizer)
    orthant = orthant ∩ basicsemialgebraicset(
        algebraicset(polynomialtype(x[1])[]), polynomial.(monomials(x, 3)))
    cref = @constraint(model, sum(x) * x' * H * x in cone, domain = orthant)
    optimize!(model)

    if cone isa SOSCone
        @test termination_status(model) == MOI.OPTIMAL

        @test primal_status(model) == MOI.FEASIBLE_POINT
        @test isempty(certificate_monomials(cref))
        @test isempty(getmat(gram_matrix(cref)))
    else
        @test termination_status(model) == MOI.INFEASIBLE
    end
end
sos_horn_test(optimizer, config)   = horn_test(optimizer, config, SOSCone())
sdsos_horn_test(optimizer, config) = horn_test(optimizer, config, SDSOSCone())
dsos_horn_test(optimizer, config)  = horn_test(optimizer, config, DSOSCone())


using GameTheory: pure_nash
using QuantEcon: MVNSampler
using GameTheory.Generators

include("test_bimatrix_generators.jl")


using Pkg
Pkg.activate(joinpath(Pkg.devdir(), "MLCourse"))
using Plots, CSV, MLJ, DataFrames, Random, NearestNeighborModels
include("./Data_Processing.jl")

#Load the data
drop = generate(option = "drop", std = "false", valid = "false", test = "true");
med = generate(option = "med", std = "false", valid = "false", test = "true");

#Function to tune the model with values of K between 1 and 50
function TunedModel_KNN(data)
    Random.seed!(2711)
    model = KNNClassifier()
    self_tuning_model = TunedModel(model = model,
                                resampling = CV(nfolds = 10),
                                tuning = Grid(),
                                range = range(model, :K, values = 1:50),
                                measure = auc)
    self_tuning_mach = machine(self_tuning_model,
                            select(data, Not(:precipitation_nextday)),
                            data.precipitation_nextday) |> fit!
    self_tuning_mach
end

#Test the tuned model on the drop data set and med data set:

rep1 = report(TunedModel_KNN(drop.train));
scatter(reshape(rep1.plotting.parameter_values, :), rep1.plotting.measurements, xlabel = "K", ylabel = "AUC")
# K = 27 with drop

rep2 = report(TunedModel_KNN(med.train));
scatter(reshape(rep2.plotting.parameter_values, :), rep2.plotting.measurements, xlabel = "K", ylabel = "AUC")
# K = 25 with med

# Test the best model over the test set with the best hyper-parameters to pick the best one
best_mach1 = machine(KNNClassifier(K = rep1.best_model.K), select(drop.train[:,:], Not(:precipitation_nextday)), drop.train.precipitation_nextday)|> fit!;
best_mach2 = machine(KNNClassifier(K = rep2.best_model.K), select(med.train[:,:], Not(:precipitation_nextday)), med.train.precipitation_nextday)|> fit!;
AUC1 = losses(best_mach1, select(drop.test, Not(:precipitation_nextday)), drop.test.precipitation_nextday)
AUC2 = losses(best_mach2, select(med.test, Not(:precipitation_nextday)), med.test.precipitation_nextday)
# AUC_drop = 0.904 , AUC_med = 0.898

#Write in the submission file with a machine trained on all the data
train, test = generate(option = "drop", std = "false", valid = "false", test = "false");
if AUC1.auc > AUC2.auc
    best_mach = machine(KNNClassifier(K = rep1.best_model.K), select(train[:,:], Not(:precipitation_nextday)), train.precipitation_nextday)|> fit!;
else
    best_mach = machine(KNNClassifier(K = rep2.best_model.K), select(train[:,:], Not(:precipitation_nextday)), train.precipitation_nextday)|> fit!;
end
sample = CSV.read(joinpath(@__DIR__, "data", "sample_submission.csv"), DataFrame);
pred = pdf.(predict(best_mach, test), true);
sample.precipitation_nextday = pred;
CSV.write(joinpath(@__DIR__, "results", "KNN_submission.csv"), sample);


"""
    tfisingMPO(h, N[; T = ComplexF64])
returns a Z2-symmetric MPO for the transverse field Ising model:
```
    H = -∑_i σ^i_x σ^(i+1)_x - ∑_i h σ^i_z
```
on N sites.
"""
function tfisingMPO(h, N; T = ComplexF64)
    bulk = DASTensor{T,4}(
        Z2(),
        ntuple(i -> Z2Charges(),4),
        ([2,2], [1,1], [1,1], [2,2]),
        InOut(1,1,-1,-1))
    #this construction is dangerous insofar as we deliberately don't
    #populate certain sectors to avoid having to introduce more logic
    #for a more robust setting:
    #vd = 2
    #track whether operation happened or not

    preapp  = T[0 0; 0 1] # 1 -> 1
    postapp = T[1 0; 0 0] # 2 -> 2
    app     = T[0 0; 1 0] # 1 -> 2
    peye= Matrix{T}(I, 1, 1)

    for ch in Z2Charges()
        #transverse field term
        hz = ifelse(iszero(ch), h, -h)
        @tensor deg[1,2,3,4] := (preapp[1,4] + postapp[1,4] - hz*app[1,4]) * peye[2,3]
        bulk[DASSector(Z2Charge(0), ch, ch, Z2Charge(0))] = deg

        #XX term
        @tensor deg[1, 2, 3, 4] := -preapp[1, 4] * peye[2, 3]
        bulk[DASSector(Z2Charge(0), ch ⊕ Z2Charge(1), ch, Z2Charge(1))] = deg

        @tensor deg[1, 2, 3, 4] :=  app[1, 4] * peye[2, 3]
        bulk[DASSector(Z2Charge(1), ch, ch ⊕ Z2Charge(1), Z2Charge(0))] = deg
    end

    # (3)        (4)        (3)
    # [lb](2) (2)[b](3) (2)[rb]
    # (1)        (1)        (1)
    lb = DASTensor{T,1}(Z2(), (Z2Charges(),), ([2,2],), InOut(-1))
    rb = DASTensor{T,1}(Z2(), (Z2Charges(),), ([2,2],), InOut( 1))
    lb[DASSector(Z2Charge(0))] = [0, 1]
    rb[DASSector(Z2Charge(0))] = [1, 0]

    @tensor bulk[1,2,3,4]       := bulk[2,1,4,3]
    @tensor leftbound[1, 2, 3]  := lb[-1] *  bulk[1, -1, 2, 3]
    @tensor rightbound[1, 2, 3] := bulk[1, 2, -1, 3] * rb[-1]

    return MPO{N}(leftbound, bulk,  rightbound)
end

"""
    sigmaZop([T = ComplexF64])
return the Z2-symmetric σz operator.
"""
function sigmaZop(T = ComplexF64)
    a = DASTensor{T,2}(Z2(),
        (Z2Charges(), Z2Charges()),
        ([1,1],[1,1]),
        InOut(1,-1))
    for ch in Z2Charges()
        a[DASSector(ch, ch)] = reshape([ifelse(iszero(ch), 1, -1)],1,1)
    end
    return a
end


using ConstrainedDynamics
using LinearAlgebra

include("examples/chain_in_chain.jl")

ConstrainedDynamics.setentries!(mech)
A1, _, b1 = ConstrainedDynamics.densesystem(mech)
x1 = A1\b1

ConstrainedDynamics.factor!(mech.graph,mech.ldu)
ConstrainedDynamics.solve!(mech)

A2, x2, b2 = ConstrainedDynamics.densesystem(mech)

@test isapprox(norm(x1-x2), 0.0; atol = 2e-5)


function padding(p)

    θlim, slim = size(p);
    θ_left = Int(θlim/2+1);
    θ_right = Int(θlim/2+θlim);
    s_left = Int(slim/2+1);
    s_right = Int(slim/2+slim);

    p_Padding = zeros(θlim*2, slim*2);
    p_Padding[θ_left:θ_right, s_left:s_right] = p;

    return p_Padding;
end

function conical_radon(p, ϕ, μ)

    ratio = 2;
    p = padding(p);
    θlim, slim = size(p);
    θ = LinRange(-2, 2, θlim);
    U = LinRange(-2, 2, θlim);
    V = LinRange(2, -2, slim);
    C = zeros(length(V), length(U));
    interval = 2*ratio / (θlim - 1);

    for Cx = 1:length(U), Cy=1:length(V)

        U_left = 1:Cx;
        y_left = tan(π / 2 + ϕ) * (θ[U_left] .- U[Cx]) .+ V[Cy];
        bit_left = abs.(y_left) .<= 1;

        if sum(bit_left)>0
            y_left = y_left[bit_left];
            height_left = ratio .- y_left;
            nearest_row_left = round.(Int64, ((height_left) - mod.(height_left, interval)) / interval) .+ 1;
            U_left = U_left[bit_left];
            for i=1:length(nearest_row_left)
                value_left = p[nearest_row_left[i], U_left[i]];
                r_left = sqrt.((value_left-U[Cx])^2 + (value_left-V[Cy])^2);
               C[Cy, Cx]+= value_left*exp(-μ*r_left);
            end
        end

        U_right = Cx:length(U);
        y_right = tan(π / 2 - ϕ) * (θ[U_right] .- U[Cx]) .+ V[Cy];
        bit_right = abs.(y_right) .<= 1;

        if sum(bit_right)>0
            y_right = y_right[bit_right];
            height_right = ratio .- y_right;
            nearest_row_right = round.(Int64, ((height_right) - mod.(height_right, interval)) / interval) .+ 1;
            U_right = U_right[bit_right];
            for i=1:length(nearest_row_right)
                value_right = p[nearest_row_right[i], U_right[i]];
                r_right = sqrt.((value_right-U[Cx])^2 + (value_right-V[Cy])^2);
               C[Cy, Cx]+= value_right*exp(-μ*r_right);
            end
        end
    end

    return C
end


@doc doc"""

    PowerLens(ϕ, order)

`PowerLens` is a lensing operator which lenses a map with a power
series expansion in $\nabla \phi$ to any order.

```math
f(x+\nabla x) \approx f(x) + (\nabla f)(\nabla \phi) + \frac{1}{2} (\nabla \nabla f) (\nabla \phi)^2 + ...
```

The action of the operator and its adjoint can be computed.
"""
struct PowerLens{D<:Dict} <: ImplicitOp{Bottom}
    order :: Int
    ∇1ϕᵖ  :: D
    ∇2ϕᵖ  :: D
end

PowerLens(order) = x -> PowerLens(x, order)
PowerLens(ϕ::Field, order) = PowerLens(∇*ϕ, order)
function PowerLens(d::FieldVector, order)
    require_unbatched(d[1])
    ∇1ϕ,  ∇2ϕ  = Ł(d)
    ∇1ϕᵖ, ∇2ϕᵖ = (Dict(p => (p==0 ? 1 : ∇iϕ.^p) for p=0:order) for ∇iϕ=(∇1ϕ,∇2ϕ))
    PowerLens(order, ∇1ϕᵖ, ∇2ϕᵖ)
end

"""
    antilensing(L::PowerLens)

Create a `PowerLens` operator that lenses by `-ϕ` instead.
"""
function antilensing(L::PowerLens)
    PowerLens(N, (Dict(p=>∇iϕᵖ*(-1)^p for (p,∇iϕᵖ)=coeffs) for coeffs=(L.∇1ϕᵖ,L.∇1ϕᵖ))...)
end

function *(L::PowerLens, f::Field)
    require_unbatched(f)
    Ðf = Ð(f)
    f̃ = copy(Ł(f))
    for n in 1:L.order, (a,b) in zip(0:n,n:-1:0)
        @. f̃ += L.∇1ϕᵖ[a] * L.∇2ϕᵖ[b] * $(Ł(∇[1]^a * ∇[2]^b * Ðf)) / factorial(a) / factorial(b)
    end
    f̃
end

function *(Ladj::Adjoint{<:Any,<:PowerLens}, f::Field)
    require_unbatched(f)
    L = parent(Ladj)
    Łf = Ł(f)
    r = copy(Ð(f))
    for n in 1:L.order, (a,b) in zip(0:n,n:-1:0)
        r .+= (-1)^n .* (∇[1]^a * ∇[2]^b * Ð(@. L.∇1ϕᵖ[a] * L.∇2ϕᵖ[b] * Łf)) ./ factorial(a) ./ factorial(b)
    end
    r
end


using TensND, LinearAlgebra, SymPy, Tensors, OMEinsum, Rotations
sympy.init_printing(use_unicode=true)

Spherical = coorsys_spherical()
θ, ϕ, r = getcoords(Spherical)
𝐞ᶿ, 𝐞ᵠ, 𝐞ʳ = unitvec(Spherical)
ℬˢ = get_normalized_basis(Spherical)
𝐱 = getOM(Spherical)
@set_coorsys Spherical
𝐞₁, 𝐞₂, 𝐞₃ = unitvec(coorsys_cartesian())
# 𝟏, 𝟙, 𝕀, 𝕁, 𝕂 = init_isotropic(basis = ℬˢ)
𝕀, 𝕁, 𝕂 = ISO(Val(3),Val(Sym))
𝟏 = tensId2(Val(3),Val(Sym))
k, μ = symbols("k μ", positive = true)
λ = k -2μ/3

u = SymFunction("u", real = true)
𝐮ˢᵖʰ = u(r) * 𝐞ʳ
𝛆ˢᵖʰ = SYMGRAD(𝐮ˢᵖʰ, Spherical)
𝛔ˢᵖʰ = λ * tr(𝛆ˢᵖʰ) * 𝟏 + 2μ * 𝛆ˢᵖʰ
𝐓ˢᵖʰ = 𝛔ˢᵖʰ ⋅ 𝐞ʳ
div𝛔ˢᵖʰ = DIV(𝛔ˢᵖʰ, Spherical) ;
eqˢᵖʰ = factor(simplify(div𝛔ˢᵖʰ ⋅ 𝐞ʳ))
solˢᵖʰ = dsolve(eqˢᵖʰ, u(r))
ûˢᵖʰ = solˢᵖʰ.rhs()
T̂ˢᵖʰ = factor(simplify(subs(𝐓ˢᵖʰ ⋅ 𝐞ʳ, u(r) => ûˢᵖʰ)))

𝐄 = 𝟏 - 3𝐞₃⊗𝐞₃
fᶿ = simplify(𝐞ᶿ ⋅ 𝐄 ⋅ 𝐞ʳ)
fʳ = simplify(𝐞ʳ ⋅ 𝐄 ⋅ 𝐞ʳ)
# fᶿ = SymFunction("fᶿ", real = true)(θ)
# fʳ = SymFunction("fʳ", real = true)(θ)
uᶿ = SymFunction("uᶿ", real = true)
uʳ = SymFunction("uʳ", real = true)
𝐮ᵈᵉᵛ = uᶿ(r) * fᶿ * 𝐞ᶿ + uʳ(r) * fʳ * 𝐞ʳ
𝛆ᵈᵉᵛ = SYMGRAD(𝐮ᵈᵉᵛ, Spherical)
𝛔ᵈᵉᵛ = λ * tr(𝛆ᵈᵉᵛ) * 𝟏 + 2μ * 𝛆ᵈᵉᵛ
𝐓ᵈᵉᵛ = 𝛔ᵈᵉᵛ ⋅ 𝐞ʳ
div𝛔ᵈᵉᵛ = simplify(DIV(𝛔ᵈᵉᵛ, Spherical))
eqᶿᵈᵉᵛ = factor(simplify(div𝛔ᵈᵉᵛ ⋅ 𝐞ᶿ / fᶿ))
eqʳᵈᵉᵛ = factor(simplify(div𝛔ᵈᵉᵛ ⋅ 𝐞ʳ / fʳ))
α, Λ = symbols("α Λ", real = true)
eqᵈᵉᵛ = factor(simplify(subs([eqᶿᵈᵉᵛ,eqʳᵈᵉᵛ], uᶿ(r) => r^α, uʳ(r) => Λ*r^α)))
αΛ = solve([eq.doit() for eq ∈ eqᵈᵉᵛ], [α, Λ])
ûᶿᵈᵉᵛ = sum([Sym("C$(i+2)") * r^αΛ[i][1] for i ∈ 1:length(αΛ)])
ûʳᵈᵉᵛ = sum([Sym("C$(i+2)") * αΛ[i][2] * r^αΛ[i][1] for i ∈ 1:length(αΛ)])
T̂ᶿᵈᵉᵛ = factor(simplify(subs(simplify(𝐓ᵈᵉᵛ ⋅ 𝐞ᶿ / fᶿ), uᶿ(r) => ûᶿᵈᵉᵛ, uʳ(r) => ûʳᵈᵉᵛ)))
T̂ʳᵈᵉᵛ = factor(simplify(subs(simplify(𝐓ᵈᵉᵛ ⋅ 𝐞ʳ / fʳ), uᶿ(r) => ûᶿᵈᵉᵛ, uʳ(r) => ûʳᵈᵉᵛ)))


using Test
using Turing: BinomialLogit
using Distributions: Binomial, logpdf
using StatsFuns: logistic

n = 10
logitp = randn()
d1 = BinomialLogit(n, logitp)
d2 = Binomial(n, logistic(logitp))
k = 3
@test logpdf(d1, k) ≈ logpdf(d2, k)


# m[x,y,x',y']=a[x,x']* b[x',y']
function hmm2ddm!(model::DynamicDiscreteModel,a,b)
	dx,dy=size(b)
	for jy=1:dy
		for jx=1:dx
			for iy=1:dy
				for ix=1:dx
					model.m[ix,iy,jx,jy]=a[ix,jx]*b[jx,jy]
				end
			end
		end
	end
end

#multiple dispatch will detect jacobian or no jacobian
function hmm2ddm!(model::DynamicDiscreteModel,a,b,ajac,bjac)
	dx,dy,dtheta=size(bjac)
	for jy=1:dy
		for jx=1:dx
			for iy=1:dy
				for ix=1:dx
					model.m[ix,iy,jx,jy]=a[ix,jx]*b[jx,jy]
					for itheta=1:dtheta
						model.mjac[ix,iy,jx,jy,itheta]=ajac[ix,jx,itheta]*b[jx,jy]+a[ix,jx]*bjac[jx,jy,itheta]
					end
				end
			end
		end
	end
end


module spin

import Base.==

using ..bases, ..states, ..operators, ..operators_sparse

export SpinBasis, sigmax, sigmay, sigmaz, sigmap, sigmam, spinup, spindown

"""
Basis for spin-n particles.

The basis can be created for arbitrary spinnumbers by using a rational number,
e.g. SpinBasis(3//2). The Pauli operators are defined for all possible
spinnumbers.

Arguments
---------

spinnumber
    Rational number specifying the spin quantum number.
"""
type SpinBasis <: Basis
    shape::Vector{Int}
    spinnumber::Rational{Int}
    function SpinBasis(spinnumber::Rational{Int})
        @assert den(spinnumber) == 2 || den(spinnumber) == 1
        @assert den(spinnumber) > 0
        new([num(spinnumber*2 + 1)], spinnumber)
    end
end

==(b1::SpinBasis, b2::SpinBasis) = b1.spinnumber==b2.spinnumber

"""
Pauli :math:`\\sigma_x` operator for the given SpinBasis.
"""
function sigmax(b::SpinBasis)
    N = num(b.spinnumber*2 + 1)
    diag = Complex128[complex(sqrt(real((b.spinnumber + 1)*2*a - a*(a+1)))) for a=1:num(2*b.spinnumber)]
    data = spdiagm(diag, 1, N, N) + spdiagm(diag, -1, N, N)
    SparseOperator(b, data)
end

"""
Pauli :math:`\\sigma_y` operator for the given SpinBasis.
"""
function sigmay(b::SpinBasis)
    N = num(b.spinnumber*2 + 1)
    diag = Complex128[1im*complex(sqrt(real((b.spinnumber + 1)*2*a - a*(a+1)))) for a=1:num(2*b.spinnumber)]
    data = spdiagm(diag, -1, N, N) - spdiagm(diag, 1, N, N)
    SparseOperator(b, data)
end

"""
Pauli :math:`\\sigma_z` operator for the given SpinBasis.
"""
function sigmaz(b::SpinBasis)
    N = num(b.spinnumber*2 + 1)
    diag = Complex128[complex(2*m) for m=b.spinnumber:-1:-b.spinnumber]
    data = spdiagm(diag, 0, N, N)
    SparseOperator(b, data)
end

"""
Raising operator :math:`\\sigma_+` for the given SpinBasis.
"""
function sigmap(b::SpinBasis)
    N = num(b.spinnumber*2 + 1)
    S = (b.spinnumber + 1)*b.spinnumber
    diag = Complex128[complex(sqrt(float(S - m*(m+1)))) for m=b.spinnumber-1:-1:-b.spinnumber]
    data = spdiagm(diag, 1, N, N)
    SparseOperator(b, data)
end

"""
Lowering operator :math:`\\sigma_-` for the given SpinBasis.
"""
function sigmam(b::SpinBasis)
    N = num(b.spinnumber*2 + 1)
    S = (b.spinnumber + 1)*b.spinnumber
    diag = [complex(sqrt(float(S - m*(m-1)))) for m=b.spinnumber:-1:-b.spinnumber+1]
    data = spdiagm(diag, -1, N, N)
    SparseOperator(b, data)
end

"""
Spin up state for the given SpinBasis.
"""
spinup(b::SpinBasis) = basis_ket(b, 1)

"""
Spin down state for the given SpinBasis.
"""
spindown(b::SpinBasis) = basis_ket(b, b.shape[1])

end #module


include("euler/euler.jl")
using .Calculus: is_pentagonal

function compute()::Int
    pentagonals = [1]
    j = 1
    for i ∈ pentagonals
        for k ∈ pentagonals
            if is_pentagonal(i - k) && is_pentagonal(abs(i - 2k))
                return abs(i - 2k)
            end
        end
        j += 3
        append!(pentagonals, i + j)
    end
end


using MAT
cd(dirname(Base.source_path()))
include("../src/utils.jl")
include("../src/2d/cpab2dIntegration.jl")
include("../src/2d/imageWarp.jl")

file = matopen("tri_R2toR2_MINS_0_0_MAXS_512_512_nc_4_4_Xbdr_0_0_ext_0.mat")
As = read(file, "As")
pts_src = read(file, "pts_src")
tw_args = read(file, "tw_args")
nCols = tw_args["nCols"]
nRows = tw_args["nRows"]
close(file)

nC1 = 4
nC2 = 4
incX = nCols/nC1
incY = nRows/nC2
nPts = size(pts_src,1)
t=0.001
nSteps = 100
Nsteps = 10
Δt = t/Nsteps
ΔtAs = Δt*As
solutionSingleThread = Array(Float64, nPts, 2)

@time solver2d(pts_src,solutionSingleThread,t,As,nCols,nRows,incX,incY,Nsteps,Δt,nSteps,ΔtAs)

using PyPlot
PyPlot.figure("CPAB 2D Integration",figsize=(12,6))
PyPlot.subplot(1,2,1)
PyPlot.title("Original Image")
img = load("lena.jpg")
imr = reinterpret(UInt8, separate(img))
PyPlot.imshow(imr.data)

PyPlot.subplot(1,2,2)
PyPlot.title("Destination Image")
newimg1 = reinterpret(UInt8, separate(imwarp(solutionSingleThread, img, nCols, nRows)))
PyPlot.imshow(newimg1.data)


"""
    center_points_simulation(dcm_array, output, header, tmp_center, CCI_slice)

Function ...
"""
function center_points_simulation(dcm_array, output, header, tmp_center, CCI_slice)
	PixelSpacing = PhantomSegmentation.get_pixel_size(header)
	rows, cols = Int(header[(0x0028, 0x0010)]), Int(header[(0x0028, 0x0011)])
    sizes = []
    for row in eachrow(output[3])
		area = row[:area]
		append!(sizes, area)
	end

	centroids = output[4]
    largest = Dict()
    for index in 1:length(centroids)
		x = centroids[index][1]
		y = centroids[index][2]
		dist_loc = sqrt((tmp_center[2] - x)^2 + (tmp_center[1] - y)^2)
        dist_loc *= PixelSpacing[1]
        if dist_loc > 31
            largest[index] = [round(y), round(x)]
        else
            nothing
		end
	end

    max_dict = Dict()
	radius = round(2.5 / PixelSpacing[1], RoundUp)
    for key in largest
        tmp_arr = create_circular_mask(rows, cols, (key[2][2], key[2][1]), radius)
        tmp_arr = @. abs(tmp_arr * dcm_array[:,:,CCI_slice]) + abs(tmp_arr * dcm_array[:,:,CCI_slice - 1]) + abs(tmp_arr * dcm_array[:,:,CCI_slice + 1])
        tmp_arr = @. ifelse(tmp_arr == 0, missing, tmp_arr)
        max_dict[key[1]] = median(skipmissing(tmp_arr))
	end
    _, large1_key = maximum(zip(values(max_dict), keys(max_dict)))
    center1 = largest[large1_key]

	center = vec(tmp_center')
	p1 = vec(center1')
	offset = -2π/180
	offset2 = -5π/180
	p2, p3 = find_triangle_points(p1, center; offset=offset, offset2=offset2)
	cent1, cent2, cent3 = vec(p1'), vec(p2'), vec(p3')

    center = find_circle(cent1, cent2, cent3)
	return center, cent1, cent2, cent3
end

"""
    calc_centers_simulation(dcm_array, output, header, tmp_center, CCI_slice)

Function ...
"""
function calc_centers_simulation(dcm_array, output, header, tmp_center, CCI_slice)
    PixelSpacing = PhantomSegmentation.get_pixel_size(header)
    center, center1, center2, center3 = center_points(
        dcm_array, output, header, tmp_center, CCI_slice
    )
    centers = Dict()
    for center_index in (center1, center2, center3)
        side_x = abs(center[1] - center_index[1])
        side_y = abs(center[2] - center_index[2])
        angle = angle_calc(side_x, side_y)
        if (center_index[1] < center[1] && center_index[2] < center[2])
            medium_calc = [
                center_index[1] + (10.5 / PixelSpacing[1]) * sin(angle),
                (center_index[2] + (10.5 / PixelSpacing[2]) * cos(angle)),
            ]
            low_calc = [
                center_index[1] + (17 / PixelSpacing[1]) * sin(angle),
                (center_index[2] + (17 / PixelSpacing[2]) * cos(angle)),
            ]

        elseif (center_index[1] < center[1] && center_index[2] > center[2])
            medium_calc = [
                center_index[1] + (10.5 / PixelSpacing[1]) * sin(angle),
                (center_index[2] - (10.5 / PixelSpacing[2]) * cos(angle)),
            ]
            low_calc = [
                center_index[1] + (17 / PixelSpacing[1]) * sin(angle),
                (center_index[2] - (17 / PixelSpacing[2]) * cos(angle)),
            ]

        elseif (center_index[1] > center[1] && center_index[2] < center[2])
            medium_calc = [
                center_index[1] - (10.5 / PixelSpacing[1]) * sin(angle),
                (center_index[2] + (10.5 / PixelSpacing[2]) * cos(angle)),
            ]
            low_calc = [
                center_index[1] - (17 / PixelSpacing[1]) * sin(angle),
                (center_index[2] + (17 / PixelSpacing[2]) * cos(angle)),
            ]

        elseif (center_index[1] > center[1] && center_index[2] > center[2])
            medium_calc = [
                center_index[1] - (10.5 / PixelSpacing[1]) * sin(angle),
                (center_index[2] - (10.5 / PixelSpacing[2]) * cos(angle)),
            ]
            low_calc = [
                center_index[1] - (17 / PixelSpacing[1]) * sin(angle),
                (center_index[2] - (17 / PixelSpacing[2]) * cos(angle)),
            ]

        elseif (side_x == 0 && center_index[2] < center[2])
            medium_calc = [center_index[1], center_index[2] + (10.5 / PixelSpacing[2])]
            low_calc = [center_index[1], center_index[2] + (17 / PixelSpacing[2])]

        elseif (side_x == 0 && center_index[2] > center[2])
            medium_calc = [center_index[1], center_index[2] - (10.5 / PixelSpacing[2])]
            low_calc = [center_index[1], center_index[2] - (17 / PixelSpacing[2])]

        elseif (center_index[1] > center[1] && side_y == 0)
            medium_calc = [center_index[1] - (10.5 / PixelSpacing[1]), center_index[2]]
            low_calc = [center_index[1] - (17 / PixelSpacing[1]), center_index[2]]

        elseif (center_index[1] > center[1] && side_y == 0)
            medium_calc = [center_index[1] + (10.5 / PixelSpacing[1]), center_index[2]]
            low_calc = [(center_index[1] + (17 / PixelSpacing[1])), center_index[1]]

        else
            error("unknown angle")
        end

        if center_index == center1
            centers[:Large_HD] = Int.(round.(center_index))
            centers[:Medium_HD] = Int.(round.(medium_calc))
            centers[:Small_HD] = Int.(round.(low_calc))

        elseif center_index == center2
            centers[:Large_MD] = Int.(round.(center_index))
            centers[:Medium_MD] = Int.(round.(medium_calc))
            centers[:Small_MD] = Int.(round.(low_calc))

        elseif center_index == center3
            centers[:Large_LD] = Int.(round.(center_index))
            centers[:Medium_LD] = Int.(round.(medium_calc))
            centers[:Small_LD] = Int.(round.(low_calc))

        else
            nothing
        end
    end
    return centers
end

"""
    mask_inserts_simulation(
        dcm_array,
        masked_array,
        header,
        CCI_slice,
        center_insert;
        calcium_threshold=130,
        comp_connect=trues(3, 3),
    )

Function ...
"""
function mask_inserts_simulation(
    dcm_array,
    masked_array,
    header,
    CCI_slice,
    center_insert;
    calcium_threshold=130,
    comp_connect=trues(3, 3),
)
    output = calc_output(masked_array, header, CCI_slice, calcium_threshold, comp_connect)
    insert_centers = calc_centers_simulation(
        dcm_array, output, header, center_insert, CCI_slice
    )

    PixelSpacing = PhantomSegmentation.get_pixel_size(header)
    rows, cols = Int(header[(0x0028, 0x0010)]), Int(header[(0x0028, 0x0011)])

    mask_L_HD = create_circular_mask(
        cols, rows, insert_centers[:Large_HD], (round(5 / PixelSpacing[1], RoundUp) / 2) + 1
    )
    mask_L_MD = create_circular_mask(
        cols, rows, insert_centers[:Large_MD], (round(5 / PixelSpacing[1], RoundUp) / 2) + 1
    )
    mask_L_LD = create_circular_mask(
        cols, rows, insert_centers[:Large_LD], (round(5 / PixelSpacing[1], RoundUp) / 2) + 1
    )
    mask_M_HD = create_circular_mask(
        cols,
        rows,
        insert_centers[:Medium_HD],
        (round(3 / PixelSpacing[1], RoundUp) / 2) + 1,
    )
    mask_M_MD = create_circular_mask(
        cols,
        rows,
        insert_centers[:Medium_MD],
        (round(3 / PixelSpacing[1], RoundUp) / 2) + 1,
    )
    mask_M_LD = create_circular_mask(
        cols,
        rows,
        insert_centers[:Medium_LD],
        (round(3 / PixelSpacing[1], RoundUp) / 2) + 1,
    )
    mask_S_HD = create_circular_mask(
        cols, rows, insert_centers[:Small_HD], (round(1 / PixelSpacing[1], RoundUp) / 2) + 1
    )
    mask_S_MD = create_circular_mask(
        cols, rows, insert_centers[:Small_MD], (round(1 / PixelSpacing[1], RoundUp) / 2) + 1
    )
    mask_S_LD = create_circular_mask(
        cols, rows, insert_centers[:Small_LD], (round(1 / PixelSpacing[1], RoundUp) / 2) + 1
    )

    return transpose(mask_L_HD),
    transpose(mask_M_HD),
    transpose(mask_S_HD),
    transpose(mask_L_MD),
    transpose(mask_M_MD),
    transpose(mask_S_MD),
    transpose(mask_L_LD),
    transpose(mask_M_LD),
    transpose(mask_S_LD)
end


module TestEntropy
  using Test
  using EcologicalNetworks

# DIRAC NETWORK
# only a single interaction

Ndirac = UnipartiteProbabilisticNetwork([0 0 1; 0 0 0; 0 0 0] / 2)

inf_decomp_dirac = information_decomposition(Ndirac, norm=true)

@test inf_decomp_dirac[:D] ≈ 1.0
@test inf_decomp_dirac[:I] ≈ 0.0
@test inf_decomp_dirac[:V] ≈ 0.0

# IDENTITY NETWORK
# one partner per species, only mutual information is different from 0

Nid = BipartiteQuantitativeNetwork([1 0 0 0; 0 1 0 0; 0 0 1 0; 0 0 0 1])

inf_decomp_id = information_decomposition(Nid)

@test inf_decomp_id[:D] ≈ 0.0
@test inf_decomp_id[:I] ≈ 2.0log2(4.0)
@test inf_decomp_id[:V] ≈ 0.0

# UNIFORM NETWORK
# every species interacts with every other species

Nunif = BipartiteNetwork(ones(Bool, 10, 6))

inf_decomp_unif = information_decomposition(Nunif)

@test convert2effective(inf_decomp_unif[:D]) ≈ 1.0
@test convert2effective(inf_decomp_unif[:I]) ≈ 1.0
@test convert2effective(inf_decomp_unif[:V]) ≈ 10 * 6

# network with assymetric information
# information is only determined by columns

Nassym = BipartiteQuantitativeNetwork(ones(5, 1) * rand(1, 7))

@test mutual_information(Nassym) < 1e-10  # no information between the two variables
@test diff_entropy_uniform(Nassym, 1) < 1e-10
@test diff_entropy_uniform(Nassym, 2) > 0.0
@test conditional_entropy(Nassym, 2) ≈ entropy(Nassym, 1)
end


"""
    Rf, m = calib_force(POSES, F, m0::Real = 0.3; offset=true)
    Rf, m = calib_force(POSES, F, g::Vector; offset=true)

`POSES` ∈ ℜ(4,4,N) are the positions of the tool frame from the robot forward kinematics, 4x4 transformation matrices
`F` ∈ ℜ(N,3) vector of forces (accepts ℜ(Nx6) matrix with torques also)
# Usage
```
Rf*force[i,1:3] + forceoffs = POSES[1:3,1:3,i]'*[0, 0, mf*-9.82]
```
If `m0` is supplied, it's assumed that the gravity vector is [0,0,-g], or in words, that the gravity is acting along the negative z axis. A custom gravity vector `g` can also be supplied.

This function can also be used to calibrate an accelerometer.

[Bagge Carlson, F.](https://www.control.lth.se/staff/fredrik-bagge-carlson/), ["Machine Learning and System Identification for Estimation in Physical Systems"](https://lup.lub.lu.se/search/publication/ffb8dc85-ce12-4f75-8f2b-0881e492f6c0) (PhD Thesis 2018).

See also `calib_force_iterative` and `calib_force_eigen`.
"""
function calib_force(POSES,F,m0::Real=0.3; kwargs...)
    g0 = [0,0,-m0*9.82]
    calib_force(POSES,F,g0; kwargs...)
end
function calib_force(POSES,F,g::AbstractVector; offset=true, verbose=true)

    N = size(POSES,3)

    local forceoffs
    I = Robotlib.I3
    A  = Array{eltype(F)}(undef, 3N, offset ? 12 : 9)
    B  = Array{eltype(F)}(undef, 3N)
    A2 = Array{eltype(F)}(undef, 3N, offset ? 4 : 1)
    B2 = Array{eltype(F)}(undef, 3N)
    At = offset ? (F,i) -> [F[i,1]*I   F[i,2]*I    F[i,3]*I   -I] : (F,i) -> [F[i,1]*I   F[i,2]*I    F[i,3]*I]

    for i = 1:N
        RA               = POSES[1:3,1:3,i]'
        b                = RA*g
        A[3(i-1)+1:3i,:] = At(F,i)
        B[3(i-1)+1:3i]   = b
    end
    verbose && @info("Condition number of Gram matrix: ", cond(A'A))
    w  = tls(A,B)
    Rf = reshape(w[1:9],3,3)
    verbose && @info("Determinant of Rf before orthonormalization: ", det(Rf), opnorm(Rf))
    if det(Rf) < 0
        verbose && @warn("det(Rf) < 0, left handed coordinate system? Trying to solve the problem with reversed gravity (-g)")
        return calib_force(POSES,F,-g; offset=offset, verbose=verbose)
    end
    orthonormal!(Rf)
    At = offset ? RA -> [RA*g/norm(g) -I] : RA -> RA*g/norm(g)

    for i = 1:N
        RA                = POSES[1:3,1:3,i]'
        b                 = Rf*F[i,1:3]
        A2[3(i-1)+1:3i,:] = At(RA)
        B2[3(i-1)+1:3i]   = b
    end
    w         = A2\B2
    m         = w[1]/9.82
    offset && (forceoffs = w[2:4])

    m < 0 && @error("Estimated mass is negative, left handed coordinate system for sensor?")
    if offset
        return Rf,m,forceoffs
    else
        return Rf, m
    end
end

"""
    Rf, g, m = calib_force_eigen(POSES, F, g)

`POSES` ∈ ℜ(4,4,N) or ℜ(3,3,N) are the orientations of the tool frame from the robot forward kinematics, 4×4 transformation matrices or 3×3 rotation matrices.
`F` ∈ ℜ(N,3) vector of forces (accepts ℜ(N×6) matrix with torques also)
´g´ is an initial guess for the gravity vector.
# Usage
```
Rf*force[i,1:3] = POSES[1:3,1:3,i]'*g
```
See also `calib_force_eigen`.
"""
function calib_force_iterative(POSES, F, g; trace = false)
    N = size(POSES, 3)
    I = Robotlib.I3
    A = Array{eltype(F)}(undef, 3N, 3)
    B = F'[:]
    local m, Rf
    trace && (Rg = [])
    for iter = 1:6
        Rf, m = calib_force(
            POSES,
            F,
            g;
            offset = false,
            verbose = false,
        )
        for i = 1:N
            A[3(i-1)+1:3i, :] = Rf'POSES[:, :, i]'
        end
        g = A \ B
        trace && push!(Rg, (Rf, g))
    end
    trace && (return Rf, g, m, Rg)
    Rf, g, m
end

function eigenR(K)
    v = real(eigen(K).vectors[:, end])
    R = reshape(v, 3, 3)
    det(R) < 0 && (R .*= -1)
    orthonormal(R)
end

"""
    calib_force_eigen(POSES, F)

`POSES` ∈ ℜ(4,4,N) or ℜ(3,3,N) are the orientations of the tool frame from the robot forward kinematics, 4×4 transformation matrices or 3×3 rotation matrices.
`F` ∈ ℜ(N,3) vector of forces (accepts ℜ(N×6) matrix with torques also)
# Usage
```
Rf*force[i,1:3] + forceoffs = POSES[1:3,1:3,i]'*[0, 0, mf*-9.82]
```
"""
function calib_force_eigen(POSES, forces)
    N = size(POSES, 3)
    D = reshape(POSES[1:3,1:3,:], 3, 3N)'
    F = kron(forces, I(3))
    K = F \ D * (D \ F)
    R̂ = eigenR(K)
    ĝ = D \ F * vec(R̂)
    R̂, ĝ
end

import Robotlib: skew

"""
    This function uses Cayley transform to solve for R. It requires a known mass so it is recommended to use `calib_force` instead.
"""
function calib_force2(POSES,F,m,g  = -9.82)
    N  = size(POSES,3)

    mg = m*g

    I  = Robotlib.I3
    A  = Array{eltype(F)}(undef, 3N, 3)
    B  = Array{eltype(F)}(undef, 3N)

    for i = 1:N
        RA               = POSES[3,1:3,i]
        A[3(i-1)+1:3i,:] = skew(F[i,1:3][:] + RA*mg)
        B[3(i-1)+1:3i]   = RA*mg - F[i,1:3]
    end
    @info("Condition number of Gram matrix: ", cond(A'A))
    w  = tls(A,B)
    S  = skew(w)
    Rf = (I+S)\(I-S)
end


## HC

function hc_score(pValues)
    validPvalues(pValues)
    n = length(pValues)
    f = competerank(pValues) / n ## ties = max
    var = f .* (1-f) / n
    var[var == 0] = min( var[var > 0] ) # just to make sure we have no zero variance
    hc = abs(f-pValues) / sqrt(var)
    hc
end


include("../matplotlib_settings.jl")

using GeometricExamples

include("guiding_center_4d_settings_firk.jl")

tableau_list = get_tableau_list_firk()

equ = problem.guiding_center_4d_ode()

include("guiding_center_4d.jl")

run_tableau_list(tableau_list, equ, Δt, ntime)

generate_html(tableau_list, figures, prefix, fig_suff, fig_dir, out_dir, html_file;
              title=coll_title, params=get_params(Δt, ntime, problem))


using StructJuMP

m = StructuredModel()

@variable(m, xh[ASSETS] >= 0)
@variable(m, xb[ASSETS] >= 0)
@variable(m, xs[ASSETS] >= 0)

# Stage 0
sub = StructuredModel(parent=m)
@constraint(sub, (1+Ct)*sum(val[j]*xb[j] for j = ASSETS) <= BUDGET)

for stage = 1:T
    sub = StructuredModel(parent=m) # return a ConstraintRef to a Stage?
    for j in ASSETS
        @constraint(sub, xh[j] == (1+Ret[j])*xh[getParent(sub),j] + xb[j] - xs[j])
        @constraint(sub, (1-Ct)*sum(val[j]*xs[j] for j=ASSETS) == LIABILITY + (1+Ct)*sum(val[j]*xb[j] for j=ASSETS))
    end
end

sub = StructuredModel(parent=m)
@variable(sub, wealth)
@constraint(sub, wealth == sum(val[j]*xh[j] for j=ASSETS))
@constraint(sub, @Exp(wealth) == mu)
@objective(sub, mu - Rho*((wealth*wealth)- mu*mu))

# Exp(x) = sum(ndinNODESET:ndincurrentstage) * Prob[nd] ∗x[nd]


using Catalyst
using PWS
using Test
using StaticArrays
using Distributions

sn = @reaction_network begin
    κ, ∅ --> S
    λ, S --> ∅
end κ λ

κ = 0.25
λ = 0.005

dist = PWS.distribution(sn, [κ, λ])

traj = PWS.Trajectory([[50.0], [51.0], [50.0]], [1.0, 2.0, 3.0], [1, 2])

p_wait(s, dt) = exp(- (κ + s * λ) * dt)

wait1 = log(p_wait(50, 1.0))
wait2 = log(p_wait(51, 1.0))
wait3 = log(p_wait(50, 1.0))
reac1 = log(κ)
reac2 = log(51 * λ)

@test PWS.logpdf(dist, traj) ≈ sum((wait1, wait2, wait3, reac1, reac2))
@test PWS.trajectory_energy(dist, traj) ≈ sum((wait1, wait2, wait3, reac1, reac2))
@test PWS.trajectory_energy(dist, traj, tspan=(1.0,2.0)) ≈ sum((wait2, reac2))

cumulative = PWS.cumulative_logpdf(dist, traj, 0:0.5:3)
@test cumulative ≈ [
    0.0,
    0.5wait1,
    wait1 + reac1,
    wait1 + reac1 + 0.5wait2,
    wait1 + reac1 + wait2 + reac2,
    wait1 + reac1 + wait2 + reac2 + 0.5wait3,
    wait1 + reac1 + wait2 + reac2 + wait3,
]

subset = PWS.cumulative_logpdf(dist, traj, 2:0.33:3)
@test subset ≈ [
    0.0,
    0.33wait3,
    0.66wait3,
    0.99wait3
]

traj2 = PWS.Trajectory([[50.0]], [1.0])
@test PWS.trajectory_energy(dist, traj2, tspan=(0.0,0.5)) == -0.25


using WaveProp
using WaveProp.Geometry
using WaveProp.Interpolation

d = ReferenceTriangle()
basis = MonomialBasis(d,1)
for el in basis
    @show el
end
xi = [(0,0),(1,0),(0,1)]
vandermond(xi,basis)

C = lagrange_basis(xi,basis)


using Yao, Yao.Blocks
using NLopt

include("maxcut_gw.jl")

HB(nbit::Int) = sum([put(nbit, i=>X) for i=1:nbit])
tb = TimeEvolution(HB(3), 0.1)
function HC(W::AbstractMatrix)
    nbit = size(W, 1)
    ab = add(nbit)
    for i=1:nbit,j=i+1:nbit
        if W[i,j] != 0
            push!(ab, 0.5*W[i,j]*repeat(nbit, Z, [i,j]))
        end
    end
    ab
end

function qaoa_circuit(W::AbstractMatrix, depth::Int; use_cache::Bool=false)
    nbit = size(W, 1)
    hb = HB(nbit)
    hc = HC(W)
	use_cache && (hb = hb |> cache; hc = hc |> cache)
    c = chain(nbit, [repeat(nbit, H, 1:nbit)])
    append!(c, [chain(nbit, [timeevolve(hc, 0.0, tol=1e-5), timeevolve(hb, 0.0, tol=1e-5)]) for i=1:depth])
end

function cobyla_optimize(circuit::MatrixBlock{N}, hc::MatrixBlock; niter::Int) where N
    function f(params, grad)
        reg = zero_state(N) |> dispatch!(circuit, params)
        loss = expect(hc, reg) |> real
        #println(loss)
        loss
    end
    opt = Opt(:LN_COBYLA, nparameters(circuit))
    min_objective!(opt, f)
    maxeval!(opt, niter)
    cost, params, info = optimize(opt, parameters(circuit))
    pl = zero_state(N) |> circuit |> probs
    cost, params, pl
end

using Random, Test
@testset "qaoa circuit" begin
	Random.seed!(2)
	nbit = 5
	W = [0 5 2 1 0;
		 5 0 3 2 0;
		 2 3 0 0 0;
		 1 2 0 0 4;
		 0 0 0 4 0]

	# the exact solution
	exact_cost, sets = goemansWilliamson(W)

	hc = HC(W)
	@test ishermitian(hc)
	# the actual loss is -<Hc> + sum(wij)/2
	@test -expect(hc, product_state(5, bmask(sets[1]...)))+sum(W)/4 ≈ exact_cost

	# build a QAOA circuit and start training.
	qc = qaoa_circuit(W, 20; use_cache=false)
	opt_pl = nothing
	opt_cost = Inf
	for i = 1:10
		dispatch!(qc, :random)
		cost, params, pl = cobyla_optimize(qc, hc; niter=2000)
		@show cost
		cost < opt_cost && (opt_pl = pl)
	end

	# check the correctness of result
	config = argmax(opt_pl)-1
	@test config in [bmask(set...) for set in sets]
end


using NCDatasets
using Formatting

using Statistics

Re  = 6371000.0 #  m
c_p  = 1004.0    #  J / kg K
g   = 9.81      #  m / s^2
Lw  = 2.26e6    #  J / kg

t_beg = 240-10*12+1
t_end = 240
t_rng = t_beg:t_end
years = (t_end - t_beg + 1) / 12.0
nc_file_dir = "extract_nc_files"
domain_nc_file_dir = "domain_nc_files"

casenames = [
    "lowres_STD_SOM",
    "lowres_SSM_SOM",
    "lowres_SSM_NK",
    "lowres_SSM_SOM_noQflux",
]

linestyle = ["k-", "r-", "g-", "b-"]

mreplace = (x,) -> replace(x, missing=>NaN)

Dataset("$domain_nc_file_dir/domain.lnd.fv4x5_gx3v7.091218.nc", "r") do ds
    global mask
    mask = ds["mask"][:] |> mreplace

    mask[mask.!=0] .= 1
end

Dataset("$nc_file_dir/$(casenames[1]).h1.nc", "r") do ds
    global lat, lon, lev, ilev, rec

    lat = ds["lat"][:] |> mreplace
    lon = ds["lon"][:] |> mreplace

    ilev = (ds["ilev"][:] |> mreplace) * 100.0

    rec = ds["time"][:] |> mreplace

    lev = (ilev[2:end] + ilev[1:end-1]) / 2.0
end

# mass
Δp = (ilev[2:end] - ilev[1:end-1])

function calMeanTransport(v)

    # Get the zonal mean
    v = mean(v, dims=(1,))[1, :, :, :]

    # Weight with mass
    for j=1:length(lat), k=1:length(lev), l=1:length(rec)
        v[j, k, l] *= Δp[k] * cos(lat[j] |>  deg2rad)
    end

    v .*= 2 * π * Re / g

    # Get the vertical integration and time mean
    v = sum(v, dims=(2,))[:, 1, :]
    v = mean(v, dims=(2,))[:, 1]

    return v
end

function integrate(x, dydx)
    y = copy(dydx) * 0.0

    for i = 2:length(x)
        y[i] = y[i-1] + (dydx[i-1] + dydx[i]) * (x[i] - x[i-1]) / 2.0
    end

    return y
end

print("Import Pyplot... ")
using PyPlot
plt[:rcParams]["font.size"] = 20

#using PyCall
#@pyimport mpl_toolkits.basemap as basemap

println("Done. ")

#include("subcode_compare_models/global_surface_temperature.jl")
#include("subcode_compare_models/map_sea_surface_pressure.jl")
include("subcode_compare_models/map_SST.jl")

#=
### Streamfunction

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

fig_compare, ax_compare = plt[:subplots](1, 1, figsize=(12,8))
ax_compare[:plot]([-90, 90], [0, 0], linewidth=2, color="#cccccc")

for i = 1:length(casenames)

    casename = casenames[i]

    fig, ax = plt[:subplots](1, 1, figsize=(12,8))

    Dataset("$nc_file_dir/$casename.h0.nc", "r") do ds

        global V = mean(replace(ds["V"][:], missing=>NaN), dims=(1, 4))[1, :, :, 1]
        ψ  = invertStreamfunction(lat * π / 180.0, ilev, V; g0=g, Re=Re) / 1e10

        println(size(ψ), ";", size(lat), ";", size(ilev))
        cs = ax[:contour](lat, ilev/100, ψ'[:,:], range(-20, step=2, stop=20), colors="k")

        ax[:clabel](cs, fmt="%d")
        ax[:set_xlabel]("Lat [deg]")
        ax[:set_ylabel]("Pressure [hPa]")

        ax[:set_title](casename)
        ax[:set_xlim](-90, 90)
        ax[:set_ylim](100, 1000)
        ax[:invert_yaxis]()

        fig[:savefig]("$(casename)_psi.png", dpi=200)

        ax_compare[:plot](lat, mean(ψ[:, 19:20], dims=(2,))[:,1], linestyle[i], label="$casename")

    end

end

ax_compare[:legend]()
ax_compare[:set_xlabel]("Lat [deg]")
ax_compare[:set_ylabel]("Streamfunction [ \$ \\times 10^{10} \$ kg / s]")

ax_compare[:set_title]("Streamfunction along 510 hPa (meridional mass transport below 510 hPa)")
ax_compare[:set_xlim](-90, 90)

fig_compare[:savefig]("compare_models_psi.png", dpi=200)

=#

#=

### Meridional Static Energy Transport

fig, ax = plt[:subplots](1, 1, figsize=(12,8))

ax[:plot]([-90, 90], [0, 0], linewidth=2, color="#cccccc")

for i = 1:length(casenames)

    casename = casenames[i]

    Dataset(casename * ".nc", "r") do ds
        global MHT, MGT, MLT
        MHT = c_p * ds["VT"][:]  |> calMeanTransport
        MGT = ds["VZ"][:]       |> calMeanTransport
        MLT = Lw * ds["VQ"][:]  |> calMeanTransport

    end

    MMSET = (MHT + MGT + MLT) / 1e15

    ax[:plot](lat, MMSET, linestyle[i], linewidth=2,  label=casename)

end

ax[:legend]()
ax[:set_xlabel]("Latitude [deg]")
ax[:set_ylabel]("[PW]")

ax[:set_title]("10 years")
ax[:set_xlim](-90, 90)
ax[:set_ylim](-15, 20)

fig[:savefig]("compare_models_MSEFLX.png", dpi=200)

### Precipitation
fig, ax = plt[:subplots](1, 1, figsize=(12,8))

ax[:plot]([-90, 90], [0, 0], linewidth=2, color="#cccccc")

for i = 1:length(casenames)

    casename = casenames[i]

    Dataset(casename * ".nc", "r") do ds
        global tot_prec = replace(ds["PRECC"][:] + ds["PRECL"][:], missing=>NaN)

        tot_prec = mean(tot_prec, dims=(1,3))[1, :, 1] # average over time and longitude
    end

    ax[:plot](lat, tot_prec * 365.0 * 86400.0 * 1000.0, linestyle[i], linewidth=2,  label=casename)
end

ax[:legend]()
ax[:set_xlabel]("Latitude [deg]")
ax[:set_ylabel]("Precipitation [mm / yr]")

ax[:set_title]("10 years")
ax[:set_xlim](-90, 90)
#ax[:set_ylim](, 20)

fig[:savefig]("compare_models_PRECIP.png", dpi=200)

### Energy convergence

fig_met, ax_met = plt[:subplots](1, 1, figsize=(12,8))
ax_met[:plot]([-90, 90], [0, 0], linewidth=2, color="#cccccc")

for i = 1:length(casenames)

    fig, ax = plt[:subplots](1, 1, figsize=(12,8))
    ax[:plot]([-90, 90], [0, 0], linewidth=2, color="#cccccc")

    casename = casenames[i]

    Dataset(casename * ".h0.nc", "r") do ds
        t_rng = (:, :,  240-5*12+1:240)
        FFLX_TOA = mean( - mean(ds["FSNT"][t_rng...], dims=(3,)) + mean(ds["FLNT"][t_rng...], dims=(3,)), dims=(1,))[1, :, 1]
        FFLX_SFC = mean( - mean(ds["FSNS"][t_rng...], dims=(3,)) + mean(ds["FLNS"][t_rng...], dims=(3,)), dims=(1,))[1, :, 1]
        HFLX_SFC = mean( mean(ds["SHFLX"][t_rng...], dims=(3,)) + mean(ds["LHFLX"][t_rng...], dims=(3,)), dims=(1,))[1, :, 1]

        FFLX_TOA .*= Re * cos.(deg2rad.(lat)) * 2π
        FFLX_SFC .*= Re * cos.(deg2rad.(lat)) * 2π
        HFLX_SFC .*= Re * cos.(deg2rad.(lat)) * 2π

        EFLX_CONV = - ( FFLX_TOA - FFLX_SFC - HFLX_SFC )

        ax[:plot](lat, FFLX_TOA, "r--", linewidth=2,  label="FTOA")
        ax[:plot](lat, FFLX_SFC, "b--", linewidth=2,  label="FSFC")
        ax[:plot](lat, HFLX_SFC, "g--", linewidth=2,  label="HFLX")
        ax[:plot](lat, EFLX_CONV, "k-", linewidth=2,  label="CONV")

        MET = integrate(Re * deg2rad.(lat), EFLX_CONV)
        ax_met[:plot](lat, MET / 1e15, linestyle[i], linewidth=2, label=casename)

    end

    ax[:legend]()
    ax[:set_xlabel]("Latitude [deg]")
    ax[:set_ylabel]("[PW / m]")

    ax[:set_title](format("10 years in case {}", casename))
    ax[:set_xlim](-90, 90)

    fig[:savefig](format("compare_models_EBM_{}.png", casename), dpi=200)

end

ax_met[:legend]()
ax_met[:set_xlabel]("Latitude [deg]")
ax_met[:set_ylabel]("[PW]")

ax_met[:set_title](format("16-20 years annual implied northward heat transport"))
ax_met[:set_xlim](-90, 90)
ax_met[:set_ylim](-6, 6)

fig_met[:savefig](format("compare_models_MET.png"), dpi=200)

#ax[:set_ylim](-15, 20)

### Temperature trend

fig, ax = plt[:subplots](1, 1, figsize=(12,8))

dθ = deg2rad(lat[2] - lat[1])
for i = 1:length(casenames)

    casename = casenames[i]

    Dataset(casename * ".h1.nc", "r") do ds

        TREFHT = sum(mean(ds["TREFHT"][:], dims=(1,)), dims=(1,) )[1, :, :]
        Tavg = zeros(Float64, length(rec))
        for l = 1:length(rec)
            TREFHT[:, l] .*= cos.(deg2rad.(lat))
            Tavg[l] = sum(TREFHT[:, l]) * dθ / 2.0
        end

        Tavg .-= 273.15

        ax[:plot](collect(Float64, 0:length(rec)-1) / 365.0, Tavg, linestyle[i], linewidth=2,  label=casename)
    end

end

ax[:legend]()
ax[:set_xlabel]("Time [month]")
ax[:set_ylabel]("Temperature")

ax[:set_title]("Global Mean surface temperature of 20 years")
ax[:set_xlim](0, 20)
ax[:set_ylim](0, 20)

fig[:savefig]("compare_models_ECNVG.png", dpi=200)

=#


using DifferentialEquations
using Plots
include("is_it.jl")
barrier=300
H=.19
h=1/(2*H)
max_time=exp(10)
a=(2+h-2*sqrt(h+1))/h; # paramter a (ratio) in terms of h

pert=1e-3
Y_0=1/sqrt(h+1)
u0= [0;0;pert;Y_0]
tspan = (0.0,max_time)

ODE1(z,w)=conj(  im * w.*( 1 ./(z.^2-w.^2)+1 ./(1+w.^2) ))
ODE2(z,w)=conj(  im * z.*( 1 ./(w.^2-z.^2)+1 ./(1+z.^2) ))
Centroid_ODE(z,w)=conj( 1 ./(1+z.^2)+1 ./(1+w.^2) )
# refernce
# z=X+iP W=Q+iY . Plane Q=0, P=0
# z=u[1]+i u[2] ; w=u[3]+i u[4]

function Eq_of_M(du,u,p,t)
du[1]=real(ODE1(u[1]+im*u[2],u[3]+im*u[4]))
du[2]=imag(ODE1(u[1]+im*u[2],u[3]+im*u[4]))
du[3]=real(ODE2(u[1]+im*u[2],u[3]+im*u[4]))
du[4]=imag(ODE2(u[1]+im*u[2],u[3]+im*u[4]))
return
end
condition(u,t,integrator)=abs(u[1])>barrier || abs(u[2])>barrier || abs(u[3])>barrier||abs(u[4])>barrier
affect!(integrator) = terminate!(integrator)
cb = DiscreteCallback(condition,affect!)

prob = ODEProblem(Eq_of_M,u0,tspan)
sol=solve(prob, Vern9(),reltol=1e-13,abstol=1e-15,maxiters=1e15,callback=cb)
A=sol[:,:]
N=length(sol.t)
Z=A[1,:]+im *A[2,:];
W=A[3,:]+im *A[4,:];
Energy=zeros(N)
for k=1:N
    Energy[k]=abs(1/(1+Z[k]^2)-1/(1+W[k]^2))
end

Centeroid=zeros(N);

lin_impulse=im*(1+a )*ones(N);
z_1pos=.5*(Z+lin_impulse+W);
z_1neg=.5*(Z-lin_impulse-W);
z_2pos=.5*(-Z+lin_impulse-W);
z_2neg=.5*(-Z-lin_impulse+W);

#1pos and 1neg bound
d_11=abs.(z_1pos-z_1neg)
#1pos and 2neg bound
d_12=abs.(z_1pos-z_2neg)

z_1posX=real(z_1pos)
z_1posY=imag(z_1pos)

z_2posX=real(z_2pos)
z_2posY=imag(z_2pos)

z_1negX=real(z_1neg)
z_1negY=imag(z_1neg)

z_2negX=real(z_2neg)
z_2negY=imag(z_2neg)
gr()

p=plot()

plot!(z_1pos,linewidth=1,seriescolor=:red)
plot!(z_1neg, linewidth=1,seriescolor=:blue)
plot!(z_2pos,linewidth=1)
plot!(z_2neg,linewidth=1)
# legend('$z_1^+$', '$z_1^-$', '$z_2^+$', '$z_2^-$','Interpreter','latex')
uf=zeros(4)
uf[1]=A[1,end] #X
uf[2]=A[2,end] #P
uf[3]=A[3,end]#Q
uf[4]=A[4,end] #Y
tol=1e-1
is_it_exit(uf,H,tol)


using ODE

function F(t,x,p)
  d1 =  p[1]*x[1] - p[2]*x[1]*x[2]
  d2 = -p[3]*x[2] + p[4]*x[1]*x[2] - p[5]*x[2]*x[3]
  d3 = -p[6]*x[3] + p[7]*x[2]*x[3]
  [d1, d2, d3]
end

pp = ones(7);		# Choose all parameters as 1.0
x0 = [0.5, 1.0, 2.0];	# Setup the initial conditions
tspan = [0.0:0.1:10.0];	# and the time range

(t,x) = ODE.ode23((t,x) -> F(t,x,pp), x0, tspan);

n = length(t);
y1 = zeros(n); [y1[i] = x[i][1] for i = 1:n];
y2 = zeros(n); [y2[i] = x[i][2] for i = 1:n];
y3 = zeros(n); [y3[i] = x[i][3] for i = 1:n];

using Winston
plot(t,y1,"b.",t,y2,"g-.",t,y3,"r--")


# generalise this into coupler-specific function

abstract type AbstractCheck end

struct ConservationCheck{A} <: AbstractCheck
    ρe_tot_atmos::A
    ρe_tot_slab::A
end

function check_conservation_callback(cs, atmos_sim, slab_sim)

    atmos_sim.domain.center_space.vertical_topology.mesh.domain
    Δz_1 =
        parent(ClimaCore.Fields.coordinate_field(atmos_sim.domain.face_space).z)[2] -
        parent(ClimaCore.Fields.coordinate_field(atmos_sim.domain.face_space).z)[1]

    atmos_field = atmos_sim.integrator.u.c.ρe  # J
    slab_field = get_slab_energy(slab_sim, slab_sim.integrator.u.T_sfc) ./ Δz_1  # J  [NB: sum of the boundary field inherits the depth from the first atmospheric layer, which ≂̸ slab depth]

    ρe_tot_atmos = sum(atmos_field)
    ρe_tot_slab = sum(slab_field)

    push!(cs.ρe_tot_atmos, ρe_tot_atmos)
    push!(cs.ρe_tot_slab, ρe_tot_slab)
end

get_slab_energy(slab_sim, T_sfc) = slab_sim.params.ρ .* slab_sim.params.c .* T_sfc .* slab_sim.params.h #[NB: upon initialisation FT(275) was subtracted from T_0]

using ClimaCorePlots
function plot(CS::ConservationCheck)
    diff_ρe_tot_atmos = CS.ρe_tot_atmos .- CS.ρe_tot_atmos[3]
    diff_ρe_tot_slab = (CS.ρe_tot_slab .- CS.ρe_tot_slab[3])
    Plots.plot(diff_ρe_tot_atmos, label = "atmos")
    Plots.plot!(diff_ρe_tot_slab, label = "slab")
    tot = diff_ρe_tot_atmos .+ diff_ρe_tot_slab
    Plots.plot!(tot .- tot[1], label = "tot")
end

function conservation_plot(atmos_sim, slab_sim, solu_atm, solu_slab, figname = "tst_c.png")

    atmos_e = [sum(u.c.ρe) for u in solu_atm] # J

    z = parent(ClimaCore.Fields.coordinate_field(atmos_sim.domain.face_space).z)
    Δz_1 = z[2] - z[1]
    slab_e = [sum(get_slab_energy(slab_sim, u)) for u in solu_slab]  # J  [NB: sum of the boundary field inherits the depth from the first atmospheric layer, which ≂̸ slab depth]

    diff_ρe_tot_atmos = atmos_e .- atmos_e[3]
    diff_ρe_tot_slab = (slab_e .- slab_e[3])
    Plots.plot(diff_ρe_tot_atmos, label = "atmos")
    Plots.plot!(diff_ρe_tot_slab, label = "slab")
    tot = diff_ρe_tot_atmos .+ diff_ρe_tot_slab
    Plots.plot!(tot .- tot[1], label = "tot", xlabel = "time [s]", ylabel = "energy(t) - energy(t=0) [s]")
    Plots.savefig(figname)

end


#=
Provides a type called LQ for solving linear quadratic control
problems.

@author : Spencer Lyon <spencer.lyon@nyu.edu>
@author : Zac Cranko <zaccranko@gmail.com>

@date : 2014-07-05

References
----------

http://quant-econ.net/jl/lqcontrol.html

=#

"""
Linear quadratic optimal control of either infinite or finite horizon

The infinite horizon problem can be written

    min E sum_{t=0}^{infty} beta^t r(x_t, u_t)

with

    r(x_t, u_t) := x_t' R x_t + u_t' Q u_t + 2 u_t' N x_t

The finite horizon form is

    min E sum_{t=0}^{T-1} beta^t r(x_t, u_t) + beta^T x_T' R_f x_T

Both are minimized subject to the law of motion

    x_{t+1} = A x_t + B u_t + C w_{t+1}

Here x is n x 1, u is k x 1, w is j x 1 and the matrices are conformable for
these dimensions.  The sequence {w_t} is assumed to be white noise, with zero
mean and E w_t w_t' = I, the j x j identity.

For this model, the time t value (i.e., cost-to-go) function V_t takes the form

    x' P_T x + d_T

and the optimal policy is of the form u_T = -F_T x_T.  In the infinite horizon
case, V, P, d and F are all stationary.

##### Fields

- `Q::ScalarOrArray` : k x k payoff coefficient for control variable u. Must be
symmetric and nonnegative definite
- `R::ScalarOrArray` : n x n payoff coefficient matrix for state variable x.
Must be symmetric and nonnegative definite
- `A::ScalarOrArray` : n x n coefficient on state in state transition
- `B::ScalarOrArray` : n x k coefficient on control in state transition
- `C::ScalarOrArray` : n x j coefficient on random shock in state transition
- `N::ScalarOrArray` : k x n cross product in payoff equation
- `bet::Real` : Discount factor in [0, 1]
- `capT::Union{Int, Void}` : Terminal period in finite horizon problem
- `rf::ScalarOrArray` : n x n terminal payoff in finite horizon problem. Must be
symmetric and nonnegative definite
- `P::ScalarOrArray` : n x n matrix in value function representation
V(x) = x'Px + d
- `d::Real` : Constant in value function representation
- `F::ScalarOrArray` : Policy rule that specifies optimal control in each period

"""
type LQ
    Q::ScalarOrArray
    R::ScalarOrArray
    A::ScalarOrArray
    B::ScalarOrArray
    C::ScalarOrArray
    N::ScalarOrArray
    bet::Real
    capT::Union{Int, Void} # terminal period
    rf::ScalarOrArray
    P::ScalarOrArray
    d::Real
    F::ScalarOrArray # policy rule
end

"""
Main constructor for LQ type

Specifies default argumets for all fields not part of the payoff function or
transition equation.

##### Arguments

- `Q::ScalarOrArray` : k x k payoff coefficient for control variable u. Must be
symmetric and nonnegative definite
- `R::ScalarOrArray` : n x n payoff coefficient matrix for state variable x.
Must be symmetric and nonnegative definite
- `A::ScalarOrArray` : n x n coefficient on state in state transition
- `B::ScalarOrArray` : n x k coefficient on control in state transition
- `;C::ScalarOrArray{zeros(size(R}(1)))` : n x j coefficient on random shock in
state transition
- `;N::ScalarOrArray{zeros(size(B,1)}(size(A, 2)))` : k x n cross product in
payoff equation
- `;bet::Real(1.0)` : Discount factor in [0, 1]
- `capT::Union{Int, Void}(Void)` : Terminal period in finite horizon
problem
- `rf::ScalarOrArray{fill(NaN}(size(R)...))` : n x n terminal payoff in finite
horizon problem. Must be symmetric and nonnegative definite.

"""
function LQ(Q::ScalarOrArray,
            R::ScalarOrArray,
            A::ScalarOrArray,
            B::ScalarOrArray,
            C::ScalarOrArray=zeros(size(R, 1)),
            N::ScalarOrArray=zero(B'A);
            bet::ScalarOrArray=1.0,
            capT::Union{Int,Void}=nothing,
            rf::ScalarOrArray=fill(NaN, size(R)...))

    k = size(Q, 1)
    n = size(R, 1)
    F = k==n==1 ? zero(Float64) : zeros(Float64, k, n)
    P = copy(rf)
    d = 0.0

    LQ(Q, R, A, B, C, N, bet, capT, rf, P, d, F)
end

"""
Update `P` and `d` from the value function representation in finite horizon case

##### Arguments

- `lq::LQ` : instance of `LQ` type

##### Returns

- `P::ScalarOrArray` : n x n matrix in value function representation
V(x) = x'Px + d
- `d::Real` : Constant in value function representation

##### Notes

This function updates the `P` and `d` fields on the `lq` instance in addition to
returning them

"""
function update_values!(lq::LQ)
    # Simplify notation
    Q, R, A, B, N, C, P, d = lq.Q, lq.R, lq.A, lq.B, lq.N, lq.C, lq.P, lq.d

    # Some useful matrices
    s1 = Q + lq.bet * (B'P*B)
    s2 = lq.bet * (B'P*A) + N
    s3 = lq.bet * (A'P*A)

    # Compute F as (Q + B'PB)^{-1} (beta B'PA)
    lq.F = s1 \ s2

    # Shift P back in time one step
    new_P = R - s2'lq.F + s3

    # Recalling that trace(AB) = trace(BA)
    new_d = lq.bet * (d + trace(P * C * C'))

    # Set new state
    lq.P, lq.d = new_P, new_d
end

"""
Computes value and policy functions in infinite horizon model

##### Arguments

- `lq::LQ` : instance of `LQ` type

##### Returns

- `P::ScalarOrArray` : n x n matrix in value function representation
V(x) = x'Px + d
- `d::Real` : Constant in value function representation
- `F::ScalarOrArray` : Policy rule that specifies optimal control in each period

##### Notes

This function updates the `P`, `d`, and `F` fields on the `lq` instance in
addition to returning them

"""
function stationary_values!(lq::LQ)
    # simplify notation
    Q, R, A, B, N, C = lq.Q, lq.R, lq.A, lq.B, lq.N, lq.C

    # solve Riccati equation, obtain P
    A0, B0 = sqrt(lq.bet) * A, sqrt(lq.bet) * B
    P = solve_discrete_riccati(A0, B0, R, Q, N)

    # Compute F
    s1 = Q + lq.bet * (B' * P * B)
    s2 = lq.bet * (B' * P * A) + N
    F = s1 \ s2

    # Compute d
    d = lq.bet * trace(P * C * C') / (1 - lq.bet)

    # Bind states
    lq.P, lq.F, lq.d = P, F, d
end

"""
Non-mutating routine for solving for `P`, `d`, and `F` in infinite horizon model

See docstring for stationary_values! for more explanation
"""
function stationary_values(lq::LQ)
    _lq = LQ(copy(lq.Q),
             copy(lq.R),
             copy(lq.A),
             copy(lq.B),
             copy(lq.C),
             copy(lq.N),
             bet=copy(lq.bet),
             capT=lq.capT,
             rf=copy(lq.rf))

    stationary_values!(_lq)
    return _lq.P, _lq.F, _lq.d
end

"""
Private method implementing `compute_sequence` when state is a scalar
"""
function _compute_sequence{T}(lq::LQ, x0::T, policies)
    capT = length(policies)

    x_path = Array{T}(capT+1)
    u_path = Array{T}(capT)

    x_path[1] = x0
    u_path[1] = -(first(policies)*x0)
    w_path = lq.C * randn(capT+1)

    for t = 2:capT
        f = policies[t]
        x_path[t] = lq.A*x_path[t-1] + lq.B*u_path[t-1] + w_path[t]
        u_path[t] = -(f*x_path[t])
    end
    x_path[end] = lq.A*x_path[capT] + lq.B*u_path[capT] + w_path[end]

    x_path, u_path, w_path
end

"""
Private method implementing `compute_sequence` when state is a scalar
"""
function _compute_sequence{T}(lq::LQ, x0::Vector{T}, policies)
    # Ensure correct dimensionality
    n, j, k = size(lq.C, 1), size(lq.C, 2), size(lq.B, 2)
    capT = length(policies)

    A, B, C = lq.A, reshape(lq.B, n, k), reshape(lq.C, n, j)

    x_path = Array{T}(n, capT+1)
    u_path = Array{T}(k, capT)
    w_path = C*randn(j, capT+1)

    x_path[:, 1] = x0
    u_path[:, 1] = -(first(policies)*x0)

    for t = 2:capT
        f = policies[t]
        x_path[:, t] = A*x_path[: ,t-1] + B*u_path[:, t-1] + w_path[:, t]
        u_path[:, t] = -(f*x_path[:, t])
    end
    x_path[:, end] = A*x_path[:, capT] + B*u_path[:, capT] + w_path[:, end]

    x_path, u_path, w_path
end

"""
Compute and return the optimal state and control sequence, assuming innovation N(0,1)

##### Arguments

- `lq::LQ` : instance of `LQ` type
- `x0::ScalarOrArray`: initial state
- `ts_length::Integer(100)` : maximum number of periods for which to return
process. If `lq` instance is finite horizon type, the sequenes are returned
only for `min(ts_length, lq.capT)`

##### Returns

- `x_path::Matrix{Float64}` : An n x T+1 matrix, where the t-th column
represents `x_t`
- `u_path::Matrix{Float64}` : A k x T matrix, where the t-th column represents
`u_t`
- `w_path::Matrix{Float64}` : A n x T+1 matrix, where the t-th column represents
`lq.C*N(0,1)`

"""
function compute_sequence(lq::LQ, x0::ScalarOrArray, ts_length::Integer=100)

    # Compute and record the sequence of policies
    if isa(lq.capT, Void)
        stationary_values!(lq)
        policies = fill(lq.F, ts_length)
    else
        capT = min(ts_length, lq.capT)
        policies = Array{typeof(lq.F)}(capT)
        for t = capT:-1:1
            update_values!(lq)
            policies[t] = lq.F
        end
    end

    _compute_sequence(lq, x0, policies)
end


module CubicTests

using Base.Test
using Interpolations

for (constructor, copier) in ((interpolate, identity), (interpolate!, copy))
    f0(x) = sin((x-3)*2pi/9 - 1)
    f1(x) = 1.0 + 0.1*x + 0.01*x^2 + 0.001*x^3

    xmax = 10
    A0 = Float64[f0(x) for x in 1:xmax]
    A1 = Float64[f1(x) for x in 1:xmax]

    f2(x, y) = sin(x/10)*cos(y/6)
    xmax2, ymax2 = 30, 10
    A2 = Float64[f2(x, y) for x in 1:xmax2, y in 1:ymax2]

    for BC in (Line, Flat, Free, Periodic), GT in (OnGrid, OnCell)
        for (A, f) in ((A0, f0), (A1, f1))
            itp1 = @inferred(constructor(copier(A), BSpline(Cubic(BC())), GT()))
            @test @inferred(size(itp1)) == size(A)

            # test that inner region is close to data
            for x in 3.1:.2:8.1
                @test_approx_eq_eps f(x) itp1[x] abs(.1*f(x))
            end

            # test that we can evaluate close to, and at, boundaries
            if GT == OnGrid
                itp1[1.]
                itp1[1.0]
                itp1[1.2]
                itp1[9.8]
                itp1[10.]
                itp1[10]
            else
                itp1[0.5]
                itp1[0.6]
                itp1[10.4]
                itp1[10.5]
            end
        end

        itp2 = @inferred(constructor(copier(A2), BSpline(Cubic(BC())), GT()))
        @test @inferred(size(itp2)) == size(A2)

        for x in 3.1:.2:xmax2-3, y in 3.1:2:ymax2-3
            @test_approx_eq_eps f2(x,y) itp2[x,y] abs(.1*f2(x,y))
        end
    end
end

end

module CubicGradientTests

using Interpolations, Base.Test

ix = 1:15
f(x) = cos((x-1)*2pi/(length(ix)-1))
g(x) = -2pi/14 * sin((x-1)*2pi/(length(ix)-1))

A = f(ix)

for (constructor, copier) in ((interpolate, identity), (interpolate!, copy))

    for BC in (Line, Flat, Free, Periodic), GT in (OnGrid,OnCell)

        itp = constructor(copier(A), BSpline(Cubic(BC())), GT())
        # test that inner region is close to data
        for x in linspace(ix[5],ix[end-4],100)
            @test_approx_eq_eps g(x) gradient(itp,x)[1] cbrt(cbrt(eps(g(x))))
        end
    end
end
itp_flat_g = interpolate(A, BSpline(Cubic(Flat())), OnGrid())
@test_approx_eq_eps gradient(itp_flat_g, 1)[1] 0 eps()
@test_approx_eq_eps gradient(itp_flat_g, ix[end])[1] 0 eps()

itp_flat_c = interpolate(A, BSpline(Cubic(Flat())), OnCell())
@test_approx_eq_eps gradient(itp_flat_c, .5)[1] 0 eps()
@test_approx_eq_eps gradient(itp_flat_c, ix[end]+.5)[1] 0 eps()

end


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

using Markdown
using InteractiveUtils

# ╔═╡ 265ed188-5bf6-11eb-0d29-51ea74e6c7d1
using Pkg, DrWatson

# ╔═╡ ac6f2dfc-5bf6-11eb-3dc0-b5368b466ff7
begin
	using Distributions
	using StatsPlots
	using StatsBase
	using LaTeXStrings
	using CSV
	using DataFrames
	using LinearAlgebra
	using Random
	using ParetoSmoothedImportanceSampling
	using StanQuap
	using StatisticalRethinking
	using StatisticalRethinkingPlots
end

# ╔═╡ a4042486-5bf5-11eb-0183-33fd00d868e4
md" ## Clip-07-16-18.2s.jl"

# ╔═╡ 5f749200-6757-11eb-0345-6b094802c649
function logprob(post_df::DataFrame, x::Matrix, y::Vector, k=k)
	b = Matrix(hcat(post_df[:, [Symbol("b.$i") for i in 1:k]]))
	mu = post_df.a .+ b * x[:, 1:k]'
    logpdf.(Normal.(mu , post_df.sigma),  y')
end

# ╔═╡ ac8227b2-5bf6-11eb-27dd-09924b99e6c2
stan7_9 = "
data {
    int<lower=1> K;
    int<lower=0> N;
    matrix[N, K] x;
    vector[N] y;
}
parameters {
    real a;
    vector[K] b;
    real<lower=0> sigma;
}
transformed parameters {
    vector[N] mu;
    mu = a + x * b;
}
model {
	a ~ normal(0, 100);
	b ~ normal(0, 10);
	sigma ~ exponential(1);
    y ~ normal(mu, sigma);          // observed model
}
generated quantities {
	vector[N] log_lik;
	for (i in 1:N)
		log_lik[i] = normal_lpdf(y[i] | mu[i], sigma);
}
";

# ╔═╡ 00ee9dc0-5c23-11eb-10c3-1dfb4f486e6f
begin
	tmpdir = joinpath(projectdir(), "tmp")
	m7_9s = SampleModel("m7.9s", stan7_9; tmpdir)
end;

# ╔═╡ 5d5b8982-64b2-11eb-0ca9-b94e3197ff31
begin
	Ns = [20, 100] 				# Number of observations
	rho = [0.15, -0.4]			# Covariance between x1 and x2
	L = 100 					    # Number of simulations
	K = 5 						# Number of slopes
	dev_is = zeros(L, K)
	dev_os = zeros(L, K)
	res = Vector{NamedTuple}(undef, length(Ns))
end;

# ╔═╡ f253787c-64c3-11eb-2cd4-9322707100b7
begin
	for (ind, N) in enumerate(Ns)
		for i in 1:L
			for j = 1:K
				println("N = $(Ns[ind]), run = $i, no of b parms = $j")
				y, x_train, x_test = sim_train_test(;N, K, rho)
				data = (N = size(x_train, 1), K = size(x_train, 2),
					y = y, x = x_train,
					N_new = size(x_test, 1), x_new = x_test)
				rc7_9s = stan_sample(m7_9s; data)
				if success(rc7_9s)

					# use `logprob()`

					post7_9s_df = read_samples(m7_9s, :dataframe)
					lp_train = logprob(post7_9s_df, x_train, y, j)
					loo_lp_is, _, _ = psisloo(lp_train)
					dev_is[i, j] = -2loo_lp_is
					lp_test = logprob(post7_9s_df, x_test, y, j)
					loo_lp_is, _, _ = psisloo(lp_test)
					dev_os[i, j] = -2loo_lp_is

				end
			end
		end
		res[ind] = (
			mean_dev_is = mean(dev_is, dims=1),
			std_dev_is = std(dev_is, dims=1),
			mean_dev_os = mean(dev_os, dims=1),
			std_dev_os =std(dev_os, dims=1)
		)
	end
end;

# ╔═╡ 7eb8b028-657f-11eb-0abf-37c3819aea78
begin
	fig = Vector{Plots.Plot{Plots.GRBackend}}(undef, 2)
	for i in 1:2
		xcoord = collect(1:K) .- 0.05
		ycoord = res[i].mean_dev_is[1,:]
		ylims = i == 1 ? (0,170) : (200, 400)
		fig[i] = scatter(xcoord, ycoord, xlab = "No of b parameters",
			ylab = "Deviance", ylims = ylims, leg=:bottomleft,
			lab = "In sample", markersize=2)
		for j in 1:K
			plot!([xcoord[j], xcoord[j]],
				[ycoord[j]-res[i].std_dev_is[j], ycoord[j]+res[i].std_dev_is[j]],
				lab=false, color=:darkblue)
		end
		title!("N = $(Ns[i])")
		xcoord = collect(1:K) .+ 0.05
		ycoord = res[i].mean_dev_os[1,:]
		scatter!(xcoord, ycoord, lab = "Out of sample", markersize=2)
		for j in 1:K
			plot!([xcoord[j], xcoord[j]],
				[ycoord[j]-res[i].std_dev_os[j], ycoord[j]+res[i].std_dev_os[j]],
				lab=false, color=:red)
		end
		title!("N = $(Ns[i])")
	end
	plot(fig..., layout=(1,2))
end

# ╔═╡ 67d8698e-6faf-11eb-18f7-318ec55c1c7a
begin
	devs = Vector{Float64}(undef, 2)
	y, x_train, x_test = sim_train_test(;N=100, K=5, rho)
	data = (N = size(x_train, 1), K = size(x_train, 2),
		y = y, x = x_train,
		N_new = size(x_test, 1), x_new = x_test)
	rc7_9s = stan_sample(m7_9s; data)
	if success(rc7_9s)

		# use `logprob()`

		post7_9s_df = read_samples(m7_9s, :dataframe)
		lp_train = logprob(post7_9s_df, x_train, y, 2)
		devs[1] = -2sum(lppd(lp_train))
		lp_test = logprob(post7_9s_df, x_test, y, 2)
		devs[2] = -2sum(lppd(lp_test))
	end
end;

# ╔═╡ 08703112-6fb0-11eb-28d7-83dcab6b1998
md"
!!! note
	In some runs, out-of-sample deviance can be much larger than PSIS and WAIC estimates. In-sample deviance is always lower than PSIS and WAIC estimates. Re-run above cell several times to see this.
"

# ╔═╡ 8ffcd37a-6faf-11eb-329f-fb9ddfe5247c
Text("In-sample deviance = $(devs[1]), out-of-sample deviance = $(devs[2])")

# ╔═╡ 8ffd04ee-6faf-11eb-1e4b-f13d4d3bfe9c
if success(rc7_9s)
	waic(m7_9s)
end

# ╔═╡ 8ffdaed0-6faf-11eb-36fc-47790fcae672
if success(rc7_9s)
	loo7_9s, loos7_9s, pk7_9s = psisloo(m7_9s)
	-2loo7_9s
end

# ╔═╡ 74d4016e-6563-11eb-0c3b-b7c4f81baca6
md" ## End of clip-07-16-18.2s.jl"

# ╔═╡ Cell order:
# ╟─a4042486-5bf5-11eb-0183-33fd00d868e4
# ╠═265ed188-5bf6-11eb-0d29-51ea74e6c7d1
# ╠═ac6f2dfc-5bf6-11eb-3dc0-b5368b466ff7
# ╠═5f749200-6757-11eb-0345-6b094802c649
# ╠═ac8227b2-5bf6-11eb-27dd-09924b99e6c2
# ╠═00ee9dc0-5c23-11eb-10c3-1dfb4f486e6f
# ╠═5d5b8982-64b2-11eb-0ca9-b94e3197ff31
# ╠═f253787c-64c3-11eb-2cd4-9322707100b7
# ╠═7eb8b028-657f-11eb-0abf-37c3819aea78
# ╠═67d8698e-6faf-11eb-18f7-318ec55c1c7a
# ╟─08703112-6fb0-11eb-28d7-83dcab6b1998
# ╠═8ffcd37a-6faf-11eb-329f-fb9ddfe5247c
# ╠═8ffd04ee-6faf-11eb-1e4b-f13d4d3bfe9c
# ╠═8ffdaed0-6faf-11eb-36fc-47790fcae672
# ╟─74d4016e-6563-11eb-0c3b-b7c4f81baca6


"""The standard 2x2 block matrix form with square diagonal blocks."""
type StandardBlockMatrix{T,TA<:AbMat,TB<:AbMat,TC<:AbMat,TD<:AbMat} <: AbstractBlockMatrix{T}
    A::TA
    B::TB
    C::TC
    D::TD

    function StandardBlockMatrix(A::AbMat,B::AbMat,C::AbMat,D::AbMat)
      sA,sB,sC,sD= map(size,Any[A,B,C,D])
      if sA[1]==sA[2] && sD[1]==sD[2]
        sA[1]==sB[2]==sC[1] && sD[1]==sB[1]==sC[2] ? new(A,B,C,D) : error("off-diagonal blocks must have compatible dimensions")
      else
        error("block diagonals must be square")
      end
    end
end
function StandardBlockMatrix(A::AbMat,B::AbMat,C::AbMat,D::AbMat)
    T = promote_type(map(eltype,(A,B,C,D))...)
    return StandardBlockMatrix{T,typeof(A),typeof(B),typeof(C),typeof(D)}(A,B,C,D)
end
StandardBlockMatrix(blocks::Matrix) = StandardBlockMatrix(blocks...)

"""Returns the four blocks in a 2x2 matrix"""
function blocks(M::StandardBlockMatrix)
    theblocks = Array{Any,2}(2,2)
    theblocks[1,1] = M.A
    theblocks[2,1] = M.B
    theblocks[1,2] = M.C
    theblocks[2,2] = M.D
    return theblocks
end

#needs to be moved into abstract
#function Base.getindex(M::StandardBlockMatrix,i,j)
#  mA,nA = size(M.A)
#  if     i<=mA && j<=nA
#      return M.A[i,j]
#  elseif i>mA  && j<=nA
#      return M.C[i-mA,j]
#  elseif i<=mA && j>nA
#      return M.B[i,j-nA]
#  else
#      return M.D[i-mA,j-nA]
#  end
#end