"""
    μ = center!(X[, μ, obsdim])

or

    μ = center!(D[, colnames, μ])

where `X` is of type Matrix or Vector and `D` of type DataFrame.

Center `X` along `obsdim` around the corresponding entry in the
vector `μ`. If `μ` is not specified then it defaults to the
feature specific means.

For DataFrames, `obsdim` is obsolete and centering is done column wise.
Instead the vector `colnames` allows to specify which columns to center.
If `colnames` is not provided all columns of type T<:Real are centered.

Example:

    X = rand(4, 100)
    D = DataFrame(A=rand(10), B=collect(1:10), C=[string(x) for x in 1:10])

    μ = center!(X, obsdim=2)
    μ = center!(X, ObsDim.First())
    μ = center!(D)
    μ = center!(D, [:A, :B])

"""
function center!(X, μ; obsdim=LearnBase.default_obsdim(X))
    center!(X, μ, convert(ObsDimension, obsdim))
end

function center!(X; obsdim=LearnBase.default_obsdim(X))
    center!(X, convert(ObsDimension, obsdim))
end

function center!(X::AbstractArray{T,N}, μ::AbstractVector, ::ObsDim.Last) where {T,N}
    center!(X, μ, ObsDim.Constant{N}())
end

function center!(X::AbstractArray{T,N}, ::ObsDim.Last) where {T,N}
    center!(X, ObsDim.Constant{N}())
end

function center!(X::AbstractArray{T,N}, obsdim::ObsDim.Constant{M}) where {T,N,M}
    μ = vec(mean(X, dims=M))
    center!(X, μ, obsdim)
end

function center!(X::AbstractVector{T}, ::ObsDim.Constant{1}) where T
    μ = mean(X)
    for i in 1:length(X)
        X[i] = X[i] - μ
    end
    μ
end

function center!(X::AbstractVector, μ::AbstractVector, ::ObsDim.Constant{1})
    @inbounds for i in 1:length(X)
        X[i] = X[i] - μ[i]
    end
    μ
end

function center!(X::AbstractVector, μ::AbstractFloat, ::ObsDim.Constant{1})
    @inbounds for i in 1:length(X)
        X[i] = X[i] - μ
    end
    μ
end

function center!(X::AbstractMatrix, μ::AbstractVector, ::ObsDim.Constant{1})
    nObs, nVars = size(X)
    for iVar in 1:nVars
        @inbounds for iObs in 1:nObs
            X[iObs, iVar] = X[iObs, iVar] - μ[iVar]
        end
    end
    μ
end

function center!(X::AbstractMatrix, μ::AbstractVector, ::ObsDim.Constant{2})
    nVars, nObs = size(X)
    for iObs in 1:nObs
        @inbounds for iVar in 1:nVars
            X[iVar, iObs] = X[iVar, iObs] - μ[iVar]
        end
    end
    μ
end

function center!(D::AbstractDataFrame)
    μ_vec = Float64[]

    flt = Bool[T <: Real for T in eltype.(eachcol(D))]
    for colname in propertynames(D)[flt]
        μ = mean(D[!, colname])
        center!(D, colname, μ)
        push!(μ_vec, μ)
    end
    μ_vec
end

function center!(D::AbstractDataFrame, colnames::AbstractVector{Symbol})
    μ_vec = Float64[]
    for colname in colnames
        if eltype(D[!, colname]) <: Real
            μ = mean(D[!, colname])
            if ismissing(μ)
                @warn("Column \"$colname\" contains missing values, skipping rescaling of this column!")
                continue
            end
            center!(D, colname, μ)
            push!(μ_vec, μ)
        else
            @warn("Skipping \"$colname\", centering only valid for columns of type T <: Real.")
        end
    end
    μ_vec
end

function center!(D::AbstractDataFrame, colnames::AbstractVector{Symbol}, μ::AbstractVector)
    for (icol, colname) in enumerate(colnames)
        if eltype(D[!, colname]) <: Real
            center!(D, colname, μ[icol])
        else
            @warn("Skipping \"$colname\", centering only valid for columns of type T <: Real.")
        end
    end
    μ
end

function center!(D::AbstractDataFrame, colname::Symbol, μ)
    if any(ismissing, D[!, colname])
        @warn("Column \"$colname\" contains missing values, skipping centering on this column!")
    else
        newcol::Vector{Float64} = convert(Vector{Float64}, D[!, colname])
        nobs = length(newcol)
        @inbounds for i in eachindex(newcol)
            newcol[i] -= μ
        end
        D[!, colname] = newcol
    end
    μ
end

"""
    μ, σ = rescale!(X[, μ, σ, obsdim])

or

    μ, σ = rescale!(D[, colnames, μ, σ])

where `X` is of type Matrix or Vector and `D` of type DataFrame.

Center `X` along `obsdim` around the corresponding entry in the
vector `μ` and then rescale each feature using the corresponding
entry in the vector `σ`.

For DataFrames, `obsdim` is obsolete and centering is done column wise.
The vector `colnames` allows to specify which columns to center.
If `colnames` is not provided all columns of type T<:Real are centered.

Example:

    X = rand(4, 100)
    D = DataFrame(A=rand(10), B=collect(1:10), C=[string(x) for x in 1:10])

    μ, σ = rescale!(X, obsdim=2)
    μ, σ = rescale!(X, ObsDim.First())
    μ, σ = rescale!(D)
    μ, σ = rescale!(D, [:A, :B])

"""
function rescale!(X, μ, σ; obsdim=LearnBase.default_obsdim(X))
    rescale!(X, μ, σ, convert(ObsDimension, obsdim))
end

function rescale!(X::AbstractArray{T,N}, μ, σ, ::ObsDim.Last) where {T,N}
    rescale!(X, μ, σ, ObsDim.Constant{N}())
end

function rescale!(X; obsdim=LearnBase.default_obsdim(X))
    rescale!(X, convert(ObsDimension, obsdim))
end

function rescale!(X::AbstractArray{T,N}, ::ObsDim.Last) where {T,N}
    rescale!(X, ObsDim.Constant{N}())
end

function rescale!(X::AbstractArray{T,N}, obsdim::ObsDim.Constant{M}) where {T,N,M}
    μ = vec(mean(X, dims=M))
    σ = vec(std(X, dims=M))
    rescale!(X, μ, σ, obsdim)
end

function rescale!(X::AbstractVector, ::ObsDim.Constant{1})
    μ = mean(X)
    σ = std(X)
    @inbounds for i in 1:length(X)
        X[i] = (X[i] - μ) / σ
    end
    μ, σ
end

function rescale!(X::AbstractMatrix, μ::AbstractVector, σ::AbstractVector, ::ObsDim.Constant{2})
    σ[σ .== 0] .= 1
    nVars, nObs = size(X)
    for iObs in 1:nObs
        @inbounds for iVar in 1:nVars
            X[iVar, iObs] = (X[iVar, iObs] - μ[iVar]) / σ[iVar]
        end
    end
    μ, σ
end

function rescale!(X::AbstractMatrix, μ::AbstractVector, σ::AbstractVector, ::ObsDim.Constant{1})
    σ[σ .== 0] .= 1
    nObs, nVars = size(X)
    for iVar in 1:nVars
        @inbounds for iObs in 1:nObs
            X[iObs, iVar] = (X[iObs, iVar] - μ[iVar]) / σ[iVar]
        end
    end
    μ, σ
end

function rescale!(X::AbstractVector, μ::AbstractVector, σ::AbstractVector, ::ObsDim.Constant{1})
    @inbounds for i in 1:length(X)
        X[i] = (X[i] - μ[i]) / σ[i]
    end
    μ, σ
end

function rescale!(X::AbstractVector, μ::AbstractFloat, σ::AbstractFloat, ::ObsDim.Constant{1})
    @inbounds for i in 1:length(X)
        X[i] = (X[i] - μ) / σ
    end
    μ, σ
end

function rescale!(D::AbstractDataFrame)
    μ_vec = Float64[]
    σ_vec = Float64[]

    flt = Bool[T <: Real for T in eltype.(eachcol(D))]
    for colname in propertynames(D)[flt]
        μ = mean(D[!, colname])
        σ = std(D[!, colname])
        rescale!(D, colname, μ, σ)
        push!(μ_vec, μ)
        push!(σ_vec, σ)
    end
    μ_vec, σ_vec
end

function rescale!(D::AbstractDataFrame, colnames::Vector{Symbol})
    μ_vec = Float64[]
    σ_vec = Float64[]
    for colname in colnames
        if eltype(D[!, colname]) <: Real
            μ = mean(D[!, colname])
            σ = std(D[!, colname])
            if ismissing(μ)
                @warn("Column \"$colname\" contains missing values, skipping rescaling of this column!")
                continue
            end
            rescale!(D, colname, μ, σ)
            push!(μ_vec, μ)
            push!(σ_vec, σ)
        else
            @warn("Skipping \"$colname\", rescaling only valid for columns of type T <: Real.")
        end
    end
    μ_vec, σ_vec
end

function rescale!(D::AbstractDataFrame, colnames::Vector{Symbol}, μ::AbstractVector, σ::AbstractVector)
    for (icol, colname) in enumerate(colnames)
        if eltype(D[!, colname]) <: Real
            rescale!(D, colname, μ[icol], σ[icol])
        else
            @warn("Skipping \"$colname\", rescaling only valid for columns of type T <: Real.")
        end
    end
    μ, σ
end

function rescale!(D::AbstractDataFrame, colname::Symbol, μ, σ)
    if any(ismissing, D[!, colname])
        @warn("Column \"$colname\" contains missing values, skipping rescaling of this column!")
    else
        σ_div = σ == 0 ? one(σ) : σ
        newcol::Vector{Float64} = convert(Vector{Float64}, D[!, colname])
        nobs = length(newcol)
        @inbounds for i in eachindex(newcol)
            newcol[i] = (newcol[i] - μ) / σ_div
        end
        D[!, colname] = newcol
    end
    μ, σ
end

struct FeatureNormalizer
    offset::Vector{Float64}
    scale::Vector{Float64}

    function FeatureNormalizer(offset::Vector{Float64}, scale::Vector{Float64})
        @assert length(offset) == length(scale)
        new(offset, scale)
    end
end

function FeatureNormalizer(X::AbstractMatrix{T}) where T<:Real
    FeatureNormalizer(vec(mean(X, dims=2)), vec(std(X, dims=2)))
end

function StatsBase.fit(::Type{FeatureNormalizer}, X::AbstractMatrix{T}) where T<:Real
    FeatureNormalizer(X)
end

function StatsBase.predict!(cs::FeatureNormalizer, X::AbstractMatrix{T}) where T<:Real
    @assert length(cs.offset) == size(X, 1)
    rescale!(X, cs.offset, cs.scale)
    X
end

function StatsBase.predict(cs::FeatureNormalizer, X::AbstractMatrix{T}) where T<:AbstractFloat
    Xnew = copy(X)
    StatsBase.predict!(cs, Xnew)
end

function StatsBase.predict(cs::FeatureNormalizer, X::AbstractMatrix{T}) where T<:Real
    X = convert(AbstractMatrix{AbstractFloat}, X)
    StatsBase.predict!(cs, X)
end


### A Pluto.jl notebook ###
# v0.19.2

using Markdown
using InteractiveUtils

# ╔═╡ d9f9bd70-508d-11ec-1d5d-b3592d653992
# ╠═╡ show_logs = false
begin
	let
		using Pkg
		Pkg.activate(mktempdir())
		Pkg.Registry.update()
		Pkg.add("PlutoUI")
		Pkg.add(url="https://github.com/Dale-Black/DistanceTransforms.jl")
		Pkg.add(url="https://github.com/Dale-Black/Losers.jl")
	end

	using PlutoUI
	using DistanceTransforms
	using Losers
end

# ╔═╡ 5398d829-ad81-4971-94b5-f87eeaaa381e
md"""
## Import packages
"""

# ╔═╡ 44119707-2397-4e7f-83f7-f362eae6004f
TableOfContents()

# ╔═╡ 6e03547c-d9b6-4972-9574-7cb4f30948f9
md"""
## Distance transforms + deep learning
One area of recent interest in deep learning has been distance transform-based loss functions to directly reduce Hausdorff loss and other boundary-specific loss metrics [D. Karimi and S. E. Salcudean](http://arxiv.org/abs/1904.10030).

This package can be utilized directly in various training loops, along with [Losers.jl](https://github.com/Dale-Black/Losers.jl) to train segmentation networks, like the common UNet architecture, to directly reduce Hausdorff distance using the `hausdorff()` loss function.

In practice, it is common to combine `hausdorff` with the `dice` loss function to help stabalize during training.

Below we will show how one might set this up, without actually training on anything.
"""

# ╔═╡ 2810025c-9940-4da2-a587-938c425f511e
y1 = [1 1 1 0; 1 1 1 0; 1 1 1 0; 1 1 1 0]

# ╔═╡ dfc81bb5-29fa-46f2-a335-66f541205d22
y = cat(y1, y1; dims=3)

# ╔═╡ 89347efa-ee13-4c61-a4e6-3d4d1a77955d
ŷ = copy(y)

# ╔═╡ 96d6cb3f-a500-490b-9bc9-500731af2651
md"""
Given some fake three dimensional array, lets apply a distance transform to this array for use in the `hausdorff` loss function
"""

# ╔═╡ 646d3c5b-e72b-49ac-b1e1-74d2d89c57c2
y_dtm = euclidean(y)

# ╔═╡ 10f64f57-e052-4895-9fe3-58e1c4fb2d88
ŷ_dtm = euclidean(ŷ)

# ╔═╡ 6e52eae8-e1cd-4d07-a000-d591b7bfa3ff
hausdorff(ŷ, y, ŷ_dtm, y_dtm)

# ╔═╡ b7c539d1-0e65-4716-ad23-cf4400bf4bf7
md"""
As one would expect, the difference between two identical arrays is zero
"""

