### A Pluto.jl notebook ###
# v0.17.5

using Markdown
using InteractiveUtils

# ╔═╡ a4b92179-59f1-4406-a6d8-e2c844f63e95
# hideall
let
    docs_dir = dirname(dirname(@__DIR__))
    pkg_dir = dirname(docs_dir)

    using Pkg: Pkg
    Pkg.activate(docs_dir)
    Pkg.develop(; path=pkg_dir)
    Pkg.instantiate()

    # Putting the include here to avoid Pluto getting confused about cell order.
    include(joinpath(docs_dir, "src", "tutorials_utils.jl"))
end;

# ╔═╡ 79cac787-9b01-4c0e-bb7d-f4555fcf6216
using CSV

# ╔═╡ 2d99b5d8-761c-489c-9bbe-bb48fd471ea5
using DataFrames

# ╔═╡ 4c99ba26-0041-4aaa-b224-44819dc72a87
using TuringGLM

# ╔═╡ 1346efe0-72f4-11ec-0251-5fac34b1eec8
md"""
Let's cover **Linear Regression** with a famous dataset called `kidiq` (Gelman & Hill, 2007), which is data from a survey of adult American women and their respective children.
Dated from 2007, it has 434 observations and 4 variables:

* `kid_score`: child's IQ
* `mom_hs`: binary/dummy (0 or 1) if the child's mother has a high school diploma
* `mom_iq`: mother's IQ
* `mom_age`: mother's age

For the purposes of this tutorial, we download the dataset from the TuringGLM repository:
"""

# ╔═╡ e0a8ebec-3b29-42c2-a406-47d99df64b68
url = "https://github.com/TuringLang/TuringGLM.jl/raw/main/data/kidiq.csv";

# ╔═╡ 65b092a5-e6c5-48a0-9058-39a5c0094afa
kidiq = CSV.read(download(url), DataFrame)

# ╔═╡ 2eb258a0-854b-4ea6-80f9-f88d868ec839
md"""
Using `kid_score` as dependent variable and `mom_hs` along with `mom_iq` as independent variables with a moderation (interaction) effect:
"""

# ╔═╡ a1031c0e-2237-409e-bf7d-65efb071f1da
fm = @formula(kid_score ~ mom_hs * mom_iq)

# ╔═╡ 43d63761-adf5-4a52-b996-4ad3adfb35d0
md"""
Next, we instantiate our model with `turing_model` without specifying any model, thus the default model will be used (`model=Normal`):
"""

# ╔═╡ 55b91963-001e-4753-93a6-2fa64190f353
model = turing_model(fm, kidiq);

# ╔═╡ 92bec7a8-3524-4daa-9ef9-b58f2e59dea4
n_samples = 2_000;

# ╔═╡ 7935378b-9f7a-418c-9d51-e77e4a551bde
md"""
This model is a valid Turing model, which we can pass to the default `sample` function from Turing to get our parameter estimates.
We use the `NUTS` sampler with $n_samples samples.
"""

# ╔═╡ 3f26f821-0985-4e15-845d-4791c623a799
chns = sample(model, NUTS(), n_samples);

# ╔═╡ 4adf2c8d-3f05-4ceb-897c-ee2e1b156474
# hide
plot_chains(chns)

# ╔═╡ 28ef5f85-7645-4b26-ad71-b043ef141a50
md"""
## References

Gelman, A., & Hill, J. (2007). Data analysis using regression and multilevel/hierarchical models. Cambridge university press.
"""

# ╔═╡ Cell order:
# ╠═a4b92179-59f1-4406-a6d8-e2c844f63e95
# ╠═1346efe0-72f4-11ec-0251-5fac34b1eec8
# ╠═79cac787-9b01-4c0e-bb7d-f4555fcf6216
# ╠═2d99b5d8-761c-489c-9bbe-bb48fd471ea5
# ╠═4c99ba26-0041-4aaa-b224-44819dc72a87
# ╠═e0a8ebec-3b29-42c2-a406-47d99df64b68
# ╠═65b092a5-e6c5-48a0-9058-39a5c0094afa
# ╠═2eb258a0-854b-4ea6-80f9-f88d868ec839
# ╠═a1031c0e-2237-409e-bf7d-65efb071f1da
# ╠═43d63761-adf5-4a52-b996-4ad3adfb35d0
# ╠═55b91963-001e-4753-93a6-2fa64190f353
# ╠═92bec7a8-3524-4daa-9ef9-b58f2e59dea4
# ╠═7935378b-9f7a-418c-9d51-e77e4a551bde
# ╠═3f26f821-0985-4e15-845d-4791c623a799
# ╠═4adf2c8d-3f05-4ceb-897c-ee2e1b156474
# ╠═28ef5f85-7645-4b26-ad71-b043ef141a50


@info "Loading packages"
using CPUTime
using Plots
using Printf

include("src/grids.jl")
include("src/util.jl")

## Parameters / Configuration
# Size Distribtion
r̅ = (10.0) * 1e-6 # mode radius of initial distribution, input micron convert to m
x̅ = mass_from_r(r̅) # mean initial droplet mass (kg)
L = (1.0)*1e-3 # total water content, g/m3 convert to kg/m3

# Grid Spacing
n = 4
α = 2^(1/n) # Mass bin scaling ratio, 2^(1/n) where 'n' is the number of bins
            # between mass doubling

m = ceil(Integer, 1 + 3*log(rm / r₁)/log(α)) # number of mass bins to populate
Δy = Δlnr = log(α) / 3  # constant grid distance of logarithmic grid

# Collision kernel
# Options:
#   1) :golovin - Golovin (1963)
#   2) :hydro - basic hydrodynamic kernel with unit coal/coll efficiency
#   3) :long - Long collision efficiency
kernel = :long

# Time Integration
tmax = 3601 # seconds
Δt = 5.0 # s
Δt_plot = 10 # minutes
nt = ceil(Integer, tmax / Δt)

# Other configs
debug = false
do_plots = true

# Report configuration to terminal
println("""
1D Stochastic Collection Equation Solver
========================================

Initial Conditions
------------------
L = $(L*1e3) g/m³ - Total LWC
r̅ = $(r̅*1e6) μm - Mode radius
x̅ = $x̅ kg - Mean droplet mass

Grid Setup
----------
$m bins spanning radii ($(r₁*1e6) μm - $(rm*1e6) μm)
α factor = $α (mass doubling every $n bins)

Time Integration
----------------
Integrating to $tmax seconds by Δt=$Δt

- DEBUG is $debug
- DO_PLOTS is $do_plots

""")

## Model Setup

# Mass and Radius Grid
#=
I find it a bit easier to reason through things in droplet-radius-space. We can
adopt the same convention of constructing a mass grid with an α scaling factor
(e.g. for α=√2 every mass bin is twice as big as the bin twice before it) and
write simple, similar recurrence relationship in radius space:

  rᵢ = α rᵢ₋₁

Given r₀ as the smallest bin, we can then solve that

  rᵢ = r₁ α^{(i-1)/3} s.t. i ∈ ℤ > 0

Computing the number of grid cells...
  ln rm > ln(r1) + ln(α) (m-1)/3
  ln rm - ln r1 > ln α (m-1)/3
  3 ln(rm / r1) / (ln α) > m - 1

=#
rᵢ = r₁*(α.^((collect(1:m) .- 1)./3)) # meter
xᵢ = mass_from_r.(rᵢ) # kg

# Initial droplet distribution; g(y, t) = 3x² * n(x, t), eq. 2 from B98
gᵢ = (3 .* xᵢ.^2) .* nc.(xᵢ, L=L, x̅=x̅) # kg / m3

# Courant Numbers
#=
The final take-home here is that these courant limits are purely a function of
the mass grid discretization. On the one hand, they can be pre-computed and then
cached for retrieval as needed. This fact is used in combination with the fact
that the same holds true for the collision kernel to pre-compute all the limits
in the original code, as well as pre-estimate where collisions will land in the
actual coad subroutine
=#

println("""
COMPUTE COURANT LIMITS ON GRID""")
CPUtic()
c, ima = courant(xᵢ)
elapsed = CPUtoq()
@printf "%3.1f seconds elapsed\n" elapsed

# Collision Kernel

println("""
COMPUTE COLLISION KERNELS ON GRID""")
CPUtic()
ck = kernels(xᵢ, kernel)
elapsed = CPUtoq()
@printf "%3.1f seconds elapsed\n" elapsed

## Plot Initial Check

# Plot initial conditions and then begin the time loop
if do_plots
  println("Plotting initial conditions")
  p = plot(
    rᵢ*1e6, gᵢ*1e3, label="t = 0 min",
    xaxis=:log, xlabel="r (μm)",
    # xlim=(0.5, 5000), xticks=[1, 10, 100, 1000],
    # ylim=(0, 0.9), yticks=0:0.06:0.9,
    ylabel="g (g / m³)",
  )
  display(p)
end

## TIME LOOP

#=
All of the time looping logic here is super old-school and can be totally
re-written for simplicity.
=#
# time integration
tlmin = 1e-6
t = 0.0
lmin = 0.0
total_runtime = 0.0

@inline function coad!(g, x, c, ck, ima, Δt, Δy; debug=false, gmin=gmin)

  # Lower and Upper integration limit i0, i1
  i0, i1 = find_bounds(g, gmin=gmin)

  if debug
    @printf "DEBUG bnds_check %6d %8d %8d\n" t i0 i1
  end

  # Main collision/coalescence loop
  for i ∈ i0:i1
    for j ∈ i:i1
      k = ima[i, j]  # Get pre-computed index of coalescence bin edge
      kp = k + 1

      # PORT - handle a weird edge condition in the initialization of ima?
      if k < 1
        continue
      end

      # We did not pre-scale the collison kernels LUT by Δt and Δy in this
      # version so we have to account for them here.
      x0 = ck[i, j] * g[i] * g[j] * Δt * Δy
      x0 = min(x0, g[i] * x[j])
      if j != k # Not sure what's going on here.
        x0 = min(x0, g[j] * x[i])
      end

      gsi = x0 / x[j]
      gsj = x0 / x[i]
      gsk = gsi + gsj
      g[i] = g[i] - gsi
      g[j] = g[j] - gsj
      gk = g[k] + gsk

      if gk > gmin

        # ORIGINAL
        # x1 = log(gᵢ[kp] / gk + 1e-60)

        # MODIFIED - Apply a limiter to avoid negative args to log
        #   We note that this may not strictly obey the formulation of the flux
        #   algorithm, but this tends to work okay in practice.
        log_arg = g[kp] / gk + 1e-60
        log_arg = max(1e-60, log_arg)
        x1 = log(log_arg)

        # We did not pre-scale the cached Courant limits for this version so we
        # have to apply them when estimating the flux directly.
        flux = gsk / x1 * (exp(0.5 * x1) - exp(x1 * (0.5 - (c[i, j]/3/Δy))))
        flux = min(flux, gk)
        g[k] = gk - flux
        g[kp] = g[kp] + flux
      end

    end # j
  end # i

end

println("""
BEGIN TIME INTEGRATION
""")
CPUtic()
for i ∈ 1:nt
  global t = t + Δt
  global tlmin = tlmin + Δt

  # Collision
  # subroutine coad
  coad!(gᵢ, xᵢ, c, ck, ima, Δt, Δy, debug=debug, gmin=gmin)

  # Plotting
  if tlmin ≥ 60
    global tlmin = tlmin - 60
    global lmin = lmin + 1

    if (lmin % Δt_plot) < 1 & do_plots
      display(plot!(p, rᵢ*1e6, gᵢ*1e3, label = "t = $lmin min"))
    end

    # Mass balance? Not sure what's going on here. Maybe numerical checking?
    x0 = 0.0
    x1 = 1.0
    imax = 0
    for i ∈ 1:m
      x0 = x0 + gᵢ[i] * Δy
      x1 = max(x1, gᵢ[i])
      if abs(x1 - gᵢ[i]) < 1e-9
        imax = i
      end
    end
    @printf "    %4d mins |" lmin
    @printf " mass %10.3e  max %10.3e imax %3d" x0 x1 imax
    @printf "\n"
  end

  # Mass balance? Not sure what's going on here. Maybe numerical checking?
  x0 = 0.0
  x1 = 1.0
  imax = 0
  for i ∈ 1:m
    x0 = x0 + gᵢ[i] * Δy
    x1 = max(x1, gᵢ[i])
    if abs(x1 - gᵢ[i]) < 1e-9
      imax = i
    end
  end
  @printf "    %4d s |" t
  @printf " mass %10.3e  max %10.3e  imax %3d" x0 x1 imax
  @printf "\n"

  local elapsed = CPUtoq()
  global total_runtime += elapsed
  CPUtic()
end
@printf "Total time - %5.2f seconds\n" total_runtime

println("End; press any key to close.")
xxx = readline()


module argminUniProb
  export argminProb
  using Statistics

  """
      argminProb( r::Matrix{Float64} )

      compute minimum value per row, if several cols hold the minimum value, randomly choose among them.
      #Arguments
        - `r::Matrix{Float64}`: risk matrix each column represents one hypothesis.
      #Return
        - `Vector{Float64}`:    indices of maximal values.
  """
  function argminProb( r::Matrix{Float64} )
    index = Vector{Int64}( undef, size( r, 1 ) );
    for i in 1:1:size( index, 1 )
      index[i] = argminProb( r[i,:] );
    end
    return index;
  end

  """
      argminProb!( r::Matrix{Float64}, v::Vector{Int64} )

      compute minimum value per row, if several cols have maximum value randomly choose among them.
      #Arguments
        - `r::Matrix{Float64}`: risk matrix each column represents one hypothesis.
        - `v::Vector{Int64} `:  resulting vector.
      #Return
        - `nothing`:            nothing.
  """
  function argminProb!( r::Matrix{Float64}, v::AbstractVector )
    @assert size( v,1 ) == size( r, 2 );
    for i in 1:1:size( v, 1 )
      v[i] = argminProb( r[i,:] );
    end
  end

  """
      argminProb( values::Vector{Float64} )

      compute minimum, if several indices have minimum value randomly choose among them.
      #Arguments
        - `values::Vector{Float64}`: values to compute argmin on.
      #Return
        - `Vector{Float64}`:         indices of minimal values.
  """
  function argminProb( values::AbstractVector )
    minPositions = collect( 1:1:size( values, 1 ) )[ values .== minimum( values ) ];
    if size( minPositions, 1 ) > 1
      res = minPositions[ rand( 1:size( minPositions, 1 ) ) ];
    else
      res = minPositions[1];
    end
    return res;
  end
end


"""
     sys = series(sys1, sys2)
     sys = sys2*sys1

Connect the descriptor systems `sys1` and `sys2` in series such that `sys = sys2*sys1`.
This coupling corresponds to the multiplication of their transfer function matrices.
Series coupling of systems with constant matrices and vectors having suitable dimensions
or with UniformScalings is also supported.
Series coupling with a constant is equivalent to elementwise multiplication of
the transfer function matrix with the constant.
"""
series(sys1::DescriptorStateSpace, sys2::DescriptorStateSpace) = sys2*sys1
series(sys1::DescriptorStateSpace, sys2::Union{AbstractNumOrArray,UniformScaling}) = sys2*sys1
series(sys1::Union{AbstractNumOrArray,UniformScaling}, sys2::DescriptorStateSpace) = sys2*sys1

