#################### Geweke Diagnostic ####################

function gewekediag(x::Vector{T}; first::Real=0.1, last::Real=0.5,
                             etype=:imse, args...) where {T<:Real}
  if !(0.0 < first < 1.0)
    throw(ArgumentError("first is not in (0, 1)"))
  elseif !(0.0 < last < 1.0)
    throw(ArgumentError("last is not in (0, 1)"))
  elseif first + last > 1.0
    throw(ArgumentError("first and last proportions overlap"))
  end
  n = length(x)
  x1 = x[1:round(Int, first * n)]
  x2 = x[round(Int, n - last * n + 1):n]
  z = (mean(x1) - mean(x2)) /
      sqrt(mcse(x1, etype; args...)^2 + mcse(x2, etype; args...)^2)
  [round(z, digits = 3), round(1.0 - erf(abs(z) / sqrt(2.0)), digits = 4)]
end

function gewekediag(c::AbstractChains; first::Real=0.1, last::Real=0.5,
                    etype=:imse, args...)
  _, p, m = size(c.value)
  vals = Array{Float64}(undef, p, 2, m)
  for j in 1:p, k in 1:m
    vals[j, :, k] = gewekediag(c.value[:, j, k], first=first, last=last,
                               etype=etype; args...)
  end
  hdr = header(c) * "\nGeweke Diagnostic:\nFirst Window Fraction = $first\n" *
        "Second Window Fraction = $last\n"
  ChainSummary(vals, c.names, ["Z-score", "p-value"], hdr)
end


#Construct correlator from umbrella window time series.
function corr_i(W,mean)
    dt = 2 #2 fs timestep
    tCorr = 3000 #time over which correlations expected to die off.
    z =  W[:,2] .- mean # z time series (centered)
    nSamp = length(z)
    corr = zeros(tCorr+1,2) #initilize array that will contain correlation averges for a given t.
    norm = zeros(tCorr+1,1) #array of normalizations for different t.
    corr[1,1] = 0.0
    var0 = var(z)
    corr[1,2] = var0
    for i in 1:(nSamp) #loop over all samples (as separate time origins)
        tMax = nSamp
        if i+tCorr<nSamp
            tMax = i + tCorr
        end
        idx = 2 #shift index to account for  <δz(0)δ(0)> term.
        for j in (i+1):tMax
            t = (j - i)
            corr[idx,2] = corr[idx,2] + z[j]*z[i]
            corr[idx,1] = t*dt/1000 #conv to picoseconds.
            norm[idx] = norm[idx] + 1
            idx = idx + 1
        end
        #if i % tCorr == 0
        #    println(norm[2])
        #end
    end
    for k = 2:length(corr[:,1])
        corr[k,2] = corr[k,2]/norm[k]
    end

    #println(corr[length(corr[:,2]),1]," ",corr[length(corr[:,2]),2])

    ##################################
    #Below I make the resolution for longer times more coarse, to get rid of some of the noise
    #that might impact numerical integration.
    corr_cut = zeros(trunc(Int,1+750+2250/150),2)
    idx = 1
    for kk = 1:750
        corr_cut[kk,2] = corr[kk,2]
        corr_cut[kk,1] = corr[kk,1]
        idx = idx + 1
    end
    for jj in 751:150:(tCorr+1)
        corr_cut[idx,2] = corr[jj,2]
        corr_cut[idx,1] = corr[jj,1]
        idx = idx + 1
    end
    ###################################

    return corr, corr_cut, var0, tCorr
end

#Integrate correlation function from 0 → "∞". Trapezoid rule.
function Nint_corr(δδ)
    l = length(δδ[:,2])
    s = 0
    for i in 1:(l-1)
        s = s + (δδ[i+1,1] - δδ[i,1])*0.5*(δδ[i+1,2] + δδ[i,2])
    end
    return s
end


#=
block_diag.jl
Like SparseArrays:blockdiag but slightly different name, just to be safe,
because SparseArrays are also an AbstractArray type.

Motivated by compressed sensing dynamic MRI
where each frame has different k-space sampling
2020-06-08 Jeff Fessler
=#

export block_diag
import LinearMaps
using SparseArrays: blockdiag, sparse

same_idim(As::LinearMapAX...) = all(map(A -> A._idim == As[1]._idim, As))
same_odim(As::LinearMapAX...) = all(map(A -> A._odim == As[1]._odim, As))

"""
    B = block_diag(As::LinearMapAX... ; tryop::Bool)
Make block diagonal `LinearMapAX` object from blocks.

Return a `LinearMapAM` unless `tryop` and all blocks have same `idim` and `odim`.

Default for `tryop` is `true` if all blocks are type `LinearMapAO`.
"""
block_diag(As::LinearMapAO... ; tryop::Bool = true) = block_diag(tryop, As...)

block_diag(As::LinearMapAX... ; tryop::Bool = false) = block_diag(tryop, As...)

function block_diag(tryop::Bool, As::LinearMapAX...)
    B = LinearMaps.blockdiag(map(A -> A._lmap, As)...)
    nblock = length(As)
    prop = (nblock = nblock,)
    if (tryop && nblock > 1 && same_idim(As...) && same_odim(As...)) # operator version
        return LinearMapAA(B ; prop=prop,
            idim = (As[1]._idim..., nblock),
            odim = (As[1]._odim..., nblock),
        )
    end
    return LinearMapAA(B, prop)
end


# # Tutorial 🐧

# This is a gentle and lighthearted tutorial on how to use tools from AlgebraOfGraphics,
# using as example dataset a collection of measurements on penguins[^1]. See
# the [Palmer penguins website](https://allisonhorst.github.io/palmerpenguins/index.html)
# for more information.
#
# [^1]: Gorman KB, Williams TD, Fraser WR (2014) Ecological Sexual Dimorphism and Environmental Variability within a Community of Antarctic Penguins (Genus Pygoscelis). PLoS ONE 9(3): e90081. [DOI](https://doi.org/10.1371/journal.pone.0090081)

using PalmerPenguins, DataFrames

penguins = dropmissing(DataFrame(PalmerPenguins.load()))
first(penguins, 6)

# ## Frequency plots
#
# Let us start by getting a rough idea of how the data is distributed

using AlgebraOfGraphics, CairoMakie
set_aog_theme!()

axis = (width = 225, height = 225)
penguin_frequency = data(penguins) * frequency() * mapping(:species)
draw(penguin_frequency; axis)

# Next, let us see whether the distribution is the same across islands.

plt = penguin_frequency * mapping(color = :island)
draw(plt; axis)

# Ups! The bars are in the same spot and are hiding each other. We need to specify
# how we want to fix this. Bars can either `dodge` each other, or be `stack`ed on top
# of each other.

plt = penguin_frequency * mapping(color = :island, dodge = :island)
draw(plt; axis)

# This is our first finding. `Adelie` is the only species of penguins that can be
# found on all three islands. To be able to see both which species is more numerous
# and how different species are distributed across islands in a unique plot,
# we could have used `stack`.

plt = penguin_frequency * mapping(color = :island, stack = :island)
draw(plt; axis)

# ## Correlating two variables
#
# Now that we have understood the distribution of these three penguin species, we can
# start analyzing their features.

penguin_bill = data(penguins) * mapping(:bill_length_mm, :bill_depth_mm)
draw(penguin_bill; axis)

# We would actually prefer to visualize these measures in centimeters, and to have
# cleaner axes labels. As we want this setting to be preserved in all of our `bill`
# visualizations, let us save it in the variable `penguin_bill`, to be reused in
# subsequent plots.

penguin_bill = data(penguins) * mapping(
    :bill_length_mm => (t -> t / 10) => "bill length (cm)",
    :bill_depth_mm => (t -> t / 10) => "bill depth (cm)",
)
draw(penguin_bill; axis)

# Much better! Note the parentheses around the function `t -> t / 10`. They are
# necessary to specify that the function maps `t` to `t / 10`, and not to
# `t / 10 => "bill length (cm)"`.

# There does not seem to be a strong correlation between the two dimensions, which
# is odd. Maybe dividing the data by species will help.

plt = penguin_bill * mapping(color = :species)
draw(plt; axis)

# Ha! Within each species, penguins with a longer bill also have a deeper bill.
# We can confirm that with a linear regression

plt = penguin_bill * linear() * mapping(color = :species)
draw(plt; axis)

# This unfortunately no longer shows our data!
# We can use `+` to plot both things on top of each other:

plt = penguin_bill * linear() * mapping(color = :species) + penguin_bill * mapping(color = :species)
draw(plt; axis)

# Note that the above expression seems a bit redundant, as we wrote the same thing twice.
# We can "factor it out" as follows

plt = penguin_bill * (linear() + mapping()) * mapping(color = :species)
draw(plt; axis)

# where `mapping()` is a neutral multiplicative element.
# Of course, the above could be refactored as

layers = linear() + mapping()
plt = penguin_bill * layers * mapping(color = :species)
draw(plt; axis)

# We could actually take advantage of the spare `mapping()` and use it to pass some
# extra info to the scatter, while still using all the species members to compute
# the linear fit.

layers = linear() + mapping(marker = :sex)
plt = penguin_bill * layers * mapping(color = :species)
draw(plt; axis)

# This plot is getting a little bit crowded. We could instead analyze female and
# male penguins in separate subplots.

layers = linear() + mapping(col = :sex)
plt = penguin_bill * layers * mapping(color = :species)
draw(plt; axis)

# ## Smooth density plots
#
# An alternative approach to understanding how two variables interact is to consider
# their joint probability density distribution (pdf).

using AlgebraOfGraphics: density
plt = penguin_bill * density(npoints=50) * mapping(col = :species)
draw(plt; axis)

# The default colormap is multi-hue, but it is possible to pass single-hue colormaps as well:

plt *= visual(colormap = :grayC)
draw(plt; axis)

# A `Heatmap` (the default visualization for a 2D density) is a bit unfortunate if
# we want to mark species by color. In that case, one can use `visual` to change
# the default visualization and, optionally, fine tune some arguments.
# In this case, a `Wireframe` with thin lines looks quite nice. (Note that, for the
# time being, we must specify explicitly that we require a 3D axis.)

axis = (type = Axis3, width = 300, height = 300)
layer = density() * visual(Wireframe, linewidth=0.05)
plt = penguin_bill * layer * mapping(color = :species)
draw(plt; axis)

# Of course, a more traditional approach would be to use a `Contour` plot instead:

axis = (width = 225, height = 225)
layer = density() * visual(Contour)
plt = penguin_bill * layer * mapping(color = :species)
draw(plt; axis)

# The data and the linear fit can also be added back to the plot:

layers = density() * visual(Contour) + linear() + mapping()
plt = penguin_bill * layers * mapping(color = :species)
draw(plt; axis)

# In the case of many layers (contour, density and scatter) it is important to think
# about balance. In the above plot, the markers are quite heavy and can obscure the linear
# fit and the contour lines.
# We can lighten the markers using alpha transparency.

layers = density() * visual(Contour) + linear() + visual(alpha = 0.5)
plt = penguin_bill * layers * mapping(color = :species)
draw(plt; axis)

# ## Correlating three variables
#
# We are now mostly up to speed with `bill` size, but we have not consider how
# it relates to other penguin features, such as their weight.
# For that, a possible approach is to use a continuous color
# on a gradient to denote weight and different marker shapes to denote species.
# Here we use `group` to split the data for the linear regression without adding
# any additional style.

body_mass = :body_mass_g => (t -> t / 1000) => "body mass (kg)"
layers = linear() * mapping(group = :species) + mapping(color = body_mass, marker = :species)
plt = penguin_bill * layers
draw(plt; axis)

# Naturally, within each species, heavier penguins have bigger bills, but perhaps
# counter-intuitively the species with the shallowest bills features the heaviest penguins.
# We could also try and see the interplay of these three variables in a 3D plot.

axis = (type = Axis3, width = 300, height = 300)
plt = penguin_bill * mapping(body_mass, color = :species)
draw(plt; axis)

#

plt = penguin_bill * mapping(body_mass, color = :species, layout = :sex)
draw(plt; axis)

# Note that static 3D plot can be misleading, as they only show one projection
# of 3D data. They are mostly useful when shown interactively.


@testset "kernelpdmat" begin
    rng = MersenneTwister(123456)
    A = rand(rng, 10, 5)
    k = SqExponentialKernel()
    for obsdim in [1,2]
        @test all(Matrix(kernelpdmat(k,A,obsdim=obsdim)) .≈ Matrix(PDMat(kernelmatrix(k,A,obsdim=obsdim))))
        # @test_throws ErrorException kernelpdmat(k,ones(100,100),obsdim=obsdim)
    end
end


using Pumas.NCA

using Pumas, CSV

file = Pumas.example_nmtran_data("nca_test_data/dapa_IV")
data = CSV.read(file);

first(data, 6) # take first 6 rows

NCA.auc(data[:CObs][1:16], data[:TIME][1:16])

NCA.auc(data[:CObs][1:16], data[:TIME][1:16], method=:linuplogdown)

timeu = u"hr"
concu = u"mg/L"
amtu  = u"mg"
data.id = data.ID
data.time = data.TIME
data.conc = data.CObs
data.amt = data.AMT_IV
data.route = "iv"
pop = read_nca(data, llq=0concu, timeu=timeu, concu=concu, amtu=amtu)

NCA.auc(pop)

NCA.auc(pop[2], auctype=:last)

NCA.auc(pop, auctype=:last)

NCA.auc(pop[1], interval=(10,Inf).*timeu)

NCA.auc(pop[1], interval=[(10,Inf).*timeu, (10, 15).*timeu])

NCA.auc_extrap_percent(pop[1])

NCA.aumc_extrap_percent(pop[1])
NCA.aumc(pop[1])

NCA.lambdaz(pop[1])

NCA.lambdazr2(pop)
NCA.lambdazadjr2(pop)
NCA.lambdazintercept(pop)
NCA.lambdaztimefirst(pop)
NCA.lambdaznpoints(pop)

NCA.lambdaz(pop[1], threshold=3)

NCA.lambdaz(pop[1], idxs=[10, 15, 16])

NCA.lambdaz(pop[1], slopetimes=[1,2,3].*timeu)

NCA.tmax(pop[1])
NCA.cmax(pop[1])
NCA.cmax(pop[1], interval=(20, 24).*timeu)
NCA.cmax(pop[1], interval=[(20, 24).*timeu, (10, 15).*timeu])

NCA.tlast(pop[1])
NCA.clast(pop[1])

NCA.thalf(pop[1])

NCA.interpextrapconc(pop[1], 12timeu, method=:linear)

using Plots # load the plotting library
plot(pop)

plot(pop, loglinear=false)

plot(pop, linear=false)

report = NCAReport(pop)