# ╔═╡ 2780485d-ce45-4d5f-a88c-d4823bcf4273
begin
	y2 = rand([0, 1], 10, 10, 10)
	ŷ2 = rand([0, 1], 10, 10, 10)

	y2_dtm = euclidean(y2)
	ŷ2_dtm = euclidean(ŷ2)
end

# ╔═╡ 423704a8-926e-4fd3-98c4-7da8374a35a5
hausdorff(ŷ2, y2, ŷ2_dtm, y2_dtm)

# ╔═╡ 5d1ac21a-8073-4817-869b-5500649e486f
md"""
## GPU considerations
When using distance transform-based loss functions for deep learning, one must consider GPUs for many real-world deep learning tasks. The benefit of DistanceTransforms.jl is that GPU compatible distance transform algorithms are ready to use, out-of-the-box.

We will look at how one would go about this, without actually calling a CUDA array from the CUDA.jl library.
"""

# ╔═╡ cfd9c96d-1ce7-4edc-954c-4d844aff6629
begin
	y3 = rand([0, 1], 10, 10, 10)
	ŷ3 = rand([0, 1], 10, 10, 10)
end

# ╔═╡ 6345df57-dba7-4608-9e08-9c9d981810e5
begin
	y3_bool = boolean_indicator(y3)
	ŷ3_bool = boolean_indicator(ŷ3)
end

# ╔═╡ dde2505d-1a69-4201-a75a-e29a62197163
begin
	tfm1 = SquaredEuclidean(y3_bool)
	tfm2 = SquaredEuclidean(ŷ3_bool)
end

# ╔═╡ d2a987a3-f44c-4239-9e96-d83ae16c265e
begin
	y3_dtm = transform(y3_bool, tfm1)
	ŷ3_dtm = transform(ŷ3_bool, tfm1)
end

# ╔═╡ 90467d7a-785d-44d3-92bc-9d05bcb5f5b4
hausdorff(ŷ3, y3, ŷ3_dtm, y3_dtm)

# ╔═╡ 43e694a4-bbd5-4072-a4de-0ea65b16c906
md"""
If we were to be running an actual training loop it might look something like this...

```julia
# Begin training loop
for epoch in 1:max_epochs
	step = 0
    α = ...
	@info epoch

	# Loop through training data
	for (xs, ys) in train_loader
		step += 1
		@info step

		# Send data to GPU
		xs, ys = xs |> gpu, ys |> gpu
		gs = Flux.gradient(ps) do
			ŷs = model(xs)

			# Apply distance transform using GPU compatible `SquaredEuclidean`
			# Data will usually be 4D or 5D [x, y, (z), channel, batch]

			ys_dtm = CuArray{Float32}(undef, size(ys))
			ŷs_dtm = CuArray{Float32}(undef, size(ŷs))
			for b in size(ys, 5)
				for c in size(ys, 4)
					for z in size(ys, 3)
						bool_arr_gt = boolean_indicator(ys[:, :, z, c, b])
						bool_arr_pred = boolean_indicator(ŷs[:, :, z, c, b])
						tfm_gt = SquaredEuclidean(bool_arr_gt)
						tfm_pred = SquaredEuclidean(bool_arr_pred)
						ys_dtm[:, :, z, c, b] = transform!(bool_arr_gt, tfm_gt, threads)
						ŷs_dtm[:, :, z, c, b] = transform!(bool_arr_pred, tfm_pred, threads)
					end
				end
			end
			hd_loss = hausdorff(ŷs, ys, ŷs_dtm, ys_dtm)
			dice_loss = dice(ŷs, ys)
			loss = α*dice_loss + (1-α)*hausdorff_loss
			return loss
		end
		Flux.update!(optimizer, ps, gs)
	end
end
```
"""

# ╔═╡ Cell order:
# ╟─5398d829-ad81-4971-94b5-f87eeaaa381e
# ╠═d9f9bd70-508d-11ec-1d5d-b3592d653992
# ╠═44119707-2397-4e7f-83f7-f362eae6004f
# ╟─6e03547c-d9b6-4972-9574-7cb4f30948f9
# ╠═2810025c-9940-4da2-a587-938c425f511e
# ╠═dfc81bb5-29fa-46f2-a335-66f541205d22
# ╠═89347efa-ee13-4c61-a4e6-3d4d1a77955d
# ╟─96d6cb3f-a500-490b-9bc9-500731af2651
# ╠═646d3c5b-e72b-49ac-b1e1-74d2d89c57c2
# ╠═10f64f57-e052-4895-9fe3-58e1c4fb2d88
# ╠═6e52eae8-e1cd-4d07-a000-d591b7bfa3ff
# ╟─b7c539d1-0e65-4716-ad23-cf4400bf4bf7
# ╠═2780485d-ce45-4d5f-a88c-d4823bcf4273
# ╠═423704a8-926e-4fd3-98c4-7da8374a35a5
# ╟─5d1ac21a-8073-4817-869b-5500649e486f
# ╠═cfd9c96d-1ce7-4edc-954c-4d844aff6629
# ╠═6345df57-dba7-4608-9e08-9c9d981810e5
# ╠═dde2505d-1a69-4201-a75a-e29a62197163
# ╠═d2a987a3-f44c-4239-9e96-d83ae16c265e
# ╠═90467d7a-785d-44d3-92bc-9d05bcb5f5b4
# ╟─43e694a4-bbd5-4072-a4de-0ea65b16c906


using LoopVectorization, Test

function multiassign!(y, x)
  @assert length(y) + 3 == length(x)
  @inbounds for i ∈ eachindex(y)
    x₁, ((x₂, x₃), (x₄, x₅)) = x[i], (sincos(x[i+1]), (x[i+2], x[i+3]))
    y[i] = x₁ * x₄ - x₂ * x₃
  end
  y
end
multiassign(x) = multiassign!(similar(x, length(x) - 3), x)
function multiassign_turbo!(y, x)
  @assert length(y) + 3 == length(x)
  @turbo for i ∈ eachindex(y)
    x₁, ((x₂, x₃), (x₄, x₅)) = x[i], (sincos(x[i+1]), (x[i+2], x[i+3]))
    y[i] = x₁ * x₄ - x₂ * x₃
  end
  y
end
multiassign_turbo(x) = multiassign_turbo!(similar(x, length(x) - 3), x)

multistorefunc(x) = exp(x), sincos(x)
function multistore!(x, y, z, a)
  @inbounds for i ∈ eachindex(x, y, z, a)
    x[i], (y[i], z[i]) = multistorefunc(a[i])
  end
end
function multistore_turbo!(x, y, z, a)
  @turbo for i ∈ eachindex(x, y, z, a)
    x[i], (y[i], z[i]) = multistorefunc(a[i])
  end
end

@testset "Multiple assignments" begin
  @show @__LINE__
  x = rand(111)
  @test multiassign(x) ≈ multiassign_turbo(x)
  a0 = similar(x)
  b0 = similar(x)
  c0 = similar(x)
  a1 = similar(x)
  b1 = similar(x)
  c1 = similar(x)
  multistore!(a0, b0, c0, x)
  multistore_turbo!(a1, b1, c1, x)
  @test a0 ≈ a1
  @test b0 ≈ b1
  @test c0 ≈ c1
end


using Plots, LaTeXStrings, Statistics

function Deposition_2DG(; len, tot_time, dep_rate, color_step)
    surface = InitialSurf(len)
    for n in 0:tot_time
        randsurf = rand(1:len, dep_rate)
        for i in randsurf
            surface = check_and_depose(surface, i, len, (n%color_step) + 1)
        end
    end
    return surface
end

function check_and_depose(surf, index, len, addval) #checking down and sides and first layer
    i1 , i2 = sides(index, len)
    if surf[end][index] != 0
        newsurf = [0 for j = 1:len]
        newsurf[index] = addval
        push!(surf, newsurf)
    end
    for num in 0:length(surf) -1
        if surf[end - num][index] == 0 && (surf[end - num][i1] != 0 || surf[end - num][i2] != 0)
            surf[end - num][index] = addval
            break
        elseif length(surf) > num + 1 && surf[end - 1 - num][index] != 0
            surf[end - num][index] = addval
            break
        end
    end
    return surf
end

function sides(n, L)
    if n == L
        return n-1 , 1
    elseif n == 1
        return L , n+1
    else
        return n-1, n+1
    end
end

function InitialSurf(len)
    pos1 = [0 for i=1:len]
    pos1[floor(Int, len/2)] = 1
    return [pos1]
end

Parameters = Dict(:len => 300,
                    :tot_time => 1000,
                        :dep_rate => 50,
                            :color_step => 100)

surf = Deposition_2DG(;Parameters...)

heatmap(transpose(hcat(surf...)), c = :solar, legend = false)
savefig("../../Figs/Q3/Graphic.pdf")


### A Pluto.jl notebook ###
# v0.12.21

using Markdown
using InteractiveUtils

# ╔═╡ 52d41aa0-769c-11eb-2eb1-e78eeed9d5c8
begin
	using Plots
	using LsqFit
	using Statistics
	using PlutoUI
	md"_hidden code cell_"
end

# ╔═╡ c65eeff0-76b9-11eb-3543-3739a6484c24
html"""

Term Paper 2

This work is a submission for the course Modelling Complex Systems(IDC-631) for the year 2020-21 Spring (Even) Semester by Dhruva Sambrani, MS18163
"""

# ╔═╡ d846746e-76bb-11eb-2938-adb0dc210f71
md"""# Abstract
This paper laregly recreates the paper[^soc] by Bak, Per and Tang, Chao and Wiesenfeld, Kurt on models that exhibit self-organized criticality. These models show certain structures repetitive structures and show distinctive power laws in certain properties.

The specific model I have looked at is the Sandpile model, which has a relaxation time frequency having a log-log slope of between 1 to 2. This paper first introduces some working definitions and their equivalent definition in Julia[^julia] code.

First a pretty, but small simulation is shown which also verifies that the model is infact correctly translated. Post this, more rigourous analysis is done, and finally results and observations are mentioned
"""

# ╔═╡ a3a53a50-7830-11eb-380b-ed771520d060
md"## How to read the paper

This paper is made using the Pluto[^pluto] Package which provides an interface which can show both code and theory side by side, which will hopefully enhance the understanding of the work and also provide a reproducable implementation without having to go into the details of it. Those who wish to read the paper for the theory and results can easily skip over the code and still understand everything. Those who are attempting to recreate the work can use the code as a guide. However, in the interest of keeping the code aligned to the topic at hand, certain auxillary functions have been hidden. The content of these hidden cells can be found in the executable version of this notebook [here]()
"

# ╔═╡ f02e0ed0-784c-11eb-08b5-4746af933e1a
TableOfContents()

# ╔═╡ 2ed745de-76cf-11eb-101c-89f9d2ba170b
md"# Code Setup"

# ╔═╡ 1f57fcde-76cf-11eb-23c1-f5895772e792
md"## Packages used
The Plots.jl[^plots] package is used for plotting and the Statistics[^stats] Package for some statistical analysis. LsqFit.jl[^lsqfit] is used to fit the log-log plots to nice straight lines. PlutoUI.jl[^plutoui] is used for some UI elements such as the ToC
"

# ╔═╡ b239c390-76cf-11eb-1c95-4bb63c153f87
md"## Define Structure and Functions

A Sandpile grid is nothing but an Integer Array and doesn't need any special definitions.
We just need to define the functions that act on a 2 dimensional grid which can simulate the rules of a Sandpile model
"

# ╔═╡ 1f872dbe-76d0-11eb-055c-cd0487ce223a
md"### The Sandpile model
In the most abstract sense, a Sandpile is just an array of integers which evolve in a specific way. However, more intuitively, it can be thought of as a distributions of equally sized grains of sand on a discrete grid. Then, the array of integers represent the height _or_ the slope of the pile in that location. Note that in the following report, we think of the values as the _slope_ of the pile. However,
!!! warning \"The use of Height\"
    We use **height** to simply refer to the value of the cell, not the actual height of the pile. This poor choice of words is a common tactic to confuse readers so that they don't read the paper too carefully and find mistakes.

### Method of evolution

The model evolves with the following algorithm -

$$z_{i,j}(t+1) = z_{i,j}(t) - h_c(z_{i,j}(t)>h_c) +
\text{count}(n\rightarrow h_c>n,\ \text{neighborhood})$$

Basically, we define a **neighborhood** of $h_c$ cells. If the height of the sand in a cell is greater than $h_c$, it is in a **critical** state, $h_c$ grains **topple** into the neighboring cells equally. Hence, we can simply subtract $h_c$ and count the number of **critical** neighbors to calculate the value of the cell in the next time step. Note that the sand is conserved in this form of evolution.

!!! info \"Relevant Quote\"
    Waste not the smallest thing for grains of sand make mountains

    **\- Eric Knight**
"

# ╔═╡ 5221fe40-7784-11eb-2d14-b7ca964ba546
md"""!!! tip "Our neighborhood and $h_c$"
    In this work, we use the VonNeumann Neighborhood in 2D and radius 1. Hence, the neighbors are $(-1,0), (+1,0), (0,-1), (0,+1)$ and $h_c$ = 4
"""

# ╔═╡ b0797e20-769c-11eb-2711-7587647ce0c5
hc = 4

# ╔═╡ 7b2d6380-769c-11eb-12f6-35402dc5e96e
@inline iscritical(i::Int, hc::Int) = i >= hc

# ╔═╡ b1c3cc10-76d6-11eb-38e1-75a2c56e0044
md"#### Boundary Conditions

We use a fixed boundary condition of 0. This is like the edge of a table, because all sand keeps falling and is lost forever.

!!! info \"Relevant Quote\"
    Summer Vacation is slipping through our hands like grains of sand!

    **\- Calvin from Calvin and Hobbes**
"