"""
    sys = parallel(sys1, sys2)
    sys = sys1 + sys2

Connect the descriptor systems `sys1` and `sys2` in parallel such that `sys = sys1 + sys2`.
This coupling corresponds to the addition of their transfer function matrices.
Parallel coupling of systems with constant matrices or vectors having the same row and column dimensions
or with UniformScalings is also supported.
Parallel coupling with a constant is equivalent to elementwise parallel coupling of
the transfer function matrix with the constant.
"""
parallel(sys1::DescriptorStateSpace, sys2::DescriptorStateSpace) = sys1 + sys2
parallel(sys1::DescriptorStateSpace, sys2::Union{AbstractNumOrArray,UniformScaling}) = sys1 + sys2
parallel(sys1::Union{AbstractNumOrArray,UniformScaling},sys2::DescriptorStateSpace) = sys1 + sys2
"""
    sys = append(systems...)

Append the descriptor systems `systems` by concatenating the input and output vectors
of individual systems. This corresponds to the block diagonal concatenation of
their transfer function matrices.
Appending systems with constant matrices, vectors or scalars or with UniformScalings is also supported.
"""
function append(systems::(DST where DST<:DescriptorStateSpace)...)
    T = promote_type(eltype.(systems)...)
    Ts = systems[1].Ts
    if !all(s.Ts == Ts for s in systems)
        error("All systems must have same sampling time")
    end
    A = blockdiag([s.A for s in systems]...)
    if all(s.E == I for s in systems)
        E = I
    elseif any(s.E == I for s in systems)
        blockdims = Int[size(s.A,1) for s in systems]
        E = sblockdiag(blockdims, [s.E for s in systems]...)
    else
        E = blockdiag([s.E for s in systems]...)
    end

    B = blockdiag([s.B for s in systems]...)
    C = blockdiag([s.C for s in systems]...)
    D = blockdiag([s.D for s in systems]...)
    return DescriptorStateSpace{T}(A, E, B, C, D, Ts)
end

function append(A::Union{DescriptorStateSpace,AbstractNumOrArray,UniformScaling}...)
    for a in A
        isa(a, UniformScaling) && @warn "All UniformScaling objects in append are set to scalars"
        require_one_based_indexing(a)
    end
    Ts = promote_system_SamplingTime(A...)
    return append(promote_to_systems(Ts, fill(1,length(A)), 1, promote_type(eltype.(A)...), A...)...)
end

function hcat(SYS1 :: DescriptorStateSpace, SYS2 :: DescriptorStateSpace)
    ny = SYS1.ny
    ny == size(SYS2, 1) ||  error("The systems must have the same output dimension")
    T = promote_type(eltype(SYS1), eltype(SYS2))
    Ts = promote_Ts(SYS1.Ts, SYS2.Ts)

    A = blockdiag(T.(SYS1.A), T.(SYS2.A))

    if SYS1.E == I && SYS2.E == I
        E = I
    elseif SYS1.E == I || SYS2.E == I
        blockdims = [size(SYS1.A,1), size(SYS2.A,1)]
        E = sblockdiag(blockdims, SYS1.E == I ? SYS1.E : T.(SYS1.E), SYS2.E == I ? SYS2.E : T.(SYS2.E))
    else
        E = blockdiag(T.(SYS1.E), T.(SYS2.E))
    end
    B = blockdiag(T.(SYS1.B), T.(SYS2.B))
    C = [ T.(SYS1.C) T.(SYS2.C)]
    D = [ T.(SYS1.D) T.(SYS2.D)]
    return DescriptorStateSpace{T}(A, E, B, C, D, Ts)
end
#hcat(SYS :: DescriptorStateSpace, MAT :: AbstractNumOrArray) = hcat(SYS,dss(MAT,Ts=SYS.Ts))
hcat(SYS :: DescriptorStateSpace{T} where T, MAT :: Union{Number, AbstractVecOrMat{<:Number}}) = hcat(SYS,dss(MAT,Ts=SYS.Ts))
#hcat(MAT :: AbstractNumOrArray, SYS :: DescriptorStateSpace) = hcat(dss(MAT,Ts=SYS.Ts),SYS)
hcat(MAT :: Union{Number, AbstractVecOrMat{<:Number}}, SYS :: DescriptorStateSpace) = hcat(dss(MAT,Ts=SYS.Ts),SYS)
hcat(SYS :: DescriptorStateSpace, MAT :: UniformScaling) = hcat(SYS,dss(Matrix{promote_type(eltype(SYS),eltype(MAT))}(MAT,SYS.ny,SYS.ny),Ts=SYS.Ts))
hcat(MAT :: UniformScaling, SYS :: DescriptorStateSpace) = hcat(dss(Matrix{promote_type(eltype(SYS),eltype(MAT))}(MAT,SYS.ny,SYS.ny),Ts=SYS.Ts),SYS)

function isadss(DST :: Union{DescriptorStateSpace,AbstractNumOrArray,UniformScaling}...)
    # pick the index i of first system
    for i = 1:length(DST)
        DST[i] isa DescriptorStateSpace && (return i)
    end
    # set index to 0 if no system is present
    return 0
end

"""
    sys = horzcat(sys1,sys2)
    sys = [sys1 sys2]
    sys = horzcat(systems...)

Concatenate horizontally two systems `sys1` and `sys2` or several descriptor systems `systems...`
by concatenating the input vectors of individual systems. This corresponds to the horizontal
concatenation of their transfer function matrices.
Concatenation of systems with constant matrices, vectors, or scalars having the same row dimensions
or with UniformScalings is also supported.
"""
horzcat(systems::Union{DST,AbstractNumOrArray,UniformScaling}...) where DST <: DescriptorStateSpace = hcat(systems...)
function Base.hcat(systems::DST...) where DST <: DescriptorStateSpace
    # Perform checks
    T = promote_type(eltype.(systems)...)
    Ts = systems[1].Ts
    !all(s.Ts == Ts for s in systems) && error("All systems must have the same sampling time")
    ny = systems[1].ny
    !all(s.ny == ny for s in systems) && error("All systems must have the same output dimension")
    A = blockdiag([s.A for s in systems]...)
    if all(s.E == I for s in systems)
        E = I
    elseif any(s.E == I for s in systems)
        blockdims = Int[size(s.A,1) for s in systems]
        E = sblockdiag(blockdims, [s.E for s in systems]...)
    else
        E = blockdiag([s.E for s in systems]...)
    end
    B = blockdiag([s.B for s in systems]...)
    C = hcat([s.C for s in systems]...)
    D = hcat([s.D for s in systems]...)

    return DescriptorStateSpace{T}(A, E, B, C, D, Ts)
end

"""
    sys = vertcat(sys1,sys2)
    sys = [sys1; sys2]
    sys = vertcat(systems...)

Concatenate vertically two descriptor systems `sys1` and `sys2` or several descriptor systems `systems...`
by concatenating the output vectors of individual systems. This corresponds to the vertical
concatenation of their transfer function matrices.
Concatenation of systems with constant matrices, vectors, or scalars having the same column dimensions
or with UniformScalings is also supported.
"""
vertcat(systems::Union{DST,AbstractNumOrArray,UniformScaling}...) where DST <: DescriptorStateSpace = vcat(systems...)

function Base.vcat(systems::DST...) where DST <: DescriptorStateSpace
    # Perform checks
    T = promote_type(eltype.(systems)...)
    Ts = systems[1].Ts
    !all(s.Ts == Ts for s in systems) && error("All systems must have the same sampling time")
    nu = systems[1].nu
    !all(s.nu == nu for s in systems) && error("All systems must have the same input dimension")
    A = blockdiag([s.A for s in systems]...)
    if all(s.E == I for s in systems)
        E = I
    elseif any(s.E == I for s in systems)
        blockdims = Int[size(s.A,1) for s in systems]
        E = sblockdiag(blockdims, [s.E for s in systems]...)
    else
        E = blockdiag([s.E for s in systems]...)
    end
    B = vcat([s.B for s in systems]...)
    C = blockdiag([s.C for s in systems]...)
    D = vcat([s.D for s in systems]...)
    return DescriptorStateSpace{T}(A, E, B, C, D, Ts)
end

function vcat(SYS1 :: DescriptorStateSpace, SYS2 :: DescriptorStateSpace)
    nu = SYS1.nu
    nu == size(SYS2, 2) ||  error("The systems must have the same input dimension")
    T = promote_type(eltype(SYS1), eltype(SYS2))
    Ts = promote_Ts(SYS1.Ts, SYS2.Ts)

    A = blockdiag(T.(SYS1.A), T.(SYS2.A))

    if SYS1.E == I && SYS2.E == I
        E = I
    elseif SYS1.E == I || SYS2.E == I
        blockdims = [size(SYS1.A,1), size(SYS2.A,1)]
        E = sblockdiag(blockdims, SYS1.E == I ? SYS1.E : T.(SYS1.E), SYS2.E == I ? SYS2.E : T.(SYS2.E))
    else
        E = blockdiag(T.(SYS1.E), T.(SYS2.E))
    end
    C = blockdiag(T.(SYS1.C), T.(SYS2.C))
    B = [ T.(SYS1.B); T.(SYS2.B)]
    D = [ T.(SYS1.D); T.(SYS2.D)]
    return DescriptorStateSpace{T}(A, E, B, C, D, Ts)
end

vcat(SYS :: DescriptorStateSpace, MAT :: Union{Number, AbstractVecOrMat{<:Number}}) = vcat(SYS,dss(MAT,Ts=SYS.Ts))
#vcat(MAT :: AbstractNumOrArray, SYS :: DescriptorStateSpace) = vcat(dss(MAT,Ts=SYS.Ts),SYS)
vcat(MAT :: Union{Number, AbstractVecOrMat{<:Number}}, SYS :: DescriptorStateSpace) = vcat(dss(MAT,Ts=SYS.Ts),SYS)
vcat(SYS :: DescriptorStateSpace, MAT :: UniformScaling) = vcat(SYS,dss(Matrix{promote_type(eltype(SYS),eltype(MAT))}(MAT,SYS.nu,SYS.nu),Ts=SYS.Ts))
vcat(MAT :: UniformScaling, SYS :: DescriptorStateSpace) = vcat(dss(Matrix{promote_type(eltype(SYS),eltype(MAT))}(MAT,SYS.nu,SYS.nu),Ts=SYS.Ts),SYS)

for (f, _f, dim, name) in ((:hcat, :_hcat, 1, "rows"), (:vcat, :_vcat, 2, "cols"))
    @eval begin
        @inline $f(A::Union{DescriptorStateSpace, AbstractVecOrMat,UniformScaling,Number}...) = $_f(A...)
        function $_f(A::Union{DescriptorStateSpace, AbstractVecOrMat,UniformScaling,Number}...)
            n = -1
            for a in A
                if !isa(a, UniformScaling)
                    require_one_based_indexing(a)
                    na = size(a,$dim)
                    n >= 0 && n != na &&
                        throw(DimensionMismatch(string("number of ", $name,
                            " of each array must match (got ", n, " and ", na, ")")))
                    n = na
                end
            end
            n == -1 && throw(ArgumentError($("$f of only UniformScaling objects cannot determine the matrix size")))
            if isnothing(findfirst(a -> isa(a, DescriptorStateSpace), A))
                @warn "Type piracy in $(string($f)): to be fixed in Julia 1.8 (make an issue otherwise)"
                # alleviate type piracy problematic
                n == 1 && (A = promote_to_arrays(A...))  # convert all scalars to vectors
                return cat(LinearAlgebra.promote_to_arrays(fill(n, length(A)), 1, Matrix, A...)..., dims=Val(3-$dim))
            else
                Ts = promote_system_SamplingTime(A...)
                return $f(promote_to_systems(Ts, fill(n,length(A)), 1, promote_type(eltype.(A)...), A...)...)
            end
        end
    end
end

function Base.hvcat(rows :: Tuple{Vararg{Int}}, DST :: DescriptorStateSpace...)
    j2 = rows[1]
    sys = hcat(DST[1:j2]...)
    for i = 2:length(rows)
        j1 = j2+1
        j2 = j2+rows[i]
        sys = [sys; hcat(DST[j1:j2]...)]
    end
    return sys
end

#function Base.hvcat(rows::Tuple{Vararg{Int}}, A::Union{DescriptorStateSpace, AbstractVecOrMat{<:Number}, Number, UniformScaling}...)
Base.hvcat(rows::Tuple{Vararg{Int}}, A::Union{DescriptorStateSpace, AbstractVecOrMat, UniformScaling, Number}...) = _hvcat(rows, A...)
function _hvcat(rows::Tuple{Vararg{Int}}, A::Union{DescriptorStateSpace, AbstractVecOrMat, UniformScaling, Number}...)
    require_one_based_indexing(A...)
    nr = length(rows)
    sum(rows) == length(A) || throw(ArgumentError("mismatch between row sizes and number of arguments"))
    n = fill(-1, length(A))
    needcols = false # whether we also need to infer some sizes from the column count
    j = 0
    for i = 1:nr # infer UniformScaling sizes from row counts, if possible:
        ni = -1 # number of rows in this block-row, -1 indicates unknown
        for k = 1:rows[i]
            if !isa(A[j+k], UniformScaling)
                isa(A[j+k],Number) ? na = 1 : na = size(A[j+k], 1)
                ni >= 0 && ni != na &&
                    throw(DimensionMismatch("mismatch in number of rows"))
                ni = na
            end
        end
        if ni >= 0
            for k = 1:rows[i]
                n[j+k] = ni
            end
        else # row consisted only of UniformScaling objects
            needcols = true
        end
        j += rows[i]
    end
    if needcols # some sizes still unknown, try to infer from column count
        nc = -1
        j = 0
        for i = 1:nr
            nci = 0
            rows[i] > 0 && n[j+1] == -1 && (j += rows[i]; continue)
            for k = 1:rows[i]
                nci += isa(A[j+k], UniformScaling) ? n[j+k] : (isa(A[j+k],Number) ? 1 : size(A[j+k], 2))
            end
            nc >= 0 && nc != nci && throw(DimensionMismatch("mismatch in number of columns"))
            nc = nci
            j += rows[i]
        end
        nc == -1 && throw(ArgumentError("sizes of UniformScalings could not be inferred"))
        j = 0
        for i = 1:nr
            if rows[i] > 0 && n[j+1] == -1 # this row consists entirely of UniformScalings
                # nci = nc ÷ rows[i]
                # nci * rows[i] != nc && throw(DimensionMismatch("indivisible UniformScaling sizes"))
                nci, r = divrem(nc, rows[i])
                r != 0 && throw(DimensionMismatch("indivisible UniformScaling sizes"))
                for k = 1:rows[i]
                    n[j+k] = nci
                end
            end
            j += rows[i]
        end
    end
    if isnothing(findfirst(a -> isa(a, DescriptorStateSpace), A))
        # alleviate type piracy problematic
        # convert all scalars to vectors: not needed in Julia 1.8
        @warn "Type piracy in hvcat: to be fixed in Julia 1.8 (make an issue otherwise)"
        any(n .== 1) && (A = promote_to_arrays(A...))
        Amat = LinearAlgebra.promote_to_arrays(n, 1, Matrix, A...)
        return Base.typed_hvcat(promote_type(eltype.(Amat)...), rows, Amat...)
    else
        Ts = promote_system_SamplingTime(A...)
        return hvcat(rows, promote_to_systems(Ts, n, 1, promote_type(eltype.(A)...), A...)...)
    end