reportdf = NCA.to_dataframe(report)
first(reportdf,6) # Print only the first 6 rows

multiple_doses_file = Pumas.example_nmtran_data("nca_test_data/dapa_IV_ORAL")
mdata = CSV.read(multiple_doses_file)

timeu = u"hr"
concu = u"mg/L"
amtu  = u"mg"
mdata.id = mdata.ID
mdata.time = mdata.TIME
mdata.conc = mdata.COBS
mdata.amt = mdata.AMT
mdata.route = replace(mdata.FORMULATION, "IV"=>"iv", "ORAL"=>"ev")
mdata.occasion = mdata.OCC
mpop = read_nca(mdata, timeu=timeu, concu=concu, amtu=amtu)

plot(mpop)

df = NCA.auc(mpop)
first(df,6) # Print only the first 6 rows

rep = NCAReport(mpop, ithdose=1)
reportdf = NCA.to_dataframe(rep)
first(reportdf,6) # Print only the first 6 rows

using PumasTutorials
PumasTutorials.tutorial_footer(WEAVE_ARGS[:folder],WEAVE_ARGS[:file])


# "in this version, set data and smoothness terms using arrays.
#  grid neighborhood is set up "manually"."
#
using GCoptimization

width = 10
height = 5
num_pixels = width * height
num_labels = 7
result = zeros(num_pixels)

# "first set up data costs individually"
data = Vector{Int}(num_pixels * num_labels)
for i = 0:num_pixels-1, l = 0:num_labels-1
    if i < 25
        if l == 0
            data[i*num_labels+l+1] = 0
        else
            data[i*num_labels+l+1] = 10
        end
    else
        if l == 5
            data[i*num_labels+l+1] = 0
        else
            data[i*num_labels+l+1] = 10
        end
    end
end

# "next set up the array for smooth costs"
smooth = Vector{Int}(num_labels * num_labels)
for l1 = 0:num_labels-1, l2 = 0:num_labels-1
    smooth[l1+l2*num_labels+1] = (l1-l2)*(l1-l2) <= 4  ? (l1-l2)*(l1-l2) : 4;
end

gco = GCoptimizationGeneralGraph(num_pixels, num_labels)
setDataCost(gco, data)
setSmoothCost(gco, smooth)

# "now set up a grid neighborhood system.
#  first set up horizontal neighbors"
for y = 0:height-1, x = 1:width-1
    setNeighbors(gco, x+y*width, x-1+y*width)
end

# "next set up vertical neighbors"
for y = 1:height-1, x = 0:width-1
    setNeighbors(gco, x+y*width, x+(y-1)*width)
end

println("Before optimization energy is ", compute_energy(gco))
expansion(gco, 2)    # "run expansion for 2 iterations."
println("After optimization energy is ", compute_energy(gco))

for i = 0:num_pixels-1
    result[i+1] = whatLabel(gco, i)
end

showall(result)


using Stheno, Random, Plots, Statistics
using Stheno: @model, EQ, Noise

###########################  Define and inspect our model  ###########################

rng = MersenneTwister(123456);

#=
In this example, `f` is an unknown real-valued function that we wish to infer. To achieve
this, we have access to two sensors. The first returns noisy estimates of `f`, where we have
been reliably informed by whoever designed the sensor that the mean of the noise is given by
`sin(x) - 5 + sqrt(abs(x))`, and that it's variance is low (1e-2). (how the designer
estimated this function, is why a sensor might possibly have such a weird mean error is
beyond the scope of this example) The second returns biased measurements of `f`, where the
bias is known to be 3.5. The model below specifies a model for this scenario.
=#
@model function model()

    # Define a smooth latent process that we wish to infer.
    f = GP(EQ())

    # Define the two noise processes described.
    noise1 = sqrt(1e-2) * GP(Noise()) + (x->sin.(x) .- 5.0 .+ sqrt.(abs.(x)))
    noise2 = sqrt(1e-1) * GP(3.5, Noise())

    # Define the processes that we get to observe.
    y1 = f + noise1
    y2 = f + noise2

    return f, noise1, noise2, y1, y2
end
f, noise₁, noise₂, y₁, y₂ = model();

# Generate some toy observations of `y1` and `y2`.
X₁, X₂ = sort(rand(rng, 3) * 10), sort(rand(rng, 10) * 10);
ŷ₁, ŷ₂ = rand(rng, [y₁(X₁), y₂(X₂)]);

# Compute the posterior processes.
(f′, y₁′, y₂′) = (f, y₁, y₂) | (y₁(X₁)←ŷ₁, y₂(X₂)←ŷ₂);

# Sample jointly from the posterior processes and compute posterior marginals.
Xp = range(-2.5, stop=12.5, length=500);
f′Xp, y₁′Xp, y₂′Xp = rand(rng, [f′(Xp, 1e-9), y₁′(Xp, 1e-9), y₂′(Xp, 1e-9)], 100);

ms1 = marginals(f′(Xp));
ms2 = marginals(y₁′(Xp));
ms3 = marginals(y₂′(Xp));

μf′, σf′ = mean.(ms1), std.(ms1);
μy₁′, σy₁′ = mean.(ms2), std.(ms2);
μy₂′, σy₂′ = mean.(ms3), std.(ms3);

###########################  Plot results  ###########################

plotly();

posterior_plot = plot();

# Plot posterior marginal std. dev.
plot!(posterior_plot, Xp, [μy₁′ μy₁′];
    linewidth=0.0,
    fillrange=[μy₁′ .- 3 .* σy₁′, μy₁′ .+ 3 * σy₁′],
    fillalpha=0.3,
    fillcolor=:red,
    label="");
plot!(posterior_plot, Xp, [μy₂′ μy₂′];
    linewidth=0.0,
    fillrange=[μy₂′ .- 3 .* σy₂′, μy₂′ .+ 3 * σy₂′],
    fillalpha=0.3,
    fillcolor=:green,
    label="");
plot!(posterior_plot, Xp, [μf′ μf′];
    linewidth=0.0,
    fillrange=[μf′.- 3  .* σf′ μf′ .+ 3 .* σf′],
    fillalpha=0.5,
    fillcolor=:blue,
    label="");

# Plot posterior marginal samples.
scatter!(posterior_plot, Xp, y₁′Xp,
    markercolor=:red,
    markershape=:circle,
    markerstrokewidth=0.0,
    markersize=0.5,
    markeralpha=0.3,
    label="");
scatter!(posterior_plot, Xp, y₂′Xp,
    markercolor=:green,
    markershape=:circle,
    markerstrokewidth=0.0,
    markersize=0.5,
    markeralpha=0.3,
    label="");
plot!(posterior_plot, Xp, f′Xp;
    linecolor=:blue,
    linealpha=0.2,
    label="");

# Plot posterior means
plot!(posterior_plot, Xp, μy₁′;
    linecolor=:red,
    linewidth=2.0,
    label="");
plot!(posterior_plot, Xp, μy₂′;
    linecolor=:green,
    linewidth=2.0,
    label="");
plot!(posterior_plot, Xp, μf′;
    linecolor=:blue,
    linewidth=2.0,
    label="Latent Function");

# Plot posterior of first noise process.
scatter!(posterior_plot, X₁, ŷ₁;
    markercolor=:red,
    markershape=:circle,
    markerstrokewidth=0.0,
    markersize=4,
    markeralpha=0.8,
    label="Sensor 1");
scatter!(posterior_plot, X₂, ŷ₂;
    markercolor=:green,
    markershape=:circle,
    markerstrokewidth=0.0,
    markersize=4,
    markeralpha=0.8,
    label="Sensor 2");

display(posterior_plot);


# # Eady turbulence example
#
# In this example, we initialize a random velocity field and observe its viscous,
# turbulent decay in a two-dimensional domain. This example demonstrates:
#
#   * How to use a tuple of turbulence closures
#   * How to use hyperdiffusivity
#   * How to implement background velocity and tracer distributions
#   * How to use `ComputedField`s for output
#
# ## Install dependencies
#
# First let's make sure we have all required packages installed.

# ```julia
# using Pkg
# pkg"add Oceananigans, JLD2, Plots"
# ```

# ## The Eady problem
#
# The "Eady problem" simulates the baroclinic instability problem proposed by Eric Eady in
# the classic paper
# ["Long waves and cyclone waves," Tellus (1949)](https://doi.org/10.3402/tellusa.v1i3.8507).
# The Eady problem is a simple, canonical model for the generation of mid-latitude
# atmospheric storms and the ocean eddies that permeate the world sea.
#
# In the Eady problem, baroclinic motion and turublence is generated by the interaction
# between turbulent motions and a stationary, geostrophically-balanced basic state that
# is unstable to baroclinic instability. In this example, the baroclinic generation of
# turbulence due to extraction of energy from the geostrophic basic state
# is balanced by a bottom boundary condition that extracts momentum from turbulent motions
# and serves as a crude model for the drag associated with an unresolved and small-scale
# turbulent bottom boundary layer.
#
# ### The geostrophic basic state
#
# The geostrophic basic state in the Eady problem is represented by the streamfunction,
#
# ```math
# ψ(y, z) = - α y (z + L_z) \, ,
# ```
#
# where ``α`` is the geostrophic shear and ``L_z`` is the depth of the domain.
# The background buoyancy includes both the geostrophic flow component,
# ``f ∂_z ψ``, where ``f`` is the Coriolis parameter, and a background stable stratification
# component, ``N^2 z``, where ``N`` is the buoyancy frequency:
#
# ```math
# B(y, z) = f ∂_z ψ + N^2 z = - α f y + N^2 z \, .
# ```
#
# The background velocity field is related to the geostrophic streamfunction via
# ``U = - ∂_y ψ`` such that
#
# ```math
# U(z) = α (z + L_z) \, .
# ```
#
# ### Boundary conditions
#
# All fields are periodic in the horizontal directions.
# We use "insulating", or zero-flux boundary conditions on the buoyancy perturbation
# at the top and bottom. We thus implicitly assume that the background vertical density
# gradient, ``N^2 z``, is maintained by a process external to our simulation.
# We use free-slip, or zero-flux boundary conditions on ``u`` and ``v`` at the surface
# where ``z=0``. At the bottom, we impose a momentum flux that extracts momentum and
# energy from the flow.
#
# #### Bottom boundary condition: quadratic bottom drag
#
# We model the effects of a turbulent bottom boundary layer on the eddy momentum budget
# with quadratic bottom drag. A quadratic cottom drag is introduced by imposing a vertical flux
# of horizontal momentum that removes momentum from the layer immediately above: in other words,
# the flux is negative (downwards) when the velocity at the bottom boundary is positive, and
# positive (upwards) with the velocity at the bottom boundary is negative.
# This drag term is "quadratic" because the rate at which momentum is removed is proportional
# to ``\boldsymbol{u} |\boldsymbol{u}|``, where ``\boldsymbol{u} = u \boldsymbol{\hat{x}} +
# v \boldsymbol{\hat{y}}`` is the horizontal velocity.
#
# The ``x``-component of the quadratic bottom drag is thus
#
# ```math
# \tau_{xz}(z=L_z) = - c^D u \sqrt{u^2 + v^2} \, ,
# ```
#
# while the ``y``-component is
#
# ```math
# \tau_{yz}(z=L_z) = - c^D v \sqrt{u^2 + v^2} \, ,
# ```
#
# where ``c^D`` is a dimensionless drag coefficient and ``\tau_{xz}(z=L_z)`` and ``\tau_{yz}(z=L_z)``
# denote the flux of ``u`` and ``v`` momentum at ``z = L_z``, the bottom of the domain.
#
# ### Vertical and horizontal viscosity and diffusivity
#
# Vertical and horizontal viscosities and diffusivities are required
# to stabilize the Eady problem and can be idealized as modeling the effect of
# turbulent mixing below the grid scale. For both tracers and velocities we use
# a Laplacian vertical diffusivity ``κ_z ∂_z^2 c`` and a horizontal
# hyperdiffusivity ``ϰ_h (∂_x^4 + ∂_y^4) c``.
#
# ### Eady problem summary and parameters
#
# To summarize, the Eady problem parameters along with the values we use in this example are
#
# | Parameter name | Description | Value | Units |
# |:--------------:|:-----------:|:-----:|:-----:|
# | ``f``          | Coriolis parameter | ``10^{-4}`` | ``\mathrm{s^{-1}}`` |
# | ``N``          | Buoyancy frequency (square root of ``\partial_z B``) | ``10^{-3}`` | ``\mathrm{s^{-1}}`` |
# | ``\alpha``     | Background vertical shear ``\partial_z U`` | ``10^{-3}`` | ``\mathrm{s^{-1}}`` |
# | ``c^D``        | Bottom quadratic drag coefficient | ``10^{-4}`` | none |
# | ``κ_z``        | Laplacian vertical diffusivity | ``10^{-2}`` | ``\mathrm{m^2 s^{-1}}`` |
# | ``ϰ_h``        | Biharmonic horizontal diffusivity | ``10^{-2} \times \Delta x^4 / \mathrm{day}`` | ``\mathrm{m^4 s^{-1}}`` |
#
# We start off by importing `Oceananigans`, `Printf`, and some convenient constants
# for specifying dimensional units:

using Printf
using Oceananigans
using Oceananigans.Units: hours, day, days

# ## The grid
#
# We use a three-dimensional grid with a depth of 4000 m and a
# horizontal extent of 1000 km, appropriate for mesoscale ocean dynamics
# with characteristic scales of 50-200 km.

grid = RectilinearGrid(size=(48, 48, 16), x=(0, 1e6), y=(0, 1e6), z=(-4e3, 0))

# ## Rotation
#
# The classical Eady problem is posed on an ``f``-plane. We use a Coriolis parameter
# typical to mid-latitudes on Earth,

coriolis = FPlane(f=1e-4) # [s⁻¹]

# ## The background flow
#
# We build a `NamedTuple` of parameters that describe the background flow,

basic_state_parameters = ( α = 10 * coriolis.f, # s⁻¹, geostrophic shear
                           f = coriolis.f,      # s⁻¹, Coriolis parameter
                           N = 1e-3,            # s⁻¹, buoyancy frequency
                          Lz = grid.Lz)         # m, ocean depth

# and then construct the background fields ``U`` and ``B``

## Background fields are defined via functions of x, y, z, t, and optional parameters
U(x, y, z, t, p) = + p.α * (z + p.Lz)
B(x, y, z, t, p) = - p.α * p.f * y + p.N^2 * z

U_field = BackgroundField(U, parameters=basic_state_parameters)
B_field = BackgroundField(B, parameters=basic_state_parameters)

# ## Boundary conditions
#
# The boundary conditions prescribe a quadratic drag at the bottom as a flux
# condition.

cᴰ = 1e-4 # quadratic drag coefficient