# ╔═╡ 75a73440-769c-11eb-2001-a5345c6e3732
begin
	@inline function getindexbc(grid, i)
	    if checkbounds(Bool, grid, i)
	        @inbounds return grid[i]
	    else
	        return 0
	    end
	end
	getindexbc(grid) = i -> getindexbc(grid)
end

# ╔═╡ 8806b9d0-769c-11eb-28d9-bdc5b3870377
function evolveat(
        grid::Array{Int,2},
        hood::Array{CartesianIndex{2},1},
        index::CartesianIndex,
        hc::Int,
        topple::Array{Int},
		affected_grid::BitArray{2},
        step::Int
    )
    res = grid[index]
    if res >= hc
        res -= 4
        @inbounds topple[step] += 1 # we'll talk about this later
		affected_grid[index] = true # and this
    end
    c = count(h -> iscritical(getindexbc(grid, h + index), hc), hood)
    res = res + c
    return res
end

# ╔═╡ 8dddef40-769c-11eb-2656-91b79e7af1cf
function evolve!(
        grid::Array{Int,2},
        simgrid::Array{Int,2},
		affected_grid::BitArray{2},
        hood::Array{CartesianIndex{2},1},
        hc::Int,
        topple::Array{Int,1},
        step::Int
    )
    for index in CartesianIndices(grid)
        @inbounds simgrid[index] = evolveat(grid, hood, index, hc, topple, affected_grid, step)
    end
	return simgrid, grid
end

# ╔═╡ 1c132690-76d9-11eb-0edf-45dda119d525
md"## Sanity Check and Pretty Animations

Let's run a small simulation on a 10x10 grid to see if we've coding in everything correctly. We start with an empty Pile and add a grain in the center of the grid. We wait for the pile to relax, and we add another grain of sand. We do this until we add  100 grains of sand.
"

# ╔═╡ 7db36160-76b8-11eb-22ec-bd23ba1800d6
let
	# This code is only to make the animation. The exact method of simulation is described later
	function anim_ready_next!(grid::Array{Int,2}, hc::Int, step::Int, max_steps::Int)
	    if count(iscritical.(grid, hc)) == 0
	        if step == max_steps
	            return false, step
	        else
	            grid[5, 5] += 1
	            step += 1
	        end
	    end
	    return true, step
	end
	function anim_simulate(grid::Array{Int}, hc::Int, max_steps::Int, res = false)
	    hood = CartesianIndex.([(-1, 0), (1, 0), (0, -1), (0, 1)])
	    step = 0
	    st = true
	    simgrid = similar(grid)
		results = Array{Int,2}[]
	    while st
	        st, step = anim_ready_next!(grid, hc, step, max_steps)
			grid, simgrid = evolve!(grid, simgrid, falses(size(grid)), hood, hc, [0], 0)
			push!(results, copy(grid))
	    end
	    return results
	end
	res = anim_simulate(zeros(Int, 10, 10), 4, 100);
	anim = @animate for i in res
		plot(i, st=:wireframe, fill=true, camera=(43,60), zlims=(0,4), clims=(0,4))
	end
	gif(anim, fps=5)
end

# ╔═╡ e0f4aab0-7797-11eb-0ecc-910d7b9ded81
md"_hidden code cell_"

# ╔═╡ f052c870-76d9-11eb-1b82-bb53ba216dc8
md"### Observations
1. The pile builds up to a height of 4 and then topples, as expected
2. The pile falls off at the edge, as expected
3. Some topplings are quick where as some topplings set off another topple and so on. However, these _chained_ topples are not as common as the single or double topples. We'll call an entire chain of topplings until a sand grain is added as a **avalanche**
4. This looks NOTHING like an actual sandpile. What were the authors thinking?\

Regarding point 4, this is because we are keeping a track of the _slope_, not the _height_ of the Sandpile. And yet we called it the _height_. I warned you, this is super confusing.
"

# ╔═╡ 0e96d220-76dc-11eb-2b86-f5708fe4803f
md"# Power law analysis

According to the paper mentioned before, there are certain properties which show a log-log frequency plot, which means that

$$\log(F_x) = m\cdot\log(x)+c$$

for certain properties of the model, where x is the value of the property and F(x) is the frequency of the value. These properties are -

1. **Number of sites affected** during an avalanche
2. **Total number of topples** induced during the avalanche
3. **Relaxation time** of the avalanche

Lets elucidate this further. Given a random condition, a single perturbation may lead to the pertubed cell to topple.
This may, or may not, set off an avalanche which affects neighboring cells, has multiple topples, and takes a while to settle down.
If we count these properties, and plot them on a log-log plot, it was shown that they form a straight line.
"

# ╔═╡ 08191482-7781-11eb-37bb-1fb9abbceee0
md"""!!! tip "Short note on Implementation"
    The properties are measured in the following fashion
    -
        - Toggle count - The $i^{th}$ element of an array is updated everytime a topple is simulated in the $i^{th}$ perturbation.
        - Relaxation Time - The $i^{th}$ element of an array is updated after every `evolve` cycle for the $i^{th}$ perturbation.
        - Number of sites affected - The $(i,j)$ element in a grid of `false`s is set to `true` when the $(i,j)$ cell topples in a perturbation (not `evolve`) cycle. At the end of the $k^{th}$ perturbation, the total number of `true`s are counted and stored in the $k^{th}$ element of an array.
"""

# ╔═╡ 9484347e-769c-11eb-19c7-0961b3cdf2ef
function ready_next!(
		grid::Array{Int,2},
		affected_grid::BitArray{2},
		hc::Int,
		affect_count::Array{Int, 1},
		relax_time::Array{Int, 1},
		step::Int, max_steps::Int
	)
    if count(iscritical.(grid, hc)) == 0
		affect_count[step] = count(affected_grid)
		affected_grid[:] .= false
        if step == max_steps
            return false, step
        else
			# perturb
            xrand, yrand = rand.(Base.OneTo.(size(grid)))
            grid[xrand, yrand] += 1
            step += 1
        end
	else
		relax_time[step] += 1
	end
    return true, step
end

# ╔═╡ 9b1c21e0-769c-11eb-162e-35dbc2aba3f2
function simulate(grid::Array{Int}, hc::Int, max_steps::Int)
    hood = CartesianIndex.([(-1, 0), (1, 0), (0, -1), (0, 1)])
    topple_count = fill(0, max_steps)
	relax_time = fill(0, max_steps)
	affect_count = fill(0, max_steps)
    step = 1
    st = true
    simgrid = similar(grid)
	affected_grid = falses(size(grid))
    while st
        st, step = ready_next!(grid, affected_grid, hc, affect_count, relax_time, step, max_steps)
        grid, simgrid = evolve!(grid, simgrid, affected_grid, hood, hc, topple_count, step)
    end
    return topple_count, relax_time, affect_count
end

# ╔═╡ 36049990-778c-11eb-3007-4dc0bb6b2615
md"""## Method of analysis

Once we've simulated and collect the data, the analysis is done in the following way-
1. Create a frequency table of `val => freq(val)` for each of the properties.
2. Take a log for both values and frequencies.
3. Fit the data to $\text{freq} = m\cdot\text{val}+c$
4. Smooth the data and plot it against the fit

!!! warning "Things to keep in mind while fitting"
    The data does not follow a linear curve completely. The linearity is lost for reasons listed below. Hence, we should fit the data to only the first, linear part of the data, which usually is around 60%
"""

# ╔═╡ 668d2040-769d-11eb-3943-51fb559fded3
function freqtable(k::Array{T}) where T
    d = Dict{T,Int}()
    for i ∈ k
        if i != 0
            d[i] = get(d, i, 0) + 1
        end
    end
    return d
end

# ╔═╡ 6a004e02-769d-11eb-01c4-e5d635d9aa91
function getlogdata(topple_freq)
    x = log.(keys(topple_freq))
    perm = sortperm(x)
    y = log.(values(topple_freq))[perm]
    x = x[perm]
    return x, y
end

# ╔═╡ 6a009c20-769d-11eb-2d53-534435f34669
function smooth(xs, ys, length)
    binx = range(xs[1], stop=xs[end], length=length)
    _step = step(binx)
    biny = [
        mean(ys[findall(x -> val <= x < (val + _step), xs)])
        for val in binx
    ]
    binx = (binx .+ _step)[1:end - 1]
    return binx, biny
end

# ╔═╡ 4fd4eef0-769d-11eb-023e-21941d0cb482
function loglogfit(property_data)
	x, y = getlogdata(freqtable(property_data))
	smoothx, smoothy = smooth(x, y, 100)
	plot(smoothx,smoothy, st=:scatter)
	lnsp = floor(Int, 0.4*length(x))
	@. model(x, p) = p[1] * x + p[2]
	fit = curve_fit(model, x[1:lnsp], y[1:lnsp], [0.,0.])
	xp = 0:0.01:maximum(x)
	yp = model(xp, fit.param)
	return plot!(xp, yp, width=2), fit
end

# ╔═╡ 361d0af0-7789-11eb-29ca-d768d81ccc39
function stringifyfit(fit)
	val = fit.param[1]
	err = stderror(fit)[1]/val * 100
	"$(round(val, sigdigits=4))±$(round(err, sigdigits=4))%"
end

# ╔═╡ ef5f7640-7772-11eb-3d9e-4d7a092011df
function makeresults(all_data, gs, samples)
	tpl, rel, aff = loglogfit.(all_data)
	title!(tpl[1],
		"Number of Topples\nm = $(stringifyfit(tpl[2]))",
		titlefontsize=10)
	title!(rel[1],
		"Relaxation Time\nm = $(stringifyfit(rel[2]))",
		titlefontsize=10)
	title!(aff[1],
		"Sites affected\nm = $(stringifyfit(aff[2]))",
		titlefontsize=10)
	plot(tpl[1], rel[1], aff[1], layout=(1,3), size=(3,1).*225)
	md" ## Grid size = $(gs), Samples = $(samples) $(current())"
end

# ╔═╡ 6224d450-7772-11eb-0efd-a96d742fa1b3
data_50_1_5 = simulate(rand(1:hc - 1, 50, 50), 4, 100000);

# ╔═╡ 0c34ab90-76a2-11eb-2612-071caa9a3824
makeresults(data_50_1_5, 50, "1e5")

# ╔═╡ 934a6d60-778b-11eb-2e1c-0d125b4b09df
data_50_1_4 = simulate(rand(1:hc - 1, 50, 50), 4, 10000);

# ╔═╡ bcaba3b0-76b7-11eb-1c8a-ff1e69e0e16e
makeresults(data_50_1_4, 50, "1e4")

# ╔═╡ abd6c270-778b-11eb-3d22-41ba454ab170
data_50_5_4 = simulate(rand(1:hc - 1, 50, 50), 4, 50000);

# ╔═╡ eca0f07e-7788-11eb-174d-fb38d141adb7
makeresults(data_50_5_4, 50, "5e4")

# ╔═╡ f4b73920-778b-11eb-3987-13806219df5e
data_20_5_4 = simulate(rand(1:hc - 1, 20, 20), 4, 50000);

# ╔═╡ 24b4ac10-7788-11eb-2b05-d3a7f4244703
makeresults(data_20_5_4, 20, "5e4")

# ╔═╡ 6f953350-779a-11eb-29f1-ab01cd638438
data_80_5_4 = simulate(rand(1:hc - 1, 80, 80), 4, 50000);

# ╔═╡ 5ab6ae90-779b-11eb-1727-9d2cda89e743
makeresults(data_80_5_4, 80, "5e4")

# ╔═╡ 795f8790-782c-11eb-2143-b9d7220990d3
begin
	f(dat) = loglogfit.(dat) .|> x->stringifyfit(x[2])
	mktable(i) = "| Sample Size | 20x20                | 50x50                | 80x80                |
| ----------- | -------------------- | -------------------- | -------------------- |
| 1e4         | No Data              | $(f(data_50_1_4)[i]) | No Data              |
| 5e4         | $(f(data_20_5_4)[i]) | $(f(data_50_5_4)[i]) | $(f(data_80_5_4)[i]) |
| 1e5         | No Data              | $(f(data_50_1_5)[i]) | No Data              |"
end

# ╔═╡ 21db9b60-784c-11eb-1a10-29a5f75821a4
md"## Result Summary"

# ╔═╡ 93a7faa0-7823-11eb-36e1-512899012714
Markdown.parse(
"### Topple Count
$(mktable(1))"
)

# ╔═╡ 7149d3ce-782c-11eb-2074-055c558ee198
Markdown.parse(
"### Relaxation time
$(mktable(2))"
)

# ╔═╡ d2393230-782c-11eb-1d53-dfaddfeb0ed2
Markdown.parse(
"### Affected grids
$(mktable(3))"
)

# ╔═╡ a6502560-7788-11eb-150a-cdf40bd5cbbd
md"## Observations

1. All the values of slope lie within the range $[-1.1, -0.95]$
2. As the sampling size increases, the error reduces
3. As the grid size increases, the error reduces
4. The linear domain
    - increases with sample size until a point, after which it doesn't increase further
    - increases with grid size until a point, after which it doesn't increase further
"

# ╔═╡ 64b89820-782e-11eb-27df-c136babb8926
md"""# References
[^soc]: Self-organized criticality. Bak, Per and Tang, Chao and Wiesenfeld, Kurt. (Jul, 1988) Phys. Rev. A, 38: 364--374. doi: [10.1103/PhysRevA.38.364](https://link.aps.org/doi/10.1103/PhysRevA.38.364)
[^julia]: Julia: A Fresh Approach to Numerical Computing. Jeff Bezanson, Alan Edelman, Stefan Karpinski, Viral B. Shah. (2017) SIAM Review, 59: 65–98. doi: [10.1137/141000671](https://dx.doi.org/10.1137/141000671)
[^plots]: Plots - powerful convenience for visualization in Julia. Thomas Breloff. [docs.juliaplots.org](docs.juliaplots.org)
[^pluto]: Pluto - Reactive Notebooks for Julia. Fons van der Plas and Mikołaj Bochenski. [https://github.com/fonsp/Pluto.jl](https://github.com/fonsp/Pluto.jl)
[^plutoui]: PlutoUI. Fons van der Plas. [https://github.com/fonsp/PlutoUI.jl](https://github.com/fonsp/PlutoUI.jl)
[^stats]: Statistics - Julia stdlib for Statistics. [https://docs.julialang.org/en/v1/stdlib/Statistics/](https://docs.julialang.org/en/v1/stdlib/Statistics/)
"""