end

# promotion to systems of constant matrices, vectors, scalars and UniformScalings
promote_to_system_(n::Int, ::Type{T}, J::UniformScaling, Ts::Real) where {T} = dss(copyto!(Matrix{T}(undef, n,n), J), Ts = Ts)
promote_to_system_(n::Int, ::Type{T}, A::AbstractNumOrArray, Ts::Real) where {T} = dss(to_matrix(T,A), Ts = Ts)
promote_to_system_(n::Int, ::Type{T}, A::DescriptorStateSpace, Ts::Real) where {T} = T == eltype(A) ? A : dss(dssdata(T,A)..., Ts = Ts)
promote_to_systems(Ts::Real, n, k, ::Type) = ()
promote_to_systems(Ts::Real, n, k, ::Type{T}, A) where {T} = (promote_to_system_(n[k], T, A, Ts),)
promote_to_systems(Ts::Real, n, k, ::Type{T}, A, B) where {T} =
    (promote_to_system_(n[k], T, A, Ts), promote_to_system_(n[k+1], T, B, Ts))
promote_to_systems(Ts::Real, n, k, ::Type{T}, A, B, C) where {T} =
    (promote_to_system_(n[k], T, A, Ts), promote_to_system_(n[k+1], T, B, Ts), promote_to_system_(n[k+2], T, C, Ts))
promote_to_systems(Ts::Real, n, k, ::Type{T}, A, B, Cs...) where {T} =
    (promote_to_system_(n[k], T, A, Ts), promote_to_system_(n[k+1], T, B, Ts), promote_to_systems(Ts, n, k+2, T, Cs...)...)
promote_to_system_type(A::Tuple{Vararg{Union{DescriptorStateSpace,AbstractNumOrArray,UniformScaling}}}) = DescriptorStateSpace
promote_to_systems(Ts::Union{Real,Nothing}, var::Symbol, n, k, ::Type) = ()
# promote_to_array_type(A::Tuple{Vararg{Union{AbstractVecOrMat,UniformScaling,Number}}}) = Matrix

function promote_to_arrays(A...)
    N = length(A)
    B = (Vector{T} where T)(undef, N)
    for i = 1:N
        isa(A[i],Number) ? (B[i] = [A[i]]) : B[i] = A[i]
    end
    return tuple(B...)
end

function promote_system_SamplingTime(A::Union{DescriptorStateSpace,AbstractNumOrArray,UniformScaling}...)
    # pick and check the common sampling time
    Ts = nothing
    for a in A
        typeof(a) <: DescriptorStateSpace  && (isnothing(Ts) ? Ts = a.Ts : Ts = promote_Ts(Ts,a.Ts))
    end
    return isnothing(Ts) ? (return 0) : (return Ts) # for systems use Ts = 0 as defualt
end


s2int(s) = parse(Int, s)

function parse_board(s)
    lines = split(s, "\n")
    grid = map(s2int, hcat(map(split, lines)...))
    [map(Set, eachrow(grid)); map(Set, eachcol(grid))]
end

function slurp_bingo(filename)
    lines = readlines(filename)
    numbers = map(s2int, split(lines[1], ','))
    boards_string = join(lines[3:end], "\n")
    boards = split(boards_string, "\n\n")
    numbers, map(parse_board, boards)
end

function solve(filename)
    numbers, boards = slurp_bingo(filename)
    for number in numbers
        for board in boards
            broadcast(row -> delete!(row, number), board)
            if min(map(length, board)...) == 0
                return reduce(+, union(board...)) * number
            end
        end
    end
end

solve("4eg.txt")
solve("4.txt")

score(board, number) = reduce(+, union(board...)) * number
completed(board) = min(length.(board)...) == 0
remove(board, number) = broadcast(row -> delete!(row, number), board)

function solve2(filename)
    numbers, boards = slurp_bingo(filename)
    for number in numbers
        for board in boards
            remove(board, number)
            if all(completed, boards)
                return score(board, number)
            end
        end
    end
end

solve2("4eg.txt")
solve2("4.txt") == 2980


module sptests

using Base.Test
using TypedPolynomials: StaticPolynomials

macro wrappedallocs(expr)
    @assert expr.head == :call
    f = expr.args[1]
    args = expr.args[2:end]
    argnames = [gensym() for a in args]
    quote
        function g($(argnames...))
            @allocated $(Expr(:call, esc(f), argnames...))
        end
        $(Expr(:call, :g, [esc(a) for a in args]...))
    end
end

@testset "variables" begin
    @polyvar x y

    @test x * x == x^2
    @test x^1 == x
    @test x^0 == 1
    @test (@wrappedallocs x * x) == 0
    @test (@wrappedallocs x^2) == 0
    @test (@wrappedallocs x^1) == 0
#     @test (@wrappedallocs x^0) == 0
end

@testset "orderings" begin
    @polyvar x y z

#     @test Monomial(y) < Monomial(x)
#     @test Monomial(x) > Monomial(y)
#     @test Monomial(y) < x
#     @test y < Monomial(x)
    @test y < x
    @test x > y
#     @test x > Monomial(y)
#     @test Monomial(x) > y
end

@testset "monomials" begin
    @polyvar x y z

    @test typeof(@inferred x * x) == Monomial{(x,), (2,)}
    @test typeof(@inferred x * y) == Monomial{(x, y), (1, 1)}
    @test typeof(@inferred y * x) == Monomial{(x, y), (1, 1)}
    @test typeof(@inferred x * z * y) == Monomial{(x, y, z), (1, 1, 1)}
end

# include("substitution.jl")

end


"""
    function_to_interpolated_function(f_xy; physwidth = 10.0m, physheight = 4.0m)

Input: f: (Q, Q)  → (Q, Q) Function taking a tuple quantities, outputs a tuple of quantities
    physwidth:  Q
    physheight: Q

Output: f: (Q, Q)  → (Q, Q) Function taking a tuple quantities, outputs a tuple of quantities

Convert the function f_xy into an interpolated function. The output function is linearly interpolated
between pixels, based on the currently defined sketch scale.

The output of the function can be otherwise; this is just the originally intended use.
"""
function function_to_interpolated_function(f_xy; physwidth = 10.0m, physheight = 4.0m, centered = true)
    # Position iterators
    xs, ys = x_y_iterators_at_pixels(;physwidth, physheight, centered)
    tuplemat = [f_xy(x, y) for x in xs, y in reverse(ys)]
    fxy_inter = interpolate((xs, reverse(ys)),
        tuplemat,
        Gridded(Linear()));
    # Extend the domain, using the same values as on the domain border. Don't know how the closest point
    # on domain boundary is found.
    extrapolate(fxy_inter, Flat())
end

"""
    matrix_to_function(matrix::Matrix, xmin, xmax, ymin, ymax, nx, ny)
    -> f: (Q, Q)  → (Q, Q)

Make a linear interpolation function from (complex) matrix.
This could be useful for e.g. a flow field, where the matrix contains complex
quantity values.

The imaginary component is mapped to the second parameter.
"""
function matrix_to_function(matrix::Matrix, xmin, xmax, ymin, ymax, nx, ny)
    xs = range(xmin, xmax, length = nx)
    ys = range(ymin, ymax, length = ny)
    # Adapt to Interpolations
    tuplemat = map( cmplx -> (real(cmplx), imag(cmplx)), transpose(matrix)[ : , end:-1:1])
    # Function that can interpolate between nodes.
    fxy_inter = interpolate((xs, reverse(ys)),
        tuplemat,
        Gridded(Linear()));
    # Extend the domain, using the same values as on the domain border. Don't know how the closest point
    # on domain boundary is found.
    extrapolate(fxy_inter, Flat())
end


# ---
# title: 1000. Minimum Cost to Merge Stones
# id: problem1000
# author: Tian Jun
# date: 2020-10-31
# difficulty: Hard
# categories: Dynamic Programming
# link: <https://leetcode.com/problems/minimum-cost-to-merge-stones/description/>
# hidden: true
# ---
#
# There are `N` piles of stones arranged in a row.  The `i`-th pile has
# `stones[i]` stones.
#
# A _move_ consists of merging **exactly  `K` consecutive** piles into one pile,
# and the cost of this move is equal to the total number of stones in these `K`
# piles.
#
# Find the minimum cost to merge all piles of stones into one pile.  If it is
# impossible, return `-1`.
#
#
#
# **Example 1:**
#
#
#
#     Input: stones = [3,2,4,1], K = 2
#     Output: 20
#     Explanation:
#     We start with [3, 2, 4, 1].
#     We merge [3, 2] for a cost of 5, and we are left with [5, 4, 1].
#     We merge [4, 1] for a cost of 5, and we are left with [5, 5].
#     We merge [5, 5] for a cost of 10, and we are left with [10].
#     The total cost was 20, and this is the minimum possible.
#
#
# **Example 2:**
#
#
#
#     Input: stones = [3,2,4,1], K = 3
#     Output: -1
#     Explanation: After any merge operation, there are 2 piles left, and we can't merge anymore.  So the task is impossible.
#
#
# **Example 3:**
#
#
#
#     Input: stones = [3,5,1,2,6], K = 3
#     Output: 25
#     Explanation:
#     We start with [3, 5, 1, 2, 6].
#     We merge [5, 1, 2] for a cost of 8, and we are left with [3, 8, 6].
#     We merge [3, 8, 6] for a cost of 17, and we are left with [17].
#     The total cost was 25, and this is the minimum possible.
#
#
#
#
# **Note:**
#
#   * `1 <= stones.length <= 30`
#   * `2 <= K <= 30`
#   * `1 <= stones[i] <= 100`
#
#
## @lc code=start
using LeetCode

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


# ---
# title: 40. Combination Sum II
# id: problem40
# author: Tian Jun
# date: 2020-10-31
# difficulty: Medium
# categories: Array, Backtracking
# link: <https://leetcode.com/problems/combination-sum-ii/description/>
# hidden: true
# ---
#
# Given a collection of candidate numbers (`candidates`) and a target number
# (`target`), find all unique combinations in `candidates` where the candidate
# numbers sum to `target`.
#
# Each number in `candidates` may only be used **once** in the combination.
#
# **Note:**  The solution set must not contain duplicate combinations.
#
#
#
# **Example 1:**
#
#
#
#     Input: candidates = [10,1,2,7,6,1,5], target = 8
#     Output:
#     [
#     [1,1,6],
#     [1,2,5],
#     [1,7],
#     [2,6]
#     ]
#
#
# **Example 2:**
#
#
#
#     Input: candidates = [2,5,2,1,2], target = 5
#     Output:
#     [
#     [1,2,2],
#     [5]
#     ]
#
#
#
#
# **Constraints:**
#
#   * `1 <= candidates.length <= 100`
#   * `1 <= candidates[i] <= 50`
#   * `1 <= target <= 30`
#
#
## @lc code=start
using LeetCode

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


using DynamicalSystems
using Base.Test

println("\nTesting continuous system evolution...")

@testset "Lorenz System" begin

  lo11 = Systems.lorenz()
  lo22 = ContinuousDS(lo11.state, lo11.eom)
  lo33 = Systems.lorenz(big.([0.0, 10.0, 0.0]))
  @testset "Construction" begin
    @test typeof(lo11) <: ContinuousDS
    @test typeof(lo33) <: ContinuousDS
  end
  @testset "Evolve & kwargs" begin
    # test evolve(state):
    st1 = evolve(lo11.state, lo11, 1.0)
    st3 = evolve(lo33.state, lo33, 1.0)
    @test st1 ≈ st3
    # Test evolve(state,keywords):
    st1 = evolve(lo11.state, lo11, 1.0; diff_eq_kwargs=Dict(:abstol=>1e-9, :reltol=>1e-9))
    st3 = evolve(lo33.state, lo33, 1.0, diff_eq_kwargs=Dict(:abstol=>1e-9, :reltol=>1e-9))
    @test st1 ≈ st3
  end

  @testset "Timeseries" begin
    # timeseries pure:
    ts1 = timeseries(lo11, 2.0; mutate = false)
    ts3 = timeseries(lo33, 2.0; mutate = false)
    @test eltype(ts3) == BigFloat
    @test size(ts1) == size(ts3)
    @test ts1[end, :] ≈ ts3[end,:]
    # timeseries with diff_eq_kwargs and dt:
    ts1 = timeseries(lo11, 2.0; dt=0.1,
    diff_eq_kwargs=Dict(:abstol=>1e-9, :reltol=>1e-9), mutate = false)
    ts3 = timeseries(lo33, 2.0; dt=0.1,
    diff_eq_kwargs=Dict(:abstol=>1e-9, :reltol=>1e-9), mutate = false)
    @test size(ts1) == size(ts3)
    @test ts1[end, :] ≈ ts3[end,:]
  end

  @testset "Jacobians" begin
    j1 = jacobian(lo11); j2 = jacobian(lo22); j3 = jacobian(lo33)
    @test eltype(j3) == BigFloat
    @test j1 ≈ j2
    @test j1 ≈ j3
    s1 = evolve(lo11, 1.0)
    s2 = evolve(lo22, 1.0)
    s3 = evolve(lo33, 1.0)
    j1 = lo11.jacob(s1); j2 = lo22.jacob(s2); j3 = lo33.jacob(s3)
    @test eltype(j3) == BigFloat
    @test j1 ≈ j2
    @test j1 ≈ j3
  end

end


using DistributionVectors
using Test, Distributions, LinearAlgebra, Missings
using BenchmarkTools: @btime
using Plots, StatsPlots

function benchmarkSums()
    nrep = 30
    mu = log.(rand(Normal(120, 10), nrep));
    sigma = log.(rand(Normal(1.2, 0.1), nrep));
    acf1 = [0.8,0.3,0.1];
    dv = SimpleDistributionVector(LogNormal{eltype(mu)}, mu, sigma)
    @btime dsum_v = sum($dv, $acf1, SkipMissing(); method = Val(:vector) )
    @btime dsum_m = sum($dv, $acf1, SkipMissing(); method = Val(:bandedmatrix) )
    # the Bandematrix based slower after optimizing the vector for stability in missings
    #
    # repeat with missings
    mum = allowmissing(copy(mu)); mum[1] = missing
    dv = SimpleDistributionVector(LogNormal{nonmissingtype(eltype(mum))}, mum, sigma)
    @btime dsum_v = sum($dv, $acf1, SkipMissing(); method = Val(:vector) )
    @btime dsum_m = sum($dv, $acf1, SkipMissing(); method = Val(:bandedmatrix) )
    #
    #S = similar(mum);
    @code_warntype sum_lognormals(S, dv, acf1, SkipMissing())
    f() = sum_lognormals(S, dv, acf1, SkipMissing())
    @time f(); @time f()

    # # try allocating instead of view (replace line of view_nonmissing)
    # # allocating is faster
    # function sum_lognormals2!(S, dv, corr::AbstractMatrix,
    # ms::MissingStrategy=PassMissing())
    #     parms = params(dv)
    #     μ = @view parms[1,:]
    #     σ = @view parms[2,:]
    #     # S = allowmissing(similar(μ))
    #     @. S = exp(μ + abs2(σ)/2)
    #     nmissing = count(ismissing, S)
    #     anymissing = nmissing != 0
    #     !(isa(ms, HandleMissingStrategy)) && anymissing && error(
    #          "Found missing values. Use argument 'SkipMmissing()' to sum over nonmissing.")
    #     Ssum::nonmissingtype(eltype(S)) = sum(skipmissing(S))
    #     @. S = σ * S  # do only after Ssum
    #     # setting S to zero results in summing zero for missing records
    #     # which is the same as filtering both S and corr
    #     anymissing && replace!(S, missing => 0.0)
    #     #s = transpose(disallowmissing(S)) * corr * disallowmissing(S)
    #     #Spure = view_nonmissing(S) # non-allocating
    #     Spure = disallowmissing(S) # allocating
    #     s = transpose(Spure) * corr * Spure
    #     σ2eff = s/abs2(Ssum)
    #     μ_sum = log(Ssum) - σ2eff/2
    #     #@show Ssum, s, length(S) - nmissing
    #     LogNormal(μ_sum, √σ2eff)
    # end
    # @btime sum_lognormals($storage, $dv, $corMa, SkipMissing())
    # @btime sum_lognormals2!($storage, $dv, $corMa, SkipMissing())