@inline drag_u(x, y, t, u, v, cᴰ) = - cᴰ * u * sqrt(u^2 + v^2)
@inline drag_v(x, y, t, u, v, cᴰ) = - cᴰ * v * sqrt(u^2 + v^2)

drag_bc_u = FluxBoundaryCondition(drag_u, field_dependencies=(:u, :v), parameters=cᴰ)
drag_bc_v = FluxBoundaryCondition(drag_v, field_dependencies=(:u, :v), parameters=cᴰ)

u_bcs = FieldBoundaryConditions(bottom = drag_bc_u)
v_bcs = FieldBoundaryConditions(bottom = drag_bc_v)

# ## Turbulence closures
#
# We use a horizontal hyperdiffusivity and a Laplacian vertical diffusivity
# to dissipate energy in the Eady problem.
# To use both of these closures at the same time, we set the keyword argument
# `closure` to a tuple of two closures.

κ₂z = 1e-2 # [m² s⁻¹] Laplacian vertical viscosity and diffusivity
κ₄h = 1e-1 / day * grid.Δxᶜᵃᵃ^4 # [m⁴ s⁻¹] horizontal hyperviscosity and hyperdiffusivity

Laplacian_vertical_diffusivity = AnisotropicDiffusivity(νh=0, κh=0, νz=κ₂z, κz=κ₂z)
biharmonic_horizontal_diffusivity = AnisotropicBiharmonicDiffusivity(νh=κ₄h, κh=κ₄h)

# ## Model instantiation
#
# We instantiate the model with the fifth-order WENO advection scheme, a 3rd order
# Runge-Kutta time-stepping scheme, and a `BuoyancyTracer`.

model = NonhydrostaticModel(
                   grid = grid,
              advection = WENO5(),
            timestepper = :RungeKutta3,
               coriolis = coriolis,
                tracers = :b,
               buoyancy = BuoyancyTracer(),
      background_fields = (b=B_field, u=U_field),
                closure = (Laplacian_vertical_diffusivity, biharmonic_horizontal_diffusivity),
    boundary_conditions = (u=u_bcs, v=v_bcs)
)

# ## Initial conditions
#
# We seed our initial conditions with random noise stimulate the growth of
# baroclinic instability.

## A noise function, damped at the top and bottom
Ξ(z) = randn() * z/grid.Lz * (z/grid.Lz + 1)

## Scales for the initial velocity and buoyancy
Ũ = 1e-1 * basic_state_parameters.α * grid.Lz
B̃ = 1e-2 * basic_state_parameters.α * coriolis.f

uᵢ(x, y, z) = Ũ * Ξ(z)
vᵢ(x, y, z) = Ũ * Ξ(z)
bᵢ(x, y, z) = B̃ * Ξ(z)

set!(model, u=uᵢ, v=vᵢ, b=bᵢ)

# We subtract off any residual mean velocity to avoid exciting domain-scale
# inertial oscillations. We use a `sum` over the entire `parent` arrays or data
# to ensure this operation is efficient on the GPU (set `architecture = GPU()`
# in `NonhydrostaticModel` constructor to run this problem on the GPU if one
# is available).

ū = sum(model.velocities.u.data.parent) / (grid.Nx * grid.Ny * grid.Nz)
v̄ = sum(model.velocities.v.data.parent) / (grid.Nx * grid.Ny * grid.Nz)

model.velocities.u.data.parent .-= ū
model.velocities.v.data.parent .-= v̄
nothing # hide

# ## Simulation set-up
#
# We set up a simulation that runs for 10 days with a `JLD2OutputWriter` that saves the
# vertical vorticity and divergence every 2 hours. We limit the time-step to
# the maximum allowable due either to diffusion, internal waves, or advection by the background flow.

## Calculate absolute limit on time-step using diffusivities and
## background velocity.
Ū = basic_state_parameters.α * grid.Lz

max_Δt = min(grid.Δxᶜᵃᵃ / Ū, grid.Δxᶜᵃᵃ^4 / κ₄h, grid.Δxᶜᵃᵃ^2 / κ₂z, 1 / basic_state_parameters.N)

simulation = Simulation(model, Δt = max_Δt, stop_time = 8days)

# ### The `TimeStepWizard`
#
# The TimeStepWizard manages the time-step adaptively, keeping the
# Courant-Freidrichs-Lewy (CFL) number close to `1.0` while ensuring
# the time-step does not increase beyond the maximum allowable value
# for numerical stability given the specified background flow, internal wave
# time scales, and diffusion time scales.

wizard = TimeStepWizard(cfl=0.85, max_change=1.1, max_Δt=max_Δt)

simulation.callbacks[:wizard] = Callback(wizard, IterationInterval(10))

# ### A progress messenger
#
# We add a callback that prints out a helpful progress message while the simulation runs.

start_time = time_ns()

progress(sim) = @printf("i: % 6d, sim time: % 10s, wall time: % 10s, Δt: % 10s, CFL: %.2e\n",
                        sim.model.clock.iteration,
                        prettytime(sim.model.clock.time),
                        prettytime(1e-9 * (time_ns() - start_time)),
                        prettytime(sim.Δt),
                        AdvectiveCFL(sim.Δt)(sim.model))

simulation.callbacks[:progress] = Callback(progress, IterationInterval(10))

# ### Output
#
# To visualize the baroclinic turbulence ensuing in the Eady problem,
# we use `ComputedField`s to diagnose and output vertical vorticity and divergence.
# Note that `ComputedField`s take "AbstractOperations" on `Field`s as input:

u, v, w = model.velocities # unpack velocity `Field`s

## Vertical vorticity [s⁻¹]
ζ = ComputedField(∂x(v) - ∂y(u))

## Horizontal divergence, or ∂x(u) + ∂y(v) [s⁻¹]
δ = ComputedField(-∂z(w))

# With the vertical vorticity, `ζ`, and the horizontal divergence, `δ` in hand,
# we create a `JLD2OutputWriter` that saves `ζ` and `δ` and add them to
# `simulation`.

simulation.output_writers[:fields] = JLD2OutputWriter(model, (ζ=ζ, δ=δ),
                                                      schedule = TimeInterval(4hours),
                                                        prefix = "eady_turbulence",
                                                         force = true)
nothing # hide

# All that's left is to press the big red button:

run!(simulation)

# ## Visualizing Eady turbulence
#
# We animate the results by opening the JLD2 file, extracting data for
# the iterations we ended up saving at, and ploting slices of the saved
# fields. We prepare for animating the flow by creating coordinate arrays,
# opening the file, building a vector of the iterations that we saved
# data at, and defining a function for computing colorbar limits:

using JLD2, Plots

## Coordinate arrays
xζ, yζ, zζ = nodes(ζ)
xδ, yδ, zδ = nodes(δ)

## Open the file with our data
file = jldopen(simulation.output_writers[:fields].filepath)

## Extract a vector of iterations
iterations = parse.(Int, keys(file["timeseries/t"]))

# This utility is handy for calculating nice contour intervals:

function nice_divergent_levels(c, clim, nlevels=31)
    levels = range(-clim, stop=clim, length=nlevels)
    cmax = maximum(abs, c)
    clim < cmax && (levels = vcat([-cmax], levels, [cmax]))
    return levels
end
nothing # hide

# Now we're ready to animate.

@info "Making an animation from saved data..."