# ╔═╡ Cell order:
# ╟─c65eeff0-76b9-11eb-3543-3739a6484c24
# ╟─d846746e-76bb-11eb-2938-adb0dc210f71
# ╟─a3a53a50-7830-11eb-380b-ed771520d060
# ╟─f02e0ed0-784c-11eb-08b5-4746af933e1a
# ╟─2ed745de-76cf-11eb-101c-89f9d2ba170b
# ╟─1f57fcde-76cf-11eb-23c1-f5895772e792
# ╟─52d41aa0-769c-11eb-2eb1-e78eeed9d5c8
# ╟─b239c390-76cf-11eb-1c95-4bb63c153f87
# ╟─1f872dbe-76d0-11eb-055c-cd0487ce223a
# ╟─5221fe40-7784-11eb-2d14-b7ca964ba546
# ╠═b0797e20-769c-11eb-2711-7587647ce0c5
# ╠═7b2d6380-769c-11eb-12f6-35402dc5e96e
# ╠═8806b9d0-769c-11eb-28d9-bdc5b3870377
# ╟─b1c3cc10-76d6-11eb-38e1-75a2c56e0044
# ╠═75a73440-769c-11eb-2001-a5345c6e3732
# ╠═8dddef40-769c-11eb-2656-91b79e7af1cf
# ╟─1c132690-76d9-11eb-0edf-45dda119d525
# ╟─7db36160-76b8-11eb-22ec-bd23ba1800d6
# ╟─e0f4aab0-7797-11eb-0ecc-910d7b9ded81
# ╟─f052c870-76d9-11eb-1b82-bb53ba216dc8
# ╟─0e96d220-76dc-11eb-2b86-f5708fe4803f
# ╟─08191482-7781-11eb-37bb-1fb9abbceee0
# ╠═9484347e-769c-11eb-19c7-0961b3cdf2ef
# ╠═9b1c21e0-769c-11eb-162e-35dbc2aba3f2
# ╟─36049990-778c-11eb-3007-4dc0bb6b2615
# ╟─668d2040-769d-11eb-3943-51fb559fded3
# ╟─6a004e02-769d-11eb-01c4-e5d635d9aa91
# ╟─6a009c20-769d-11eb-2d53-534435f34669
# ╟─4fd4eef0-769d-11eb-023e-21941d0cb482
# ╟─361d0af0-7789-11eb-29ca-d768d81ccc39
# ╠═ef5f7640-7772-11eb-3d9e-4d7a092011df
# ╠═6224d450-7772-11eb-0efd-a96d742fa1b3
# ╟─0c34ab90-76a2-11eb-2612-071caa9a3824
# ╠═934a6d60-778b-11eb-2e1c-0d125b4b09df
# ╟─bcaba3b0-76b7-11eb-1c8a-ff1e69e0e16e
# ╠═abd6c270-778b-11eb-3d22-41ba454ab170
# ╟─eca0f07e-7788-11eb-174d-fb38d141adb7
# ╠═f4b73920-778b-11eb-3987-13806219df5e
# ╟─24b4ac10-7788-11eb-2b05-d3a7f4244703
# ╠═6f953350-779a-11eb-29f1-ab01cd638438
# ╟─5ab6ae90-779b-11eb-1727-9d2cda89e743
# ╟─795f8790-782c-11eb-2143-b9d7220990d3
# ╟─21db9b60-784c-11eb-1a10-29a5f75821a4
# ╟─93a7faa0-7823-11eb-36e1-512899012714
# ╟─7149d3ce-782c-11eb-2074-055c558ee198
# ╟─d2393230-782c-11eb-1d53-dfaddfeb0ed2
# ╟─a6502560-7788-11eb-150a-cdf40bd5cbbd
# ╟─64b89820-782e-11eb-27df-c136babb8926


#========================================================================================#
#	Laboratory 6
#
# Data visualisation
#
# Author: Niall Palfreyman, 27/01/2022
#========================================================================================#
[
	Activity(
		"""
		In this laboratory we learn how to use Julia's Makie plotting package to create
		publication-quality graphics for data visualisation. Makie is the front-end for
		several different backend graphics packages - for example, Cairo, GL and WebGL. We
		shall focus here on Cairo. First, we add and load the CairoMakie package. You already
		know how to do this using the Julia PackageManager, but like everything in Julia, we
		can also do it ourselves from within a program using method calls:

		using Pkg
		Pkg.add("CairoMakie")
		using CairoMakie

		This will take a while - as will also your first graphics call. Julia uses JiT (Just-
		in-time) compilation, so the first time you call a method, it will always take longer
		than the second time you call it. Enter the following command, then tell me the type of
		your return value fig:

		fig = scatterlines(0:10,(0:10).^2);
		""",
		"Enter the command exactly as I have done here",
		x -> x==Main.Makie.FigureAxisPlot
	),
	Activity(
		"""
		If you enter fig at the Julia prompt, you will see that a Makie figure contains a whole
		bunch of fields necessary for producing a graphic image. We will look at this image by
		saving it to a pdf file, but first we need to study the filesystem underlying the Julia
		REPL. Enter now the following command:

		pwd()

		The returned string tells you the path of your Present Working Directory. It should look
		something like this:

		"C:\\\\Users\\\\hswt136nia\\\\...\\\\Projects\\\\Ingolstadt\\\\src"

		In the chapter "Filesystem" of the Julia user manual you will find many functions for
		exploring the filesystem - we shall just use one or two here. Look at the return vale of
		the following command and tell me its type:

		readdir()
		""",
		"It should be a Vector of filenames in your PWD",
		x -> x <: Vector{String}
	),
	Activity(
		"""
		We want to create a Temp directory to save our files into. First, check whether there
		already exists such a directory in the PWD by entering:

		"Temp" in readdir()

		If the result is false, create the Temp subdirectory by entering:

		mkdir("Temp")

		What happens if you use this command when Temp already exists?
		""",
		"",
		x -> occursin( "error", lowercase(x)) || occursin( "exist", lowercase(x))
	),
	Activity(
		"""
		Now change directory into your new Temp directory, then check that you've changed PWD:

		cd("Temp")
		pwd()

		Now we're ready to save our Makie figure. Enter the following command, go and look at
		the resulting file using Adobe Acrobat, then come back here and tell me the highest
		number on the vertical axis:

		save("myfig.png", fig)
		""",
		"",
		x -> x==100
	),
	Activity(
		"""
		Didn't it look pretty? :) Wouldn't it be nice to view the graphic from within the Julia
		REPL? The REPL is a very basic environment that cannot display graphics, but we can easily
		include this functionality by adding and loading the package ElectronDisplay. When you have
		done this, just enter fig at the Julia prompt and ElectronDisplay will automatically
		display() your image.

		What is the type of the return value from this display?
		""",
		"",
		x -> x <: Main.Makie.FigureAxisPlot
	),
	Activity(
		"""
		Our figure is still a little primitive - let's customise its attributes. First notice
		that fig has a structure that includes a figure, an axis and a plot object:

		fg,ax,plt = fig;

		Check out the attributes field of plt to find out how many attributes it has:
		""",
		"",
		x -> x >= 15
	),
	Activity(
		"""
		We can ask for help on any function in the REPL by using the help() function. Try:

		help(lines)

		and then tell me how many different ways there are of calling the lines() function:
		""",
		"The signature of a function is the pattern of arguments that defines which method to call",
		x -> x==3
	),
	Activity(
		"""
		OK, now let's get fancy. We don't always want the scatter points on our plot, so we'll try
		out the lines() function. Also, it would be nice to plot something a little more exciting,
		like maybe the Hill function. The Hill function describes the activation and inhibition of
		DNA expression in biological cells by a transcription factor (TF):

		function hill( tf, K::Real=2, n::Real=1)
			abs_n = abs(n)
			if abs_n!=1 K = K^abs_n; tf = tf.^abs_n end
			n>=0 ? tf./(K.+tf) : K./(K.+tf)
		end
		trfactor = 0:20
		lines(trfactor,hill(trfactor))

		We can brighten up this figure by adding some interesting line attributes:

		lines( trfactor, hill(trfactor); color=:blue, linewidth=3, linestyle=:dash)

		The ; symbol marks the beginning of keyword arguments of a function. Experiment with this
		plot by changing the keyword arguments of lines(). What argument value displays a line
		consisting of a sequence of two dots and a dash (-..-..-..-)?
		""",
		"Search for keyword arguments on the site: https://makie.juliaplots.org/",
		x -> x==:dashdotdot
	),
	Activity(
		"""
		The lines() function returns a Plot object: the line that is to be plotted. This line is
		positioned within a set of axes: an Axis object. We can design the structure of this Axis
		object by adding an Axis specification to lines():

		lines( trfactor, hill(trfactor,2,1); color=:blue, label="activation",
			linewidth=3, linestyle=:dash,
			axis=(;
				xlabel="TF concentration", ylabel="Expression rate",
				title="Transcription regulation",
				xgridstyle=:dash, ygridstyle=:dash
			)
		)

		Notice carefully how I have used indentation to make it clear where each structure or
		function begins and ends. It's not easy to do this in the REPL, but soon we will be writing
		our own code, so it's good to practise making our code easily readable by other programmers.

		Experiment now with the axis parameters ...
		""",
		"",
		x -> true
	),
	Activity(
		"""
		This Axis object is then displayed within a Figure object. We can design the appearance of
		this figure by adding a Figure specification to lines():

		lines( trfactor, hill(trfactor); color=:blue, label="activation",
			linewidth=3, linestyle=:dash,
			axis=(;
				xlabel="TF concentration", ylabel="Expression rate",
				title="Transcription regulation",
				xgridstyle=:dash, ygridstyle=:dash
			),
			figure=(;
				figure_padding=5, resolution=(600,400), font="sans",
				backgroundcolor=:green, fontsize=20
			)
		)

		Experiment now with the figure parameters ...
		""",
		"",
		x -> true
	),
	Activity(
		"""
		Once we have set up a Plot object inside an Axis, inside a Figure, we can add extra Plot
		objects using bang! methods that change the current state of an object, for example:

		scatterlines!( trfactor, hill(trfactor,2,-1); color=:red, label="repression", linewidth=3)
		axislegend( "Legend"; position=:rc)

		Once we have done this, we can redisplay the current figure like this:

		current_figure()

		Take a look at the point where the curves of the two Hill-functions cross. You can see that
		this TF concentration correspondes to the value of K that we pass to the Hill function.
		Check that this is generally true, and tell me what value of Expression rate corresponds
		in general to the TF concentration value K:
		""",
		"Create a new graphic, and set the new value of K in BOTH calls to hill()",
		x -> x==0.5
	),
	Activity(
		"""
		Let's use our new-found knowledge of plotting to display a bubble-plot:

		data = randn(50,3)
		figr, ax, pltobj = scatter( data[:,1], data[:,2];
			color=data[:,3], label="Bubbles", colormap=:plasma, markersize=15*abs.(data[:,3]),
			axis=(; aspect=DataAspect()),
			figure=(; resolution=(600,400))
		)
		limits!(-3,3,-3,3)
		Legend( fig[1,2], ax, valign=:top)
		Colorbar( fig[1,2], pltobj, height=Relative(3/4))
		figr

		Do you like my pretty bubble-plot?
		""",
		"I'm very insecure: I won't let you go on until you say you like my art!",
		x -> occursin("yes",lowercase(x))
	),
	Activity(
		"""
		As you see, to create a nice graphic, we often need to build it up in several steps, so it
		is a idea to collect these steps together inside a single function. For example, suppose
		we want to display several different 2-d functions as heatmaps - we might do like this:

		function prettyheatmap( f::Function)
			figure = (; resolution=(600,400), font="CMU Serif")
			axis = (; xlabel="x", ylabel="y", aspect=DataAspect())
			xs = range(-2, 2, length = 25)
			ys = range(-2, 2, length = 25)
			zs = [f(x,y) for x in xs, y in ys]

			global fig, ax, pltobj = heatmap( xs, ys, zs, axis=axis, figure=figure)
			Colorbar( fig[1,2], pltobj, label="A colour bar")
			fig
		end

		Now call prettyheatmap() as follows, then tell me how many yellow rings you see:

		prettyheatmap( (x,y) -> cos(5*hypot(x,y)) )
		""",
		"Just count the complete rings - not the circle",
		x -> x==2
	),
	Activity(
		"""
		Now that we have wrapped our plotting code inside the function prettyheatmap(), we can
		easily reuse this encapsulated code to display a completely different function:

		mountains(x,y) =
			3 * (1-x)^2 *				exp(-(x^2+(y+1)^2))
			-(1/3) *					exp(-((x+1)^2 + y^2))
			-10 * (x/5 - x^3 - y^5) *	exp(-(x^2+y^2))
		prettyheatmap( mountains)

		How many mountains can you see?
		""",
		"Use the colour bar to count the number of patches where z > 0",
		x -> x==3
	),
	Activity(
		"""
		You see how useful it is to encapsulate (i.e.: hide) code? As soon as we have hidden the
		complicated graphics code inside the function myheatmap(), we can easily reuse this
		encapsulated code, by simply passing to it the name of the function we wish to plot:

		Encapsulation is THE CENTRAL tool for abstraction in modern computer languages!

		However, encapsulation can get broken! Before we move on to the next laboratory, I want
		you to notice something important. Do you remember our graph of the quadratic function
		y=x^2 that we saved in a file at the beginning of this laboratory? If you remember, we
		stored this graphic in a global variable fig, so we should be able to look at it again
		now, yes?

		OK, so please now display the quadratic graph contained in the global variable fig, and
		then, without doing anything else, move on to the next activity in this laboratory:

		fig
		""",
		"Your result may puzzle you, but please just move on to the next activity",
		x -> true
	),
	Activity(
		"""
		You should have found that the variable fig no longer contains the quadratic graph, but
		instead contains our most recent graph of the mountains() function.

		I want you to think now about this question: Why has the value of fig changed, and what
		single word can you delete in this laboratory to fix this problem? Don't worry if you
		can't immediately work it out: we will return to this important issue of encapsulation
		in laboratory 7.
		""",
		"",
		x -> true
	),
]