end

function bootstrap_sums_lognormal()
    #using StatsBase
    nObs = 5
    xTrue = fill(10.0, nObs)
    sigmaStar = fill(1.5, nObs) # multiplicative stddev of 1.5
    dv = SimpleDistributionVector(fit.(LogNormal, xTrue, Σstar.(sigmaStar))...)
    dsum = sum(dv)
    nboot = 100_000
    x = rand(dv, nboot);
    sums = vec(sum(x, dims = 2));
    #stds = std.(Ref(x));
    cdf_sums = ecdf(sums);
    @test cdf_sums(median(dsum)) ≈ 0.5 rtol = 0.02
    @test cdf_sums(quantile(dsum, [0.025,0.975])) ≈ [0.025, 0.975] rtol = 0.02
end

function fplot(dsum, sums)
    plot(dsum, lab="computed", xlabel="sum", ylabel="density")
    density!(sums, lab="random sample")
    vline!([mean(dsum)], lab="mean computed")
    vline!([mean(sums)], lab="mean random")
    vline!(quantile(dsum, [0.025, 0.975]), lab="cf computed")
    vline!(quantile(sums, [0.025, 0.975]), lab="cf random")
end

function bootstrap_sums_lognormal_acf()
    #using StatsBase
    mu = log.([110,100,80,120,160.0])
    sigma = log.([1.2,1.5,1.1,1.3,1.1])
    acf1 = [0.4,0.1]
    dv = SimpleDistributionVector(LogNormal{eltype(mu)}, mu, sigma);
    #ps = 0.025:0.025:0.975; qs = quantile.(dv[1], ps); plot(qs, pdf(dv[1],qs))
    drsum = boot_dvsums_acf(dv, acf1)
    dsum = sum(dv, acf1)
    @test dsum ≈ drsum rtol = 0.02
    # see plotting function in bootstrap_sums_lognormal above
end


# Copyright (c) 2022 Code Komali
#
# This software is released under the MIT License.
# https://opensource.org/licenses/MIT

module HKOpinionDynamicsExample

export hk_model_run_and_plot!

using Agents
using Statistics: mean
using Plots
using Random

mutable struct HKAgent <: AbstractAgent
    id::Int
    old_opinion::Float64
    new_opinion::Float64
    previous_opinion::Float64
end

function hk_model(; numagents = 100, ϵ = 0.2)
    model = ABM(HKAgent, scheduler = Schedulers.fastest, properties= Dict(:ϵ => ϵ))
    Random.seed!(1234)
    for i in 1:numagents
        o = rand()
        add_agent!(model, o, o, -1)
    end
    return model
end

function boundfilter(agent, model)
    filter(
        j -> abs(agent.old_opinion -j) < model.ϵ,
        [a.old_opinion for a in allagents(model)]
    )
end

function agent_step!(agent, model)
    agent.previous_opinion = agent.old_opinion
    agent.new_opinion = mean(boundfilter(agent, model))
end

function model_step!(model)
    for a in allagents(model)
        a.old_opinion = a.new_opinion
    end
end

function terminate(model, s)
    if any(
        !isapprox(a.previous_opinion, a.new_opinion; rtol= 1e-12) for a in allagents(model)
    )
        return false
    else
        return true
    end
end

function model_run(; kwargs...)
    model = hk_model(; kwargs...)
    agent_data,_ = run!(
        model,
        agent_step!,
        model_step!,
        terminate;
        adata = [:new_opinion],
        # make sure to comment the 'when' otherwise you won't get the graph as expected
        #when = terminate # when to collect data, the default is to collect on every step
    )
    return agent_data
end

#data = model_run()

plotsim(data, ϵ) = plot(
    data.step,
    data.new_opinion,
    leg = false,
    group = data.id,
    title = "epsilon = $(ϵ)",
)

function hk_model_run_and_plot!(eps_list=[0.05, 0.15, 0.3] ; kwargs...)
    plots = []
    alldata = []
    for eps in eps_list
        data = model_run(; ϵ = eps)
        push!(alldata, data)
        plt = plotsim(data, eps)
        push!(plots, plt)
    end
    plot(plots..., layout=(3,1))
    savefig("HKOpinionDynamicsPlot.png")
    return alldata
end

#hk_model_run_and_plot!()

end


#
# http://en.wikipedia.org/wiki/Allen%27s_interval_algebra
#

function before{T1, T2, N}(b1::HyperRectangle{N, T1}, b2::HyperRectangle{N, T2})
    for i = 1:N
        maximum(b1)[i] < minimum(b2)[i] || return false
    end
    true
end

@inline function meets{T1, T2, N}(b1::HyperRectangle{N, T1}, b2::HyperRectangle{N, T2})
    maximum(b1) == minimum(b2)
end

function overlaps{T1, T2, N}(b1::HyperRectangle{N, T1}, b2::HyperRectangle{N, T2})
    for i = 1:N
        maximum(b2)[i] > maximum(b1)[i] > minimum(b2)[i] &&
        minimum(b1)[i] < minimum(b2)[i] || return false
    end
    true
end

function starts{T1, T2, N}(b1::HyperRectangle{N, T1}, b2::HyperRectangle{N, T2})
    if minimum(b1) == minimum(b2)
        for i = 1:N
            maximum(b1)[i] < maximum(b2)[i] || return false
        end
        return true
    else
        return false
    end
end

function during{T1, T2, N}(b1::HyperRectangle{N, T1}, b2::HyperRectangle{N, T2})
    for i = 1:N
        maximum(b1)[i] < maximum(b2)[i] && minimum(b1)[i] > minimum(b2)[i] ||
        return false
    end
    true
end

function finishes{T1, T2, N}(b1::HyperRectangle{N, T1}, b2::HyperRectangle{N, T2})
    if maximum(b1) == maximum(b2)
        for i = 1:N
            minimum(b1)[i] > minimum(b2)[i] || return false
        end
        return true
    else
        return false
    end
end

#
# Containment
#

function isinside(rect::SimpleRectangle, x::Real, y::Real)
    rect.x <= x && rect.y <= y && rect.x + rect.w >= x && rect.y + rect.h >= y
end

function isinside(c::Circle, x::Real, y::Real)
    @inbounds ox,oy = origin(c)
    xD = abs(ox - x)
    yD = abs(oy - y)
    xD <= c.r && yD <= c.r
end

"""
contains(s::Simplex, pt; atol=0., rtol=defaulrtol(eltype(pt)))

Check if a point is contained inside a Simplex. If the intrinsic dimension
of the simplex is smaller then the dimension of the surrounding space
(for example a triangle in 3d) one needs the atol, rtol parameters.
"""
function contains(s::Simplex, pt; atol=0., rtol=Base.rtoldefault(eltype(pt)))
    w = weights(pt, s)
    all(w .>= -atol) || return false  # projection lies outside of s
    pt_proj = vertexmat(s) * w
    return isapprox(pt_proj, pt, atol=atol, rtol=rtol)
end

"""
Check if HyperRectangles are contained in each other. This does not use
strict inequality, so HyperRectangles may share faces and this will still
return true.
"""
function contains{T1, T2, N}(b1::HyperRectangle{N, T1}, b2::HyperRectangle{N, T2})
    for i = 1:N
        maximum(b2)[i] <= maximum(b1)[i] && minimum(b2)[i] >= minimum(b1)[i] ||
        return false
    end
    true
end

"""
Check if a point is contained in a HyperRectangle. This will return true if
the point is on a face of the HyperRectangle.
"""
function contains{T, N}(b1::HyperRectangle{N, T}, pt::Union{FixedVector, AbstractVector})
    for i = 1:N
        pt[i] <= maximum(b1)[i] && pt[i] >= minimum(b1)[i] || return false
    end
    true
end

"""
Check if HyperRectangles are contained in each other. This does not use
strict inequality, so HyperRectangles may share faces and this will still
return true.
"""
in(b1::HyperRectangle, b2::HyperRectangle) = contains(b2, b1)

"""
Check if a point is contained in a HyperRectangle. This will return true if
the point is on a face of the HyperRectangle.
"""
in(pt::Union{FixedVector, AbstractVector}, b1::HyperRectangle) = contains(b1, pt)

#
# Equality
#

==(a::AbstractMesh, b::AbstractMesh) = false
function =={M <: AbstractMesh}(a::M, b::M)
    for ((ka, va), (kb, vb)) in zip(all_attributes(a), all_attributes(b))
        va != vb && return false
    end
    true
end

(==){T1, T2, N}(b1::HyperRectangle{N, T1}, b2::HyperRectangle{N, T2}) =
    minimum(b1) == minimum(b2) && widths(b1) == widths(b2)

isequal(b1::HyperRectangle, b2::HyperRectangle) = b1 == b2

isless(a::SimpleRectangle, b::SimpleRectangle) = isless(area(a), area(b))


"""
    mae(ŷ, y; agg=mean)

Return the loss corresponding to mean absolute error:

    agg(abs.(ŷ .- y))
"""
mae(ŷ, y; agg=mean) = agg(abs.(ŷ .- y))

"""
    mse(ŷ, y; agg=mean)

Return the loss corresponding to mean square error:

    agg((ŷ .- y).^2)
"""
mse(ŷ, y; agg=mean) = agg((ŷ .- y).^2)

"""
    msle(ŷ, y; agg=mean, ϵ=eps(ŷ))

The loss corresponding to mean squared logarithmic errors, calculated as

    agg((log.(ŷ .+ ϵ) .- log.(y .+ ϵ)).^2)

The `ϵ` term provides numerical stability.
Penalizes an under-estimation more than an over-estimatation.
"""
msle(ŷ, y; agg=mean, ϵ=epseltype(ŷ)) = agg((log.((ŷ .+ ϵ) ./ (y .+ ϵ))).^2)

"""
    huber_loss(ŷ, y; δ=1, agg=mean)

Return the mean of the [Huber loss](https://en.wikipedia.org/wiki/Huber_loss)
given the prediction `ŷ` and true values `y`.

                 | 0.5 * |ŷ - y|^2,            for |ŷ - y| <= δ
    Huber loss = |
                 |  δ * (|ŷ - y| - 0.5 * δ), otherwise
"""
function huber_loss(ŷ, y; agg=mean, δ=ofeltype(ŷ, 1))
   abs_error = abs.(ŷ .- y)
   #TODO: remove dropgrad when Zygote can handle this function with CuArrays
   temp = Zygote.dropgrad(abs_error .<  δ)
   x = ofeltype(ŷ, 0.5)
   agg(((abs_error.^2) .* temp) .* x .+ δ*(abs_error .- x*δ) .* (1 .- temp))
end

"""
    crossentropy(ŷ, y; dims=1, ϵ=eps(ŷ), agg=mean)

Return the cross entropy between the given probability distributions;
calculated as

    agg(-sum(y .* log.(ŷ .+ ϵ); dims=dims))

Cross entropy is tipically used as a loss in multi-class classification,
in which case the labels `y` are given in a one-hot format.
`dims` specifies the dimension (or the dimensions) containing the class probabilities.
The prediction `ŷ` is supposed to sum to one across `dims`,
as would be the case with the output of a [`softmax`](@ref) operation.

Use of [`logitcrossentropy`](@ref) is recomended over `crossentropy` for
numerical stability.

See also: [`Flux.logitcrossentropy`](@ref), [`Flux.binarycrossentropy`](@ref), [`Flux.logitbinarycrossentropy`](@ref)
"""
function crossentropy(ŷ, y; dims=1, agg=mean, ϵ=epseltype(ŷ))
    agg(.-sum(xlogy.(y, ŷ .+ ϵ); dims=dims))
end

"""
    logitcrossentropy(ŷ, y; dims=1, ϵ=eps(ŷ), agg=mean)

Return the crossentropy computed after a [`Flux.logsoftmax`](@ref) operation;
calculated as

    agg(.-sum(y .* logsoftmax(ŷ; dims=dims); dims=dims))

`logitcrossentropy(ŷ, y)` is mathematically equivalent to
[`Flux.Losses.crossentropy(softmax(ŷ), y)`](@ref) but it is more numerically stable.

See also: [`Flux.Losses.crossentropy`](@ref), [`Flux.Losses.binarycrossentropy`](@ref), [`Flux.Losses.logitbinarycrossentropy`](@ref)
"""
function logitcrossentropy(ŷ, y; dims=1, agg=mean)
    agg(.-sum(y .* logsoftmax(ŷ; dims=dims); dims=dims))
end

"""
    binarycrossentropy(ŷ, y; agg=mean, ϵ=eps(ŷ))

Return the binary cross-entropy loss, computed as

    agg(@.(-y*log(ŷ + ϵ) - (1-y)*log(1-ŷ + ϵ)))

The `ϵ` term provides numerical stability.

Typically, the prediction `ŷ` is given by the output of a [`sigmoid`](@ref) activation.

Use of `logitbinarycrossentropy` is recomended over `binarycrossentropy` for numerical stability.

See also: [`Flux.Losses.crossentropy`](@ref), [`Flux.Losses.logitcrossentropy`](@ref), [`Flux.Losses.logitbinarycrossentropy`](@ref)
"""
function binarycrossentropy(ŷ, y; agg=mean, ϵ=epseltype(ŷ))
    agg(@.(-xlogy(y, ŷ+ϵ) - xlogy(1-y, 1-ŷ+ϵ)))
end
# Re-definition to fix interaction with CuArrays.
# CUDA.@cufunc binarycrossentropy(ŷ, y; ϵ=eps(ŷ)) = -y*log(ŷ + ϵ) - (1 - y)*log(1 - ŷ + ϵ)

"""
    logitbinarycrossentropy(ŷ, y; agg=mean)

Mathematically equivalent to
[`Flux.binarycrossentropy(σ(ŷ), y)`](@ref) but is more numerically stable.

See also: [`Flux.Losses.crossentropy`](@ref), [`Flux.Losses.logitcrossentropy`](@ref), [`Flux.Losses.binarycrossentropy`](@ref)
```
"""
function logitbinarycrossentropy(ŷ, y; agg=mean)
    agg(@.((1-y)*ŷ - logσ(ŷ)))