anim = @animate for (i, iter) in enumerate(iterations)

    ## Load 3D fields from file
    t = file["timeseries/t/$iter"]
    R_snapshot = file["timeseries/ζ/$iter"] ./ coriolis.f
    δ_snapshot = file["timeseries/δ/$iter"]

    surface_R = R_snapshot[:, :, grid.Nz]
    surface_δ = δ_snapshot[:, :, grid.Nz]

    slice_R = R_snapshot[:, 1, :]
    slice_δ = δ_snapshot[:, 1, :]

    Rlim = 0.5 * maximum(abs, R_snapshot) + 1e-9
    δlim = 0.5 * maximum(abs, δ_snapshot) + 1e-9

    Rlevels = nice_divergent_levels(R_snapshot, Rlim)
    δlevels = nice_divergent_levels(δ_snapshot, δlim)

    @info @sprintf("Drawing frame %d from iteration %d: max(ζ̃ / f) = %.3f \n",
                   i, iter, maximum(abs, surface_R))

    xy_kwargs = (xlims = (0, grid.Lx), ylims = (0, grid.Lx),
                 xlabel = "x (m)", ylabel = "y (m)",
                 aspectratio = 1,
                 linewidth = 0, color = :balance, legend = false)

    xz_kwargs = (xlims = (0, grid.Lx), ylims = (-grid.Lz, 0),
                 xlabel = "x (m)", ylabel = "z (m)",
                 aspectratio = grid.Lx / grid.Lz * 0.5,
                 linewidth = 0, color = :balance, legend = false)

    R_xy = contourf(xζ, yζ, surface_R'; clims=(-Rlim, Rlim), levels=Rlevels, xy_kwargs...)
    δ_xy = contourf(xδ, yδ, surface_δ'; clims=(-δlim, δlim), levels=δlevels, xy_kwargs...)
    R_xz = contourf(xζ, zζ, slice_R'; clims=(-Rlim, Rlim), levels=Rlevels, xz_kwargs...)
    δ_xz = contourf(xδ, zδ, slice_δ'; clims=(-δlim, δlim), levels=δlevels, xz_kwargs...)

    plot(R_xy, δ_xy, R_xz, δ_xz,
           size = (1000, 800),
           link = :x,
         layout = Plots.grid(2, 2, heights=[0.5, 0.5, 0.2, 0.2]),
          title = [@sprintf("ζ(t=%s) / f", prettytime(t)) @sprintf("δ(t=%s) (s⁻¹)", prettytime(t)) "" ""])

    iter == iterations[end] && close(file)
end

mp4(anim, "eady_turbulence.mp4", fps = 8) # hide


using Plots, DifferentialEquations
gr()

function f′(du, u, p, t)
    du[1] = p.input(t) - p.α*u[1]
    du[2] = -p.β * u[2]
    du[3] = p.β*(p.input(t) - u[3])
end

spike_condition(u,t,integrator) = u[1]-integrator.p.δ

function spike!(integrator)
    integrator.u[1] = 0.0
    integrator.u[2] += 1.0
    push!(integrator.p.spiketimes, integrator.t)
end

u₀ = [0.0,1.0,1.0]
tspan = (0.0,10.0)
p = (
    input=t-> 0.5*(sin(t)+1) + 0.3*(cos(0.7+1.5*t)+1)-0.3,
    α = 0.5,
    β = 0.5,
    δ = 0.2,
    spiketimes = Float64[0.0]
)

cb = ContinuousCallback(spike_condition, spike!, nothing)
sol = solve(ODEProblem(f′, u₀, tspan, p); callback=cb)

#ŝ = make_analytical_solution(p.spiketimes; α=p.α, δ=p.δ, basis=:sample_and_hold)
mean_dt = (tspan[2]-tspan[1])/length(p.spiketimes)
mean_scale = 2*p.δ*p.β/(1+exp(-p.α*mean_dt))

sample_times = LinRange(tspan..., length(p.spiketimes))
sample_val = p.input.(sample_times)
digitized_signal(t) = get(sample_val,searchsortedlast(sample_times,t),0.0)
q = (
    input=digitized_signal,
    α = 0.5,
    β = 0.5,
    δ = 0.2,
    spiketimes = Float64[0.0]
)
digital_sol = solve(ODEProblem(f′, u₀, tspan, q))

let last_spike=tspan[1], next_spike=tspan[1], spike_idx=0
    ts = LinRange(tspan...,300)[2:end]
    anim = @animate for (i,t) in enumerate(ts)
        if t >= next_spike
            last_spike = next_spike
            if spike_idx < length(p.spiketimes)
                spike_idx += 1
                next_spike = p.spiketimes[spike_idx]
            else
                next_spike = Inf
            end
        end

        sample_idx=searchsortedlast(sample_times,t)

        p_in_analog=plot(p.input, tspan[1],t, xlims=tspan, ylims=(-0.2,1.3), title="input", ylabel="analog")
        p_state_analog=plot(p.input, tspan[1],t, xlims=tspan, ylims=(-0.2,1.3), title="encoding")
        p_trans_analog=plot(p.input, tspan[1],t, xlims=tspan, ylims=(-0.2,1.3), title="transmission")
        p_out_analog=plot(t->sol(t,idxs=3), tspan[1],t, xlims=tspan, ylims=(-0.2,1.3), title="decoding")

        p_in_digital=plot(p.input, tspan[1],t, xlims=tspan, ylims=(-0.2,1.3), ylabel="digital")
        plot!(sample_times[1:sample_idx], p.input.(sample_times[1:sample_idx]), line=:stem, xlims=tspan, ylims=(-0.5,1.6))
        p_state_digital=plot(sample_times[1:sample_idx], p.input.(sample_times[1:sample_idx]), seriestype=:steppre, xlims=tspan, ylims=(-0.2,1.3))
        p_trans_digital=plot(sample_times[1:sample_idx], p.input.(sample_times[1:sample_idx]), marker=:o, line=:stem, xlims=tspan, ylims=(-0.2,1.3))
        t_samp = [sample_times[1:sample_idx]; t]
        val_samp= [sample_val[1:sample_idx];sample_val[sample_idx]]
        p_out_digital=plot(t->digital_sol(t,idxs=3), tspan[1], t, xlims=tspan, ylims=(-0.2,1.3))

        p_in_spike=plot(p.input, tspan[1],t, xlims=tspan, ylims=(-0.2,1.3), fill=(0,:gray), ylabel="spiking")
        plot!(p.input, last_spike,t, xlims=tspan, ylims=(-0.2,1.3), fill=(0,"#fcaf3eff"))
        p_state_spike=plot(t->sol(t,idxs=1), tspan[1],t, xlims=tspan, ylims=(-0.05  ,1.2*p.δ), color="#fcaf3eff")
        hline!([p.δ], color=:red, linestyle=:dash)
        p_trans_spike=vline(p.spiketimes[1:spike_idx], xlims=tspan)
        p_out_spike = plot(t->sol(t,idxs=2)*mean_scale+exp(-p.β*t)*(1-mean_scale), tspan[1],t, xlims=tspan, ylims=(-0.2,1.3))

        plot(p_in_analog, p_state_analog, p_trans_analog, p_out_analog,
             p_in_digital, p_state_digital, p_trans_digital, p_out_digital,
             p_in_spike, p_state_spike, p_trans_spike, p_out_spike, layout=grid(3,4), legend=false, grid=false, size=(1600,800), linewidth=2, fontsize=20)
    end
    gif(anim, "figures/lif_animation.gif", fps = 30)
end


@testset "sparse_constructor" begin
    A = sprand(10,10,0.1)
    s = SparseTensor(A)
    @test run(sess, s)≈A
    I = [1;2;4;2;3;5]
    J = [1;3;2;2;2;1]
    V = rand(6)
    A = sparse(I,J,V,6,5)
    s = SparseTensor(I,J,V,6,5)
    @test run(sess, s)≈A
    indices = [I J]
    s = SparseTensor(I,J,V,6,5)
    @test run(sess, s)≈A

    S = Array(s)
    @test run(sess, S)≈Array(A)

    @test size(s)==(6,5)
    @test size(s,1)==6
    @test size(s,2)==5
end

@testset "sparse_arithmetic" begin
    A1 = rand(10,5); B1 = sprand(10,5,0.3)
    A = constant(A1)
    B = SparseTensor(B1)
    @test run(sess, -B)≈-B1
    for op in [+,-]
        C = op(A, B)
        C1 = op(A1, B1)
        @test run(sess, C)≈C1
    end
    @test run(sess, B-A)≈B1-A1
end

@testset "sparse_adjoint" begin
    A = sprand(10,5,0.3)
    A1 = SparseTensor(A)
    @test run(sess, A1')≈sparse(A')
end

@testset "sparse_mul" begin
    A1 = rand(10,10); B1 = sprand(10,10,0.3)
    C1 = rand(10)
    A = constant(A1)
    B = SparseTensor(B1)
    C = constant(C1)
    @test run(sess, B*A) ≈ B1*A1
    @test run(sess, A*B) ≈ A1*B1
    @test run(sess, B*A1) ≈ B1*A1
    @test run(sess, A1*B) ≈ A1*B1
    @test run(sess, B*C) ≈ B1*C1
    @test run(sess, B*C1) ≈ B1*C1
end

@testset "sparse_vcat_hcat" begin
    B1 = sprand(10,3,0.3)
    B = SparseTensor(B1)
    @test run(sess, [B;B])≈[B1;B1]
    @test run(sess, [B B])≈[B1 B1]
end

@testset "sparse_indexing" begin
    B1 = sprand(10,10,0.3)
    B = SparseTensor(B1)
    @test run(sess, B[2:3,2:3])≈B1[2:3,2:3]
    @test run(sess, B[2:3,:])≈B1[2:3,:]
    @test run(sess, B[:,2:3])≈B1[:,2:3]
end

@testset "sparse_solve" begin
    A = sparse(I, 10,10) + sprand(10,10,0.1)
    b = rand(10)
    A1 = SparseTensor(A)
    b1 = constant(b)
    u = A1\b1
    @test run(sess, u) ≈ A\b
end

@testset "sparse_assembler" begin

    m = 20
    n = 100
    handle = SparseAssembler(100, m, 0.0)
    op = PyObject[]
    A = zeros(m, n)
    for i = 1:1
        ncol = rand(1:n, 10)
        row = rand(1:m)
        v = rand(10)
        for (k,val) in enumerate(v)
            @show k
            A[row, ncol[k]] += val
        end
        @show v
        push!(op, accumulate(handle, row, ncol, v))
    end
    op = vcat(op...)
    J = assemble(m, n, op)
    B = run(sess, J)
    @test norm(A-B)<1e-8

    handle = SparseAssembler(100, 5, 1.0)
    op1 = accumulate(handle, 1, [1;2;3], [2.0;0.5;0.5])
    J = assemble(5, 5, op1)
    B = run(sess, J)
    @test norm(B-[2.0  0.0  0.0  0.0  0.0
        0.0  0.0  0.0  0.0  0.0
        0.0  0.0  0.0  0.0  0.0
        0.0  0.0  0.0  0.0  0.0
        0.0  0.0  0.0  0.0  0.0])<1e-8

    handle = SparseAssembler(100, 5, 0.0)
    op1 = accumulate(handle, 1, [1;1], [1.0;1.0])
    op2 = accumulate(handle, 1, [1;2], [1.0;1.0])

    J = assemble(5, 5, [op1;op2])
    B = run(sess, J)
    @test norm(B-[3.0  1.0  0.0  0.0  0.0
        0.0  0.0  0.0  0.0  0.0
        0.0  0.0  0.0  0.0  0.0
        0.0  0.0  0.0  0.0  0.0
        0.0  0.0  0.0  0.0  0.0])<1e-8

end

@testset "sparse_least_square" begin
    ii = Int32[1;1;2;2;3;3]
    jj = Int32[1;2;1;2;1;2]
    vv = Float64[1;2;3;4;5;6]
    ff = Float64[1;1;1]
    A = SparseTensor(ii, jj, vv, 3, 2)
    o = A\ff
    @test norm(run(sess, o)-[-1;1])<1e-6
end

@testset "sparse mat mul" begin
    A = sprand(10,5,0.3)
    B = sprand(5,20,0.3)
    C = A*B
    CC = SparseTensor(A)*SparseTensor(B)
    C_ = run(sess, CC)
    @test C_≈C

    A = spdiagm(0=>[1.;2.;3;4;5])
    B = sprand(5,20,0.3)
    C = A*B
    CC = SparseTensor(A)*SparseTensor(B)
    C_ = run(sess, CC)
    @test C_≈C

    A = sprand(10,5,0.5)
    B = spdiagm(0=>[1.;2.;3;4;5])
    C = A*B
    CC = SparseTensor(A)*SparseTensor(B)
    C_ = run(sess, CC)
    @test C_≈C
end

@testset "spdiag" begin
    p = rand(10)
    A = spdiagm(0=>p)
    B = spdiag(constant(p))
    C = spdiag(10)
    @test run(sess, B)≈A
    @test B._diag
    @test run(sess, C)≈spdiagm(0=>ones(10))
end

@testset "spzero" begin
    q = spzero(10)
    @test run(sess, q)≈sparse(zeros(10,10))
    q = spzero(10,20)
    @test run(sess, q)≈sparse(zeros(10,20))
end

@testset "sparse indexing" begin
    i1 = unique(rand(1:20,3))
    j1 = unique(rand(1:30,3))
    A = sprand(20,30,0.3)
    Ad = Array(A[i1, j1])
    B = SparseTensor(A)
    Bd = Array(B[i1, j1])
    Bd_ = run(sess, Bd)
    @test Ad≈Bd_
end

@testset "sum" begin
    s = sprand(10,20,0.2)
    S = SparseTensor(s)
    @test run(sess, sum(S)) ≈ sum(s)
    @test run(sess, sum(S,dims=1)) ≈ sum(Array(s),dims=1)[:]
    @test run(sess, sum(S,dims=2)) ≈ sum(Array(s),dims=2)[:]
end

@testset "dense_to_sparse" begin
    A = sprand(10,20,0.3)
    B = Array(A)
    @test run(sess, dense_to_sparse((B))) ≈ A
    @test run(sess, dense_to_sparse(constant(B))) ≈ A
end

@testset "spdiagm" begin
    a = rand(10)
    b = rand(9)
    A = spdiag(
        10,
        0=>a,
        -1=>b
    )
    B = diagm(
        0=>a,
        -1=>b
    )
    @test Array(run(sess, A)) ≈ B

    b = rand(7)
    A = spdiag(
        10,
        0=>a,
        -3=>b
    )
    B = diagm(
        0=>a,
        -3=>b
    )
    @test Array(run(sess, A)) ≈ B

    b = rand(7)
    A = spdiag(
        10,
        0=>a,
        -3=>b,
        3=>4b
    )
    B = diagm(
        0=>a,
        -3=>b,
        3=>4b
    )
    @test Array(run(sess, A)) ≈ B

    b = rand(7)
    A = spdiag(
        10,
        0=>a,
        -3=>b
    )
    B = diagm(
        0=>a,
        -3=>b
    )
    @test Array(run(sess, A)) ≈ B
end

@testset "hvcat" begin
    A = sprand(10,5,0.3)
    B = sprand(10,5,0.2)
    C = sprand(5,10,0.4)
    D = [A B;C]
    D_ = [SparseTensor(A) SparseTensor(B); SparseTensor(C)]
    @test run(sess, D_)≈D
end

@testset "find" begin
    A = sprand(10,10, 0.3)
    ii = Int64[]
    jj = Int64[]
    vv = Float64[]
    for i = 1:10
        for j = 1:10
            if A[i,j]!=0
                push!(ii, i)
                push!(jj, j)
                push!(vv, A[i,j])
            end
        end
    end
    a = SparseTensor(A)
    i, j, v = find(a)
    @test run(sess, i)≈ii
    @test run(sess, j)≈jj
    @test run(sess, v)≈vv

    @test run(sess, rows(a))≈ii
    @test run(sess, cols(a))≈jj
    @test run(sess, values(a))≈vv
end

@testset "sparse scatter update add" begin
    A = sprand(10,10,0.3)
    B = sprand(3,3,0.6)
    ii = [1;4;5]
    jj = [2;4;6]
    u = scatter_update(A, ii, jj, B)
    C = copy(A)
    C[ii,jj] = B
    @test run(sess, u)≈C

    u = scatter_add(A, ii, jj, B)
    C = copy(A)
    C[ii,jj] += B
    @test run(sess, u)≈C
end

@testset "constant sparse" begin
    A = sprand(10,10,0.3)
    B = constant(A)
    @test run(sess, B)≈A
end

@testset "get index" begin
    idof = [false;true]
    M = spdiag(constant(ones(2)))
    Md = M[idof, idof]
    @test run(sess, Md) ≈ sparse(reshape([1.0],1,1))
end

@testset "sparse_factorization_and_solve" begin
    A = sprand(10,10,0.7)
    rhs1 = rand(10)
    rhs2 = rand(10)
    Afac = factorize(constant(A))
    v1 = Afac\rhs1
    v2 = Afac\rhs2

    @test norm(run(sess, v1) - A\rhs1)<1e-10
    @test norm(run(sess, v2) - A\rhs2)<1e-10
end


using Test
using TestSetExtensions
using LinearAlgebra
using Qaintessent

##==----------------------------------------------------------------------------------------------------------------------

@testset ExtendedTestSet "density matrix" begin

    pauli = [
        Matrix{Float64}(I, 2, 2),
        Qaintessent.matrix(X),
        Qaintessent.matrix(Y),
        Qaintessent.matrix(Z),
    ]
    @testset "density matrix pauli group matrix" begin
        @test Qaintessent.matrix(pauli_group_matrix("XIZY")) ≈
            kron(pauli[3], pauli[4], pauli[1], pauli[2])
    end

    @testset "density matrix constructor" begin
        N = 3

        v = randn(Float64, 4^N)
        d1 = DensityMatrix(v)
        d2 = DensityMatrix(v, N)
        @test d1.v ≈ d2.v
        @test d1.N == d2.N
    end

    @testset "density matrix constructor helper function" begin
        N = 4
        ψ = randn(ComplexF64, 2^N)
        # matrix representation of density matrix |ψ⟩⟨ψ|
        ρmat = reshape(kron(conj(ψ), ψ), 2^N, 2^N)
        # density matrix represented with respect to Pauli basis
        ρ = density_from_statevector(ψ)
        @test Qaintessent.matrix(ρ) ≈ ρmat
        # compute Pauli representation from matrix form
        ρalt = density_from_matrix(ρmat)
        @test ρalt.v ≈ ρ.v
    end

end


using StructuresKit

#Test bimoment calculation for a single span beam.
#Example II-11, AISI D100-17, 9CS2.5x059, 25 ft. span, load is 10 lbs/ft
#AISC assumes vertical location of load is at center of web, so do that here also

#Ix Iy Ixy J Cw
sectionProperties = [(10.2911, 0.69678, 0.0, 0.001022, 11.9)]

#E  ν
materialProperties = [(29500, 0.30)]

#ax ay
loadLocation = [(-1.05, 0.0)]

#kx kϕ
springStiffness = [(0.0, 0.0)]

#member information
#L dL SectionProperties MaterialProperties LoadLocation SpringStiffness
memberDefinitions = [(25*12,3.0,1,1,1,1)]

#end boundary conditions
#type=1 u''=v''=ϕ''=0 (simply supported), type=2 u'=v'=ϕ'=0  (fixed), type=3 u''=v''=ϕ''=u'''=v'''=ϕ'''=0 (free end, e.g., a cantilever)
endBoundaryConditions = [1 1]

#supports
#location where u=v=ϕ=0
supports = [0.0 25.0*12]

#load  qx   qy
uniformLoad = (0.0, 10/12/1000)  #10 lbs/ft to kips/in.

z, u, v, ϕ, beamProperties = Beam.solve(memberDefinitions, sectionProperties, materialProperties, loadLocation, springStiffness, endBoundaryConditions, supports, uniformLoad)

B = InternalForces.bimoment(z, beamProperties.dm, ϕ, beamProperties.E, beamProperties.Cw)

Bmax = maximum(B)

B_D100 = 7.44   #from Bob Glauz's ballot example, see /testfiles/InternalForcesTest5/495A_Examples_v2.pdf

Error= abs((Bmax - B_D100)/ B_D100)

#accept 1% error from numerical solution
@test Error <= 0.01


# Solution to Q4e (implement a CrossEntropyLog)
import Base

type CrossEntropyLog{T}
    # Buffer of the targets, needed during fprop and bprop.
    t::Matrix{T}

    CrossEntropyLog() = new(zeros(T, 0, 0))
end

target!{T}(l::CrossEntropyLog{T}, t::Matrix{T}) = l.t = t

function fprop!{T}(l::CrossEntropyLog{T}, x::Matrix{T})
    # No need to remember anything here.
    return -sum(l.t .* x, 2)
end

function bprop!{T}(l::CrossEntropyLog{T}, ∂y::Matrix{T})
    # This expression comes from Q4d. Notice the absence of division,
    # which means there's no potential for instability!
    return ∂y .* -l.t
end

params(::CrossEntropyLog) = ()
grads(::CrossEntropyLog) = ()

Base.show(io::IO, l::CrossEntropyLog) = print(io, "CrossEntropy(Log)")


using CVChannel
using Test
using DelimitedFiles

"""
This obtains data about the Werner-Holevo channel when tensored
with the Siddhu channel.
"""

print("\nThis script shows that there is non-multiplicativity of")
print("cvPPT for the Werner-Holevo channel tensored with the Siddhu channel.")

@testset "Werner-Holevo with Siddhu is Non-Multiplicative" begin
    println("Now we actually verify non-multiplicativity.")
    λ_range = [0:0.025:0.325;]
    s_range = [0:0.1:0.5;]
    λ_ctr, s_ctr = 1,1
    data_table = zeros(length(s_range),length(λ_range))
    for λ in λ_range
        println("---Now scanning λ=",λ,"---")
        A,B,g,a = generalWHLPConstraints(1,3,λ*ones(1))
        cvWH, v1 =wernerHolevoCVPPT(1,3,A,B,g,a)
        whChan(X) = wernerHolevoChannel(X,λ)
        wh_chan = Choi(whChan,3,3)
        for s in s_range
            println("Now scanning s=",s)
            target_val = 2*cvWH
            sidchan(X) = siddhuChannel(X,s)
            sid_chan= Choi(sidchan,3,3)

            par_cv, = pptCV(parChoi(wh_chan, sid_chan), :dual)

            non_mult = par_cv - target_val
            isapprox(non_mult,0,atol=3e-6) ?
                data_table[s_ctr,λ_ctr] = 0 :
                    data_table[s_ctr,λ_ctr] = non_mult

            s_ctr += 1
            end
        λ_ctr += 1
        s_ctr = 1
    end

    println("\n Here is the data.")
    show(stdout, "text/plain", data_table)
    println("Note that the non-multiplicativity is the same for all s.")
    println("Also note that as λ increases, the non-multiplicativity decreases.")

    header, s_label = hcat("s|λ:",λ_range'), s_range
    data_to_save = vcat(header,hcat(s_label, data_table))

    info_vec = Vector{Union{Nothing,String}}(nothing, size(data_to_save)[1])
    info_vec[1] = "INFO:"
    info_vec[2] = "Generated by werner-holevo-with-siddhu.jl"
    info_vec[4] = "λ_range ="*string(λ_range)
    info_vec[5] = "s_range="*string(s_range)
    data_to_save = hcat(data_to_save,info_vec)

    println("Please name the file you'd like to write the results to: \n")
    file_name = readline()
    file_to_open = string(file_name,".csv")
    writedlm(file_to_open, data_to_save, ',')

    #This is sufficient for checking that there is non-multiplicativity somewhere
    @test data_table[1:1:1]>[0.1]
end


function part1(dataset, wvar)
	display(cor(wvar););
	#= +++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++
	[1.3 Figure 2]
		In this part we are just interested in some graphics, then
		I will just consider a light version of each model
	+++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++ =#
	p = 2;
	h = 40;
	nc = nfields(dataset)
	tt = size(dataset[1])[1] - p
	U = fill(NaN,tt,nc);
	j = 1
	for y in dataset
		_a1,U[:,j], _a2= GShock.GSstimation(y, p, h, VarGS = 2, NF = false, nmodls = 1000);
		j+=1
	end
	umean = mean(U, dims=2)
	UB = maximum(U, dims=2) - umean
	LB = umean - minimum(U, dims=2)
	umean = umean .- mean(umean);
	gd1 = wvar[end-tt+1:end,1]
	gd1 = gd1 .- mean(gd1);
	pp = plot(1:tt, umean, ribbon = (LB, UB),c = :teal, fillalpha=0.1, w = 1.5,
	    	fg_legend = :transparent, label = "News-augmented comm. shocks",
	    	title = "Figure 2: News-augmented commodity shocks vs G20 growth" ,
	    	titlefontsize = 10,xticks = (1:12:tt, qlabel[end+1-tt:12:end]),
	    	legendfontsize=7, legend=:bottomleft);
	plot!(1:tt,gd1, c=:red, w = 1.25, alpha = 0.85, label= "G20 output");
	savefig(".//Figures//intro2.svg");

	#= +++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++
	[1.5 Correlations]
	+++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++ =#
	display(cor([umean gd1]));
	comprice = Array{Float64,2}(undef,tt+p,nc)
	[comprice[:,i] = dataset[i][:,2] for i in 1:10]
	corvec_a = vec(cor(comprice));
	corvec   = corvec_a[corvec_a .<1]
	display(sum(corvec .> 0.9))
	pp= histogram(corvec, fillalpha = 0.4, legend = :false,
		xlabel = "Correlation", lc = :black);
	savefig(".//Figures//histogram.svg");
end
#=	for i in 1:nc
	    start = 1 + nv*(i-1);
	    y = dataset[:, start:nv*i];
		model1, Uaux, _tup1 =
	    GShock.GSstimation(y, p, h, VarGS = 2, NF = false, nmodls = 1000);
	    tt = length(Uaux);
	    U[end+1-tt:end,i] = Uaux;
	end

	umean = [mean(U[i,:][.!isnan.(U[i,:])]) for i in 4:nq];
	umax  = [maximum(U[i,:][.!isnan.(U[i,:])]) for i in 4:nq];
	umin  = [minimum(U[i,:][.!isnan.(U[i,:])]) for i in 4:nq];
=#


using Permutations
using DelimitedFiles
using Plots
using DeferredAcceptance
using StatsBase

students = readdlm("examples/sysopt/students.dat", Int)
schools = readdlm("examples/sysopt/schools.dat", Int)
n, m = size(schools)
capacities = ones(Int64, m)

descr = "Hybrid market with $n students, $m schools"
println(descr)

schools_STB = singletiebreaking(schools)
schools_MTB = multipletiebreaking(schools)
schools_XTB = hybridtiebreaking(schools, 0.5)
schools_WXTB = welfaretiebreaking(students, schools, 0.5)

# Update rank risks
println("Starting STB")
cdf_STB = DA_rank_dist(students, schools_STB, capacities)
println("Starting MTB")
cdf_MTB = DA_rank_dist(students, schools_MTB, capacities)
println("Starting XTB")
cdf_XTB = DA_rank_dist(students, schools_XTB, capacities)
println("Starting WXTB")
cdf_WXTB = DA_rank_dist(students, schools_WXTB, capacities)

TTC_results = TTC_match(students, capacities)[2]
TTC_cmap = countmap(TTC_results)
cdf_TTC = cumsum([get(TTC_cmap, i, 0) for i in 1:m])

LPso_results = vec([students[c, s] for (s, c) in enumerate(readdlm("examples/sysopt/LPSystemOpt.dat", Int))])
LPso_cmap = countmap(LPso_results)
cdf_LPso = cumsum([get(LPso_cmap, i, 0) for i in 1:m])

IPso_results = vec([students[c, s] for (s, c) in enumerate(readdlm("examples/sysopt/IPStableOpt.dat", Int))])
IPso_cmap = countmap(IPso_results)
cdf_IPso = cumsum([get(IPso_cmap, i, 0) for i in 1:m])

p = plot([cdf_STB, cdf_MTB, cdf_XTB, cdf_WXTB, cdf_TTC, cdf_LPso, cdf_IPso],
     label = ["DA-STB" "DA-MTB" "DA-XTB, λ=0.5" "DA-WXTB, λ=0.5" "RSD to top trading cycles" "System welfare opt (TUM LP)" "Stable welfare opt (IP)"],
     lc = [:teal :rebeccapurple :firebrick :olivedrab :orangered :gold :black],
        ls = [:dash :dot :dashdot :dash :dot :solid :dashdot],
        legend = :bottomright,
        title = descr, titlefontsize=11,
        xlabel = "rank", ylabel= "number of students")

display(p)

# savefig(p, string("examples/sysopt/plot.pdf"))
# savefig(p, string("examples/sysopt/plot.png"))


struct WaveletParams{TF, Tw}
    s::TF
    N::Int
    J::Int
    L::Int
    K::Int
    σ::TF

    ws::Tw

    Ψs::Vector{Matrix{TF}}

    Γ_H::Vector{NamedTuple{
        (:j, :l, :θ, :k, :j_, :l_, :θ_, :k_, :τ_),
        Tuple{Int, Int, TF, Int, Int,
        Int, TF, Int, Tuple{Int, Int}}}
        }

    ix::Vector{Tuple{Int, Int}}
    ix_shift::Vector{Tuple{Int, Int, Tuple{Int, Int}}}
    ix_0::Vector{Tuple{Int, Int}}
end

function WaveletParams(s, N, J, L, K, σ, Γ_H = default_Γ_H(J, L, K, θtype=typeof(s)))
    dx = 2s / N
    xs = range(-s; length=N, step=dx)
    ws = fftfreq(N, 1 / dx)

    g = MvNormal([0, 0], σ^2 * I)
    gaus = [pdf(g, [x, y]) for x in xs, y in xs]
    Gaus = gaus |> ifftshift |> fft |> real

    float_type = typeof(s)

    # wavelet_matrices(ws, J, L; ξ0=maximum(ws) / 2, c=1) =
    ξ0 = maximum(ws) / 2
    c = 1
    wavelet_matrices = [wavelet_matrix(ws, j, θ; ξ0, L, c) for j in 0:J - 1, θ in range(0., length=L, step=2π / L)]
    Ψs = map(Ψ -> float_type.(Ψ .* Gaus), reshape(wavelet_matrices, :))

    ix, ix_shift, ix_0, perm_ix, perm_ix_shift = get_ix_subsets_from_Γ_H(Γ_H, J, L, K)

    WaveletParams(s, N, J, L, K, σ, ws, Ψs, Γ_H, ix, ix_shift, ix_0)
end

e_θ(l, L) = (cos(π / 2 + 2π * l / L), sin(π / 2 + 2π * l / L))
τ_θ_j(l, j_, L) = round.(Int, 2^j_ .* e_θ(l, L))

function default_Γ_H(J, L, K; θtype=Float64)
    Γ_H = NamedTuple{
        (:j, :l, :θ, :k, :j_, :l_, :θ_, :k_, :τ_),
        Tuple{Int, Int, θtype, Int, Int,
        Int, θtype, Int, Tuple{Int, Int}}
    }[]
    p = Iterators.product(
        0:L - 1,
        0:L - 1,
        0:J - 1,
        0:J - 1,
        [false, true],
        0:K - 1,
        0:K - 1
    )

    for (l, l_, j, j_, isshift, k, k_) in p
        l, l_, j, j_, isshift, k, k_
        if l_ > l
            continue
        end
        if (j > j_) || (j < j_ - 2)
            continue
        end
        if k > min(k_, 1)
            continue
        end

        if isshift
            τ_ = τ_θ_j(l, j_, L)
        else
            τ_ = (0, 0)
        end

        θ = 2π * l / L
        θ_ = 2π * l_ / L
        dl = abs(l - l_)

        if j == j_
            if k == 0 && k_ > 1
                continue
            end
            if k == 1
                if k_ != 1
                    continue
                elseif dl > 2
                    continue
                end
            end
        else
            if k == 0 && k_ == 4
                continue
            end
            if k == 1
                if k_ != 2^(j_ - j)
                    continue
                elseif dl > 2
                    continue
                end
            end
        end

        push!(Γ_H, (;j, l, θ, k, j_, l_, θ_, k_, τ_))
    end
    Γ_H
end

function get_ix_from_Γ_H(p, J, L)
    ix1 = (p.j + 1) + J * (p.l) + J * L * (p.k)
    ix2 = (p.j_ + 1) + J * (p.l_) + J * L * (p.k_)
    (p.τ_ == (0, 0), ix1, ix2, p.τ_)
end

function get_ix_subsets_from_Γ_H(Γ_H, J, L, K)
    ix_all = map(p -> get_ix_from_Γ_H(p, J, L), Γ_H)

    perm_ix = findall(s -> s[1] == true, ix_all)
    ix = map(s -> s[2:3], ix_all[perm_ix])

    perm_ix_shift = findall(s -> s[1] == false, ix_all)
    ix_shift = map(s -> s[2:4], ix_all[perm_ix_shift])

    ix_0 = map(i -> (1, i), 1:J * L * K)
    ix, ix_shift, ix_0, perm_ix, perm_ix_shift
end


# In this tutorial, we will learn:
#
#   * How to formulate the weak form for a scalar time-harmonic electromagnetic problem
#   * How to implement a perfectly matched layer (PML) to absorb outgoing waves
#   * How to impose periodic boundary conditions in Gridap
#   * How to discretize PDEs with complex-valued solutions
#
# ## Problem statement
#
# We are going to solve a scalar electromagnetic wave scattering problem: a plane wave (Hz-polarized $H_{inc}$) scattering of a dielectric cylinder (of radius $R$ and permittivity $\varepsilon$), as illustrated below. The computational cell is of height $H$ and length $L$, and we employ a perfectly matched layer (PML) thickness of $d_{pml}$ to implement outgoing (radiation) boundary conditions for this finite domain.
#
# ![](../assets/emscatter/Illustration.png)
#
# From Maxwell's equations, considering a time-harmonic electromagnetic field, we can derive the governing equation of this problem in 2D (Helmholtz equation) [1]:
#
# ```math
# \left[-\nabla\cdot\frac{1}{\varepsilon(x)}\nabla -k^2\mu(x)\right] H = f(x),
# ```
#
# where $k=\omega/c$ is the wave number in free space and $f(x)$ is the source term (which corresponds to a magnetic current density in Maxwell's equations).
#
# In order to simulate this scattering problem in a finite computation cell, we need outgoing (radiation) boundary conditions such that all waves at the boundary would not be reflected back since we are simulating an infinite space.
# One commonly used technique to simulate such infinite space is through the so called "perfectly matched layers" (PML) [2]. Actually, PML is not a boundary condition but an artificial absorbing "layer" that absorbs waves with minimal reflections (going to zero as the resolution increases).
# There are many formulations of PML. Here, we use one of the most flexible formulations, the "stretched-coordinate" formulation, which takes the following replace in the PDE [3]:
#
# ```math
# \frac{\partial}{\partial x}\rightarrow \frac{1}{1+\mathrm{i}\sigma(u_x)/\omega}\frac{\partial}{\partial x},
# ```
#
# ```math
# \frac{\partial}{\partial y}\rightarrow \frac{1}{1+\mathrm{i}\sigma(u_y)/\omega}\frac{\partial}{\partial y},
# ```
#
# where $u_{x/y}$ is the depth into the PML, $\sigma$ is a profile function (here we chose $\sigma(u)=\sigma_0(u/d_{pml})^2$) and different derivative corresponds to different absorption directions.
# Note that at a finite mesh resolution, PML reflects some waves, and the standard technique to mitigate this is to "turn on" the PML absorption gradually—in this case we use a quadratic profile. The amplitude $\sigma_0$ is chosen so that in the limit of infinite resolution the "round-trip" normal-incidence is some small number.
#
# Since PML absorbs all waves in $x/y$ direction, the associated boundary condition is then usually the zero Dirichlet boundary condition. Here, the boundary conditions are zero Dirichlet boundary on the top and bottom side $\Gamma_D$ but periodic boundary condition on the left ($\Gamma_L$) and right side ($\Gamma_R$).
# The reason that we use a periodic boundary condition for the left and right side instead of zero Dirichlet boundary condition is that we want to simulate a plane wave exicitation, which then requires a periodic boundary condition.
#
# Consider $\mu(x)=1$ (which is mostly the case in electromagnetic problems) and denote $\Lambda=\operatorname{diagm}(\Lambda_x,\Lambda_y)$ where $\Lambda_{x/y}=\frac{1}{1+\mathrm{i}\sigma(u_{x/y})/\omega}$, we can formulate the problem as
#
# ```math
# \left\{ \begin{aligned}
# \left[-\Lambda\nabla\cdot\frac{1}{\varepsilon(x)}\Lambda\nabla -k^2\right] H &= f(x) & \text{ in } \Omega,\\
# H&=0 & \text{ on } \Gamma_D,\\
# H|_{\Gamma_L}&=H|_{\Gamma_R},&\\
# \end{aligned}\right.
# ```
# For convenience, in the weak form and Julia implementation below we represent $\Lambda$ as a vector instead of a diagonal $2 \times 2$ matrix, in which case $\Lambda\nabla$ becomes the elementwise product.
#
# ## Numerical scheme
#
# Similar to the previous tutorials, we need to construct the weak form for the above PDEs. After integral by part and removing the zero boundary integral term, we get:
#
# ```math
# a(u,v) = \int_\Omega \left[\nabla(\Lambda v)\cdot\frac{1}{\varepsilon(x)}\Lambda\nabla u-k^2uv\right]\mathrm{d}\Omega
# ```
#
# ```math
# b(v) = \int_\Omega vf\mathrm{d}\Omega
# ```
# Notice that the $\nabla(\Lambda v)$ is also a element-wise "product" of two vectors $\nabla$ and $\Lambda v$.
#
# ## Setup
#
# We import the packages that will be used, define the geometry and physics parameters.
#

using Gridap
using GridapGmsh
using Gridap.Fields
using Gridap.Geometry

λ = 1.0          # Wavelength (arbitrary unit)
L = 4.0          # Width of the area
H = 6.0          # Height of the area
xc = [0 -1.0]    # Center of the cylinder
r = 1.0          # Radius of the cylinder
d_pml = 0.8      # Thickness of the PML
k = 2*π/λ        # Wave number
const ϵ₁ = 3.0   # Relative electric permittivity for cylinder

# ## Discrete Model
#
# We import the model from the `geometry.msh` mesh file using the `GmshDiscreteModel` function defined in `GridapGmsh`. The mesh file is created with GMSH in Julia (see the file ../assets/emscatter/MeshGenerator.jl). Note that this mesh file already contains periodic boundary information for the left and right side, and that is enough for gridap to realize a periodic boundary condition should be implemented.
#

model = GmshDiscreteModel("../models/geometry.msh")

# ## FE spaces
#
# We use the first-order lagrangian as the finite element function space basis. The dirihlet edges are labeld with `DirichletEdges` in the mesh file. Since our problem involves complex numbers (because of PML), we need to assign the `vector_type` to be `Vector{ComplexF64}`.
#

# ### Test and trial finite element function space
order = 1
reffe = ReferenceFE(lagrangian,Float64,order)
V = TestFESpace(model,reffe,dirichlet_tags="DirichletEdges",vector_type=Vector{ComplexF64})
U = V # mathematically equivalent to TrialFESpace(V,0)

# ## Numerical integration
#
# We generate the triangulation and a second-order Gaussian quadrature for the numerial integration. Note that we create a boundary triangulation from a `Source` tag for the line excitation. Generally, we do not need such additional mesh tags for the source, we can use a delta function to approximate such line source excitation. However, by generating a line mesh, we can increase the accuracy of this source excitation.
#

# ### Generate triangulation and quadrature from model
degree = 2
Ω = Triangulation(model)
dΩ = Measure(Ω,degree)
# ### Source triangulation
Γ = BoundaryTriangulation(model;tags="Source")
dΓ = Measure(Γ,degree)

# ## PML formulation
#
# Here we first define a `s_PML` function: $s(x)=1+\mathrm{i}\sigma(u)/\omega,$ and its derivative `ds_PML`. The parameter `LH` indicates the size of the inner boundary of the PML regions. Finally, we create a function-like object `Λ` that returns the PML factors and define its derivative in gridap.
# Note that here we are defining a "callable object" of type `Λ` that encapsulates all of the PML parameters. This is convenient, both because we can pass lots of parameters around easily and also because we can define additional methods on `Λ`, e.g. to express the `∇(Λv)` operation.
#

# ### PML parameters
Rpml = 1e-12      # Tolerence for PML reflection
σ = -3/4*log(Rpml)/d_pml # σ_0
LH = [L,H] # Size of the PML inner boundary (a rectangular centere at (0,0))

# ### PML coordinate streching functions
function s_PML(x,σ,k,LH,d_pml)
    u = abs.(Tuple(x)).-LH./2  # get the depth into PML
    return @. ifelse(u > 0,  1+(1im*σ/k)*(u/d_pml)^2, $(1.0+0im))
end

function ds_PML(x,σ,k,LH,d_pml)
    u = abs.(Tuple(x)).-LH./2 # get the depth into PML
    ds = @. ifelse(u > 0, (2im*σ/k)*(1/d_pml)^2*u, $(0.0+0im))
    return ds.*sign.(Tuple(x))
end

struct Λ<:Function
    σ::Float64
    k::Float64
    LH::Vector{Float64}
    d_pml::Float64
end

function (Λf::Λ)(x)
    s_x,s_y = s_PML(x,Λf.σ,Λf.k,Λf.LH,Λf.d_pml)
    return VectorValue(1/s_x,1/s_y)
end

Fields.∇(Λf::Λ) = x->TensorValue{2,2,ComplexF64}(-(Λf(x)[1])^2*ds_PML(x,Λf.σ,Λf.k,Λf.LH,Λf.d_pml)[1],0,0,-(Λf(x)[2])^2*ds_PML(x,Λf.σ,Λf.k,Λf.LH,Λf.d_pml)[2])

# ## Weak form
#
# In the mesh file, we labeled the cylinder region with `Cylinder` to distinguish it from other regions. Using this tag, we can assign material properties correspondingly (basically a function with different value in different regions). The weak form is very similar to its mathematical form in gridap.
#

# ### Intermediate variables
labels = get_face_labeling(model)
dimension = num_cell_dims(model)
tags = get_face_tag(labels,dimension)
const cylinder_tag = get_tag_from_name(labels,"Cylinder")

function ξ(tag)
    if tag == cylinder_tag
        return 1/ϵ₁
    else
        return 1.0
    end
end

τ = CellField(tags,Ω)
Λf = Λ(σ,k,LH,d_pml)

# ### Bi-linear term (from weak form)
# Note that we use a element-wise product .* here for the vector-vector product $\Lambda \nabla$
a(u,v) = ∫(  (∇.*(Λf*v))⊙((ξ∘τ)*(Λf.*∇(u))) - (k^2*(v*u))  )dΩ

# ### Source term (uniform line source)
b(v) = ∫(v)*dΓ

# ## Solver phase
#
# We can assemble the finite element operator in Gridap with the bi-linear and linear form, then solve for the field.
#

op = AffineFEOperator(a,b,U,V)
uh = solve(op)

# ## Analytical solution
# ### Theoretical analysis
# In this section, we construct the semi-analytical solution to this scattering problem, for comparison to the numerical solution. This is possible because of the symmetry of the cylinder, which allows us to expand the solutions of the Helmoltz equation in Bessel functions and match boundary conditions at the cylinder interface. (In 3d, the analogous process with spherical harmonics is known as "Mie scattering".) For more information on this technique, see Ref [4].
# In 2D cylinder coordinates, we can expand the plane wave in terms of Bessel functions (this is the Jacobi–Anger identity [5]):
#
# ```math
# H_0=\sum_m i^mJ_m(kr)e^{im\theta},
# ```
#
# where $m=0,\pm 1,\pm 2,\dots$ and $J_m(z)$ is the Bessel function of the fisrt kind.
#
# For simplicity, we start with only the $m$-th component and take it as the incident part:
#
# ```math
# H_{inc}=J_m(kr).
# ```
#
# For the scattered field, since the scattered wave should be going out, we can then expand it in terms of the Hankel function of the first kind (outgoing and incoming cylindrical waves are Hankel functions of the first and second kind [6]):
#
# ```math
# H_1=\alpha_mH_m^1(kr).
# ```
#
# For the fields inside the cylinder, we require the field to be finite at $r=0$, which then constrains the field to be only the expansion of the Bessel functions:
#
# ```math
# H_2=\beta_mJ_m(nkr),
# ```
#
# where $n=\sqrt{\varepsilon}$ is the refractive index.
#
# Applying the boundary conditions (tangential part of the electric and magnetic field to be continuous):
#
# ```math
# H_{inc}+H_1=H_2|_{r=R},
# ```
#
# ```math
# \frac{\partial H_{inc}}{\partial r}+\frac{\partial H_1}{\partial r}=\frac{1}{\epsilon}\frac{\partial H_2}{\partial r}|_{r=R}.
# ```
#
# After some math, we get:
#
# ```math
# \alpha_m=\frac{J_m(nkR)J_m(kR)^\prime-\frac{1}{n}J_m(kR)J_m(nkR)^\prime}{\frac{1}{n}H_m^1(kR)J_m(nkr)^\prime-J_m(nkr)H_m^1(kr)^\prime},
# ```
#
# ```math
# \beta_m = \frac{H_m^1(kR)J_m(kR)^\prime-J_m(kR)H_m^1(kR)^\prime}{\frac{1}{n}J_m(nkR)^\prime H_m^1(kR)-J_m(nkR)H_m^1(kR)^\prime},
# ```
#
# where $^\prime$ denotes the derivative, and the derivatives of the Bessel functions are obtained with the recurrent relations:
#
# ```math
# Y_m(z)^\prime=\frac{Y_{m-1}(z)-Y_{m+1}(z)}{2}
# ```
#
# where $Y_m$ denotes any Bessel functions (Hankel functions).
#
#
# Finally, the analytical field is ($1/2k$ is the amplitude that comes from the unit line source excitation):
# ```math
# H(r>R)=\frac{1}{2k}\sum_m\left[\alpha_mi^mH_m^1(kr)+J_m(kr)\right]e^{im\theta}
# ```
#
# ```math
# H(r\leq R)=\frac{1}{2k}\sum_m\beta_mi^mJ_m(nkr)e^{im\theta}
# ```
#

# ### Define the analytical functions
using SpecialFunctions
dbesselj(m,z) = (besselj(m-1,z)-besselj(m+1,z))/2
dhankelh1(m,z)= (hankelh1(m-1,z)-hankelh1(m+1,z))/2
α(m,n,z) = (besselj(m,n*z)*dbesselj(m,z)-1/n*besselj(m,z)*dbesselj(m,n*z))/(1/n*hankelh1(m,z)*dbesselj(m,n*z)-besselj(m,n*z)*dhankelh1(m,z))
β(m,n,z) = (hankelh1(m,z)*dbesselj(m,z)-besselj(m,z)*dhankelh1(m,z))/(1/n*dbesselj(m,n*z)*hankelh1(m,z)-besselj(m,n*z)*dhankelh1(m,z))

function H_t(x,xc,r,ϵ,λ)
    n = √ϵ
    k = 2*π/λ
    θ = angle(x[1]-xc[1]+1im*(x[2]-xc[2]))+π
    M = 40 # Number of Bessel function basis used
    H0 = 0
    if norm([x[1]-xc[1],x[2]-xc[2]])<=r
        for m=-M:M
            H0 += β(m,n,k*r)*cis(m*θ)*besselj(m,n*k*norm([x[1]-xc[1],x[2]-xc[2]]))
        end
    else
        for m=-M:M
            H0 += α(m,n,k*r)*cis(m*θ)*hankelh1(m,k*norm([x[1]-xc[1],x[2]-xc[2]]))+cis(m*θ)*besselj(m,k*norm([x[1]-xc[1],x[2]-xc[2]]))
        end
    end
    return 1im/(2*k)*H0
end
#
# ### Construct the analytical solution in finite element basis
uh_t = CellField(x->H_t(x,xc,r,ϵ₁,λ),Ω)

# ## Output and compare results
#
# The simulated field is shown below. We can see that the simulated fields and the analytical solution matched closed except for the top and PML regions. This is because the simulated source generate plane waves in two directions but we only consider the downward propagating wave in the analytical solution and the PML effect is also not considered in the analytical solution. Therefore, we just need to focus on the "center" regions which excludes the PML and top region above the source, the difference is within 6% of the field amplitude integral. As we increase the resolution, this difference should decrease (until it becomes limited by the PML reflection coefficient from $\sigma_0$, the number of Bessel function basis $M$ or by floating-point error.)
# ![](../assets/emscatter/Results.png)

# ### Save to file and view
writevtk(Ω,"demo",cellfields=["Real"=>real(uh),
        "Imag"=>imag(uh),
        "Norm"=>abs2(uh),
        "Real_t"=>real(uh_t),
        "Imag_t"=>imag(uh_t),
        "Norm_t"=>abs2(uh_t),
        "Difference"=>abs(uh_t-uh)])
# ### Compare the difference in the "center" region
function AnalyticalBox(x) # Get the "center" region
    if abs(x[1])<L/2 && abs(x[2]+0.5)<2.5
        return 1
    else
        return 0
    end
end

Difference=sqrt(sum(∫(abs2(uh_t-uh)*AnalyticalBox)*dΩ)/sum(∫(abs2(uh_t)*AnalyticalBox)*dΩ))

@assert Difference < 0.1
#
# ## References
# [1] [Wikipedia: Electromagnetic wave equation](https://en.wikipedia.org/wiki/Electromagnetic_wave_equation)
#
# [2] [Wikipedia: Perfectly matched layer](https://en.wikipedia.org/wiki/Perfectly_matched_layer)
#
# [3] A. Oskooi and S. G. Johnson, “[Distinguishing correct from incorrect PML proposals and a corrected unsplit PML for anisotropic, dispersive media](http://math.mit.edu/~stevenj/papers/OskooiJo11.pdf),” Journal of Computational Physics, vol. 230, pp. 2369–2377, April 2011.
#
# [4] Stratton, J. A. (1941). Electromagnetic Theory. New York: McGraw-Hill.
#
# [5] [Wikipedia: Jacobi–Anger expansion](https://en.wikipedia.org/wiki/Jacobi%E2%80%93Anger_expansion)
#
# [6] [Wikipedia: Bessel function](https://en.wikipedia.org/wiki/Bessel_function)


function zad38(A,B, el1)
    if size(A) != size(B)
        return NaN
    end
    n = size(A)
    C = zeros(n)
    s = 0
    br = 0
    for i in 1:n[1]
        for j in 1:n[1]
            if A[i,j] > el1
                s +=A[i,j]
                br +=1
            end
            if i == j
                C[i,j] = el1
            elseif i < j
                C[i,j] = A[i,j]
            else
                C[i,j] = B[i,j]
            end
        end
    end
    return C, s/br
end


using FiniteVolumeMethod
using BenchmarkTools

operators = ((gradient, grid -> CellVar(x -> sum(x), grid)), (divergence, grid -> FaceVar(x -> x, grid)))

suite = BenchmarkGroup()
refinement_steps = 5

for dim = 1
    suite[dim] = BenchmarkGroup()

    # Create a uniform grid
    uniform = TreeGrid(zeros(dim))
    for step=1:refinement_steps
        FiniteVolumeMethod.refine!(uniform, [uniform.tree], recurse=true)
    end

    # Create a locally refined grid
    refined = TreeGrid(zeros(dim))
    child = refined.tree
    for step=1:refinement_steps
        FiniteVolumeMethod.refine!(refined, [child])
        child = child.children[1+mod(step-1, 1<<dim)]
    end

    cartesian = CartesianGrid(1 << refinement_steps, dim)

    grids = [refined, uniform, cartesian]
    gridNames = ["refined", "uniform", "cartesian"]

    for (oper, constructor) ∈ operators
        suite[dim][string(oper)] = BenchmarkGroup()

        for (i, grid) ∈ enumerate(grids)
            var = constructor(grid)
            suite[dim][string(oper)][gridNames[i]] = @benchmarkable $(oper)($var, $grid)
        end
    end
end


using IterTools
using Plots
using EllZeroTrendFiltering
using DelimitedFiles

data = readdlm(joinpath(joinpath(dirname(@__FILE__),"data","snp500.txt")))
data = data[1:300]

N = length(data)

@time ℓ = compute_discrete_transition_costs(data);
cost_last = QuadraticPolynomial(1.0, -2*data[end], data[end]^2)

start_time = time()

Λ = pwq_dp_constrained(ℓ, cost_last, 10, 1.65)
println("Time: ", time()-start_time)

for k=3:10
    println(k, " : ", recover_optimal_index_set(Λ[k, 1])[3])
end
###

@time I2, y2, f2 = recover_optimal_index_set(Λ[7, 1])
println(I2)
Y2, f2_2 = EllZeroTrendFiltering.find_optimal_y_values(ℓ, cost_last, I2)

I2, Y2, f3 = recover_solution(Λ[7, 1], ℓ, cost_last, )
#println("Comparison: ", sqrt(f), " --- ", sqrt(f2))

#find_optimal_y_values

plot(layout=(1,2))
##
l = [(isassigned(Λ,i,j) ? length(Λ[i,j]) : 0) for i = 1:size(Λ,1), j = 1:size(Λ,2)]
sum(l)
##

plot!(l', subplot=2)

sum(l)
##

I2, y2, f2 = EllZeroTrendFiltering.recover_solution(Λ[6, 1], ℓ, cost_last)
println(I2)
Y2, _ = EllZeroTrendFiltering.find_optimal_y_values(ℓ, cost_last, I2)

plot(data)
plot!(I2, Y2)


function two_scales_correlation(prices::DataFrame, times::Vector{<:Real}, asseti::Symbol, assetj::Symbol, gamma::Real, num_grids::Real)
    vol_plus_version  = two_scales_volatility(prices[:,asseti] * gamma + (1-gamma) * prices[:,assetj] , times, Symbol("CompoundAsset__",asseti,"__+__",assetj,"__",gamma), num_grids)[1]
    vol_minus_version = two_scales_volatility(prices[:,asseti] * gamma - (1-gamma) * prices[:,assetj] , times, Symbol("CompoundAsset__",asseti,"__-__",assetj,"__",gamma), num_grids)[1]
    covv = (1/(4*gamma*(1-gamma))) * ( vol_plus_version^2 -  vol_minus_version^2 )
    # Note that we calculate these again rather than feeding in the ones calculated earlier because it is improtant the vols are calculated with the same data as the sums/differences.
    asseti_vol = two_scales_volatility(prices[:,asseti], times, Symbol("CompoundAsset__",asseti,"__ONLY"), num_grids)[1]
    assetj_vol = two_scales_volatility(prices[:,assetj], times, Symbol("CompoundAsset__",assetj,"__ONLY"), num_grids)[1]
    correl = covv/(asseti_vol * assetj_vol)
    return correl
end

function get_refresh_times_and_prices(ts::SortedDataFrame, asset1::Symbol, asset2::Symbol)
    assets = [asset1, asset2]
    at_times    = get_all_refresh_times(ts, assets)
    prices = random_value_in_interval(ts, at_times; assets = assets)
    times = prices[:,:Time]
    return prices, times
end

"""
    two_scales_covariance(ts::SortedDataFrame, assets::Vector{Symbol} = get_assets(ts);
                          regularisation::Union{Missing,Symbol} = :correlation_default,
                          regularisation_params::Dict = Dict(), only_regulise_if_not_PSD::Bool = false,
                          equalweight::Bool = false, num_grids::Real = default_num_grids(ts))

Estimation of a CovarianceMatrix using the two scale covariance method.
### Inputs
* `ts` - The tick data.
* `assets` - The assets you want to estimate volatilities for.
* `regularisation` - A symbol representing what regularisation technique should be used. If missing no regularisation is performed.
* `regularisation_params` - keyword arguments to be consumed in the regularisation algorithm.
* `only_regulise_if_not_PSD` - Should regularisation only be attempted if the matrix is not psd already.
* `equalweight` - Should we use equal weight for the two different linear combinations of assets. If false then an optimal weight is calculated (from volatilities).
* `num_grids` - Number of grids used in order in two scales estimation.
### Returns
* A `CovarianceMatrix`.
"""
function two_scales_covariance(ts::SortedDataFrame, assets::Vector{Symbol} = get_assets(ts);
                               regularisation::Union{Missing,Symbol} = :correlation_default,
                               regularisation_params::Dict = Dict(), only_regulise_if_not_PSD::Bool = false,
                               equalweight::Bool = false, num_grids::Real = default_num_grids(ts))

    two_scales_vol, micro_noise = two_scales_volatility(ts, assets; num_grids = num_grids)

    N = length(assets)
    mat = zeros(N,N)
    for i in 1:N
        asseti = assets[i]
        for j in i:N
            assetj = assets[j]
            gamma = min(max(0.1, equalweight ? 0.5 : two_scales_vol[assetj] / (two_scales_vol[asseti] + two_scales_vol[assetj]))   , 0.9)
            if (i == j)
                mat[i,j] = 1
            else
                zero_vol = (two_scales_vol[asseti] < 2*eps()) |  (two_scales_vol[assetj] < 2*eps())
                if zero_vol
                    mat[i,j] = 0.0
                else
                    prices, times = get_refresh_times_and_prices(ts, asseti, assetj)
                    mat[i,j] = two_scales_correlation(prices, times, asseti, assetj, gamma, num_grids)
                end
            end
        end
    end
    mat = Hermitian(mat)

    # Regularisation - It is done on the correlation matrix for this algo rather than the covariance matrix.
    dont_regulise = ismissing(regularisation) || (only_regulise_if_not_PSD && is_psd_matrix(mat))
    covmat = dont_regulise ? mat : regularise(mat, ts, assets, regularisation; regularisation_params... )

    vols = map(a -> two_scales_vol[a], assets)
    covmat = CovarianceMatrix(covmat, vols, assets, ts.time_period_per_unit)

    return covmat
end


## indices
import ..ParallelStencil: INDICES, WITHIN_DOC
ix, iy, iz = INDICES[1], INDICES[2], INDICES[3]
ixi, iyi, izi = :($ix+1), :($iy+1), :($iz+1);
## 2nd-order accuracy for first-order derivatives
ix_1=:($ix+1);
iy_1=:($iy+1);
iz_1=:($iz+1);

macro   d_xi_1(A::Symbol)  esc(:( $A[$ix_1+1,$iyi ,$izi ] - $A[$ix_1  ,$iyi ,$izi ] )) end
macro   d_yi_1(A::Symbol)  esc(:( $A[$ixi ,$iy_1+1,$izi ] - $A[$ixi ,$iy_1  ,$izi ] )) end
macro   d_zi_1(A::Symbol)  esc(:( $A[$ixi ,$iyi ,$iz_1+1] - $A[$ixi ,$iyi ,$iz_1  ] )) end
## 8th-order accuracy for first-order derivatives
ix_8=:($ix+3);
iy_8=:($iy+3);
iz_8=:($iz+3);

macro   d_xi_8(A::Symbol)
    esc(:(1225/1024*($A[$ix_8+1,$iyi,$izi]-$A[$ix_8,$iyi,$izi])-
    245/3072*($A[$ix_8+2,$iyi,$izi]-$A[$ix_8-1,$iyi,$izi])+
    49/5120*($A[$ix_8+3,$iyi,$izi]-$A[$ix_8-2,$iyi,$izi])-
    5/7186*($A[$ix_8+4,$iyi,$izi]-$A[$ix_8-3,$iyi,$izi])));
end

macro   d_yi_8(A::Symbol)
    esc(:(1225/1024*($A[$ixi,$iy_8+1,$izi]-$A[$ixi,$iy_8,$izi])-
    245/3072*($A[$ixi,$iy_8+2,$izi]-$A[$ixi,$iy_8-1,$izi])+
    49/5120*($A[$ixi,$iy_8+3,$izi]-$A[$ixi,$iy_8-2,$izi])-
    5/7168*($A[$ixi,$iy_8+4,$izi]-$A[$ixi,$iy_8-3,$izi])));
end
macro   d_zi_8(A::Symbol)
    esc(:(1225/1024*($A[$ixi,$iyi,$iz_8+1]-$A[$ixi,$iyi,$iz_8])-
    245/3072*($A[$ixi,$iyi,$iz_8+2]-$A[$ixi,$iyi,$iz_8-1])+
    49/5120*($A[$ixi,$iyi,$iz_8+3]-$A[$ixi,$iyi,$iz_8-2])-
    5/7168*($A[$ixi,$iyi,$iz_8+4]-$A[$ixi,$iyi,$iz_8-3])));
end
## 12th-order accuracy for first-order derivatives
ix_12=:($ix+5);
iy_12=:($iy+5);
iz_12=:($iz+5);

macro   d_xi_12(A::Symbol)
    esc(:(160083/131072*($A[$ix_12+1,$iyi,$izi]-$A[$ix_12,$iyi,$izi])-
    12705/131072*($A[$ix_12+2,$iyi,$izi]-$A[$ix_12-1,$iyi,$izi])+
    22869/1310720*($A[$ix_12+3,$iyi,$izi]-$A[$ix_12-2,$iyi,$izi])-
    5445/1835008*($A[$ix_12+4,$iyi,$izi]-$A[$ix_12-3,$iyi,$izi])+
    847/2359296*($A[$ix_12+5,$iyi,$izi]-$A[$ix_12-4,$iyi,$izi])-
    63/2883584*($A[$ix_12+6,$iyi,$izi]-$A[$ix_12-5,$iyi,$izi])
    ));
end

macro   d_yi_12(A::Symbol)
    esc(:(160083/131072*($A[$ixi,$iy_12+1,$izi]-$A[$ixi,$iy_12,$izi])-
    12705/131072*($A[$ixi,$iy_12+2,$izi]-$A[$ixi,$iy_12-1,$izi])+
    22869/1310720*($A[$ixi,$iy_12+3,$izi]-$A[$ixi,$iy_12-2,$izi])-
    5445/1835008*($A[$ixi,$iy_12+4,$izi]-$A[$ixi,$iy_12-3,$izi])+
    847/2359296*($A[$ixi,$iy_12+5,$izi]-$A[$ixi,$iy_12-4,$izi])-
    63/2883584*($A[$ixi,$iy_12+6,$izi]-$A[$ixi,$iy_12-5,$izi])
    ));
end
macro   d_zi_12(A::Symbol)
    esc(:(160083/131072*($A[$ixi,$iyi,$iz_12+1]-$A[$ixi,$iyi,$iz_12])-
    12705/131072*($A[$ixi,$iyi,$iz_12+2]-$A[$ixi,$iyi,$iz_12-1])+
    22869/1310720*($A[$ixi,$iyi,$iz_12+3]-$A[$ixi,$iyi,$iz_12-2])-
    5445/1835008*($A[$ixi,$iyi,$iz_12+4]-$A[$ixi,$iyi,$iz_12-3])+
    847/2359296*($A[$ixi,$iyi,$iz_12+5]-$A[$ixi,$iyi,$iz_12-4])-
    63/2883584*($A[$ixi,$iyi,$iz_12+6]-$A[$ixi,$iyi,$iz_12-5])
    ));
end


w = [1.0, 0.1, 0.3, 1.0]
cosa = cos(0.5)
sina = sin(0.5)
KitBase.global_frame(randn(2), cosa, sina)
KitBase.local_frame(randn(2), cosa, sina)
KitBase.global_frame(w, cosa, sina)
KitBase.local_frame(w, cosa, sina)

KitBase.global_frame(rand(3), rand(3, 3))
KitBase.local_frame(rand(3), rand(3, 3))
KitBase.global_frame(rand(5), rand(3, 3))
KitBase.local_frame(rand(5), rand(3, 3))

KitBase.unit_normal(rand(2), rand(2))
KitBase.unit_normal(rand(3), rand(3), rand(3))
KitBase.point_distance(rand(2), rand(2), rand(2))

#--- structure mesh ---#
KitBase.PSpace1D()
KitBase.PSpace1D(0.0, 1.0) |> show
KitBase.PSpace2D()
KitBase.PSpace2D(0.0, 1.0, 0.0, 1.0) |> show
KitBase.CSpace2D(0.0, 1.0, 10, 0.0, π, 10) |> show

x0 = 0.0
nx = 15
dx = 0.1
KitBase.uniform_mesh(x0, nx, dx)

x = collect(0:0.1:1)
y = collect(0:0.1:1)
z = collect(0:0.1:1)
KitBase.ndgrid(x)
KitBase.ndgrid(x, y)
KitBase.ndgrid(x, y, z)
KitBase.meshgrid(x)
KitBase.meshgrid(x, y)
KitBase.meshgrid(x, y, z)

KitBase.find_idx(randn(20), 0.13, mode = :uniform)
KitBase.find_idx(randn(20), 0.13, mode = :nonuniform)

#--- unstructure mesh ---#
cd(@__DIR__)
ps = KitBase.UnstructPSpace("t1.msh")
KitBase.write_vtk(ps.points, ps.cellid, randn(size(ps.cellid, 1)))


struct BudgetedSpanningTreeInstance{T, U} <: CombinatorialInstance
  graph::AbstractGraph{T}
  rewards::Dict{Edge{T}, Float64}
  weights::Dict{Edge{T}, U}
  budget::U
end

function _budgeted_spanning_tree_compute_value(i::Union{SpanningTreeInstance{T}, BudgetedSpanningTreeInstance{T, U}}, tree::Vector{Edge{T}}) where {T, U}
  return sum(i.rewards[(e in keys(i.rewards)) ? e : reverse(e)] for e in tree)
end

abstract type BudgetedSpanningTreeSolution{T, U} <: CombinatorialSolution
  # instance::BudgetedSpanningTreeInstance{T, U}
  # tree::Vector{Edge{T}}
  # value::Float64 # TODO: remove me, only useful for Lagrangian.
end

struct SimpleBudgetedSpanningTreeSolution{T, U} <: BudgetedSpanningTreeSolution{T, U}
  instance::BudgetedSpanningTreeInstance{T, U}
  tree::Vector{Edge{T}}
  value::Float64

  function SimpleBudgetedSpanningTreeSolution(instance::BudgetedSpanningTreeInstance{T, U}, tree::Vector{Edge{T}}) where {T, U}
    if length(tree) > 0
      return new{T, U}(instance, tree, _budgeted_spanning_tree_compute_value(instance, tree))
    else
      return new{T, U}(instance, tree, -Inf)
    end
  end

  function SimpleBudgetedSpanningTreeSolution(instance::BudgetedSpanningTreeInstance{T, U}) where {T, U}
    # No feasible solution.
    return new{T, U}(instance, edgetype(instance.graph)[], -Inf)
  end
end


#
# This function computes the center of mass of a selection
#

function cm(selection,simulation,x,y,z)

  cm = [ 0., 0., 0. ]
  totmass = 0.
  for i in 1:length(selection)
    cm[1] = cm[1] + simulation.atoms[selection[i]].mass*x[selection[i]]
    cm[2] = cm[2] + simulation.atoms[selection[i]].mass*y[selection[i]]
    cm[3] = cm[3] + simulation.atoms[selection[i]].mass*z[selection[i]]
    totmass = totmass + simulation.atoms[selection[i]].mass
  end
  cm[1] = cm[1] / totmass
  cm[2] = cm[2] / totmass
  cm[3] = cm[3] / totmass

  return cm
end


import ForneyLab: ProbabilityDistribution, Univariate, GaussianMeanVariance
import Plots: plot, plot!, xlabel!, ylabel!

function plot_messages(state_prediction::ProbabilityDistribution{Univariate, GaussianMeanVariance},
                       measurement_likelihood::ProbabilityDistribution{Univariate, GaussianMeanVariance},
                       corrected_state_prediction::ProbabilityDistribution{Univariate, GaussianMeanVariance})

    # Extract parameters
    m_p = mean(state_prediction)
    v_p = var(state_prediction)
    m_l = mean(measurement_likelihood)
    v_l = var(measurement_likelihood)
    m_c = mean(corrected_state_prediction)
    v_c = var(corrected_state_prediction)

    # Set plot range
    range_p = [m_p - 8*sqrt(v_p), m_p + 8*sqrt(v_p)]
    range_l = [m_l - 8*sqrt(v_l), m_l + 8*sqrt(v_l)]
    range_c = [m_c - 8*sqrt(v_c), m_c + 8*sqrt(v_c)]
    plot_range = range(minimum([range_p[1], range_l[1], range_c[1]]),
                       maximum([range_p[2], range_l[2], range_c[2]]), length=300)

    # Probabilities over range
    probs_p = [pdf(state_prediction, a) for a in plot_range]
    probs_l = [pdf(measurement_likelihood, a) for a in plot_range]
    probs_c = [pdf(corrected_state_prediction, a) for a in plot_range]

    # Start plotting
    plot(plot_range, probs_p, color="red", label="state prediction", size=(800,500))
    plot!(plot_range, probs_l, color="blue", label="observation likelihood")
    plot!(plot_range, probs_c, color="green", label="state estimate")
    xlabel!("x-coordinate")
    ylabel!("Probability")
end

function pdf(dist::ProbabilityDistribution{Univariate, GaussianMeanVariance}, x::Number)
    return exp(-0.5*(log(2*pi) + log(dist.params[:v]) + (x - dist.params[:m])^2 / dist.params[:v]))
end;


using Random
using WindingNumbers
using WindingNumbers: angle, Point
Random.seed!(10)

using Test

@testset "Winding Numbers tests" begin

@testset "Basic angle tests" begin
  a = Point(1.0, 0.0)
  b = Point(0.0, 1.0)
  @test π/2 ≈ angle(b) - angle(a)

  a = Point(-1.0, 0.0)
  b = Point(1.0, 0.0)
  @test -π ≈ angle(b) - angle(a)

  a = Point(1.0, 0.0)
  b = Point(-1.0, 0.0)
  @test π ≈ angle(b) - angle(a)

  a = Point(1.0, 0.0)
  b = Point(1.0, 1.0)
  @test π/4 ≈ angle(b) - angle(a)
 end

function mock(x::Vector{Float64}, n, r=zeros(2))
  θ = atan(x[2] - r[2], x[1] - r[1])
  return [cos(n*θ), sin(n*θ)]
end

turnings(n::Int) = -n:n

@testset "Basic understanding - angles" begin
  a = Point(1.0, 1.0)
  b = Point(-1.0, 1.0)
  c = Point(-1.0, -1.0)
  d = Point(1.0, -1.0)
  @test angle(a, b) ≈ π/2
  @test angle(b, c) ≈ π/2
  @test angle(c, d) ≈ π/2
  @test angle(d, a) ≈ π/2
  @test angle(b, a) ≈ -π/2
  @test angle(c, b) ≈ -π/2
  @test angle(d, c) ≈ -π/2
  @test angle(a, d) ≈ -π/2
  ϕ = angle(Point(-1.0, 2*eps()), Point(-1.0, -2*eps()))
  @test ϕ ≈ 4*eps()
  @test ϕ > 0.0
  ϕ = angle(Point(-1.0, -2*eps()), Point(-1.0, 2*eps()))
  @test ϕ ≈ -4*eps()
  @test ϕ < 0.0
end

@testset "Basic understanding - square with pole" begin
  a = Point(1.0, 1.0)
  b = Point(-1.0, 1.0)
  c = Point(-1.0, -1.0)
  d = Point(1.0, -1.0)
  ϕ = angle(a, b) + angle(b, c) + angle(c, d) + angle(d, a)
  @test ϕ ≈ 2π
  ϕ = angle(b, a) + angle(c, b) + angle(d, c) + angle(a, d)
  @test ϕ ≈ -2π
end

@testset "Basic understanding - square without pole" begin
  a = Point(1.0, 1.0)
  b = Point(2.0, 1.0)
  c = Point(2.0, 2.0)
  d = Point(1.0, 2.0)
  ϕ = angle(a, b) + angle(b, c) + angle(c, d) + angle(d, a)
  @test ϕ ≈ 0.0
  ϕ = angle(b, a) + angle(c, b) + angle(d, c) + angle(a, d)
  @test ϕ ≈ 0.0
end

@testset "Pole at origin in grid" begin
  for i in turnings(5)
    n = WindingNumbers.windingnumber(x->mock(x, i), -ones(2), ones(2))
    @test Int(round(n)) == i
  end
end

@testset "Random location of pole inside grid" begin
  for i in turnings(2)
    r = rand(2) # solution
    a, b = r - rand(2), r + rand(2)
    @assert all(a .< r .<b)
    n = WindingNumbers.windingnumber(x->mock(x, i, r), a, b)
    @test Int(round(n)) == i
  end
end

@testset "Unphysical pole" begin
  for i in turnings(5)
    s = rand(2) # square
    @assert s[1] > 0.0
    @assert s[2] > 0.0
    n = WindingNumbers.windingnumber(x->[0.0, 0.0], s, 2*s)
    @test Int(round(n)) == 0
  end
end

@testset "Physical pole not in grid" begin
  for i in turnings(5)
    s = rand(2) # square
    @assert s[1] > 0.0
    @assert s[2] > 0.0
    n = WindingNumbers.windingnumber(x->mock(x, i), s, 2*s)
    @test Int(round(n)) == 0
  end
end

@testset "solve" begin
  for i ∈ 1:100
    rtol = 10^Float64(-rand(1:14))
    root = 2 .* rand(2) .- 1
    lower = root .- rand(2)
    upper = root .+ rand(2)
    allpointsinside = true
    function solvemock(x)
      allpointsinside &= all(lower .<= x .<= upper)
      y = mock(x, 1, root)
      return y[1] + im * y[2]
    end
    solutions = WindingNumbers.boundingboxes(solvemock, lower, upper,
                                             xtol_rel=rtol, stopval=eps())
    @test length(solutions) == 1
    @test allpointsinside
    length(solutions) == 1 || continue
    A, B = solutions[1]
    @test all(A .<= root .<= B)
    @test all((B .- A) ./ (upper .- lower) .< rtol)
  end
end

end


# we could just store intervals as a rational number
# but there's more than one way to represent a given rational
# so we don't want to lose user information
mutable struct Interval
    numerator::Int
    denominator::Int
    octave::Int
end

function Interval(; numerator = 1, denominator = 1, octave = 0)
    Interval(numerator, denominator, octave)
end

function Rational(interval::Interval)
    interval.numerator // interval.denominator * (2 // 1)^interval.octave
end

const INTERVAL_REGEX = r"(?<numerator>[^/o]+)(?:/(?<denominator>[^o]+))?(?:o(?<octave>.+))?"

# TODO: reuse from AudioSchedules
function parse(::Type{Interval}, text::AbstractString; line_number)
    a_match = match(INTERVAL_REGEX, text)
    if a_match === nothing
        throw_parse_error(text, "interval", line_number)
    else
        numerator_string = a_match["numerator"]
        numerator = tryparse(Int, numerator_string)
        denominator_string = a_match["denominator"]
        octave_string = a_match["octave"]
        Interval(
            if numerator === nothing
                throw_parse_error(numerator_string, "numerator", line_number)
            else
                numerator
            end,
            if denominator_string === nothing
                1
            else
                denominator = tryparse(Int, denominator_string)
                if denominator === nothing
                    throw_parse_error(denominator_string, "denominator", line_number)
                else
                    denominator
                end
            end,
            if octave_string === nothing
                0
            else
                octave = tryparse(Int, octave_string)
                if octave === nothing
                    throw_parse_error(octave_string, "octave", line_number)
                else
                    octave
                end
            end,
        )
    end
end

function print(io::IO, interval::Interval)
    denominator = interval.denominator
    octave = interval.octave
    print(io, interval.numerator)
    if denominator != 1
        print(io, '/')
        print(io, denominator)
    end
    if octave != 0
        print(io, 'o')
        print(io, octave)
    end
end


import DFWannier: Point3, SVector

import DFWannier: WannierFunction

x_range = -1:0.01:1
y_range = -1:0.01:1
z_range = -1:0.01:1

wfc_grid = [Point3(x, y, z) for x in x_range, y in y_range, z in z_range]

px_orb = normalize(WannierFunction(wfc_grid, [SVector(((p[1] + 0im) * ℯ^(-norm(p)^2)),) for p in wfc_grid]))

px_orb2 = px_orb'
LinearAlgebra.adjoint!(px_orb2, px_orb2)

@test values(px_orb2) == values(px_orb)

@test norm(px_orb) ≈ norm(px_orb2) ≈ 1.0

py_orb = normalize(WannierFunction(wfc_grid, [SVector(((p[2] + 0im) * ℯ^(-norm(p)^2)),) for p in wfc_grid]))

pz_orb = normalize(WannierFunction(wfc_grid, [SVector(((p[3] + 0im) * ℯ^(-norm(p)^2)),) for p in wfc_grid]))

@test dot(px_orb, py_orb) <= 1.0e-15
@test dot(px_orb, px_orb) ≈ 1.0

Lx = zeros(ComplexF64, 3, 3)
Ly = zeros(ComplexF64, 3, 3)
Lz = zeros(ComplexF64, 3, 3)
for (i1, p1) in enumerate((px_orb, py_orb, pz_orb)), (i2, p2) in enumerate((px_orb, py_orb, pz_orb))
	Lx[i1, i2], Ly[i1, i2], Lz[i1, i2] = DFW.calc_angmom(p1, p2, zero(Point3))
end

@test norm(sum(Lx .- [0 0 0; 0 0 -im; 0 im 0])) < 1e-4
@test norm(sum(Ly .- [0 0 im; 0 0 0; -im 0 0])) < 1e-4
@test norm(sum(Lz .- [0 -im 0; im 0 0; 0 0 0])) < 1e-4

px_orb_up = normalize(WannierFunction(wfc_grid, [SVector((p[1] + 0im, zero(ComplexF64)) .* ℯ^(-norm(p)^2)) for p in wfc_grid]))
px_orb_dn = normalize(WannierFunction(wfc_grid, [SVector((zero(ComplexF64), p[1] + 0im) .* ℯ^(-norm(p)^2)) for p in wfc_grid]))

@test dot(px_orb_dn, px_orb_up) ≈ 0.0
@test DFW.calc_spin(px_orb_up, px_orb_up) ≈ DFW.Point3(0.0 + 0im, 0.0+0.0im, 0.5 + 0.0im)

@test norm(DFW.calc_dip(px_orb, py_orb)) < 1e-17
@test norm(DFW.calc_dip(px_orb_up, px_orb_dn)) < 1e-17


# include("pascaltriangle.jl") # define pascaltriangle(m,n)

function in2bEne(jj::Int64, Msize::Int64, Np::Int64, matp::Matrix{Int64})

    # this function return the index of Fock state at the index of many body state
    indfk = copy(jj)

    # matp = zeros(Int,Msize+1,Np+1);
    # pascaltriangle!(Msize,Np,matp) # the size is Msize+1 times Np+1
    # note the indices are m+1 and n+1 for N^m_n

    # unlike Ponomarev's way, I count the index of Fock state in the backwards
    # so that "1" means |Np,0,0,...>.
    indfk = matp[Msize+1,Np+1] + 1 - indfk # the indices are m+1 and n+1 for N^m_n

    if indfk > matp[Msize+1,Np+1] # the indices are m+1 and n+1 for N^m_n
       error("indfk larger than the maximum of N^M_N.")
    end

    indfk = indfk - 1
    # vecmb = spzeros(Int64,Msize+1) #spzeros(Int64,Msize+1)

    indM = Msize-1
    indN = Np

    Eneminumhalf = 0

    while true

          if indN == 0
             break
          end

          if indfk >= matp[indM+1,indN+1]

             indfk = indfk - matp[indM+1,indN+1] # the indices are m+1 and n+1 for N^m_n
             # vecmb[Msize-indM] = vecmb[Msize-indM] + 1
             indN = indN - 1

             Eneminumhalf = Eneminumhalf + ceil(Int64,(Msize-indM)/2)-1

          else
             indM = indM - 1
          end

    end

    # the last element for the coefficient
    # vecmb[Msize+1] = 1

    # return vecmb
    return Eneminumhalf #+1/2

end


my_f(x,y) = 2x + 3y
my_g(x,y) = x^2 + y^2