# ---
# title: 1589. Maximum Sum Obtained of Any Permutation
# id: problem1589
# author: Tian Jun
# date: 2020-10-31
# difficulty: Medium
# categories: Greedy
# link: <https://leetcode.com/problems/maximum-sum-obtained-of-any-permutation/description/>
# hidden: true
# ---
#
# We have an array of integers, `nums`, and an array of `requests` where
# `requests[i] = [starti, endi]`. The `ith` request asks for the sum of
# `nums[starti] + nums[starti + 1] + ... + nums[endi - 1] + nums[endi]`. Both
# `starti` and `endi` are _0-indexed_.
#
# Return _the maximum total sum of all requests **among all permutations** of_
# `nums`.
#
# Since the answer may be too large, return it **modulo** `109 + 7`.
#
#
#
# **Example 1:**
#
#
#
#     Input: nums = [1,2,3,4,5], requests = [[1,3],[0,1]]
#     Output: 19
#     Explanation: One permutation of nums is [2,1,3,4,5] with the following result:
#     requests[0] -> nums[1] + nums[2] + nums[3] = 1 + 3 + 4 = 8
#     requests[1] -> nums[0] + nums[1] = 2 + 1 = 3
#     Total sum: 8 + 3 = 11.
#     A permutation with a higher total sum is [3,5,4,2,1] with the following result:
#     requests[0] -> nums[1] + nums[2] + nums[3] = 5 + 4 + 2 = 11
#     requests[1] -> nums[0] + nums[1] = 3 + 5  = 8
#     Total sum: 11 + 8 = 19, which is the best that you can do.
#
#
# **Example 2:**
#
#
#
#     Input: nums = [1,2,3,4,5,6], requests = [[0,1]]
#     Output: 11
#     Explanation: A permutation with the max total sum is [6,5,4,3,2,1] with request sums [11].
#
# **Example 3:**
#
#
#
#     Input: nums = [1,2,3,4,5,10], requests = [[0,2],[1,3],[1,1]]
#     Output: 47
#     Explanation: A permutation with the max total sum is [4,10,5,3,2,1] with request sums [19,18,10].
#
#
#
# **Constraints:**
#
#   * `n == nums.length`
#   * `1 <= n <= 105`
#   * `0 <= nums[i] <= 105`
#   * `1 <= requests.length <= 105`
#   * `requests[i].length == 2`
#   * `0 <= starti <= endi < n`
#
#
## @lc code=start
using LeetCode

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


# This file is part of the ValidatedNumerics.jl package; MIT licensed

## Promotion

## Promotion rules

# Avoid ambiguity with ForwardDiff:

promote_rule{T<:Real, N, R<:Real}(::Type{Interval{T}},
    ::Type{ForwardDiff.Dual{N,R}}) = ForwardDiff.Dual{N, Interval{promote_type(T,R)}}

promote_rule{T<:Real, S<:Real}(::Type{Interval{T}}, ::Type{Interval{S}}) =
    Interval{promote_type(T, S)}

promote_rule{T<:Real, S<:Real}(::Type{Interval{T}}, ::Type{S}) =
    Interval{promote_type(T, S)}

promote_rule{T<:Real}(::Type{BigFloat}, ::Type{Interval{T}}) =
    Interval{promote_type(T, BigFloat)}

## Conversion rules

# convert{T<:Real}(::Type{Interval}, x::T) = convert(Interval{Float64}, x)

doc"""`split_interval_string` deals with strings of the form
\"[3.5, 7.2]\""""

function split_interval_string(T, x::AbstractString)
    if !(contains(x, "["))  # string like "3.1"
        return @thin_round(T, parse(T, x))
    end

    m = match(r"\[(.*),(.*)\]", x)  # string like "[1, 2]"

    if m == nothing
        throw(ArgumentError("Unable to process string $x as interval"))
    end

    @round(T, parse(T, m.captures[1]), parse(T, m.captures[2]))
end

# Floating point intervals:

convert{T<:AbstractFloat}(::Type{Interval{T}}, x::AbstractString) =
    split_interval_string(T, x)

function convert{T<:AbstractFloat, S<:Real}(::Type{Interval{T}}, x::S)
    Interval{T}( T(x, RoundDown), T(x, RoundUp) )
    # the rounding up could be down as nextfloat of the rounded down one
end

function convert{T<:AbstractFloat}(::Type{Interval{T}}, x::Float64)
    II = convert(Interval{T}, rationalize(x))
    # This prevents that rationalize(x) returns a zero when x is very small
    if x != zero(x) && II == zero(Interval{T})
        II = Interval{T}(string(x))
    end
    II
end

function convert{T<:AbstractFloat}(::Type{Interval{T}}, x::Interval)
    Interval{T}( T(x.lo, RoundDown), T(x.hi, RoundUp) )
end

# Complex numbers:
convert{T<:AbstractFloat}(::Type{Interval{T}}, x::Complex{Bool}) = (x == im) ?
    one(T)*im : throw(ArgumentError("Complex{Bool} not equal to im"))

# Rational intervals
function convert(::Type{Interval{Rational{Int}}}, x::Irrational)
    a = float(convert(Interval{BigFloat}, x))
    convert(Interval{Rational{Int}}, a)
end

function convert(::Type{Interval{Rational{BigInt}}}, x::Irrational)
    a = convert(Interval{BigFloat}, x)
    convert(Interval{Rational{BigInt}}, a)
end

convert{T<:Integer, S<:Integer}(::Type{Interval{Rational{T}}}, x::S) =
    Interval(x*one(Rational{T}))

convert{T<:Integer, S<:Integer}(::Type{Interval{Rational{T}}}, x::Rational{S}) =
    Interval(x*one(Rational{T}))

convert{T<:Integer, S<:Float64}(::Type{Interval{Rational{T}}}, x::S) =
    Interval(rationalize(T, x))

convert{T<:Integer, S<:BigFloat}(::Type{Interval{Rational{T}}}, x::S) =
    Interval(rationalize(T, x))


"A triangle defined by three points in a plane."
abstract type Triangle{T} <: Domain{SVector{2,T}} end
# This is just a conceptual type that acts as a data container. No functionality is implemented
# to manipulate triangles.

"The triangle below the diagonal of a square `[a,b] × [a,b]`."
struct LowerRightTriangle{T} <: Triangle{T}
   a  :: T
   b  :: T
end

DomainSets.indomain(x, d::LowerRightTriangle) =
   (x[1] >= d.a) && (x[1] <= d.b) && (x[2] >= d.a) && (x[2] <= d.b) && (x[1] <= x[2])

"The triangle above the diagonal of a square `[a,b] × [a,b]`."
struct UpperRightTriangle{T} <: Triangle{T}
   a  :: T
   b  :: T
end

DomainSets.indomain(x, d::UpperRightTriangle) =
   (x[1] >= d.a) && (x[1] <= d.b) && (x[2] >= d.a) && (x[2] <= d.b) && (x[1] >= x[2])

const StaticTypes = DomainSets.StaticTypes


import Base.isless

struct PixelKey{CT}
    val::CT
    time_step::Int
end
isless(a::PixelKey{T}, b::PixelKey{T}) where {T} = (a.val < b.val) || (a.val == b.val && a.time_step < b.time_step)

"""
```
segments                = watershed(img, markers)
```
Segments the image using watershed transform. Each basin formed by watershed transform corresponds to a segment.
If you are using image local minimas as markers, consider using [`hmin_transform`](@ref) to avoid oversegmentation.

Parameters:
-    img            = input grayscale image
-    markers        = An array (same size as img) with each region's marker assigned a index starting from 1. Zero means not a marker.
                      If two markers have the same index, their regions will be merged into a single region.
                      If you have markers as a boolean array, use `label_components`.

"""
function watershed(img::AbstractArray{T, N}, markers::AbstractArray{S,N}) where {T<:Images.NumberLike, S<:Integer, N}

    if axes(img) != axes(markers)
        error("image size doesn't match marker image size")
    end

    segments = copy(markers)
    pq = PriorityQueue{CartesianIndex{N}, PixelKey{T}}()
    time_step = 0

    R = CartesianIndices(axes(img))
    Istart, Iend = first(R), last(R)
    for i in R
        if markers[i] > 0
            for j in CartesianIndices(_colon(max(Istart,i-one(i)), min(i+one(i),Iend)))
                if segments[j] == 0
                    segments[j] = markers[i]
                    enqueue!(pq, j, PixelKey(img[i], time_step))
                    time_step += 1
                end
            end
        end
    end

    while !isempty(pq)
        current = dequeue!(pq)
        segments_current = segments[current]
        img_current = img[current]
        for j in CartesianIndices(_colon(max(Istart,current-one(current)), min(current+one(current),Iend)))
            if segments[j] == 0
                segments[j] = segments_current
                enqueue!(pq, j, PixelKey(img_current, time_step))
                time_step += 1
            end
        end
    end

    TM = meantype(T)
    num_segments        = Int(maximum(segments))
    labels              = Array(1:num_segments)
    region_means        = Dict{Int, TM}()
    region_pix_count    = Dict{Int, Int}()

    for i in R
        region_pix_count[segments[i]] = get(region_pix_count, segments[i], 0) + 1
        region_means[segments[i]] = get(region_means, segments[i], zero(TM)) + (img[i] - get(region_means, segments[i], zero(TM)))/region_pix_count[segments[i]]
    end
    return SegmentedImage(segments, labels, region_means, region_pix_count)
end

"""
```
out = hmin_transform(img, h)
```
Suppresses all minima in grayscale image whose depth is less than h.

H-minima transform is defined as the reconstruction by erosion of (img + h) by img. See Morphological image analysis by Soille pg 170-172.
"""
function hmin_transform(img::Array{T, N}, h::Real) where {T<:Images.NumberLike, N}
    out = img.+h
    while true
        temp = max.(img, erode(out))
        if temp == out
            break
        end
        out = temp
    end
    return out
end


using ADCME
using PyPlot
using ForwardDiff
using NiLang, NiLang.AD

# ADCME
function myloss_adcme(b; n=101)
    h = 1/(n-1)
    x = LinRange(0,1,n)[2:end-1]
    f = @. (4*(2 + x - x^2))

    u = trisolve(-b/h^2 * ones(n-3), 2b/h^2+1 * ones(n-2), -b/h^2 * ones(n-3), f)
    ue = u[div(n+1,2)] # extract values at x=0.5

    return (ue-1.0)^2
end

function benchmark_adcme(n)
    b = Variable(1.0)
    loss = myloss_adcme(b; n = n)
    g = gradients(loss, b)
    sess = Session(); init(sess)
    run(sess, g)
    ts_fwd = 0.0
    ts_bwd = 0.0
    for i = 1:11
        init(sess)
        d1 = @timed run(sess, loss)
        (i>1) && (ts_fwd += d1[2])
        init(sess)
        d2 = @timed run(sess, g)
        (i>1) && (ts_bwd += d2[2])
        @info "adcme", n, d1[2], d2[2]
    end
    ts_fwd /= 10
    ts_bwd /= 10
    return ts_fwd, ts_bwd
end

# ForwardDiff
function trisolve!(a::T, b::T, c::T, D::AbstractVector, X::AbstractVector) where T
    N = length(X)
    D = copy(D)
    B = zeros(T, N)
    @inbounds B[1] = b
    @inbounds for i = 2:N
        w = a / B[i-1]
        B[i] = b - w * c
        D[i] = D[i] - w * D[i - 1]
    end
    @inbounds X[N] = D[N] / B[N]
    @inbounds for i = N-1:-1:1
        X[i] = (D[i] - c * X[i + 1]) / B[i]
    end
    return X
end

function myloss(b::T; n=101) where T
    h = 1/(n-1)
    x = LinRange(0,1,n)[2:end-1]
    f = @. T(4*(2 + x - x^2))

    u = trisolve!(-b/h^2, 2b/h^2+1, -b/h^2, f, zeros(T, n-2))
    ue = u[div(n+1,2)] # extract values at x=0.5

    return (ue-1.0)^2
end

function benchmark_forwarddiff(n)
    ts_fwd = 0.0
    ts_bwd = 0.0
    for i = 1:11
        d1 = @timed myloss(10.0, n = n)
        (i>1) && (ts_fwd += d1[2])
        d2 = @timed myloss(ForwardDiff.Dual(10.0, 1.0),  n = n)
        (i>1) && (ts_bwd += d2[2])
        @info "forwarddiff", n, d1[2], d2[2]
    end
    ts_fwd /= 10
    ts_bwd /= 10
    ts_fwd, ts_bwd
end

# NiLang
@i function i_trisolve!(a::T, b::T, c::T, D!::AbstractVector, X!::AbstractVector, B!::AbstractVector) where T
    @invcheckoff @inbounds begin
    B![1] += b
    for i = 2:length(X!)
        @routine begin
            w ← zero(T)
            w += a / B![i-1]
        end
        B![i] += b
        B![i] -= w * c
        D![i] -= w * D![i - 1]
        ~@routine
    end
    X![end] += D![end] / B![end]
    for i = length(X!)-1:-1:1
        @routine begin
            anc ← zero(T)
            anc += D![i]
            anc -= c * X![i + 1]
        end
        X![i] += anc / B![i]
        ~@routine
    end
    end