end
# Re-definition to fix interaction with CuArrays.
# CUDA.@cufunc logitbinarycrossentropy(ŷ, y) = (1 - y)*ŷ - logσ(ŷ)

"""
    kldivergence(ŷ, y; agg=mean)

Return the
[Kullback-Leibler divergence](https://en.wikipedia.org/wiki/Kullback%E2%80%93Leibler_divergence)
between the given probability distributions.

KL divergence is a measure of how much one probability distribution is different
from the other.
It is always non-negative and zero only when both the distributions are equal
everywhere.
"""
function kldivergence(ŷ, y; dims=1, agg=mean, ϵ=epseltype(ŷ))
  entropy = agg(sum(xlogx.(y), dims=dims))
  cross_entropy = crossentropy(ŷ, y; dims=dims, agg=agg, ϵ=ϵ)
  return entropy + cross_entropy
end

"""
    poisson_loss(ŷ, y)

# Return how much the predicted distribution `ŷ` diverges from the expected Poisson
# distribution `y`; calculated as `sum(ŷ .- y .* log.(ŷ)) / size(y, 2)`.
REDO
[More information.](https://peltarion.com/knowledge-center/documentation/modeling-view/build-an-ai-model/loss-functions/poisson_loss).
"""
poisson_loss(ŷ, y; agg=mean) = agg(ŷ .- xlogy.(y, ŷ))

"""
    hinge_loss(ŷ, y; agg=mean)

Return the [hinge_loss loss](https://en.wikipedia.org/wiki/Hinge_loss) given the
prediction `ŷ` and true labels `y` (containing 1 or -1); calculated as
`sum(max.(0, 1 .- ŷ .* y)) / size(y, 2)`.

See also: [`squared_hinge_loss`](@ref)
"""
hinge_loss(ŷ, y; agg=mean) = agg(max.(0, 1 .-  ŷ .* y))

"""
    squared_hinge_loss(ŷ, y)

Return the squared hinge_loss loss given the prediction `ŷ` and true labels `y`
(containing 1 or -1); calculated as `sum((max.(0, 1 .- ŷ .* y)).^2) / size(y, 2)`.

See also: [`hinge_loss`](@ref)
"""
squared_hinge_loss(ŷ, y; agg=mean) = agg((max.(0, 1 .- ŷ .* y)).^2)

"""
    dice_coeff_loss(ŷ, y; smooth=1)

Return a loss based on the dice coefficient.
Used in the [V-Net](https://arxiv.org/abs/1606.04797) image segmentation
architecture.
Similar to the F1_score. Calculated as:

    1 - 2*sum(|ŷ .* y| + smooth) / (sum(ŷ.^2) + sum(y.^2) + smooth)
"""
dice_coeff_loss(ŷ, y; smooth=ofeltype(ŷ, 1.0)) = 1 - (2*sum(y .* ŷ) + smooth) / (sum(y.^2) + sum(ŷ.^2) + smooth) #TODO agg

"""
    tversky_loss(ŷ, y; β=0.7)

Return the [Tversky loss](https://arxiv.org/abs/1706.05721).
Used with imbalanced data to give more weight to false negatives.
Larger β weigh recall more than precision (by placing more emphasis on false negatives)
Calculated as:
    1 - sum(|y .* ŷ| + 1) / (sum(y .* ŷ + β*(1 .- y) .* ŷ + (1 - β)*y .* (1 .- ŷ)) + 1)
"""
function tversky_loss(ŷ, y; β=ofeltype(ŷ, 0.7))
    #TODO add agg
    num = sum(y .* ŷ) + 1
    den = sum(y .* ŷ + β*(1 .- y) .* ŷ + (1 - β)*y .* (1 .- ŷ)) + 1
    1 - num / den
end


using Oceananigans
using Oceananigans.ImmersedBoundaries: ImmersedBoundaryGrid, GridFittedBoundary
using Printf
using Plots

solid(x, y, z) = z < 0

underlying_grid = RectilinearGrid(size=128, z=(-0.5, 0.5), topology=(Flat, Flat, Bounded))
grid = ImmersedBoundaryGrid(underlying_grid, GridFittedBoundary(solid))

closure = IsotropicDiffusivity(κ = 1.0)

model_kwargs = (tracers=:c, buoyancy=nothing, velocities=PrescribedVelocityFields())

model = HydrostaticFreeSurfaceModel(; grid=underlying_grid, closure=closure, model_kwargs...)
immersed_model = HydrostaticFreeSurfaceModel(; grid=grid, closure=closure, model_kwargs...)

initial_temperature(x, y, z) = exp(-z^2 / 0.02)
[set!(m, c=initial_temperature) for m in (model, immersed_model)]

z = znodes(model.tracers.c)
c = view(interior(model.tracers.c), 1, 1, :)
c_immersed = view(interior(immersed_model.tracers.c), 1, 1, :)
c_plot = plot(c, z, linewidth = 2, label = "t = 0", xlabel = "Tracer concentration", ylabel = "z")
plot!(c_plot, c_immersed, z, linewidth = 2, label = "t = 0, immersed model", xlabel = "Tracer concentration", ylabel = "z")

diffusion_time_scale = model.grid.Δzᵃᵃᶜ^2 / model.closure.κ.c
stop_time = 100diffusion_time_scale

simulations = [simulation = Simulation(m, Δt = 1e-1 * diffusion_time_scale, stop_time = stop_time) for m in (model, immersed_model)]
[run!(sim) for sim in simulations]

plot!(c_plot, c, z, linewidth = 2, alpha = 0.6, linestyle = :dash, label = @sprintf("Ordinary model, t = %.3e", model.clock.time))
plot!(c_plot, c_immersed, z, linewidth = 3, alpha = 0.6, label = @sprintf("Immersed model, t = %.3e", immersed_model.clock.time))

display(c_plot)


using IdentityMatrix
using Test
using LinearAlgebra
using FillArrays
using MethodInSrc