end

@i function i_myloss!(loss::T, f!::AbstractVector{T}, u!::AbstractVector{T},
              b_cache!::AbstractVector{T}, b::T) where T
    @invcheckoff @inbounds begin
    n ← length(f!) + 2
    h ← zero(T)
    h += 1 / (n-1)
    for i=1:length(f!)
        @routine begin
            @zeros T xi anc
            xi += i * h
            anc += 2 + xi
            anc -= xi ^ 2
        end
        f![i] += 4 * anc
        ~@routine
    end

    @routine begin
        @zeros T h2 factor_a factor_b factor_c
        h2 += h^2
        factor_a -= b / h2
        factor_c += factor_a
        factor_b -= 2 * factor_a
        factor_b += 1
    end
    i_trisolve!(factor_a, factor_b, factor_c, f!, u!, b_cache!)
    ~@routine
    @routine begin
        ue ← zero(T)
        ue += u![div(n+1,2)] # extract values at x=0.5
        ue -= 1
    end
    loss += ue^2
    ~@routine
    h -= 1 / (n-1)
    end
end

function benchmark_nilang(n)
    ts_fwd = 0.0
    ts_bwd = 0.0
    for i = 1:11
        d1 = @timed i_myloss!(0.0, zeros(n-2), zeros(n-2), zeros(n-2), 10.0)
        (i>1) && (ts_fwd += d1[2])
        d2 = @timed Grad(i_myloss!)(Val(1), 0.0, zeros(n-2), zeros(n-2), zeros(n-2), 10.0)
        (i>1) && (ts_bwd += d2[2])
        @info "nilang", n, d1[2], d2[2]
    end
    ts_fwd /= 10
    ts_bwd /= 10
    ts_fwd, ts_bwd
end

tsf, tsb = zeros(3,6), zeros(3,6)
for (k,n) in enumerate([101, 1001,10001, 100001, 1000001, 10000001])
    tsf[1,k], tsb[1,k] = benchmark_adcme(n)
    tsf[2,k], tsb[2,k] = benchmark_forwarddiff(n)
    tsf[3,k], tsb[3,k] = benchmark_nilang(n)
end

figure(figsize=(10,4))
subplot(121)
plot([101, 1001,10001, 100001, 1000001, 10000001], tsf[1,:], "o-", label="ADCME")
plot([101, 1001,10001, 100001, 1000001, 10000001], tsf[2,:], "o-", label="Julia")
plot([101, 1001,10001, 100001, 1000001, 10000001], tsf[3,:], "o-", label="NiLang")
title("Forward Computation")
legend()
xlabel("\$n\$")
ylabel("Time (seconds)")
grid("on")
subplot(122)
plot([101, 1001,10001, 100001, 1000001, 10000001], tsb[1,:], "o-", label="ADCME")
plot([101, 1001,10001, 100001, 1000001, 10000001], tsb[2,:], "o-", label="ForwardDiff")
plot([101, 1001,10001, 100001, 1000001, 10000001], tsb[3,:], "o-", label="NiLang")
legend()
title("Gradient Back-Propagation")
xlabel("\$n\$")
ylabel("Time (seconds)")
grid("on")
savefig("benchmark.png")


# ------------------------------------------------------------------
# Licensed under the MIT License. See LICENCE in the project root.
# ------------------------------------------------------------------

module GeoEstimation

using Meshes
using GeoStatsBase
using Variography
using KrigingEstimators

using NearestNeighbors
using Distances
using LinearAlgebra

import GeoStatsBase: solve

include("idw.jl")
include("lwr.jl")
include("kriging.jl")

export IDW, LWR, Kriging

end


"""
    add_crowding_constraints(weights::Dict{Int64, Float64})

Adds a molecular crowding constraint to the optimization problem: `∑ wᵢ × vᵢ ≤ 1` where `wᵢ`
is a weight and `vᵢ` is a flux index in the model's reactions specified in `weights` as `vᵢ
=> wᵢ` pairs.

See Beg, Qasim K., et al. "Intracellular crowding defines the mode and sequence of
substrate uptake by Escherichia coli and constrains its metabolic activity." Proceedings of
the National Academy of Sciences 104.31 (2007) for more details.
"""
add_crowding_constraints(weights::Dict{Int64,Float64}) =
    (model, opt_model) -> begin
        idxs = collect(keys(weights)) # order of keys and values is the same
        ws = values(weights)
        # since fluxes can be positive or negative, need absolute value: ∑ wᵢ × |vᵢ| ≤ 1
        # introduce slack variables to handle this
        @variable(opt_model, crowding_slack[1:length(weights)])
        @constraint(opt_model, crowding_slack .>= opt_model[:x][idxs])
        @constraint(opt_model, crowding_slack .>= -opt_model[:x][idxs])
        @constraint(opt_model, sum(w * crowding_slack[i] for (i, w) in enumerate(ws)) <= 1)
    end

"""
    add_crowding_constraints(weights::Dict{String, Float64})

Variant of [`add_crowding_constraints`](@ref) that takes a dictinary of reactions `ids`
instead of reaction indices mapped to weights.
"""
add_crowding_constraints(weights::Dict{String,Float64}) =
    (model, opt_model) -> begin
        idxs = indexin(keys(weights), reactions(model))
        nothing in idxs && throw(ArgumentError("Reaction id not found in model."))
        add_crowding_constraints(Dict(zip(Int.(idxs), values(weights))))(model, opt_model)
    end


# This was run in Julia 1.5.0

using BGEN, OrdinalMultinomialModels, Random, DataFrames, CSV

bgenfile = "bgen_test.bgen"
bgenfile = "/Users/christophergerman/.julia/dev/OrdinalGWAS/data/bgen_test.bgen"
b = Bgen(BGEN.datadir("example.8bits.bgen"))

people = n_samples(b)
nsnps = n_variants(b)
sampleids = samples(b)
variants = parse_variants(b; from_bgen_start=false)
#sample 3 and 15 and 25 to be causal
v3 = minor_allele_dosage!(b, variants[3]; T=Float64, mean_impute=true)
v25 = minor_allele_dosage!(b, variants[25]; T=Float64, mean_impute=true)

Random.seed!(310)
sex = rand(0:1, people)
X = [sex v3 v25]
θ = [-1.5; 0; 1.5]
β = [0.5; 0.5; -0.5]
y = rpolr(X, β, θ, LogitLink())

df = DataFrame(y = y, sex = sex)
CSV.write("bgen_ex.csv", df)

using OrdinalGWAS

b_rsids = rsids(b)
reps1, rem1 = divrem(nsnps, 10)
repeats = [fill(10, reps1); fill(rem1, 1)]
genenames = ["gene" * string(i) for i in 1:length(repeats .!= 0)]
annotations = vcat(map((gname, reps) ->
    fill(gname, reps),
    genenames, repeats)...)
dfannote = DataFrame(gene = annotations, rsids = b_rsids)
CSV.write("bgen_snpsetfile.txt", dfannote, delim = " ",
    header = false)


using QuantumCumulants
using Test

@testset "cluster" begin

order = 4
N_c = 3 #number of clusters
N = [Parameter(Symbol(:N_, i)) for i=1:N_c]

hf = FockSpace(:cavity)
ha = [NLevelSpace(Symbol(:atoms, j),3) for j=1:N_c]
ha_c = [ClusterSpace(ha[j],N[j],order) for j=1:N_c]
h = ⊗(hf, ha_c...)
# Define the fundamental operators
a = Destroy(h,:a,1)
S(i,j,c) = Transition(h,Symbol(:σ, c),i, j, 1+c) #c=cluster

@test QuantumCumulants.has_cluster(h)
@test !(QuantumCumulants.has_cluster(hf))

@test S(2,2,1) ≠ S(2,2,2) ≠ S(2,2,3)
@test isequal(S(2,2,2)[1]*S(2,2,1)[2], S(2,2,1)[2]*S(2,2,2)[1])
@test length(S(2,2,1)) == length(S(2,2,2)) == length(S(2,2,3)) == order
@test acts_on(S(2,2,1)[1]) == QuantumCumulants.ClusterAon(2,1)
@test acts_on(S(2,2,2)[2]) == QuantumCumulants.ClusterAon(3,2)
@test acts_on(S(2,2,1)[1]) < acts_on(S(2,2,1)[2]) < acts_on(S(2,2,2)[1])

end #testset


function klamp(t, a, b)
    if t <= a
        return a
    elseif b <= t
        return b
    else
        da = abs(t - a)
        db = abs(t - b)

        if da < db
            return a
        else
            return b
        end
        # return t
    end
end

function contact_point_box(ps, xc, qc, kx, ky, kz)
    # rotate point into box frame
    ps_box_frame = vector_rotate(ps - xc, inv(qc))

    # find box point in box frame
    px = clamp(ps_box_frame[1], -0.5 * kx, 0.5 * kx)
    py = clamp(ps_box_frame[2], -0.5 * ky, 0.5 * ky)
    pz = clamp(ps_box_frame[3], -0.5 * kz, 0.5 * kz)

    # box point in world frame
    return vector_rotate([px; py; pz], qc) + xc
end

function ∂contact_point_box∂p(ps, xc, qc, kx, ky, kz)

    FiniteDiff.finite_difference_jacobian(p -> contact_point_box(p, xc, qc, kx, ky, kz), ps)
end

function ∂contact_point_box∂x(ps, xc, qc, kx, ky, kz)

    FiniteDiff.finite_difference_jacobian(x -> contact_point_box(ps, x, qc, kx, ky, kz), xc)
end

function ∂contact_point_box∂q(ps, xc, qc, kx, ky, kz)

    FiniteDiff.finite_difference_jacobian(q -> contact_point_box(ps, xc, Quaternion(q..., false), kx, ky, kz), vector(qc))
end

# x_min = -1.0
# x_max = 1.0
# z_min = -1.0
# z_max = 1.0

# p = [-1.0; 1.0] #.- 1.0e-6

# function klamp(t, a, b)
#     if t <= a
#         return a
#     elseif b <= t
#         return b
#     else
#         da = abs(t - a)
#         db = abs(t - b)

#         if da < db
#             return a
#         else
#             return b
#         end
#     end
# end

# function nearest_point(p, x_min, x_max, z_min, z_max)
#     return [klamp(p[1], x_min, x_max), klamp(p[2], z_min, z_max)]
# end

# nearest_point(p, x_min, x_max, z_min, z_max)


"""
    A, B, C, D = partitionABCD(ABCD, m)

Partition `ABCD` into `A`, `B`, `C`, `D` for an `m`-input state-space system.
"""
function partitionABCD(ABCD, m=NaN)
    if isnan(m)
        n = minimum(size(ABCD)) - 1
        m = size(ABCD, 2) - n
    else
        n = size(ABCD, 2) - m
    end
    r = size(ABCD, 1) - n

    A = ABCD[1:n, 1:n]
    B = ABCD[1:n, n+1:n+m]
    C = ABCD[n+1:n+r, 1:n]
    D = ABCD[n+1:n+r, n+1:n+m]

    return A, B, C, D
end


# compare times of direct calculation of c's compared to Crandall's recursive calculation
using Polylogarithms
using BenchmarkTools
using DataFrames, CSV
using PyPlot
using SpecialFunctions
import Base: parse
using Printf

Q = Polylogarithms.Q
c = Polylogarithms.c
max_terms = 10
max_n = 5

s = 5.0001;
r = s - n;
μ = -0.6931471805599453 + 3.141592653589793im #  = log(z) = log(-0.5)
ℒ = log(complex(-μ))  # '\u2112'

result_direct = Array{Any}(undef, max_n, max_terms)
result_crandall = Array{Any}(undef, max_n, max_terms)

for k=0:max_n-1
    for j=1:max_terms
       if j<=3
            # b_direct = @benchmarkable Polylogarithms.c($k, $j, $ℒ)
            # result_crandall[k+1, n_terms] = run(b_direct)
            result_crandall[k+1, n_terms] =  minimum(@benchmark Polylogarithms.c($k, $j, $ℒ))
        end
        # b_crandall = @benchmarkable Polylogarithms.c_crandall($k, $j, $ℒ)
        # result_crandall[k+1, n_terms] = run(b_crandall)
        result_crandall[k+1, n_terms] = minimum(@benchmark Polylogarithms.c_crandall($k, $j, $ℒ))
   end
end


import Base: iterate
import Base.Iterators.take

"""
Structure to hold the iterates of the lifted coordinate descent method.

LiftedCCIterable(M, b, A, λ; optTol=√eps)

Construct a lifted coordinate descent iterator for the problem

    minimize  ½‖Mx-b‖² + λ⋅gauge(A|x)

The absolute tolerance `optTol` is used to determine the stopping condition.

"""
mutable struct LiftedCCIterable{T1<:Float64, T2<:Vector{T1}, T3<:AbstractArray{T1}, T4<:AbstractLinearOp, T5<:AbstractAtomicSet}
    M::T4           # linear operator
    b::T2           # observation
    A::T5           # atomic set
    r::T2           # residual
    z::T3           # negative gradient
    Ma::T3          # image of the atom
    Mx::T3          # image of primal variable
    λ::T1           # current scale of atomic set
    optTol::T1      # tolerance
    function LiftedCCIterable(M::AbstractLinearOp, b::Vector{Float64}, A::AbstractAtomicSet, λ::Float64, optTol = sqrt(eps(1.0)))
        r = copy(b)
        z = M'*r
        Ma = M*expose(A, z)
        Mx = copy(Ma); fill!(Mx, 0.0)
        new{Float64, Vector{Float64}, AbstractArray{Float64}, AbstractLinearOp, AbstractAtomicSet}(M, b, A, r, z, Ma, Mx, λ, optTol)
    end
end