@testset  "norms" begin
    for ncols in (0, 1, 3, 10)
        for T in (Float64, Int)
            m = Eye{T}(ncols)
            md = Matrix(m)
            for op in (opnorm, norm, x -> norm(x, 1), x -> norm(x, 2), x -> norm(x, Inf), x -> norm(x, -Inf), x -> norm(x, 3//2))
                @test op(m) == op(md)
            end
            @test @isinsrc opnorm(m)
        end
    end
end

@testset "Eye" begin
    for ncols in (1, 2, 5)
        IM = Eye(ncols)
        IMd = Matrix(IM)
        @test isa(imag(IM), Diagonal)
        for f in (imag, isposdef, sum, prod, first, last,
                  minimum, maximum, extrema, triu, triu!, tril, tril!, inv,
                  diag, det, logdet, sqrt)
            @test @insrc(f(IM)) == f(copy(IMd))
        end
        for f in (iszero, )
            @test @ninsrc(f(IM)) == f(copy(IMd))
        end
        for p in -5:5
            @test IM == @insrc IM^p
        end
        @test isa(copy(IM), Diagonal)
        @test isa(@insrc(Matrix(IM)), Matrix)
    end
end

@testset "mul and div" begin
    for ncols in (1, 3 ,10)
        IM = Eye(ncols)
        @test IM === @insrc IM * IM
        @test IM === @insrc IM / IM
        @test IM === @insrc IM \ IM

        rm = rand(ncols, ncols)
        @test rm === @insrc IM * rm
        @test rm === @insrc rm * IM
        @test rm === @insrc rm / IM
        @test inv(rm) == @insrc IM / rm
        @test inv(rm) == @insrc rm \ IM
        @test rm == @insrc IM \ rm
    end
end

@testset "algebra with AbstractFill" begin
    for d in (1, 2, 3, 10)
       for  T1 in (Int, Float64, ComplexF64)
           for T2 in (Int, Float64, ComplexF64)
               @test @insrc(Eye{T1}(d) * Zeros{T2}(d, d)) == Zeros{typeof(one(T1)*one(T2))}(d, d)
               @test @insrc(Eye{T1}(d) * Zeros{T2}(d)) == Zeros{typeof(one(T1)*one(T2))}(d)
               @test @insrc(Eye{T1}(d) * Ones{T2}(d, d)) == Ones{typeof(one(T1)*one(T2))}(d, d)
               @test @insrc(Eye{T1}(d) * Ones{T2}(d)) == Ones{typeof(one(T1)*one(T2))}(d)
               for val in (1, 2)
                   @test @insrc(Eye{T1}(d) * Fill{T2}(val, d, d)) == Fill{typeof(one(T1)*one(T2))}(val, d, d)
                   @test @insrc(Eye{T1}(d) * Fill{T2}(val, d)) == Fill{typeof(one(T1)*one(T2))}(val, d)
               end
           end
       end
    end
end

@testset "any all iszero isone" begin
    for T in (Int, Float64, ComplexF64)
        d1 = 1
        m = Eye{T}(d1)
        @test ! @inmodule FillArrays any(iszero, m)
        @test ! @ninsrc all(iszero, m)
        @test @ninsrc any(isone, m)
        @test @ninsrc all(isone, m)

        onem = Ones{T}(d1, d1)
        @test @ninsrc isone(onem)
        @test ! iszero(onem)

        zerom = Zeros{T}(d1, d1)
        @test ! isone(zerom)
        @test  iszero(zerom)

        fillm0 = Fill(T(0), d1, d1)
        @test ! isone(fillm0)
        @test   iszero(fillm0)

        fillm1 = Fill(T(1), d1, d1)
        @test isone(fillm1)
        @test ! iszero(fillm1)

        fillm2 = Fill(T(2), d1, d1)
        @test ! isone(fillm2)
        @test ! iszero(fillm2)

        for d in (2, 3)
            m = Eye{T}(d)
            @test any(iszero, m)
            @test ! all(iszero, m)
            @test any(isone, m)
            @test ! all(isone, m)

            m1 = Ones{T}(d, d)
            @test ! isone(m1)
            @test ! iszero(m1)
            @test all(isone, m1)
            @test ! all(iszero, m1)

            m2 = Zeros{T}(d, d)
            @test ! isone(m2)
            @test iszero(m2)
            @test ! all(isone, m2)
            @test  all(iszero, m2)

            m3 = Fill(T(2), d, d)
            @test ! isone(m3)
            @test ! iszero(m3)
            @test ! all(isone, m3)
            @test ! all(iszero, m3)
        end
    end
end

@testset "reduction types" begin
    for ncols in (1, 3, 10)
        for T in (Int, Float64, Int32, BigInt) # Rational{Int})
            for fop in (sum, prod)
                mi = Eye{T}(ncols)
                midense = Matrix(mi)
                r = @insrc fop(mi)
                rdense = fop(midense)
                @test r == rdense
                @test typeof(r) == eltype(mi)
            end
        end
    end
end

@testset "transform" begin
    m = Eye(3)
    for f in (triu, triu!, tril, tril!, inv)
        @test m == @insrc f(m) # FIXME: why does === fail ?
    end
end

@testset "kron Eye" begin
    for ncols1 in (1, 2, 3, 4)
        for ncols2 in (1, 2, 3, 4)
            im = Eye(ncols1)
            imd = Matrix(im)
            md = rand(ncols2, ncols2)
            @test kron(imd, md) == @insrc kron(im, md)
            @test kron(md, imd) == @insrc kron(md, im)
            @test isone(@isinsrc(kron(im, im)))
            @test isa(kron(im, im), typeof(im))
            @test size(@insrc(kron(im, zeros(0, 0)))) == (0,0)
            @test size(kron(zeros(0, 0), im)) == (0,0)
        end
    end
end

@testset "kron Diagonal" begin
    Ntests = 10
    for m in (1, 2, 3, 4)
        for m1 in (1, 2, 3, 4, 7)
            for n1 in (1, 2, 3, 4, 7)
                for icount in 1:Ntests
                    md = Diagonal(rand(m))
                    mdd = Matrix(md)
                    m2 = rand(m1, n1)
                    @test kron(md, m2) == kron(mdd, m2)
                    @test kron(m2, md) == kron(m2, mdd)
                end
            end
        end
    end
end

@testset "eigen" begin
    for ncols in (0, 1, 3, 5)
        im = Eye(ncols)
        md = Matrix(im)
        ri = @insrc eigen(im)
        rd = eigen(md)
        @test isapprox(ri.values, rd.values)
        @test isapprox(ri.vectors, rd.vectors)
        @test isapprox(eigvals(md), @insrc(eigvals(im)))
        @test isapprox(eigvecs(im), eigvecs(md))
        @test isa(ri.vectors, typeof(im))
    end
end

@testset "one oneunit zero" begin
    for ncols in (0, 1, 3, 10)
        for T in (Float64, Int)
            m = Eye{T}(ncols)
            @test one(m) === m
            @test oneunit(m) === m
            @test zero(m) == Diagonal(Zeros{eltype(m)}(size(m, 1)))
        end
    end
end

@testset "Fill" begin
    for ncols in (1, 3, 10)
        for T in (Float64,)
            for val in (1, 2, 2.5, 10)
                fm = Diagonal(Fill{T}(val, ncols))
                fmd = Matrix(fm)
                @test inv(fm) == inv(fmd)
            end
        end
    end
end


using ScatteredInterpolation, Random

Random.seed!(2)

if VERSION < v"0.7-"
    using Base.Test
else
    using Test
end

include("rbf.jl")
include("idw.jl")
include("nearestNeighbor.jl")


using ConvOptDL

include("data_loader_tests.jl")
include("util_tests.jl")
include("qp_tests.jl")
include("adjoin_tests.jl")
include("drop_block_tests.jl")


# Various heap implementation

###########################################################
#
#   Heap interface specification
#
#   Each heap is associated with a handle type (H), and
#   a value type v.
#
#   Here, the value type must be comparable, and a handle
#   is an object through which one can refer to a specific
#   node of the heap and thus update its value.
#
#   Each heap type must implement all of the following
#   functions. Here, let h be a heap, i be a handle,
#   v be a value and s be a size.
#
#   - length(h)           returns the number of elements
#
#   - isempty(h)          returns whether the heap is
#                         empty
#
#   - push!(h, v)         add a value to the heap
#
#   - sizehint!(h, s)     set size hint to a heap
#
#   - first(h)            return the first (top) value of a heap
#
#   - pop!(h)             removes the first (top) value, and
#                         returns it
#
#  For mutable heaps, it should also support
#
#   - push!(h, v)         adds a value to the heap and
#                         returns a handle to v
#
#   - update!(h, i, v)    updates the value of an element
#                         (referred to by the handle i)
#
#   - delete!(h, i)       deletes the node with
#                         handle i from the heap
#
#   - top_with_handle(h)  return the top value of a heap
#                         and its handle
#
#
###########################################################

# HT: handle type
# VT: value type

abstract type AbstractHeap{VT} end

abstract type AbstractMutableHeap{VT,HT} <: AbstractHeap{VT} end

abstract type AbstractMinMaxHeap{VT} <: AbstractHeap{VT} end

# heap implementations

include("heaps/binary_heap.jl")
include("heaps/mutable_binary_heap.jl")
include("heaps/minmax_heap.jl")

# generic functions

Base.eltype(::Type{<:AbstractHeap{T}}) where T = T

"""
    extract_all!(h)

Return an array of heap elements in sorted order (heap head at first index).

Note that for simple heaps (not mutable or minmax)
sorting the internal array of elements in-place is faster.
"""
function extract_all!(h::AbstractHeap{VT}) where VT
    n = length(h)
    r = Vector{VT}(undef, n)
    for i in 1 : n
        r[i] = pop!(h)
    end
    return r
end

"""
    extract_all_rev!(h)

Return an array of heap elements in reverse sorted order (heap head at last index).

Note that for simple heaps (not mutable or minmax)
sorting the internal array of elements in-place is faster.
"""
function extract_all_rev!(h::AbstractHeap{VT}) where VT
    n = length(h)
    r = Vector{VT}(undef, n)
    for i in 1 : n
        r[n + 1 - i] = pop!(h)
    end
    return r
end

# Array functions using heaps

"""
    nextreme(ord, n, arr)

return an array of the first `n` values of `arr` sorted by `ord`.
"""
function nextreme(ord::Base.Ordering, n::Int, arr::AbstractVector{T}) where T
    if n <= 0
        return T[] # sort(arr)[1:n] returns [] for n <= 0
    elseif n >= length(arr)
        return sort(arr, order = ord)
    end

    rev = Base.ReverseOrdering(ord)

    buffer = heapify(arr[1:n], rev)

    for i = n + 1 : length(arr)
        @inbounds xi = arr[i]
        if Base.lt(rev, buffer[1], xi)
            buffer[1] = xi
            percolate_down!(buffer, 1, rev)
        end
    end

    return sort!(buffer, order = ord)
end

"""
    nlargest(n, arr)

Return the `n` largest elements of the array `arr`.

Equivalent to:
    sort(arr, order = Base.Reverse)[1:min(n, end)]

Note that if `arr` contains floats and is free of NaN values,
then the following alternative may be used to achieve 2x performance.
    DataStructures.nextreme(DataStructures.FasterReverse(), n, arr)
This faster version is equivalent to:
    sort(arr, lt = >)[1:min(n, end)]
"""
function nlargest(n::Int, arr::AbstractVector)
    return nextreme(Base.Reverse, n, arr)
end

"""
    nsmallest(n, arr)

Return the `n` smallest elements of the array `arr`.

Equivalent to:
    sort(arr, order = Base.Forward)[1:min(n, end)]

Note that if `arr` contains floats and is free of NaN values,
then the following alternative may be used to achieve 2x performance.
    DataStructures.nextreme(DataStructures.FasterForward(), n, arr)
This faster version is equivalent to:
    sort(arr, lt = <)[1:min(n, end)]
"""
function nsmallest(n::Int, arr::AbstractVector)
    return nextreme(Base.Forward, n, arr)
end


# Magnetic pendulum basins of attraction
using DrWatson
@quickactivate "NonlinearDynamicsTextbook"
include(srcdir("style.jl"))
using DynamicalSystems, PyPlot, Random

α = 0.2; ω = 1.0; d = 0.3
gx = gy = range(-5, 5; length = 1500)
config = @strdict(α, ω, d, gx, gy)

function magnetic_basins(config)
    @unpack α, ω, d, gx, gy = config
    ma = Systems.magnetic_pendulum(; α, ω, d)
    if haskey(config, "attractors")
        attractors = config["attractors"]
        @time basins, attractors = basins_of_attraction((gx, gy), ma;
        mx_chk_att = 1, mx_chk_fnd_att = 10, attractors)
    else
        @time basins, attractors = basins_of_attraction((gx, gy), ma;
        mx_chk_att = 1, mx_chk_fnd_att = 10, mx_chk_lost = 10)
    end
    return @strdict(gx, gy, basins, attractors)
end

# produce high quality
file, _ = produce_or_load(datadir(), config, magnetic_basins; prefix = "magnetic_basins", tag = false)
attractors = file["attractors"]

# produce zoomed version
gx = range(1.80, 1.95; length = 1000)
gy = range(0, 0.12; length = 1000)
config = @strdict(α, ω, d, gx, gy, attractors)
produce_or_load(datadir(), config, magnetic_basins; prefix = "magnetic_basins_zoomed", tag = false)

# %% Plot this
@unpack gx, gy, basins, attractors = wload(datadir(savename("magnetic_basins", config, "jld2")))
ma = Systems.magnetic_pendulum(; α, ω, d)

LC =  matplotlib.colors.ListedColormap
cmap = LC([matplotlib.colors.to_rgb("C$k") for k in 0:2])

fig = figure(figsize=(0.33figx, figy))
ax = gca()
ax.pcolormesh(gx, gy, basins'; cmap, shading = "gouraud")
ax.set_aspect("auto")
xticks([-5, 5])
yticks([-5, 5])
xlabel("\$x\$", labelpad=-32)
ylabel("\$y\$", labelpad=-32)
for m in values(attractors)
    m1, m2 = columns(m)
    scatter(m1, m2; color = "white", edgecolor = "black", zorder = 99, s = 100)
end
# wsave(plotsdir("magneticpendulum"), fig)

# %% Plot zoomed version
@unpack gx, gy, basins, attractors = wload(datadir(savename("magnetic_basins_zoomed", config, "jld2")))
# fig = figure(figsize=(figx/2, figx/2))
x0 = 0.55; wx = 0.42; y0 = 0.65; wy = 0.32;
axin = ax.inset_axes([x0, y0, wx, wy])

axin.pcolormesh(gx, gy, basins'; cmap, shading = "gouraud")
# axin.set_aspect("equal")

# xticks([1.8, 1.95], size = 20)
# yticks([0, 0.12], size = 20)
zbox = ((gx[1], gy[1]), (gx[end], gy[end]))
axis_zoomin!(axin, ax, zbox, zbox, "C3"; dir = :top, lw = 4)

axin.axis("off")

fig.tight_layout(pad=0.3)

wsave(plotsdir("5", "magneticpendulum.png"), fig)


# Simplistic kernel calculation
function kernelize!(out::AbstractArray, x::AbstractArray, y::AbstractArray, kernel::Function;
		    center::Bool=false, symmetric::Bool=false)
	m = size(x, 2)
	n = size(y, 2)

	if !symmetric   # generic case, x and y are different
		@inbounds for j in 1:n
        		yj = view(y,:,j)
        		@simd for i in 1:m
				xi = view(x,:,i)
				out[i,j] = kernel(xi, yj)
        		end
		end
	else
    		@inbounds for j = 1:n
        		yj = view(y,:,j)
        		@simd for i in 1:j
				xi = view(x,:,i)
            			out[i,j] = kernel(xi, yj)
				out[j,i] = out[i,j]
        		end
    		end
	end

	if center
		μi = mean(out, 2)	# kernel mean of individual rows
		μj = mean(out, 1)	# kernel mean of individual columns
		μ = mean(out)		# overall mean
		@inbounds @simd for i in 1:m
			for j in 1:n
				out[i,j] = out[i,j] - μi[i] - μj[j] + μ
			end
		end
	end
	return out
end

kernelize!(out::AbstractArray, x::AbstractArray, kernel::Function; center::Bool=false, symmetric::Bool=true) =
	kernelize!(out, x, x, kernel; center=center, symmetric=symmetric)

function kernelize(x::AbstractArray, y::AbstractArray, kernel::Function; center::Bool=false, symmetric::Bool=false)
	m = size(x, 2)
	n = size(y, 2)
	out = similar(x, m, n)
	kernelize!(out, x, y, kernel; center=center, symmetric=symmetric)
end

kernelize(x::AbstractArray, kernel::Function; center::Bool=false, symmetric::Bool=true) =
	kernelize(x, x, kernel; center=center, symmetric=symmetric)

##########################
# FunctionCell Interface #
##########################
"""
	kernel([f];[center=false, symmetric=false])

Constructs an untrained function cell using the kernel function `f` that accepts two vector arguments `x` and `y`,
and returns a scalar value (default `(x,y)->x'y`). When piped data into, the untrained function cell returns a
trained function cell, retaining the training data.

# Keyword arguments
  * `symmetric::Bool` specifies whether the kernel should be symmetric i.e. calculated using the same observations (default `false`)
  * `center::Bool` specifies whether the center the kernel (default `false`)

A list of common used kernels can be found below (definitions from the documentation of to MLKernels.jl):
# Kernels
  * `(x,y)-> a⋅xᵀy + c		Linear, a>0, c≥0`
  * `(x,y)-> (a⋅xᵀy+c)ᵈ		Polynomial, a>0, c≥0, d∈Ζ₊`
  * `(x,y)-> tanh(a⋅xᵀy+c)		Sigmoid, a>0, c≥0`
  * `(x,y)-> exp(a⋅xᵀy)		Exponentiated, a>0`
  * `(x,y)-> exp(-α*‖x-y‖)		Exponential, α>0`
  * `(x,y)-> exp(-α⋅‖x-y‖²)		Squared exponential, α>0`
  * `(x,y)-> exp(-α⋅‖x-y‖ᵞ)		Gamma exponential, α>0, 0<γ<1`
  * `(x,y)-> (1 + α⋅‖x-y‖²)⁻ᵝ	Rational quadratic, α>0, β>0`
  * `(x,y)-> (1 + α⋅‖x-y‖ᵞ)⁻ᵝ	Gamma-rational, α>0, β>0, 0<γ,1`
  * `(x,y)-> exp{-α⋅Σᵢsin²[p⋅(xᵢ-yᵢ)]}	Periodic kernel, p>0, α>0`
  * `(x,y)-> ‖x-y‖²ᵞ		Power, 0<γ≤1 (negative definite)`
  * `(x,y)-> log(1+α⋅‖x-y‖²ᵞ)	Log, α>0, 0<γ≤1 (negative definite)`

# Examples
```
julia> k(a,c) = (x,y)->a*x'y+c # linear kernel
k (generic function with 1 method)

julia> K = kernel(k(1.0,2.0))
Kernel: symmetric=false, varying I/O dimensions, untrained

julia> x=[1,2,3]; y=[4,5,6,7];

julia> y |> K(x)
3×4 Array{Int64,2}:
  6   7   8   9
 10  12  14  16
 14  17  20  23
```
"""
kernel(f::Function=(x,y)->x'y; center::Bool=false, symmetric::Bool=false) =
	FunctionCell(kernel, (f,), ModelProperties(), "Kernel: center=$center, symmetric=$symmetric";
	      center=center, symmetric=symmetric)

############################
# DataCell/Array Interface #
############################
"""
	kernel(x [,f] [;symetric=false])

Trains the function cell using the kernel function `f` and data `x` which will be used
for future kernel calculations.
"""
# Training
kernel(x::T where T<:CellData, f::Function=(x,y)->x'y; center::Bool=false, symmetric::Bool=false) =
	kernel(getx!(x), f; center=center, symmetric=symmetric)
kernel(x::Tuple{T,S} where T<:AbstractArray where S<:AbstractArray, f::Function=(x,y)->x'y;
       center::Bool=false, symmetric::Bool=false) = kernel(x[1], f; center=center, symmetric=symmetric)
kernel(x::T where T<:AbstractVector, f::Function=(x,y)->x'y; center::Bool=false, symmetric::Bool=false) =
	kernel(mat(x, LearnBase.ObsDim.Constant{2}()), f; center=center, symmetric=symmetric)
kernel(x::T where T<:AbstractMatrix, f::Function=(x,y)->x'y; center::Bool=false, symmetric::Bool=false) = begin

	# Build model properties
	modelprops = ModelProperties(nvars(x),nobs(x))

	# Returned trained cell
	FunctionCell(kernel, Model((f, getobs(x), center, symmetric), modelprops),
	      "Kernel: center=$center, symmetric=$symmetric")
end

# Execution
kernel(x::T where T<:CellData, model::Model{<:Tuple{<:Function,Matrix,Bool,Bool}}) =
	datacell(kernel(getx!(x), model), gety(x))
kernel(x::T where T<:AbstractVector, model::Model{<:Tuple{<:Function,Matrix,Bool,Bool}}) =
	kernel(mat(x, LearnBase.ObsDim.Constant{2}()), model)
kernel(x::T where T<:AbstractMatrix, model::Model{<:Tuple{<:Function,Matrix,Bool,Bool}}) =
	kernelize(model.data[2], getobs(x), model.data[1] ;center=model.data[3], symmetric=model.data[4])


#=
This shows basic data analysis using the Card data
on returns to education

See the github pages of the packages for documentation

=#

##
using CSV, DataFrames, DataFramesMeta, DelimitedFiles, StatPlots, Plots
pyplot(reuse=false, show=true)
card = CSV.read("../Data/card.csv") # how to read CSV data into DataFrame
size(card)
names(card)
head(card)
describe(card)

##
# select south and black (using DataFramesMeta)
sb = @where(card, :south .== 1, :black .== 1)
# make a scatter plot
@df sb scatter(:educ, :wage, leg=false)
# select north and white
nw = @where(card, :south .== 0, :black .== 0)
@df nw scatter!(:educ, :wage, leg=false)
##
# do the same thing using the all the data, with grouping
@df card scatter(:educ, :wage, group = (:black, :south), markersize=10)
##
# density plots by groups
@df card density(:wage, group=(:black), fill=true, fillalpha=0.3, legend=false)

##
# violin and box plots
@df card violin(:wage)
@df card boxplot!(:wage) # puts is beside
# violin plot of wage, by education
@df card violin(:educ, :wage)
# add a boxplot on top
@df card boxplot!(:educ, :wage)
##
# convert the dataframe to an ordinary array,
# as this is usually better for more advanced
# modeling
data = convert(Array{Float64}, card)
writedlm("card.ascii", data) # write to plain ASCII file


println(sum(1:100)^2 - sum(i^2 for i in 1:100))


using StochasticDiffEq, DiffEqDevTools, Test, LinearAlgebra
Random.seed!(100)
dts = (1/2) .^ (10:-1:2) #14->7 good plot

const σ_const = 0.87
const μ = 1.01

u0 = rand(2)
A = [-2.0 1.0;1.0 -2.0]
B = Diagonal([σ_const for i in 1:2])

function f_commute(du,u,p,t)
  mul!(du,A,u)
  du .+= 1.01u
end

function f_commute_analytic(u0,p,t,W)
 tmp = (A+1.01I-(B^2))*t + B*sum(W)
 exp(tmp)*u0
end

function σ(du,u,p,t)
  du[1,1] = σ_const*u[1]
  du[1,2] = σ_const*u[1]
  du[2,1] = σ_const*u[2]
  du[2,2] = σ_const*u[2]
end

ff_commute = SDEFunction(f_commute,σ,analytic=f_commute_analytic)

prob = SDEProblem(ff_commute,σ,u0,(0.0,1.0),noise_rate_prototype=rand(2,2))

sol = solve(prob,RKMilCommute(),dt=1/2^(8))
sol = solve(prob,EM(),dt=1/2^(10))

dts = (1/2) .^ (10:-1:3) #14->7 good plot
sim2 = test_convergence(dts,prob,EM(),numMonte=Int(1e2))
sim2 = test_convergence(dts,prob,RKMilCommute(),numMonte=Int(1e2))

sim2.𝒪est[:final] - 1 < 0.2


export NonlocalLaplacian

#############
# NonlocalLaplacian implements the nonlocal integral operator:
#
#        𝓛 u(x) := \int_{B_δ(x)} ρ_δ(|x-y|) [u(y) - u(x)] dy
#
# where B_δ(x) = {y ∈ ℝᵈ : |x-y| ≤ δ} is the closed ball of radius 0 < δ < ∞,
# and the spherically symmetric kernel ρ_δ is defined by:
#
#       ρ_δ(r) = 2*Γ(d/2+1)*(d+2-α)/(π^{d/2} δ^{d+2-α}) χ_[0,δ](r)/r^α,
#
# where χ_[0,δ](⋅) is the indicator function, and 0 ≤ α < d + 2.
#############

abstract type AbstractNonlocalCalculusOperator{S, T} <: Operator{T} end

abstract type NonlocalLaplacian{S, T} <: AbstractNonlocalCalculusOperator{S, T} end

struct ConcreteNonlocalLaplacian{S, T} <: NonlocalLaplacian{S, T}
    space::S
    α::T
    δ::T
end

ConcreteNonlocalLaplacian(sp::Space, α::T, δ::T) where T = ConcreteNonlocalLaplacian{typeof(sp), T}(sp, α, δ)

ConcreteNonlocalLaplacian(sp::Space, α, δ) = ConcreteNonlocalLaplacian{typeof(sp), promote_type(typeof(α), typeof(δ))}(sp, α, δ)

NonlocalLaplacian(sp::Space, α, δ) = ConcreteNonlocalLaplacian(sp, α, δ)
NonlocalLaplacian(d::Domain, x...) = NonlocalLaplacian(Space(d), x...)

domain(L::ConcreteNonlocalLaplacian) = domain(L.space)

domainspace(L::ConcreteNonlocalLaplacian) = L.space

getindex(::ConcreteNonlocalLaplacian{UnsetSpace,T}, k::Integer, j::Integer) where T =
    error("Spaces cannot be inferred for operator")
rangespace(L::ConcreteNonlocalLaplacian{UnsetSpace,T}) where T = UnsetSpace()

rangespace(L::ConcreteNonlocalLaplacian{Hardy{s,DD,RR}}) where {s,DD,RR} = L.space
rangespace(L::ConcreteNonlocalLaplacian{Laurent{DD,RR}}) where {DD,RR} = L.space
rangespace(L::ConcreteNonlocalLaplacian{Fourier{DD,RR}}) where {DD,RR} = L.space
rangespace(L::ConcreteNonlocalLaplacian{CosSpace{DD,RR}}) where {DD,RR} = L.space
rangespace(L::ConcreteNonlocalLaplacian{SinSpace{DD,RR}}) where {DD,RR} = L.space

bandwidths(::ConcreteNonlocalLaplacian{Hardy{s,DD,RR}}) where {s,DD,RR} = 0,0
bandwidths(::ConcreteNonlocalLaplacian{Laurent{DD,RR}}) where {DD,RR} = 0,0
bandwidths(::ConcreteNonlocalLaplacian{Fourier{DD,RR}}) where {DD,RR} = 0,0
bandwidths(::ConcreteNonlocalLaplacian{CosSpace{DD,RR}}) where {DD,RR} = 0,0
bandwidths(::ConcreteNonlocalLaplacian{SinSpace{DD,RR}}) where {DD,RR} = 0,0

##TODO: Add scale for different periods.
function getindex(L::ConcreteNonlocalLaplacian{Hardy{true, DD, RR}, T},
                  k::Integer, j::Integer) where {DD <: PeriodicSegment, RR, T}
    if k == j
        fourier_lambda(k-1, L.α, L.δ, 1)
    else
        zero(T)
    end
end

function getindex(L::ConcreteNonlocalLaplacian{Hardy{false, DD, RR}, T},
                  k::Integer, j::Integer) where {DD <: PeriodicSegment, RR, T}
    if k == j
        fourier_lambda(k, L.α, L.δ, 1)
    else
        zero(T)
    end
end

for SP in (:Laurent, :Fourier)
    @eval begin
        function getindex(L::ConcreteNonlocalLaplacian{$SP{DD, RR}, T},
                          k::Integer, j::Integer) where {DD <: PeriodicSegment, RR, T}
            if k == j
                fourier_lambda(k÷2, L.α, L.δ, 1)
            else
                zero(T)
            end
        end
    end
end

function getindex(L::ConcreteNonlocalLaplacian{CosSpace{DD, RR}, T},
                  k::Integer, j::Integer) where {DD <: PeriodicSegment, RR, T}
    if k == j
        fourier_lambda(k-1, L.α, L.δ, 1)
    else
        zero(T)
    end
end

function getindex(L::ConcreteNonlocalLaplacian{SinSpace{DD, RR}, T},
                  k::Integer, j::Integer) where {DD <: PeriodicSegment, RR, T}
    if k == j
        fourier_lambda(k, L.α, L.δ, 1)
    else
        zero(T)
    end
end

function fourier_lambda(k, α, δ, d::Integer)
    d < 1 && error("The $d-dimensional Fourier spectrum is absurd.")
    0 ≤ α < d + 2 || error("The algebraic singularity strength α = $α is not permissible.")
    if abs(k*δ) < 6
        fourier_lambda_small(k, α, δ, d)
    else
        fourier_lambda_large(k, α, δ, d)
    end
end

fourier_lambda_small(k::Real, α::Real, δ::Real, d::Integer) = fourier_lambda_small(promote(k, α, δ)..., d)

function fourier_lambda_small(k::T, α::T, δ::T, d::Integer) where T <: Real
    z = -(k*δ/2)^2
    C = (d+2-α)/2
    D = 1+T(d)/2
    S₀, S₁, err, j = one(T), one(T) + C/(C+one(T))*z/(2*D), one(T), 1
    while err > 10*abs(S₀)*eps(T)
        rⱼ = ((C+j)/(C+j+1))/((D+j)*(j+2))
        S₀, S₁ = S₁, S₁ + (S₁ - S₀)*rⱼ*z
        err = abs(S₁-S₀)
        j += 1
    end
    return -k^2*S₁
end

fourier_lambda_large(k::Real, α::Real, δ::Real, d::Integer) = fourier_lambda_large(promote(k, α, δ)..., d)

function fourier_lambda_large(k::T, α::T, δ::T, d::Integer) where T <: Real
    if norm(d+2-α) < eps(T)
        return T(-k^2)
    else
        kδ = k*δ
        Td = T(d)
        scl = (d+2-α)*gamma(Td/2+1)*2/δ^2
        cst = multif((d-α)/2, 2/kδ, Td/2)/gamma(Td/2)
        tail = 2^(Td/2)*(kδ)^(α+1-d)*((d-2-α)*besselj((Td-2)/2, kδ)*lommelS2((d-2-2α)/2, (Td-4)/2, kδ) - besselj((Td-4)/2, kδ)*lommelS2((d-2α)/2, (Td-2)/2, kδ))
        return scl*(cst+tail)
    end
end

function multif(x, y, z)
    (y^(2x)*gamma(x+1)*gamma(z)/gamma(z-x)-1)/x
end

function multif(x::Float64, y::Float64, z::Float64)
    if x ≠ 0.0
        w = 2x*log(y) + HypergeometricFunctions.lanczosapprox(1.0, x) - HypergeometricFunctions.lanczosapprox(z, -x)
        return expm1(w)/x
    else
        return 2*log(y) - MathConstants.γ + digamma(z)
    end
end

function lommelS2(μ, ν, z)
    z^(μ-1)*drummond((1-μ+ν)/2, (1-μ-ν)/2, (z/2)^2)
end

function drummond(α::T, β::T, z::T) where T
    Nlo = one(T)
    Dlo = one(T)
    Tlo = Nlo/Dlo

    cst = (α+1)*(β+1)
    if norm(cst) < eps(real(T))
        return Tlo
    end
    Nmid = (z+α+β+1)/cst
    Dmid = (z+cst)/cst
    Tmid = Nmid/Dmid

    cst = (α+2)*(β+2)
    if norm(cst) < eps(real(T))
        return Tmid
    end
    Dhi = ((z+α+β+3+cst)*Dmid - (α+β+3)*Dlo)/cst
    Nhi = ((z+α+β+3+cst)*Nmid - (α+β+3)*Nlo)/cst
    Thi = Nhi/Dhi

    k = 0
    while (abs(Thi-Tmid) > 10*abs(Thi)*eps(real(T)) || abs(Tmid-Tlo) > 10*abs(Tmid)*eps(real(T))) && k < 10_000
        cst = (α+k+3)*(β+k+3)
        Nhi, Nmid, Nlo = ((z+(k+2)*(α+β+2*k+5)+cst)*Nhi - (k+2)*(α+β+3*k+6)*Nmid + (k+2)*(k+1)*Nlo)/cst, Nhi, Nmid
        Dhi, Dmid, Dlo = ((z+(k+2)*(α+β+2*k+5)+cst)*Dhi - (k+2)*(α+β+3*k+6)*Dmid + (k+2)*(k+1)*Dlo)/cst, Dhi, Dmid
        Thi, Tmid, Tlo = Nhi/Dhi, Thi, Tmid
        k += 1
    end
    return isnan(Thi) ? isnan(Tmid) ? Tlo : Tmid : Thi
end


using Gadfly

set_default_plot_size(6inch, 3inch)

plot(x=[0.236033, 0.346517, 0.312707, 0.00790928, 0.488613, 0.210968, 0.951916, 0.999905, 0.251662, 0.986666],
     y=[0.555751, 0.437108, 0.424718, 0.773223, 0.28119, 0.209472, 0.251379, 0.0203749, 0.287702, 0.859512],
     color=[0.2774, 0.800248, 0.934636, 0.527809, 0.866783, 0.803046, 0.278974, 0.92097, 0.292635, 0.743778],
     Scale.color_sqrt)


export binhist

probabilities(x) = _non0hist(x)
function _non0hist(x)
    L = length(x)

    hist = Vector{Float64}()
    # Reserve enough space for histogram:
    sizehint!(hist, L)

    sx = sort(x, alg = QuickSort)

    # Fill the histogram by counting consecutive equal values:
    prev_val, count = sx[1], 0
    for val in sx
        if val == prev_val
            count += 1
        else
            push!(hist, count)
            prev_val = val
            count = 1
        end
    end
    push!(hist, count)

    # Shrink histogram capacity to fit its size:
    sizehint!(hist, length(hist))
    return Probabilities(hist ./ L)
end
_non0hist(x::AbstractDataset) = _non0hist(x.data)
probabilities(x::AbstractDataset) = _non0hist(x.data)


"""
Created on Mon May 22 2020

@author: Yoann Pradat

Evaluate NMF algorithms with sparsity.
"""

include("./src/VNMF.jl")
using .VNMF
using DataFrames
using Distributions
using Random

include("./demo/src/util/plot_cost.jl")
gr() # set the backend to GR

#### # 1. SIMULATE DATA
#### # ################################################################################################################

#### simulate matrix for testing NMF
F = 10
N = 25
K = 5
rng = MersenneTwister(123)

#### truncated-normal factors with normal noise
dist_W = truncated(Normal(5, 2), 0, Inf)
dist_H = truncated(Normal(5, 2), 0, Inf)
dist_E = Normal(1, 0.1)

W = rand(rng, dist_W, F, K)
H = rand(rng, dist_H, K, N)
E = rand(rng, dist_E, F, N)

V = W * H + E

#### # 2. FEVOTTE-IDIER NMF ALGORITHMS WITH KNOWN RANK
#### # ################################################################################################################

fo = "demo/plot/simulation/nmf_sparse"
mkpath(fo)

#### global params
global_params = NMFParams(
    init          = :random,
    dist          = truncated.(Normal(0, 1), 0, Inf),
    rank          = K,
    max_iter      = 10_000,
    stopping_crit = :cost,
    stopping_tol  = 1e-6,
    verbose       = false,
)

#### # 2.1 L1 Regularization on H, MM algo
list_local_params = Array{FIParams, 1}()
list_sparse_W = Array{Real, 1}()
list_sparse_H = Array{Real, 1}()
list_cur_cost = Array{Real, 1}()
list_cur_divg = Array{Real, 1}()

list_β = [-1, 0, 1, 2]
list_e = [-1, 0, 1, 2, 3, 4]

#### MM alg
for β in list_β
    for e in list_e
        println(repeat("=", 80))
        #### MM alg
        local_params = FIParams(
            β         = β,
            scale_W   = false,
            α_H       = 10.0^e,
            α_W       = 0,
            l₁ratio_H = 1,
            l₁ratio_W = 1,
            alg       = :mm,
        )
        push!(list_local_params, local_params)
        @info "params" local_params.β local_params.α_W local_params.α_H

        #### run
        res_FI_MM = nmf_FI(V, global_params, local_params)

        cur_divg = res_FI_MM.metrics.cur_divg
        cur_cost = res_FI_MM.metrics.cur_cost
        sparse_W = VNMF.matr_avg_col_sparseness(res_FI_MM.W)
        sparse_H = VNMF.matr_avg_col_sparseness(res_FI_MM.H)

        push!(list_cur_divg, cur_divg)
        push!(list_cur_cost, cur_cost)
        push!(list_sparse_W, sparse_W)
        push!(list_sparse_H, sparse_H)

        @info "at convergence" cur_divg, cur_cost
        @info "sparsity level" sparse_W sparse_H
    end
end

#### get results in a dataframe
df_metrics_MM = DataFrame(
   β = (x -> x.β).(list_local_params),
   α_W = (x -> x.α_W).(list_local_params),
   α_H = (x -> x.α_H).(list_local_params),
   cur_divg = list_cur_divg,
   cur_cost = list_cur_cost,
   sparse_W = list_sparse_W,
   sparse_H = list_sparse_H,
)

#### # 2.2 L1 Regularization on H, H algo
list_local_params = Array{FIParams, 1}()
list_sparse_W = Array{Real, 1}()
list_sparse_H = Array{Real, 1}()
list_cur_cost = Array{Real, 1}()
list_cur_divg = Array{Real, 1}()

list_β = [-1, 0, 1, 2]
list_e = [-1, 0, 1, 2, 3, 4]

#### MM alg
for β in list_β
    for e in list_e
        println(repeat("=", 80))
        #### MM alg
        local_params = FIParams(
            β         = β,
            scale_W   = false,
            α_H       = 10.0^e,
            α_W       = 0,
            l₁ratio_H = 1,
            l₁ratio_W = 1,
            alg       = :h,
        )
        push!(list_local_params, local_params)
        @info "params" local_params.β local_params.α_W local_params.α_H

        #### run
        res_FI_MM = nmf_FI(V, global_params, local_params)

        cur_divg = res_FI_MM.metrics.cur_divg
        cur_cost = res_FI_MM.metrics.cur_cost
        sparse_W = VNMF.matr_avg_col_sparseness(res_FI_MM.W)
        sparse_H = VNMF.matr_avg_col_sparseness(res_FI_MM.H)

        push!(list_cur_divg, cur_divg)
        push!(list_cur_cost, cur_cost)
        push!(list_sparse_W, sparse_W)
        push!(list_sparse_H, sparse_H)

        @info "at convergence" cur_divg, cur_cost
        @info "sparsity level" sparse_W sparse_H
    end
end

#### get results in a dataframe
df_metrics_H = DataFrame(
   β = (x -> x.β).(list_local_params),
   α_W = (x -> x.α_W).(list_local_params),
   α_H = (x -> x.α_H).(list_local_params),
   cur_divg = list_cur_divg,
   cur_cost = list_cur_cost,
   sparse_W = list_sparse_W,
   sparse_H = list_sparse_H,
)


using Test, Interpolations, DualNumbers, LinearAlgebra

@testset "Gradients" begin
    nx = 10
    f1(x) = sin((x-3)*2pi/(nx-1) - 1)
    g1gt(x) = 2pi/(nx-1) * cos((x-3)*2pi/(nx-1) - 1)
    A1 = Float64[f1(x) for x in 1:nx]
    g1 = Array{Float64}(undef, 1)
    A2 = rand(Float64, nx, nx) * 100
    g2 = Array{Float64}(undef, 2)

    for (A, g) in ((A1, g1), (A2, g2))
        # Gradient of Constant should always be 0
        itp = interpolate(A, BSpline(Constant()))
        for x in InterpolationTestUtils.thirds(axes(A))
            @test all(iszero, @inferred(Interpolations.gradient(itp, x...)))
            @test all(iszero, @inferred(Interpolations.gradient!(g, itp, x...)))
        end

        itp = interpolate(A, BSpline(Linear()))
        check_gradient(itp, g)
        i = first(eachindex(itp))
        @test Interpolations.gradient(itp, i) == Interpolations.gradient(itp, Tuple(i)...)

        for BC in (Flat,Line,Free,Periodic,Reflect,Natural), GT in (OnGrid, OnCell)
            itp = interpolate(A, BSpline(Quadratic(BC(GT()))))
            check_gradient(itp, g)
            i = first(eachindex(itp))
            @test Interpolations.gradient(itp, i) == Interpolations.gradient(itp, Tuple(i)...)
        end

        for BC in (Line, Flat, Free, Periodic), GT in (OnGrid, OnCell)
            itp = interpolate(A, BSpline(Cubic(BC(GT()))))
            check_gradient(itp, g)
            I = first(eachindex(itp))
            @test Interpolations.gradient(itp, I) == Interpolations.gradient(itp, Tuple(I)...)
        end
    end

    # Since Linear is OnGrid in the domain, check the gradients between grid points
    itp1 = interpolate(Float64[f1(x) for x in 1:nx],
                BSpline(Linear()))
    itp2 = interpolate((1:nx-1,), Float64[f1(x) for x in 1:nx-1],
                Gridded(Linear()))
    for itp in (itp1, itp2)
        for x in 2.5:nx-1.5
            @test ≈(g1gt(x),(Interpolations.gradient(itp,x))[1],atol=abs(0.1 * g1gt(x)))
            @test ≈(g1gt(x),(Interpolations.gradient!(g1,itp,x))[1],atol=abs(0.1 * g1gt(x)))
            @test ≈(g1gt(x),g1[1],atol=abs(0.1 * g1gt(x)))
        end

        for i = 1:10
            x = rand()*(nx-2)+1.5
            checkbounds(Bool, itp, x) || continue
            gtmp = Interpolations.gradient(itp, x)[1]
            xd = dual(x, 1)
            @test epsilon(itp(xd)) ≈ gtmp
        end
    end

    # test gridded on a non-uniform grid
    knots = (1.0:0.3:nx-1,)
    itp_grid = interpolate(knots, Float64[f1(x) for x in knots[1]],
                           Gridded(Linear()))

    for x in 1.5:0.5:nx-1.5
        @test ≈(g1gt(x),(Interpolations.gradient(itp_grid,x))[1],atol=abs(0.5 * g1gt(x)))
        @test ≈(g1gt(x),(Interpolations.gradient!(g1,itp_grid,x))[1],atol=abs(0.5 * g1gt(x)))
        @test ≈(g1gt(x),g1[1],atol=abs(0.5 * g1gt(x)))
    end

    # Since Quadratic is OnCell in the domain, check gradients at grid points
    itp1 = interpolate(Float64[f1(x) for x in 1:nx-1],
                       BSpline(Quadratic(Periodic(OnCell()))))
    for x in 2:nx-1
        @test ≈(g1gt(x),(Interpolations.gradient(itp1,x))[1],atol=abs(0.05 * g1gt(x)))
        @test ≈(g1gt(x),(Interpolations.gradient!(g1,itp1,x))[1],atol=abs(0.05 * g1gt(x)))
        @test ≈(g1gt(x),g1[1],atol=abs(0.1 * g1gt(x)))
    end

    for i = 1:10
        x = rand()*(nx-2)+1.5
        gtmp = Interpolations.gradient(itp1, x)[1]
        xd = dual(x, 1)
        @test epsilon(itp1(xd)) ≈ gtmp
    end

    # For a quadratic function and quadratic interpolation, we expect an
    # "exact" answer
    # 1d
    c = 2.3
    a = 8.1
    o = 1.6
    qfunc = x -> a*(x .- c).^2 .+ o
    dqfunc = x -> 2*a*(x .- c)
    xg = Float64[1:5;]
    y = qfunc(xg)

    iq = interpolate(y, BSpline(Quadratic(Free(OnCell()))))
    x = 1.8
    @test iq(x) ≈ qfunc(x)
    @test (Interpolations.gradient(iq,x))[1] ≈ dqfunc(x)

    # 2d (biquadratic)
    p = [(x-1.75)^2 for x = 1:7]
    A = p*p'
    iq = interpolate(A, BSpline(Quadratic(Free(OnCell()))))
    @test iq[4,4] ≈ (4 - 1.75) ^ 4
    @test iq[4,3] ≈ (4 - 1.75) ^ 2 * (3 - 1.75) ^ 2
    g = Interpolations.gradient(iq, 4, 3)
    @test g[1] ≈ 2 * (4 - 1.75) * (3 - 1.75) ^ 2
    @test g[2] ≈ 2 * (4 - 1.75) ^ 2 * (3 - 1.75)

    iq = interpolate!(copy(A), BSpline(Quadratic(InPlace(OnCell()))))
    @test iq[4,4] ≈ (4 - 1.75) ^ 4
    @test iq[4,3] ≈ (4 - 1.75) ^ 2 * (3 - 1.75) ^ 2
    g = Interpolations.gradient(iq, 4, 3)
    @test ≈(g[1],2 * (4 - 1.75) * (3 - 1.75) ^ 2,atol=0.03)
    @test ≈(g[2],2 * (4 - 1.75) ^ 2 * (3 - 1.75),atol=0.2)

    # InPlaceQ is exact for an underlying quadratic
    iq = interpolate!(copy(A), BSpline(Quadratic(InPlaceQ(OnCell()))))
    @test iq[4,4] ≈ (4 - 1.75) ^ 4
    @test iq[4,3] ≈ (4 - 1.75) ^ 2 * (3 - 1.75) ^ 2
    g = Interpolations.gradient(iq, 4, 3)
    @test g[1] ≈ 2 * (4 - 1.75) * (3 - 1.75) ^ 2
    @test g[2] ≈ 2 * (4 - 1.75) ^ 2 * (3 - 1.75)

    A2 = rand(Float64, nx, nx) * 100
    gni = [1.0]
    for BC in (Flat,Line,Free,Periodic,Reflect,Natural), GT in (OnGrid, OnCell)
        itp_a = interpolate(A2, (BSpline(Linear()), BSpline(Quadratic(BC(GT())))))
        itp_b = interpolate(A2, (BSpline(Quadratic(BC(GT()))), BSpline(Linear())))
        itp_c = interpolate(A2, (NoInterp(), BSpline(Quadratic(BC(GT())))))
        itp_d = interpolate(A2, (BSpline(Quadratic(BC(GT()))), NoInterp()))

        for i = 1:10
            x = rand()*(nx-2)+1.5
            y = rand()*(nx-2)+1.5
            xd = dual(x, 1)
            yd = dual(y, 1)
            gtmp = Interpolations.gradient(itp_a, x, y)
            @test length(gtmp) == 2
            @test epsilon(itp_a(xd,y)) ≈ gtmp[1]
            @test epsilon(itp_a(x,yd)) ≈ gtmp[2]
            gtmp = Interpolations.gradient(itp_b, x, y)
            @test length(gtmp) == 2
            @test epsilon(itp_b(xd,y)) ≈ gtmp[1]
            @test epsilon(itp_b(x,yd)) ≈ gtmp[2]
            ix, iy = round(Int, x), round(Int, y)
            gtmp = Interpolations.gradient(itp_c, ix, y)
            @test length(gtmp) == 1
            @test epsilon(itp_c(ix,yd)) ≈ gtmp[1]
            gni[1] = NaN
            Interpolations.gradient!(gni, itp_c, ix, y)
            @test gni[1] ≈ gtmp[1]
            gtmp = Interpolations.gradient(itp_d, x, iy)
            @test length(gtmp) == 1
            @test epsilon(itp_d(xd,iy)) ≈ gtmp[1]
            gni[1] = NaN
            Interpolations.gradient!(gni, itp_d, x, iy)
            @test gni[1] ≈ gtmp[1]
        end
    end

    # issue #239
    f(x) = log(x)
    xs = 1:0.2:5
    A = [f(x) for x in xs]
    interp_linear = LinearInterpolation(xs, A)
    g = Interpolations.gradient(interp_linear, 1.5)
    dest = Vector{Float64}(undef, 1)
    Interpolations.gradient!(dest, interp_linear, 1.5)
    @test dest == g

    @testset "Monotonic" begin
        x = [0.0, 0.2, 0.5, 0.6, 0.9, 1.0]
        ys = [[-3.0, 0.0, 5.0, 10.0, 18.0, 22.0],
              [10.0, 0.0, -5.0, 10.0, -8.0, -2.0]]
        grid = 0.0:0.1:1.0

        itypes = [LinearMonotonicInterpolation(),
            FiniteDifferenceMonotonicInterpolation(),
            CardinalMonotonicInterpolation(0.0),
            CardinalMonotonicInterpolation(0.5),
            CardinalMonotonicInterpolation(1.0),
            FritschCarlsonMonotonicInterpolation(),
            FritschButlandMonotonicInterpolation(),
            SteffenMonotonicInterpolation()]

        for y in ys
            for it in itypes
                itp = interpolate(x, y, it)
                for t in grid
                    gradval = epsilon(itp(dual(t, 1.0)))
                    @test Interpolations.gradient1(itp, t) ≈ gradval atol = 1.e-12
                    @test Interpolations.gradient(itp, t)[1] ≈ gradval atol = 1.e-12
                end
            end
        end
    end
end


# using Revise

using IncrementalInference
using Test

##

@testset "Test priors" begin

##

N=100
graphinits = [false, true]

# TEMP value when not using for loop
graphinit = true

for  graphinit = graphinits

fg = initfg()
fg.solverParams.N = N
fg.solverParams.graphinit = graphinit
fg.solverParams.treeinit = !graphinit

addVariable!(fg, :x0, ContinuousScalar, N=N)
addFactor!(fg, [:x0], Prior(Normal(-1.0, 1.0)))

addVariable!(fg, :x1, ContinuousScalar, N=N)

addVariable!(fg, :x2, ContinuousScalar, N=N)
addFactor!(fg, [:x2], Prior(Normal(+1.0, 1.0)))

addFactor!(fg, [:x0; :x1], LinearRelative(Normal(0.0, 0.01)))
addFactor!(fg, [:x1; :x2], LinearRelative(Normal(0.0, 0.01)))

#solve
tree = solveTree!(fg)
x0_m = getKDEMean(getKDE(getVariable(fg, :x0)))[1]
x1_m = getKDEMean(getKDE(getVariable(fg, :x1)))[1]
x2_m = getKDEMean(getKDE(getVariable(fg, :x2)))[1]

@info ("Testing means = 0 with 2 priors:\ngraphinit=$graphinit\nMeans: x0: $(x0_m), x1: $x1_m, x2: $x2_m")

@test_skip isapprox(x0_m, 0.0, atol = 0.1)
@test_skip isapprox(x1_m, 0.0, atol = 0.1)
@test_skip isapprox(x2_m, 0.0, atol = 0.1)

@warn "priorusetest.jl is testing with large tolerances"
@test isapprox(x0_m, 0.0, atol = 1.0)
@test isapprox(x1_m, 0.0, atol = 1.0)
@test isapprox(x2_m, 0.0, atol = 1.0)

#testing if values are close to one another
testvals = [x0_m, x1_m, x2_m]
meanval = mean(testvals)
@test all(isapprox.(testvals, meanval, atol=0.4))

end

##

for  graphinit = graphinits

fg = initfg()
fg.solverParams.N = N
fg.solverParams.graphinit = graphinit
fg.solverParams.treeinit = !graphinit

addVariable!(fg, :x0, ContinuousScalar, N=N)
addFactor!(fg, [:x0], Prior(Normal(-1.0, 1.0)))

addVariable!(fg, :l0, ContinuousScalar, N=N)
addFactor!(fg, [:l0], Prior(Normal(+1.0, 1.0)))

addVariable!(fg, :l1, ContinuousScalar, N=N)

addFactor!(fg, [:x0; :l0], LinearRelative(Normal(0, 0.01)))
addFactor!(fg, [:x0; :l1], LinearRelative(Normal(0, 0.01)))

addVariable!(fg, :x1, ContinuousScalar, N=N)
addFactor!(fg, [:x0; :x1], LinearRelative(Normal(0, 0.01)))

addVariable!(fg, :x2, ContinuousScalar, N=N)
addFactor!(fg, [:x1; :x2], LinearRelative(Normal(0, 0.01)))

addFactor!(fg, [:x2; :l0], LinearRelative(Normal(0, 0.01)))
addFactor!(fg, [:x2; :l1], LinearRelative(Normal(0, 0.01)))

#solve
tree = solveTree!(fg)

x0_m = getKDEMean(getKDE(getVariable(fg, :x0)))[1]
x1_m = getKDEMean(getKDE(getVariable(fg, :x1)))[1]
x2_m = getKDEMean(getKDE(getVariable(fg, :x2)))[1]
l0_m = getKDEMean(getKDE(getVariable(fg, :l0)))[1]
l1_m = getKDEMean(getKDE(getVariable(fg, :l1)))[1]

@info ("Testing means = 0 with 2 priors:\ngraphinit=$graphinit\nMeans: x0: $(x0_m), x1: $x1_m, x2: $x2_m, l0: $l0_m, l1: $l1_m")

@test_skip isapprox(x0_m, 0.0, atol = 0.1)
@test_skip isapprox(x1_m, 0.0, atol = 0.1)
@test_skip isapprox(x2_m, 0.0, atol = 0.1)
@test_skip isapprox(l0_m, 0.0, atol = 0.1)
@test_skip isapprox(l1_m, 0.0, atol = 0.1)

@warn "priorusetest.jl is testing with large tolerances"
@test isapprox(x0_m, 0.0, atol = 1.0)
@test isapprox(x1_m, 0.0, atol = 1.0)
@test isapprox(x2_m, 0.0, atol = 1.0)
@test isapprox(l0_m, 0.0, atol = 1.2)
@test isapprox(l1_m, 0.0, atol = 1.2)

#testing if values are close to one another
@show testvals = [x0_m, x1_m, x2_m, l0_m, l1_m]
@show meanval = mean(testvals)
@test all(isapprox.(testvals, meanval, atol=0.3))

end

##

end