"""
Defines one complete iteration of the lifted coordinate descent method.
"""
function iterate(p::LiftedCCIterable, k::Int=0)

    # Compute  M*a  where  a = expose(A, z)
    s = support(p.A, p.z)
    a = expose(p.A, p.z)
    p.Ma .= p.M*a

    # Compute g = λ - support(A|z) and possible exit
    g = p.λ - s

    # Compute step size
    δ = -g / (norm(p.Ma)^2)

    # update iterates
    p.Mx .+= δ*p.Ma
    p.r .= p.b - p.Mx
    p.z .= p.M' * p.r

    return (g, p.r), k+1
end

"""
Return an exit flag for the lifted coordinate descent algorithm.
"""
function checkExitCd(g::Float64, u::Float64, k::Int64, α::Float64, feaTol::Float64, maxIts::Int64)
    u < α + feaTol && return :feasible                  # feasible
    g > -sqrt(eps(1.0))/2 && return :optimal                    # optimal
    k ≥ maxIts && return :iterations                    # out of iterations
    return :noerror                                     # otherwise
end

"""
Lifted coordinate descent method for the problem

    min ½‖Mx-b‖² + λ⋅gauge(A|x)
"""
function coordinate_descent(M::AbstractLinearOp, b::Vector{Float64}, A::AbstractAtomicSet, λ::Float64;
                   α::Float64 = 0.0,
                   tol::Float64 = 1e-12,
                   maxIts::Int = size(M,2),
                   logger::Bool = true
                   )

    feaTol = tol*(1+norm(b))
    # Structure to hold the iterates
    lcc = LiftedCCIterable(M, b, A, λ)
    # Structure to hold the solution
    sol = Solution(M, b, A)

    # head logging
    m, n = size(M)
    logger && logger_head_cd(m, n, feaTol, α, maxIts)

    # Check whether b is zero
    if norm(b) < tol
        return sol
    end

    # ----------------------------------------------------------------
    # min loop
    # ----------------------------------------------------------------
    k = 0
    exitFlag = :noerror
    for (g, r) in lcc

        exitFlag = primalrecover!(lcc, sol, α, feaTol, exitFlag)
        exitFlag == :noerror || break
        k = k + 1
        u = norm(r)^2/2
        exitFlag = checkExitCd(g, u, k, α, feaTol, maxIts)

        logger && logger_level_cd(α, k, u, g, exitFlag, sol.feas)
    end

    primalrecover!(lcc, sol, α, feaTol, exitFlag)

    # foot logging
    logger&& logger_foot_cd(b, sol.feas, k)

    return sol
end

function logger_head_cd(m, n, feaTol, α, maxIts)
    println()
    println("  -------------------------------------------------------------------------")
    println("  Lifted Coordinate Descent Method")
    println("  -------------------------------------------------------------------------")
    @printf("  %-22s %7d %7s %-20s %7d\n","number of variables",n,"","number of constraints",m)
    @printf("  %-22s %7.2e %7s %-20s %7.2e\n","feasibility tolerance",feaTol,"","α",α)
    @printf("  %-22s %7d \n","max iterations",maxIts)
    println("  -------------------------------------------------------------------------")
    @printf("  %8s   %8s   %8s   %8s   %10s\n",
            "iteration","u-α","g","infeas-α","exitFlag")
end

function logger_foot_cd(b, feas, totIts)
    scale = 1+norm(b)
    normr = sqrt(2*feas)
    println("  -------------------------------------------------------------------------")
    @printf("  %-22s %8.1e \n","residual (abs)",normr)
    @printf("  %-22s %8.1e \n","residual (rel)",normr/scale)
    @printf("  %-22s %8d \n","total iterations",totIts)
    println("  -------------------------------------------------------------------------")
end

function logger_level_cd(α, k, u, g, exitFlag, feas)
    @printf("  %8d   %8.2e   %8.2e   %8.2e   %s\n",
            k, u-α, g, feas-α, exitFlag)
end

"""
    primal recover
"""
function primalrecover!(p::LiftedCCIterable, sol::Solution, α::Float64, feaTol::Float64, exitFlag)
    flag = exitFlag
    s = support(p.A, p.z)
    F = face(p.A, p.z/s)
    ϵ = copy(p.r); ϵ .*= sqrt(2*α)/norm(ϵ)
    c, r = face_project(p.M, F, p.b - ϵ)
    feas = norm(r + ϵ)^2/2
    # @show feas
    if feas ≤ sol.feas
        sol.F = F
        sol.c = c
        sol.feas = feas
    end
    if feas - α ≤ feaTol
        flag = :feasible
    elseif flag == :optimal
        flag = :noerror
        replace_iterate!(p, sol)
    end
    return flag
end

"""
    replace iterates
"""
function replace_iterate!(p::LiftedCCIterable, sol::Solution)
    MF = p.M * sol.F
    p.Mx .= MF * sol.c
    p.r .= p.b - p.Mx
    p.z .= p.M' * p.r
end


module PenaltyFunctions

using LinearAlgebra
using InteractiveUtils
using RecipesBase

# trait functions
include("traits.jl")

# penalty functions
include("common.jl")

# IO and plot recipes
include("printing.jl")
include("plotrecipes.jl")

export
    Penalty,
        ElementPenalty,
            ProxableElementPenalty,
                NoPenalty,
                L1Penalty,
                L2Penalty,
                ElasticNetPenalty,
            SCADPenalty,
            MCPPenalty,
            LogPenalty,
        ArrayPenalty,
            NuclearNormPenalty,
            GroupLassoPenalty,
            MahalanobisPenalty,
    addgrad, prox, prox!, deriv, value, grad, grad!, addgrad!, scaled

end


function default_in_partition(sites::Tuple{Int}, p::Integer, nparts::Integer)
  return p == mod1(sites[1], nparts)
end

# i ≥ j
_f(i, j) = (i - 1) * i ÷ 2 + j
function default_in_partition(sites::NTuple{2,Int}, p, nparts)
  i, j = sites
  if i ≤ j
    return p == mod1(_f(j, i), nparts)
  end
  return p == mod1(_f(i, j), nparts)
end

function default_in_partition_sorted(sites::NTuple{4,Int}, p, nparts)
  i, j, k, l = sites
  if j == k
    return default_in_partition((j,), p, nparts)
  end
  return default_in_partition((i, j), p, nparts)
end

function default_in_partition(sites::NTuple{4,Int}, p, nparts)
  return default_in_partition_sorted(sort(sites), p, nparts)
end

function _sites(t)
  return Tuple((only(ITensors.site(o)) for o in ITensors.terms(t)))
end

function partition(os::OpSum, nparts::Integer; in_partition=default_in_partition)
  oss = [OpSum() for _ in 1:nparts]
  for p in 1:nparts
    for n in 1:length(os)
      sites = _sites(os[n])
      if in_partition(sites, p, nparts)
        add!(oss[p], os[n])
      end
    end
  end
  return oss
end


using Test, WAVI

@testset "WAVI tests" begin
    @testset "Pos Fraction" begin
        @info "testing pos fraction"
        z=[-1.0 -1.0 -1.0;-1.0 1.0 -1.0;-1.0 -1.0 -1.0]
        pfh,pfu,pfv=WAVI.pos_fraction(z)
        @test pfh == [0.0 0.0 0.0;0.0 0.5 0.0;0.0 0.0 0.0]
        @test pfu == [0.0 0.0 0.0;0.0 0.25 0.0;0.0 0.25 0.0;0.0 0.0 0.0]
        @test pfv == [0.0 0.0 0.0 0.0;0.0 0.25 0.25 0.0;0.0 0.0 0.0 0.0]
    end

    @testset "Pos Fraction Mask" begin
        @info "testing post fraction mask"
        z=[-1.0 -1.0 -1.0;-1.0 1.0 -1.0;-1.0 -1.0 -1.0]
        mask=[false false false; false true false; false false false]
        pfh,pfu,pfv=WAVI.pos_fraction(z,mask=mask)
        @test pfh == [0.0 0.0 0.0;0.0 1.0 0.0;0.0 0.0 0.0]
        @test pfu == [0.0 0.0 0.0;0.0 0.5 0.0;0.0 0.5 0.0;0.0 0.0 0.0]
        @test pfv == [0.0 0.0 0.0 0.0;0.0 0.5 0.5 0.0;0.0 0.0 0.0 0.0]
    end

    @testset "Height and volume above floatation" begin
        @info "testing height and volume above floatation functions"
        grid = Grid(nx = 20, ny = 25, dx = 1000., dy = 1000.)
        bed = -200.0 * ones(grid.nx, grid.ny)
        initial_conditions = InitialConditions(initial_thickness = 500.0*ones(grid.nx, grid.ny))
        model = Model(grid = grid, bed_elevation = bed, initial_conditions = initial_conditions)
        update_state!(model)
        @test volume_above_floatation(model.fields.gh.h, model.fields.gh.b, Ref(model.params), model.grid) ≈ (500.0 - (1028.0/918.0)*(200.0) ) .* 25 .* 20 .* 1000.0 .^2

    end
end


"""
    Exo1Real{T <: Real} = Parabolic{T}

A 1-exo-real number is a parabolic Cayley-Dickson construct with real numbers.
"""
const Exo1Real{T <: Real} = Parabolic{T}

function Exo1Real(a::T, b::T) where T <: Real
    Exo1Real{T}(a, b)
end

function Exo1Real(a::T) where T <: Real
    Exo1Real{T}(a, zero(T))
end

function Exo1Real(a::Real, b::Real)
    Exo1Real(promote(a, b)...)
end

function selfstar(z::Exo1Real)
    (z + star(z)) / 2
end

function anti_selfstar(z::Exo1Real)
    (z - star(z)) / 2
end

"""
    unreal(z::Exo1Real)

The unreal part of a 1-exo-real number is a 1-vector.
"""
function unreal(z::Exo1Real{T}) where T <: Real
    T[z.r]
end

function show(io::IO, z::Exo1Real)
    print(io, "[1: ")
    print(io, z.l)
    print(io, ", a: ")
    print(io, z.r)
    print(io, "]")
end

function random(::Type{Exo1Real{T}}) where T <: Real
    Exo1Real{T}(rand(T), rand(T))
end

function random(::Type{Exo1Real{BigInt}})
    Exo1Real{BigInt}(randomBigInt(), randomBigInt())
end

function random(::Type{Exo1Real{BigFloat}})
    Exo1Real{BigFloat}(randomBigFloat(), randomBigFloat())
end

function random(::Type{Exo1Real{Rational{BigInt}}})
    Exo1Real{Rational{BigInt}}(randomBigRational(), randomBigRational())
end

function convert(::Type{Exo1Real{T}}, a::Real) where T <: Real
    Exo1Real{T}(a, zero(T))
end

function convert(::Type{Exo1Real{T}}, z::Exo1Real) where T <: Real
    Exo1Real{T}(z.l, z.r)
end

function promote_rule(::Type{Exo1Real{T}}, ::Type{S}) where {T <: Real, S <: Real}
    Exo1Real{promote_type(T, S)}
end

function promote_rule(::Type{Exo1Real{T}}, ::Type{Exo1Real{S}}) where {T <: Real, S <: Real}
    Exo1Real{promote_type(T, S)}
end


# Givens representation of Q matrix, memory efficient compared to dense representation
struct GivensQ{T, QT<:AbstractVector{<:Givens{T}}} <: Factorization{T}
    rotations::QT
    n::Int
    m::Int
end

# r controls how many columns of Q we instantiate
# NOTE: this also works if rotations is empty (returns identity matrix)
function Base.Matrix(F::GivensQ, r::Int = size(F, 2))
    n, m = size(F)
    Q = Matrix(one(eltype(F))*I, n, r)
    return lmul!(F, Q)
end
Base.size(F::GivensQ) = (F.n, F.n) # could be square but also (n, m)?
Base.size(F::GivensQ, i::Int) = i > 2 ? 1 : size(F)[i]
Base.copy(F::GivensQ) = GivensQ(copy(F.rotations), F.n, F.m)
Base.adjoint(F::GivensQ) = Adjoint(F)
# NOTE: this also works if rotations is empty (is identity operator)
function LinearAlgebra.lmul!(F::GivensQ, X::AbstractVecOrMat)
    m = size(X, 1)
    for G in Iterators.reverse(F.rotations) # the adjoint of a GivensQ reverses the order of the rotations
        lmul!(G', X) # ... and takes their adjoint
    end
    return X
end
function LinearAlgebra.lmul!(A::Adjoint{<:Any, <:GivensQ}, X::AbstractVecOrMat)
    F = A.parent
    for G in F.rotations # the adjoint of a GivensQ reverses the order of the rotations
        lmul!(G, X) # ... and takes their adjoint
    end
    return X
end
LinearAlgebra.rmul!(X::AbstractVecOrMat, F::GivensQ) = lmul!(F', X')'
LinearAlgebra.rmul!(X::AbstractVecOrMat, F::Adjoint{<:Any, <:GivensQ}) = lmul!(F', X')'

lmul(F, X) = lmul!(F, copy(X))
rmul(X, F) = rmul!(copy(X), F)

Base.:*(F::GivensQ, X::AbstractVector) = lmul(F, X)
Base.:*(F::GivensQ, X::AbstractMatrix) = lmul(F, X)
Base.:*(F::Adjoint{<:Any, <:GivensQ}, X::AbstractVector) = lmul(F, X)
Base.:*(F::Adjoint{<:Any, <:GivensQ}, X::AbstractMatrix) = lmul(F, X)

Base.:*(X::AbstractMatrix, F::GivensQ) = rmul(X, F)
Base.:*(X::AbstractMatrix, F::Adjoint{<:Any, <:GivensQ}) = rmul(X, F)

function Base.:*(F::GivensQ, G::Adjoint{<:Any, <:GivensQ})
    T = promote_type(eltype(F), eltype(G))
    F === G' ? (one(T) * I)(F.n) : F * Matrix(G)
end
function Base.:*(F::Adjoint{<:Any, <:GivensQ}, G::GivensQ)
    T = promote_type(eltype(F), eltype(G))
    F' === G ? (one(T) * I)(G.n) : F * Matrix(G)
end
function Base.:(==)(F::GivensQ, G::GivensQ)
    F.rotations == G.rotations && F.n == G.n && F.m == G.m
end

LinearAlgebra.ldiv!(F::GivensQ, x::AbstractVector) = lmul!(F', x)
LinearAlgebra.ldiv!(F::Adjoint{<:Any, <:GivensQ}, x::AbstractVector) = lmul!(F', x)
LinearAlgebra.ldiv!(F::GivensQ, X::AbstractMatrix) = lmul!(F', X)
LinearAlgebra.ldiv!(F::Adjoint{<:Any, <:GivensQ}, X::AbstractMatrix) = lmul!(F', X)

function LinearAlgebra.mul!(y::AbstractVector, F::GivensQ, x::AbstractVector)
    @. y = x
    lmul!(F, y)
end
function LinearAlgebra.mul!(Y::AbstractMatrix, F::GivensQ, X::AbstractMatrix)
    @. Y = X
    lmul!(F, Y)
end
function LinearAlgebra.mul!(y::AbstractVector, F::Adjoint{<:Any, <:GivensQ}, x::AbstractVector)
    @. y = x
    lmul!(F, y)
end
function LinearAlgebra.mul!(Y::AbstractMatrix, F::Adjoint{<:Any, <:GivensQ}, X::AbstractMatrix)
    @. Y = X
    lmul!(F, Y)
end


using Test
using SimplePosets
using LinearAlgebra
using SimpleGraphs

@test 1==1

P = Chain(5) + Chain(3)
@test height(P) == 5

P = StandardExample(4)
A = minimals(P)
@test length(A) == 4
B = maximals(P)
@test length(B) == 4
@test length(union(A,B))==8

P = PartitionLattice(5)
@test card(P) == 52
@test height(P) == 5
@test inv(P') == P

P = Divisors(2*3*5)
M = mobius_matrix(P)
Z = zeta_matrix(P)
n = card(P)
@test M*Z == Matrix{Int}(I,n,n)

P = BooleanLattice(3)
@test element_type(P) == String

P = RandomPoset(10,2)
@test card(P) == 10

P = Divisors(2*3*5)
@test elements(relabel(P)) == collect(1:8)
@test length(random_linear_extension(P)) == 8

d = random_average_height(P,20)
@test d[1]==0.0


using NetworkEpidemics
using LightGraphs
using Statistics
using Plots
using ColorSchemes
using Random

CorrectedMetapopulation(χ::Real, mp::Metapopulation{SI}) = Metapopulation(mp.h, mp.D, SI(χ*mp.dynamics.β))
CorrectedMetapopulation(χ::Real, mp::Metapopulation{SIS}) = Metapopulation(mp.h, mp.D, SIS(χ*mp.dynamics.β, mp.dynamics.γ))
CorrectedMetapopulation(χ::Real, mp::Metapopulation{SIR}) = Metapopulation(mp.h, mp.D, SIR(χ*mp.dynamics.β, mp.dynamics.δ))

N = 1000
M = 10

Random.seed!(2020) # for reproducibility

h = erdos_renyi(M, 0.3)
g = static_scale_free(N, Int(20*N), 2)

k = degree(g)

χ₁ = mean(closeness_centrality(g))
χ₂ = mean(k.^2)/(mean(k)*N)
χ₃ = mean(k)/N #ne(g)/(n*(n-1)/2)
χ₄ = global_clustering_coefficient(g)

β = 0.1
γ = 0.3
D = 0.1

mp = Metapopulation(h, D, SIR(β,γ))
mp_1 = CorrectedMetapopulation(χ₁, mp)
mp_2 = CorrectedMetapopulation(χ₂, mp)
mp_3 = CorrectedMetapopulation(χ₃, mp)
mp_4 = CorrectedMetapopulation(χ₄, mp)

mpx = Metaplex(g, h, D, SIR(β,γ))

x0_i = fill(1, N)
x0_i[1:10] .= 2
x0_μ = [div(i-1,div(N,M))+1 for i in 1:N]

x0_mp = Array{Int,2}(undef, M,2)
x0_mp[2:M,:] .= [100 0]
x0_mp[1, :] .= [90, 10]

# infected seed randomly dispersed over meta-nodes
x0_i = fill(1, N)
x0_i[1:10] .= 2
x0_μ = rand(1:M, N)
x0_mp = zeros(Int, M, 3)
for i in 1:N
    x0_mp[x0_μ[i], x0_i[i]] += 1
end

tmax = 60.0
nmax = 100000
nsims = 1000
nbins = 200

ts, output = gillespie(mp, x0_mp, tmax=tmax, nmax=nmax)
ts_mpx, output_mpx = gillespie(mpx, [x0_i,x0_μ], tmax=tmax, nmax=nmax)

plot(ts, output[2], label="")
plot(ts_mpx, output_mpx[2], label="")

Random.seed!(2020)

#ts_av, u_av = average(mp, x0_mp, tmax=tmax, nmax=nmax, nbins=nbins, nsims=nsims)
ts_av_mpx, u_av_mpx = average(mpx, [x0_i,x0_μ], tmax=tmax, nmax=nmax, nbins=nbins, nsims=nsims)

ts_mf, u_mf = meanfield(mp, x0_mp, tmax=tmax, saveat=ts_av_mpx)
ts_mf_1, u_mf_1 = meanfield(mp_1, x0_mp, tmax=tmax, saveat=ts_av_mpx)
ts_mf_2, u_mf_2 = meanfield(mp_2, x0_mp, tmax=tmax, saveat=ts_av_mpx)
ts_mf_3, u_mf_3 = meanfield(mp_3, x0_mp, tmax=tmax, saveat=ts_av_mpx)
ts_mf_4, u_mf_4 = meanfield(mp_4, x0_mp, tmax=tmax, saveat=ts_av_mpx)
ts_mf_mpx, u_mf_mpx = meanfield(mpx, [x0_i,x0_μ], tmax=tmax, saveat=ts_av_mpx)

colors = ColorSchemes.tab10;

plot(ts_av_mpx, sum(u_av_mpx[2], dims=2)/N + sum(u_av_mpx[3], dims=2)/N,
    label="Metaplex",
    xlabel="Time",
    ylabel="Fraction of infected individuals",
    linestyle=:dash,
    legend=:bottomright,
    color = colors[2]
    )
    plot!(ts_mf, sum(u_mf[2], dims=2)/N + sum(u_mf[3], dims=2)/N,
    #label="Meanfield (No correction)",
    label = "Metapopulation",
#    xlabel="Time",
#    ylabel="Fraction of infected individuals",
    #legend=:bottomright,
    color = colors[1]
    )
    plot!(ts_mf_mpx, sum(u_mf_mpx[2], dims=2)/N + sum(u_mf_mpx[3], dims=2)/N,
#    label="Meanfield (Metaplex)",
    label = "Metaplex IBMF",
    color = colors[2]
    )
    plot!(ts_mf_1, sum(u_mf_1[2], dims=2)/N + sum(u_mf_1[3], dims=2)/N,
#    label="Meanfield (Closeness Centrality)",
    label = "Closeness",
    color = colors[3]
    )
    plot!(ts_mf_2, sum(u_mf_2[2], dims=2)/N + sum(u_mf_2[3], dims=2)/N,
#    label="Meanfield (Second Moment)",
    label = "Second Moment",
    color = colors[4]
    )
    plot!(ts_mf_3, sum(u_mf_3[2], dims=2)/N + sum(u_mf_3[3], dims=2)/N,
#    label="Meanfield (Mean degree)",
    label = "Mean degree",
    color = colors[5]
    )
    plot!(ts_mf_4, sum(u_mf_4[2], dims=2)/N + sum(u_mf_4[3], dims=2)/N,
#    label="Meanfield (Clustering Coefficient)",
    label = "Clustering",
    color = colors[6]
    )

cols = map(x->[x.r, x.g, x.b], colors.colors)

using MATLAB;

mat"""
figure
title('SI')
xlabel('Time')
ylabel('Fraction of Infected Individuals')
hold on
plot($ts_av_mpx, $(sum(u_av_mpx[2], dims=2)/N + sum(u_av_mpx[3], dims=2)/N), 'LineStyle', '--', 'Color', $(cols[2]))

plot($ts_mf, $(sum(u_mf[2], dims=2)/N + sum(u_mf[3], dims=2)/N), 'Color', $(cols[1]))
plot($ts_mf_mpx, $(sum(u_mf_mpx[2], dims=2)/N + sum(u_mf_mpx[3], dims=2)/N), 'Color', $(cols[2]))
plot($ts_mf_1, $(sum(u_mf_1[2], dims=2)/N + sum(u_mf_1[3], dims=2)/N), 'Color', $(cols[3]))
plot($ts_mf_2, $(sum(u_mf_2[2], dims=2)/N + sum(u_mf_2[3], dims=2)/N), 'Color', $(cols[4]))
plot($ts_mf_3, $(sum(u_mf_3[2], dims=2)/N + sum(u_mf_3[3], dims=2)/N), 'Color', $(cols[5]))
plot($ts_mf_4, $(sum(u_mf_4[2], dims=2)/N + sum(u_mf_4[3], dims=2)/N), 'Color', $(cols[6]))
hold off
ylim([0.0,1.0])
legend({'Metaplex','Metapopulation', 'Metaplex IBMF', 'Closeness Centrality', '\$\\langle k^2\\rangle / \\langle k \\rangle\$', '\$\\langle k \\rangle \$', 'Clustering Coefficient'}, 'Interpreter', 'latex', 'Location', 'southeast')
""";


module TestCas
using ProjectiveGeometricAlgebra3d.CAS
using .CAS: BasisBlade, rev, dual, geomul, wedge, anti_wedge
using Accessors
using Test

function e(inds...;gendim, coeff=1)
    mask = CAS.create_mask(1:gendim, inds...)
    BasisBlade(coeff, mask)
end

@testset "BasisBlade" begin
    @testset "rev" begin
        e = BasisBlade(1, [false])
        e1 = BasisBlade(1, [true])
        e2 = BasisBlade(1, [false, true])
        e12 = BasisBlade(1, [true, true])
        e14 = BasisBlade(1, [true, false, false, true])
        e123 = BasisBlade(1, [true, false, false, true, true])
        e145 = BasisBlade(1, [true, false, false, true, true])
        e2456 = BasisBlade(1, [false, true, false, true, true, true])
        @test rev(e) == e
        @test rev(e1) == e1
        @test rev(e2) == e2
        @test rev(e12) == @set e12.coeff *= -1
        @test rev(e14) == @set e14.coeff *= -1
        @test rev(e123) == @set e123.coeff *= -1
        @test rev(e145) == @set e145.coeff *= -1
        @test rev(e2456) == e2456

    end

    b1 = BasisBlade(5, [true, true, ]) #e12
    b2 = BasisBlade(3, [true, true]) #scalar
    b = wedge(b1, b2)
    @test b.coeff == 0
    g = geomul(b1, b2, [1,1])
    @test g.coeff == -15
    @test g.mask == [false,false]
    g = geomul(b1, b2, [2,-1])
    @test g.coeff == 30
    @test g.mask == [false,false]

    b1 = BasisBlade(1, [true, false, true ]) # e13
    b2 = BasisBlade(1, [true, false, false]) # e1
    w = wedge(b1, b2)
    @test w.coeff == 0
    g = geomul(b1, b2, [1,1,1])
    @test g.coeff == -1
    @test g.mask == [false, false, true]
    g = geomul(b1, b2, [10,11,12])
    @test g.coeff == -10
    @test g.mask == [false, false, true]
    g = geomul(b1, b2, [0,11,12])
    @test g.coeff == 0

    b1 = BasisBlade(1, [true, false, true ]) #e13
    b2 = BasisBlade(1, [false, true, false]) #e2
    w = wedge(b1, b2)
    @test w.coeff == -1
    @test w.mask == [true,true,true]
    w = wedge(b2, b1)
    @test w.coeff == -1
    g = geomul(b1, b2, [1,1,1])
    @test g.coeff == -1
    @test g.mask == [true,true,true]

    b1 = BasisBlade(5, [true, true, ]) #e12
    b2 = BasisBlade(3, [false, false]) #scalar
    b = wedge(b1, b2)
    @test b.coeff == 15
    @test b.mask == [true,true]
end

@testset "dual 4d" begin
    @test dual(e(;gendim=4))    == e(1,2,3,4, gendim=4)

    @test dual(e(1;gendim=4))   == e(2,3,4, gendim=4)
    @test dual(e(2;gendim=4))   == e(1,3,4, gendim=4, coeff=-1)
    @test dual(e(3;gendim=4))   == e(1,2,4, gendim=4)
    @test dual(e(4;gendim=4))   == e(1,2,3, gendim=4, coeff=-1)

    @test dual(e(1,2;gendim=4)) == e(3,4, gendim=4)
    @test dual(e(1,3;gendim=4)) == e(2,4, gendim=4, coeff=-1)
    @test dual(e(1,4;gendim=4)) == e(2,3, gendim=4)

    @test dual(e(1,2,3,4, gendim=4))          == e(;gendim=4)

    @test dual(e(2,3,4, gendim=4))           == e(1;gendim=4)
    @test dual(e(1,3,4, gendim=4, coeff=-1))  == e(2;gendim=4)
    @test dual(e(1,2,4, gendim=4)          )  == e(3;gendim=4)
    @test dual(e(1,2,3, gendim=4, coeff=-1))  == e(4;gendim=4)

    @test dual(e(3,4, gendim=4)            )  == e(1,2;gendim=4)
    @test dual(e(2,4, gendim=4, coeff=-1)  )  == e(1,3;gendim=4)
    @test dual(e(2,3, gendim=4)            )  == e(1,4;gendim=4)
end

end#module