# # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # #
#
# Description
# ==============================================================================
#
#   Compute the satellite lighting condition considering the Earth eclipse.
#
# References
# ==============================================================================
#
# [1] Longo, C. R. O., Rickman, S. L (1995). Method for the Calculation of
#     Spacecraft Umbra and Penumbra Shadow Terminator Points. NASA Technical
#     Paper 3547.
#
# # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # #

export lighting_condition

"""
    lighting_condition(r_i::AbstractVector, s_i::AbstractVector)

Compute the lighting condition at the position `r_i` [m] considering the Sun
position vector `s_i` [m]. The possible return values are:

- `:sunlight`: The point is under direct sunlight.
- `:penumbra`: The point is in penumbra region.
- `:umbra`: The point is in umbra region.

!!! note
    The vectors `r_i` and `s_i` must be represented in the same reference frame.
"""
function lighting_condition(r_i::AbstractVector, s_i::AbstractVector)
    # Distance from the Earth to the Sun.
    norm_s_i = norm(s_i)

    # Projection of the satellite position vector on the Sun direction [1].
    rs_i = (r_i ⋅ s_i) * s_i / (norm_s_i * norm_s_i)

    # Check if the satellite is under sunlight, umbra or penumbra.

    if (rs_i ⋅ s_i / norm_s_i) < 0
        # Distance of the umbral cone and the spacecraft [1].
        Δ_i = r_i - rs_i
        norm_Δ_i = norm(Δ_i)

        # Penumbra section [1].
        xp = R0 * norm_s_i / (Rs + R0)
        αp = asin(R0 / xp)

        # Location of the penumbral cone terminator at the projected spacecraft
        # location [1].
        kp_i = (xp + norm(rs_i)) * tan(αp)

        if norm_Δ_i > kp_i
            return :sunlight
        else
            # Umbra section [1].
            xu = R0 * norm_s_i / (Rs - R0)
            αu = asin(R0 / xu)

            # Location of the umbral cone terminator at the projected spacecraft
            # location [1].
            ep_i = (xu - norm(rs_i)) * tan(αu)

            if norm_Δ_i < ep_i
                return :umbra
            else
                return :penumbra
            end
        end
    else
        return :sunlight
    end
end


cd(@__DIR__)
using Pkg; Pkg.activate("."); Pkg.instantiate()

using Revise
using OrdinaryDiffEq, Flux, DiffEqFlux, LinearAlgebra, Plots
using DiffEqSensitivity, JLD2
const EIGEN_EST = Ref(0.0f0)
const USE_GPU = Ref(false)
#USE_GPU[] = false # the network is small enough such that CPU works just great
USE_GPU[] = true

USE_GPU[] && (using CuArrays)

_gpu(arg) = USE_GPU[] ? gpu(arg) : cpu(arg)
_cu(arg) = USE_GPU[] ? cu(arg) : identity(arg)

function getops(grid, T=Float32)
    N, dz = length(grid), step(grid)
    d = ones(N-2)
    dl = ones(N-3) # super/lower diagonal
    zv = zeros(N-2) # zero diagonal used to extend D* for boundary conditions

    # D1 first order discretization of ∂_z
    D1 = diagm(-1 => -dl, 0 => d)
    D1[1, :] .= 0
    D1[end, :] .= 0

    # D2 discretization of ∂_zz
    D2 = diagm(-1 => dl, 0 => -2*d, 1 => dl)
    κ = 0.05
    D2[1, 1] = D2[end, end] = -1
    D2 = (κ/(dz^2)).*D2 #add the constant κ as the equation requires and finish the discretization
    display(D2), display(D1)

    EIGEN_EST[] = maximum(abs, eigvals(D2))
    D1 = _cu(D1)
    D2 = _cu(D2)
    return (D1=T.(D1), D2=T.(D2))
end

file = jldopen("../DataGeneration/rayleigh_taylor_instability_3d_horizontal_averages.jld2")

Is = keys(file["timeseries/t"])

N = file["grid/Nz"]
Lz = file["grid/Lz"]
Nt = length(Is)

t = Float32.(zeros(Nt))
soldata = Float32.(zeros(Nt, N))

for (i, I) in enumerate(Is)
    t[i] = file["timeseries/t/$I"]
    soldata[i, :] .= file["timeseries/b/$I"]
end

grid = range(0, 1, length = N)
tspan = (t[1], t[end])
u0 = _cu(soldata[1,2:end-1])
ops = getops(grid)

ann = Chain(Dense(N-2,N-2,tanh), Dense(N-2,N-2,tanh), Dense(N-2,N-2,tanh),
            Dense(N-2,N-2,tanh), Dense(N-2,N-2,tanh)) |> _gpu
pp = param(Flux.data(DiffEqFlux.destructure(ann)))
lyrs = Flux.params(pp)

function dudt(u::TrackedArray,p,t)
    Φ = DiffEqFlux.restructure(ann, p)
    return ops.D1*Φ(u) + ops.D2*u
end
function dudt(u::AbstractArray,p,t)
    Φ = DiffEqFlux.restructure(ann, p)
    return ops.D1*Tracker.data(Φ(u)) + ops.D2*u
end

predict_adjoint() = diffeq_adjoint(pp,
                              prob,
                              ROCK4(eigen_est = (integ)->integ.eigen_est = EIGEN_EST[]),
                              u0=u0, saveat = saveat,
                              reltol=1e-5, abstol=1e-6,
                              # no back solve
                              sensealg=SensitivityAlg(quad=false, backsolve=false))

function loss_adjoint()
    pre = predict_adjoint()
    sum(abs2, training_data - pre)
end

cb = function ()
    arr = Array(training_data)
    cur_pred = collect(Flux.data(predict_adjoint()))
    n = size(arr, 1)
    pl = scatter(1:n,arr[:,10],label="data", legend =:bottomright, title = "10th time over space")
    scatter!(pl,1:n,cur_pred[:,10],label="prediction")
    pl2 = scatter(saveat,arr[end÷2,:],label="data", legend =:bottomright, title = "middle point over time")
    scatter!(pl2,saveat,cur_pred[end÷2,:],label="prediction")
    #display(plot(pl, pl2, size=(600, 300)))
    display(loss_adjoint())
end

prob = ODEProblem{false}(dudt,u0,tspan,pp)
nn = 1
saveat = t
soldata = soldata'
soldata = soldata[2:end-1, :]
training_data = _cu(soldata)
epochs = Iterators.repeated((), 20)
learning_rate = ADAM(0.01)
Flux.train!(loss_adjoint, lyrs, epochs, learning_rate, cb=cb)
#=
learning_rate = ADAM(0.001)
epochs = Iterators.repeated((), 300)
Flux.train!(loss_adjoint, lyrs, epochs, learning_rate, cb=cb)
@time loss_adjoint()
=#


using StanSample, MCMCChains

df = filter(row -> row[:age] >= 18,
  CSV.read(joinpath(@__DIR__, "..", "..", "data", "Howell1.csv"), DataFrame))

max_height = maximum(df[!, :height])
min_height = minimum(df[!, :height])

m4_2s = "
data {
  int N;
  real h[N];
}
parameters {
  real<lower=0> sigma;
  real<lower=$(min_height),upper=$(max_height)> mu;
}
model {
  // Priors for mu and sigma
  mu ~ normal(178, 0.1);
  sigma ~ uniform( 0 , 50 );

  // Observed heights
  h ~ normal(mu, sigma);
}
";

m_4_2s = SampleModel("m4.2s", m4_2s);

m4_2_data = Dict("N" => size(df, 1), "h" => df[!, :height]);

rc = stan_sample(m_4_2s, data=m4_2_data);

if success(rc)
  chn = read_samples(m_4_2s; output_format=:mcmcchains)
  chn = replacenames(chn, Dict("mu" => "μ", "sigma" => "σ"))
  describe(chn)
end


using MultivariateSeries

r  = 4
w0 = rand(r)
A0 = rand(5,r)
B0 = rand(3,r)
C0 = rand(4,r)

T0 = tensor(w0, A0, B0, C0)
w, A, B, C = decompose(T0, eps_rkf(1.e-10), mode=2)

T = tensor(w, A, B, C)
@assert norm(T-T0)< 1.e-6


"""
        edas(decisionMat, weights, fns)

Apply EDAS (Evaluation based on Distance from Average Solution) for a given matrix and weights.

# Arguments:
 - `decisionMat::DataFrame`: n × m matrix of objective values for n alterntives and m criteria
 - `weights::Array{Float64, 1}`: m-vector of weights that sum up to 1.0. If the sum of weights is not 1.0, it is automatically normalized.
 - `fns::Array{Function, 1}`: m-vector of functions to be applied on the columns.

# Description
edas() applies the EDAS method to rank n alterntives subject to m criteria which are supposed to be
either maximized or minimized.

# Output
- `::EDASResult`: EDASResult object that holds multiple outputs including scores, rankings, and best index.

# Examples
```julia-repl
julia> decmat = [5000 5 5300 450;
       4500 5 5000 400;
       4500 4 4700 400;
       4000 4 4200 400;
       5000 4 7100 500;
       5000 5 5400 450;
       5500 5 6200 500;
       5000 4 5800 450]
8×4 Array{Int64,2}:
 5000  5  5300  450
 4500  5  5000  400
 4500  4  4700  400
 4000  4  4200  400
 5000  4  7100  500
 5000  5  5400  450
 5500  5  6200  500
 5000  4  5800  450

julia> df = makeDecisionMatrix(decmat)
8×4 DataFrame
 Row │ Crt1     Crt2     Crt3     Crt4
     │ Float64  Float64  Float64  Float64
─────┼────────────────────────────────────
   1 │  5000.0      5.0   5300.0    450.0
   2 │  4500.0      5.0   5000.0    400.0
   3 │  4500.0      4.0   4700.0    400.0
   4 │  4000.0      4.0   4200.0    400.0
   5 │  5000.0      4.0   7100.0    500.0
   6 │  5000.0      5.0   5400.0    450.0
   7 │  5500.0      5.0   6200.0    500.0
   8 │  5000.0      4.0   5800.0    450.0

julia> weights = [0.25, 0.25, 0.25, 0.25];

julia> fns = [maximum, maximum, minimum, minimum];

julia> result = edas(df, weights, fns);

julia> result.scores
8-element Array{Float64,1}:
 0.7595941163602383
 0.8860162461306114
 0.6974721951442592
 0.739657763190231
 0.05908329207449442
 0.7318326305342363
 0.6416913873322523
 0.38519414045559647

julia> result.bestIndex
2
```

# References

Keshavarz Ghorabaee, M., Zavadskas, E. K., Olfat, L., & Turskis, Z. (2015). Multi-criteria inventory classification using a new method of evaluation based on distance from average solution (EDAS). Informatica, 26(3), 435-451.

Ulutaş, A. (2017). EDAS Yöntemi Kullanılarak Bir Tekstil Atölyesi İçin Dikiş Makinesi Seçimi. İşletme Araştırmaları Dergisi, 9(2), 169-183.

"""
function edas(decisionMat::DataFrame, weights::Array{Float64,1}, fns::Array{Function,1})::EDASResult

    # df = convert(Matrix, decisionMat)
    df = Matrix(decisionMat)

    row, col = size(df)

    w = unitize(weights)

        PDAMatrix = similar(df)
    NDAMatrix = similar(df)

    AV = zeros(Float64, col)

    @inbounds for i in 1:col
        AV[i] = mean(df[:, i])
        for j in 1:row
            if fns[i] == maximum
                PDAMatrix[j, i] = max(0, df[j, i] - AV[i]) ./ AV[i]
                NDAMatrix[j, i] = max(0, AV[i] - df[j, i]) ./ AV[i]
            elseif fns[i] == minimum
                PDAMatrix[j, i] = max(0, AV[i] - df[j, i]) ./ AV[i]
                NDAMatrix[j, i] = max(0, df[j, i] - AV[i]) ./ AV[i]
            end
        end
    end

    SP = zeros(Float64, row)
    SN = zeros(Float64, row)

    for i in 1:row
        SP[i] = w .* PDAMatrix[i, :] |> sum
        SN[i] = w .* NDAMatrix[i, :] |> sum
    end

    NSP = SP ./ maximum(SP)
    NSN = 1 .- SN ./ maximum(SN)

    scores = (NSP .+ NSN) ./ 2

    rankings = sortperm(scores)

    bestIndex = rankings |> last

    result = EDASResult(
        decisionMat,
        w,
        scores,
        rankings,
        bestIndex
    )

    return result
end


"""
    ComputeAdimensionalHubbleFactor(z::Float64, w0waCDMCosmology::w0waCDMCosmology)

This function, given the value of the cosmological parameters, evaluate the
Adimensional Hubble Factor for ``w_0 w_a``CDM cosmologies.
The analitycal expression is given by:
```math
E(z)=\\sqrt{\\Omega_M(1+z)^3+\\Omega_R(1+z)^4+
\\Omega_{DE}(1+z)^{3(1+w_0+w_a)}\\exp(-3w_a \\frac{z}{1+z})+\\Omega_k(1+z)^2}
```

!!! warning
    This expression is valid only for the
    [CPL parameterization](https://arxiv.org/abs/astro-ph/0208512)
    of the Dark Energy Equation of State.
"""
function ComputeAdimensionalHubbleFactor(z::Float64, w0waCDMCosmology::w0waCDMCosmology)
    E_z = sqrt(w0waCDMCosmology.ΩM*(1+z)^3 + w0waCDMCosmology.Ωr*(1+z)^4+
    w0waCDMCosmology.Ωk*(1+z)^2
    +w0waCDMCosmology.ΩDE*(1+z)^(3*(1+w0waCDMCosmology.w0+
    w0waCDMCosmology.wa))*exp(-3*w0waCDMCosmology.wa*z/(1+z)) )
    return E_z
end

function ComputeAdimensionalHubbleFactor(z::Float64,
    flatw0waCDMCosmology::Flatw0waCDMCosmology)
    E_z = sqrt(flatw0waCDMCosmology.ΩM*(1+z)^3
    +(1-flatw0waCDMCosmology.ΩM)*(1+z)^(3*(1+flatw0waCDMCosmology.w0+
    flatw0waCDMCosmology.wa))*exp(-3*flatw0waCDMCosmology.wa*z/(1+z)) )
    return E_z
end

"""
    ComputeHubbleFactor(z::Float64, AbstractCosmology::AbstractCosmology)

This function, given the value of the cosmological parameters, evaluate the
Hubble Factor for ``w_0 w_a``CDM cosmologies, whose expression is given by
```math
H(z)=H_0\\sqrt{\\Omega_M(1+z)^3+\\Omega_R(1+z)^4+
\\Omega_{DE}(1+z)^{3(1+w_0+w_a)}\\exp(-3w_a \\frac{z}{1+z})+\\Omega_k(1+z)^2}
```

"""
function ComputeHubbleFactor(z::Float64, AbstractCosmology::AbstractCosmology)
    H_z = AbstractCosmology.H0*ComputeAdimensionalHubbleFactor(z,
    AbstractCosmology)
    return H_z
end

"""
    Computeχ(z::Float64, AbstractCosmology::AbstractCosmology)

This function, given the value of the cosmological parameters, evaluate the
Comoving Distance χ. It is evaluated as:
```math
\\chi(z)=\\frac{c}{H_0}\\int_0^z \\frac{dz'}{E(z')}
```
"""
function Computeχ(z::Float64,
    AbstractCosmology::AbstractCosmology)
    c_0 = 2.99792458e5 #TODO: find a package containing the exact value of
                       #physical constants involved in calculations
    integral, err = QuadGK.quadgk(x -> 1 /
    ComputeAdimensionalHubbleFactor(x,AbstractCosmology), 0, z, rtol=1e-12)
    return integral*c_0/AbstractCosmology.H0
end

"""
    ComputeBackgroundQuantitiesOverGrid(CosmologicalGrid::CosmologicalGrid,
    BackgroundQuantities::BackgroundQuantities,
    AbstractCosmology::AbstractCosmology)

This function evaluate the Hubble factor and the comoving distance over the
[`CosmologicalGrid`](@ref).
"""
function ComputeBackgroundQuantitiesGrid!(CosmologicalGrid::CosmologicalGrid,
    BackgroundQuantities::BackgroundQuantities,
    AbstractCosmology::AbstractCosmology)
    for idx_ZArray in 1:length(CosmologicalGrid.ZArray)
        BackgroundQuantities.HZArray[idx_ZArray] = ComputeHubbleFactor(
        CosmologicalGrid.ZArray[idx_ZArray], AbstractCosmology)
        BackgroundQuantities.rZArray[idx_ZArray] = Computeχ(
        CosmologicalGrid.ZArray[idx_ZArray], AbstractCosmology)
    end
end

function ComputeLogSpacedχGrid!(CosmologicalGrid::CosmologicalGrid,
    BackgroundQuantities::BackgroundQuantities)
    Zχ = Dierckx.Spline1D(BackgroundQuantities.rZArray, CosmologicalGrid.ZArray)
    χArray = CosmoCentral.LogSpaced(BackgroundQuantities.rZArray[1],
    last(BackgroundQuantities.rZArray), length(BackgroundQuantities.rZArray))
    CosmologicalGrid.ZArray = Zχ.(χArray)
end

function ComputeBackgroundQuantitiesOverLogSpacedχGrid(
    CosmologicalGrid::CosmologicalGrid,
    BackgroundQuantities::BackgroundQuantities,
    AbstractCosmology::AbstractCosmology)
    ComputeBackgroundQuantitiesGrid!(CosmologicalGrid, BackgroundQuantities,
    AbstractCosmology)
    ComputeLogSpacedχGrid!(CosmologicalGrid, BackgroundQuantities)
    ComputeBackgroundQuantitiesGrid!(CosmologicalGrid, BackgroundQuantities,
    AbstractCosmology)
end

function ExtractGrowthFactor!(BackgroundQuantities::BackgroundQuantities,
    PowerSpectrum::PowerSpectrum)
    BackgroundQuantities.DZArray = zeros(length(PowerSpectrum.PowerSpectrumLinArray[1,:]))
    for zidx in 1:length(BackgroundQuantities.DZArray)
        BackgroundQuantities.DZArray[zidx] =
        sqrt(PowerSpectrum.PowerSpectrumLinArray[1,zidx]./
        PowerSpectrum.PowerSpectrumLinArray[1,1])
    end
end


# Mrongowius-Rößler method
# First published by Mrongowius & Rößler, 2021

struct MronRoe <: AbstractIteratedIntegralAlgorithm end

convorder(::MronRoe) = 1//1
errcoeff(m, h, ::MronRoe, ::MaxL2) = √m*h/(√12*π)
norv(m, n, ::MronRoe) = 2*m*n+(m^2+m)÷2

"""
    levyarea(W, n, alg::MronRoe)

Simulates an approximation of the iterated Itô-integrals ``\\int_0^1W_s\\otimes dW_s``
of the given ``m``-dimensional increment of a Wiener process with step size 1.
The parameter ``n`` specifies the number of terms in the approximation and thus determines the accuracy.
This is an efficient implementation of the algorithm proposed in [Mrongowius & Rößler, 2021](@ref mr2021).
It is based on the Fourier method from Milstein but incorporates an improved tail sum approximation.
The algorithm needs approximately ``m^2+2\\cdot m\\cdot n`` Float's
and ``1/2m^2+2\\cdot m\\cdot n + 1/2m`` random numbers.
The time complexity is ``\\mathcal{O}(m^2\\cdot n)``.
"""
function levyarea(W::AbstractVector{T}, n::Integer, alg::MronRoe) where {T}
    rng = default_rng()
    m = length(W)
    # 1. Simulate Xₖ and Yₖ and approximate stochastic area integral
    X = randn(rng, T, m, n) # allocates m*n Floats
    Y = randn(rng, T, n, m) # allocates m*n Floats
    Y .= (Y .- √(T(2)).*W') ./ (1:n)
    A = X * Y # allocates m*m Floats

    # 2. Add first, simple rest approximation (a₀)
    Ψ = randn!(rng, view(X, :, 1))
    a = T(sqrt(2*trigamma(n+1)))
    A .+= a .* W .* Ψ'

    # 3. Add improved tail-sum approximation and antisymmetrize (A.-=A')
    for i = 1:m
        @inbounds A[i,i] = zero(T)
        for j = (i+1):m
            @inbounds A[i,j] = (A[i,j] + a*randn(rng, T) - A[j,i])/(2*T(π))
            @inbounds A[j,i] = -A[i,j]
        end
    end

    return A
end


using Revise

includet("src/Spring.jl")
IArray = InternalArray

# 1. Solving the equation using the euler method ------------------------------------------

# We are asked to find solutions for the time interval [t0, t0+40].
# I want to execute the following with a step of 1e-3,
# so I'll set an appropiate number of iterations.
const step = 1e-3
const nit = Int(40/step)  # round to avoid small rounding errors

prob = SpringProblem(p0 = 0.8, v0 = 0, m = 1, k = 1.5, b = 0.3, A = 0.4, ω = 2.4)

x = eulermethod(prob, step, nit)

# 2. Plotting the result ------------------------------------------------------------------

timeline = [prob.t0 + i*step for i in 0:length(x)-1]
position = plot(timeline, IArray(x, 1), label = "movement")
speed = plot(timeline, IArray(x, 2), label = "speed")
plot(position, speed)

savefig("media/spring_std.png")


function fg(d::Distributions.Logistic, x)
    scaled = (x - d.μ)/2d.θ
    core = sech( scaled )^2 / 4d.θ
    core, core * -tanh(scaled) / d.θ
end
function fgw(d::Distributions.Logistic, x)
    scaled = (x - d.μ)/2d.θ
    core = sech( scaled )^2 / 4d.θ
    weighted_core = core*exp(x^2/2)
    weighted_core, x*weighted_core - weighted_core * tanh(scaled) / d.θ
end
function fw(d::Distributions.Logistic, x)
    scaled = (x - d.μ)/2d.θ
    sech( scaled )^2*exp(x^2/2) / 4d.θ
end


using SEC, Test, GLMakie, LinearAlgebra
import SEC.Emulate: mean, cov, GaussianProcess, condition

μ(x) = 0.0
δ(x,y) = x==y ? 1 : 0
k(x,y) = 10exp(-(x-y)^2/2) + 0.1 * δ(x,y) # noise here
gp = GaussianProcess(mean = μ, covariance = k)
N = 10
X = [i for i in 1:N]
Y = [2*i for i in 1:N]
data = (; X, Y)
cgp = condition(gp, data, noise = 0.0) # noise here
# philosophical difference only
# just associates more noise at the given prediction point
##
fig = Figure(resolution = (700, 450))
ax = Axis(fig, xlabel = "x", ylabel = "y")
Xn = collect(range(0, N+1, 100))
Xn = setdiff(Xn, X)
μX = mean(cgp, Xn)
σXX =  [cgp.covariance(x,x) for x in Xn]
scatter!(ax, X, Y, color = :red)
lines!(ax, Xn, μX, color = :blue, linewidth =3)
# filled curve 1
band!(ax, Xn, μX - 3*σXX, μX + 3*σXX, color = (:blue, 0.2))
xlims!(ax, (0,11))
ylims!(ax, (-1,21))
ax.xlabelsize = 25
ax.xticklabelsize = 25
ax.ylabelsize = 25
ax.yticklabelsize = 25
update!(fig.scene)
fig[1,1] = ax
fig


module ManifoldLearning

    using MultivariateStats

    import Base: show, dump, eigvals
    import MultivariateStats: outdim, projection, transform

    export

    ## common
    outdim,             # the output dimension of the transformation
    projection,         # the projection matrix
    eigvals,            # eigenvalues from the spectral analysis
    neighbors,          # the number of nearest neighbors used for building local structure
    ccomponent,         # point indexes of the largest connected component

    transform,          # perform the manifold learning

    # Transformation types
    Isomap,             # Type: Isomap model
    HLLE,               # Type: Hessian Eigenmaps model
    LLE,                # Type: Locally Linear Embedding model
    LTSA,               # Type: Local Tangent Space Alignment model
    LEM,                # Type: Laplacian Eigenmaps model
    DiffMap,            # Type: Diffusion maps model

    # example dataset
    swiss_roll          # swiss roll dataset generator

    include("types.jl")
    include("utils.jl")
    include("transformations.jl")
    include("isomap.jl")
    include("hlle.jl")
    include("lle.jl")
    include("ltsa.jl")
    include("lem.jl")
    include("diffmaps.jl")
end


using Soss

mstep = @model pars,state begin
    # Parameters
    α = pars.α        # Daily transmission rate
    β = pars.β        # Daily recovery rate
    γ = pars.γ        # Daily case fatality rate

    # Starting counts
    s0 = state.s      # Susceptible
    i0 = state.i      # Infected
    r0 = state.r      # Recovered
    d0 = state.d      # Deceased
    n = s0 + i0 + r0  # Population

    # Transitions between states
    si ~ Binomial(s0, α * i0 / n)
    ir ~ Binomial(i0, β)
    id ~ Binomial(i0 - ir, γ)

    # Updated counts
    return ( s = s0 - si
           , i = i0 + si - ir - id
           , r = r0 + ir
           , d = d0 + id
           )
end;

m = @model s0 begin
    α ~ Uniform()
    β ~ Uniform()
    γ ~ Uniform()
    pars = (α=α, β=β, γ=γ)
    x ~ Chain(s0) do s
            mstep(pars, s)
        end
end

using CSV
covid = CSV.read("/home/chad/git/covid-19/data/countries-aggregated.csv");

using DataFramesMeta

us = @where(covid, :Country .== "US");

# Some simple data cleaning - each column should be non-decreasing
# Fix decreasing values with geometric mean
for j in [:Confirmed, :Recovered, :Deaths]
    for i in 2:(size(us,1) - 1)
        if us[i, j] < us[i-1, j]
            us[i,j] = round(Int, sqrt(us[i-1,j] * us[i+1,j]))
        end
    end
end

using DataFrames
df = DataFrame(
    :i => us.Confirmed
    , :r => us.Recovered
    , :d => us.Deaths
    )
n = 331000000

df[!, :s] .= n .- df.i .- df.r .- df.d;
df[!, :si] .= 0;
df[:si][2:end] = .-diff(df.s);
df[!, :ir] .= 0;
df[:ir][2:end] = diff(df.r);
df[!, :id] .= 0;
df[:id][2:end] = diff(df.d);

# just the last 30 days
df = df[(end-29):end,:];

using NamedTupleTools
import NamedTupleTools
NamedTupleTools.namedtuple(dfrow::DataFrameRow) = namedtuple(pairs(dfrow).iter.is...)

nts = namedtuple(names(df)).(eachrow(df));

s0 = namedtuple(df[1,:]);
post = dynamicHMC(m(s0=s0),(x=namedtuple.(eachrow(df)),));
ppost=particles(post)

# R₀
ppost.α / (ppost.β + ppost.γ)

# Case fatality rate
ppost.γ / (ppost.β + ppost.γ)

# Implied infection duration
1/(ppost.β + ppost.γ)


using MISAlgorithms, EliminateGraphs
using Test

@testset "adjacent45" begin
    eg = K_eg(4,5)
    v1, v2 = EliminateGraphs.adjacent45(eg)
    @test Set([degree(eg,v2), degree(eg, v1)]) == Set([4,5])
end

@testset "mis2" begin
    @test EliminateGraphs.minx(x->x^2, [2,-1, 3]) == -1

    g1 = EliminateGraph([0 1 1 0 0;
                         1 0 1 1 0;
                         1 1 0 1 0;
                         0 1 1 0 0;
                         0 0 0 0 0])

    @test MISAlgorithms.mis2(g1) == 3
    g2 = disconnected_cliques_eg(6, 6)
    @test MISAlgorithms.mis2(g2) == 2
    g3 = K_eg(7, 2)
    @test MISAlgorithms.mis2(g3) == 7
    g4 = K_eg(7, 7)
    @test MISAlgorithms.mis2(g4) == 7
    g5 = ring_eg(4)
    @test MISAlgorithms.mis2(g5) == 2
    @test MISAlgorithms.mis2(c60_graph) == 24
    g6 = K_eg(4, 5)
    unsafe_disconnect!(g6, 1,5)
    unsafe_connect!(g6, 2,3)
    @test MISAlgorithms.mis2(g6) == 5
    g7 = K_eg(4, 5)
    @test MISAlgorithms.mis2(g7) == 5

    for i = 1:100
        eg = rand_eg(40, 0.2)
        @test MISAlgorithms.mis1(eg) == MISAlgorithms.mis2(eg)
    end
end


#################### Mathematical Operators ####################
function cummean(x::AbstractArray)
    return mapslices(cummean, x, dims = 1)
end

function cummean(x::AbstractVector)
    y = similar(x, Float64)
    xs = 0.0
    fill!(y, xs)

    c = 0
    for i in 1:length(x)
        if !ismissing(x[i])
            c += 1
            xs += x[i]
        end
        y[i] = xs / c
    end
    return y
end

## Csorgo S and Faraway JJ. The exact and asymptotic distributions of the
## Cramer-von Mises statistic. Journal of the Royal Statistical Society,
## Series B, 58: 221-234, 1996.
function pcramer(q::Real)
    p = 0.0
    for k in 0:3
        c1 = 4.0 * k + 1.0
        c2 = c1^2 / (16.0 * q)
        p += gamma(k + 0.5) / factorial(k) * sqrt(c1) * exp(-c2) * besselk(0.25, c2)
    end
    return p / (pi^1.5 * sqrt(q))
end


using FFTW

@testset "QFT circuit tests" begin
    @test begin
        n = 2
        results = zeros(Bool, 2^n)
        for i in 1:2^n
            in_state = zeros(Float64, 2^n)
            in_state[i] = 1.0
            qreg = FullStateQuantumRegister{ComplexF64}(2, in_state)
            qc = qft_circuit(n)
            execute!(qc, qreg)
            ref_state = ifft(in_state)
            ref_state /= sqrt(sum(conj(ref_state) .* ref_state))
            qreg_ref = FullStateQuantumRegister{ComplexF64}(n, ref_state)
            results[i] = qreg ≈ qreg_ref
        end
        all(results)
    end
    @test begin
        n = 4
        in_state = rand(2^n)
        in_state /= sqrt(sum(in_state .* in_state))
        qreg = FullStateQuantumRegister{ComplexF64}(n, in_state)
        qc = qft_circuit(n)
        execute!(qc, qreg)
        ref_state = ifft(in_state)
        ref_state /= sqrt(sum(conj(ref_state) .* ref_state))
        qreg_ref = FullStateQuantumRegister{ComplexF64}(n, ref_state)
        qreg ≈ qreg_ref
    end
    @test begin
        n = 10
        in_state = rand(2^n)
        in_state /= sqrt(sum(in_state .* in_state))
        qreg = FullStateQuantumRegister{ComplexF64}(n, in_state)
        qc = qft_circuit(n)
        execute!(qc, qreg)
        ref_state = ifft(in_state)
        ref_state /= sqrt(sum(conj(ref_state) .* ref_state))
        qreg_ref = FullStateQuantumRegister{ComplexF64}(n, ref_state)
        qreg ≈ qreg_ref
    end
end

@testset "GHZ circuit tests" begin
    @test begin
        results = []
        for n in 1:5
            qc = ghz_circuit(n)
            qreg = FullStateQuantumRegister{ComplexF64}(n, "0"^n)
            execute!(qc, qreg)
            push!(results,
                  qreg.state[1] == 1/sqrt(2) && qreg.state[end] == 1/sqrt(2))
        end
        all(results)
    end
end


# Test that an example fixed point problem converges to the exact solution. Based on the pyblp code.
@testset "Iteration converges $it" for it in (SimpleFixedPointIteration(), SQUAREMIteration())
    function contraction!(x_out::AbstractVector, x::AbstractVector)
        c1 = [10, 12]
        c2 = [3, 5]

        @. x_out = sqrt(c1 / (x + c2))
        return x_out
    end

    initial = ones(2)
    result = fixed_point_iteration(it, initial, contraction!)

    @test result.status == ITERATION_CONVERGED
    @test result.delta ≈ [1.4920333, 1.37228132]
end

# Test that the solution to the fixed point problem from Hasselblad (1969) is reasonably close to the exact
# solution. This same problem is used in an original SQUAREM unit test.
@testset "SQUAREM $scheme Hasselblad" for scheme in (SQUAREM1, SQUAREM2, SQUAREM3)
    function contraction!(x_out::AbstractVector, x::AbstractVector)
        y = [162, 267, 271, 185, 111, 61, 27, 8, 3, 1]
        i = 0:length(y)-1
        z = @. (x[1] * exp(-x[2]) * x[2]^i) / (x[1] * exp(-x[2]) * x[2]^i + (1 - x[1]) * exp(-x[3]) * x[3]^i)

        x_out[:] = @. [
            $sum(y * z) / $sum(y),
            $sum(y * i * z) / $sum(y .* z),
            $sum(y * i * (1-z)) / $sum(y * (1-z))
        ]
        return x_out
    end

    it = SQUAREMIteration(scheme=scheme)

    initial = [0.2, 2.5, 1.5]

    result = fixed_point_iteration(it, initial, contraction!)
    @test result.status == ITERATION_CONVERGED
    @test result.delta ≈ [0.6401146029910, 2.6634043566619, 1.2560951012662]
end


# License for this file: MIT (expat)
# Copyright 2017-2018, DLR Institute of System Dynamics and Control
#
# This file is part of module
#   Modia3D.Composition (Modia3D/Composition/_module.jl)
#

# Measurements

"""
    angle = planarRotationAngle(e, v1, v2)
          = planarRotationAngle(frame1, frame2)

Return angle of a planar rotation, given the rotation axis e (a unit vector)
and the representations of a vector in frame 1 (v1) and frame 2 (v2).

Under the assumption that the z-axes of frame1 and frame2 coincide, return the
angle between the x-axis of frame1 and the position vector from frame1 to frame2.
"""
planarRotationAngle(e::AbstractVector, v1::AbstractVector, v2::AbstractVector) =
     atan2( dot(-cross(e,v1), v2), dot(v1,v2) - dot(e,v1)*dot(e,v2) )

function planarRotationAngle(frame1::Object3D, frame2::Object3D)
   # Extract needed vectors
   e1z = frame1.R_abs[3,:]
   e2z = frame2.R_abs[3,:]
   e1x = frame1.R_abs[1,:]
   r1  = frame1.r_abs
   r2  = frame2.r_abs

   # Check that z-axes of both frames coincide
   angle_z_axes = acos( dot(e1z,e2z) )
   if abs(angle_z_axes) > Basics.neps
      error("\nError from Modia3D.planarRotationAngle(", ModiaMath.fullName(frame1), ", ", ModiaMath.fullName(frame2), "):\n",
                 "The z-axes of the two frames do not coincide.\n",
                 "( angle(zaxis(frame1), zaxis(frame2)) = ", string(angle_z_axes*Basics.radToDeg), " > ", Basics.neps, ")")
   end

   # Angle between frame2.r_abs - frame1.r_abs and x_axis of frame1
   r12 = r2-r1
   if norm(r12) < 100*eps()
      error("\nError from Modia3D.planarRotationAngle(", ModiaMath.fullName(frame1), ", ", ModiaMath.fullName(frame2), "):\n",
              "The origins of the two frames coincide (this is not allowed for this measurement.")
   end
   r12n = r12/norm(r12)

   return atan(dot(cross(e1x,r12n), e1z), dot(e1x,r12n) )
end

"""
    d = distance(frame1, frame2)

Return the distance between the origin of frame1 and the origin of frame2
"""
distance(frame1::Object3D, frame2::Object3D) = norm(frame2.r_abs - frame1.r_abs)

function distanceAndAngles(frame1::Composition.Object3D, frame2::Composition.Object3D)
   # Extract needed vectors
   e1z = frame1.R_abs[3,:]
   e2z = frame2.R_abs[3,:]
   r1  = frame1.r_abs
   r2  = frame2.r_abs

   # Check that z-axes of both frames are pointing in the same direction
   @assert(abs(dot(e1z,e2z) - 1.0) <= 1e-10)

   # Angle between frame1 and frame2 along z-axis
   rdiff0 = r2-r1
   ey = normalize(rdiff0)

   ey1  = frame1.R_abs*ey
   ey2  = frame2.R_abs*ey
   ex1b = cross(ey1, e1z)
   ex2b = cross(ey2, e2z)
   ex1a = frame1.R_abs[:,1]
   ex2a = frame2.R_abs[:,1]

   phi1 = planarRotationAngle(e1z, ex1a, ex1b)
   phi2 = planarRotationAngle(e2z, ex2a, ex2b)
   d    = norm(rdiff0)
   return (d,phi1,phi2)
end


module ApproxFunFromFile

export ApproxFunVandermonde

using ApproxFun

#pts=points(c,100) #This is a built in function; Approxfun would then do eg. vals=cos(pts) to get the values
#show(pts) # Density looks like a 'U' with greater sampling density near the extremes of the range

# From: https://github.com/ApproxFun/ApproxFun.jl/issues/275 , courtesy of private communication with Sheehan Olver
# Least squares approximation of data on an evenly spaced grid with Chebyshev series
function vandermonde(S,n,x::AbstractVector)
    V=Array(Float64,length(x),n)
    for k=1:n
        V[:,k]=Fun([zeros(k-1);1],S)(x)
    end
    V
end

# For ...(this)... case, make sure `length(pts) >> n`.
function ApproxFunVandermonde(filename,n=20,  lower=0.0, upper=360.0)
    c=Chebyshev([lower,upper]) #Define Chebyshev domain in this range (to match data imported)

    # Standard two column data form
    df=readdlm(filename)

    pts=df[:,1] # Points
    vals=df[:,2] #Values at these points

    V=vandermonde(c,n,pts)
    # Are you ready for the magic?
    af=Fun(V\vals,c) # Approximate Function (af)
    # me is now an ApproxFun representation of the tabulated data.
    # As a Chebyshev polynomial fit we can do all sorts of differentiation + integration.
    return af,df
end

function graphFun(af,df)
    # Logscale version...
    plot(af,label="Chebyshev fn")
    plot!(df[:,1],df[:,2],label="Raw fn")
    yaxis!("Value fn(x)")
    xaxis!("x")
end

function basictest()
    U,df=ApproxFunVandermonde("test.dat")
    using Plots
    graphFun(U,df)
    printFun(U)

    println(U)
    println(df)
end

function printFun(af)
    for i=0:360
        @printf("%f %f\n",i,af(i))
    end
end

end


######### Stan batch program example  ###########

using StanSample, Test

dyes ="
data {
  int BATCHES;
  int SAMPLES;
  real y[BATCHES, SAMPLES];
  // vector[SAMPLES] y[BATCHES];
}

parameters {
  real<lower=0> tau_between;
  real<lower=0> tau_within;
  real theta;
  real mu[BATCHES];
}

transformed parameters {
  real sigma_between;
  real sigma_within;
  sigma_between <- 1/sqrt(tau_between);
  sigma_within <- 1/sqrt(tau_within);
}

model {
  theta ~ normal(0.0, 1E5);
  tau_between ~ gamma(.001, .001);
  tau_within ~ gamma(.001, .001);

  mu ~ normal(theta, sigma_between);
  for (n in 1:BATCHES)
    y[n] ~ normal(mu[n], sigma_within);
}

generated quantities {
  real sigmasq_between;
  real sigmasq_within;

  sigmasq_between <- 1 / tau_between;
  sigmasq_within <- 1 / tau_within;
}
"

dyesdata = Dict("BATCHES" => 6,
    "SAMPLES" => 5,
    "y" => reshape([
      [1545, 1540, 1595, 1445, 1595];
      [1520, 1440, 1555, 1550, 1440];
      [1630, 1455, 1440, 1490, 1605];
      [1595, 1515, 1450, 1520, 1560];
      [1510, 1465, 1635, 1480, 1580];
      [1495, 1560, 1545, 1625, 1445]
    ], 6, 5)
  )

sm = SampleModel("dyes", dyes);

rc = stan_sample(sm, data=dyesdata)

if success(rc)
  samples = read_samples(sm)

  # Fetch cmdstan summary data frame.
  df = read_summary(sm)
  @test df[df.parameters .== :theta, :mean][1] ≈ 1527.5 atol=50.0
  @test df[df.parameters .== Symbol("mu[1]"), :mean][1] ≈ 1512.5 rtol=0.1
  @test df[df.parameters .== Symbol("mu[2]"), :mean][1] ≈ 1528.1 rtol=0.1
  @test df[df.parameters .== Symbol("mu[3]"), :mean][1] ≈ 1553.8 rtol=0.1
  @test df[df.parameters .== Symbol("mu[4]"), :mean][1] ≈ 1507.2 rtol=0.1
  @test df[df.parameters .== Symbol("mu[5]"), :mean][1] ≈ 1578.7 rtol=0.1
  @test df[df.parameters .== Symbol("mu[6]"), :mean][1] ≈ 1487.3 rtol=0.1

end


sum(map(int,split(string(big(2)^1000),"")))


using Regression
using Base.Test

import Regression: loss, has_bias, pred_without_bias, pred_with_bias, initsol

d = 5
k = 3
n = 100
X = randn(d, n)
y = randn(n)
Y = randn(k, n)
c = rand(1:k, n)
b = 2.0

pb = linearreg(X, y; bias=b)
@test isa(pb, UnivariateRegression{SqrLoss})
@test has_bias(pb)
@test loss(pb) == SqrLoss()
@test pred_without_bias(pb) == LinearPred(d)
@test pred_with_bias(pb) == AffinePred(d, b)
@test initsol(pb) == zeros(d+1)

pb = linearreg(X, Y; bias=b)
@test isa(pb, MultivariateRegression{SumSqrLoss})
@test has_bias(pb)
@test loss(pb) == SumSqrLoss()
@test pred_without_bias(pb) == MvLinearPred(d, k)
@test pred_with_bias(pb) == MvAffinePred(d, k, b)
@test initsol(pb) == zeros(k, d+1)

pb = logisticreg(X, sign(y); bias=b)
@test isa(pb, UnivariateRegression{LogisticLoss})
@test has_bias(pb)
@test loss(pb) == LogisticLoss()
@test pred_without_bias(pb) == LinearPred(d)
@test pred_with_bias(pb) == AffinePred(d, b)
@test initsol(pb) == zeros(d+1)

pb = mlogisticreg(X, c, k; bias=b)
@test isa(pb, MultivariateRegression{MultiLogisticLoss})
@test has_bias(pb)
@test loss(pb) == MultiLogisticLoss()
@test pred_without_bias(pb) == MvLinearPred(d, k)
@test pred_with_bias(pb) == MvAffinePred(d, k, b)
@test initsol(pb) == zeros(k, d+1)


using LinearAlgebra

# order 4
h = [
  2.6864248295847e-01,
  1.0094667153500e+00,
  9.9312068011715e-01,
  1, 1, 1, 1, 1
  ]
e = [1//1, 0, 0, 0, 0, 0, 0, 0]
H = Diagonal(h)

q1 = [0,
  6.1697245625434e-01,
  -1.1697245625434e-01,
  0.0000000000000e+00,
  0, 0, 0, 0]'
q2 = [-q1[2], 0,
  7.0030578958767e-01,
  -8.3333333333333e-02,
  0.0000000000000e+00,
  0, 0, 0]'
q3 = [-q1[3], -q2[3], 0,
  2//3, -1//12, 0, 0, 0]'
q4 = [-q1[4], -q2[4], -q3[4], 0,
  2//3, -1//12, 0, 0]'
q5 = [-q1[5] -q2[5] -q3[5] -q4[5] 0 2//3 -1//12 0]
q6 = [0 0 0 1//12 -2//3 0 2//3 -1//12]
q7 = [0 0 0 0 1//12 -2//3 0 2//3]
q8 = [0 0 0 0 0 1//12 -2//3 0]
Q = vcat(q1, q2, q3, q4, q5, q6, q7, q8)

D1 = H \ (Q - 1//2*e*e'); display(D1)

# order 6
h = [
  1.2740260779883e-01,
  6.1820981002054e-01,
  9.4308973897679e-01,
  1.0093019060199e+00,
  9.9884825610465e-01,
  1, 1, 1, 1
  ]
e = [1//1, 0, 0, 0, 0, 0, 0, 0, 0]
H = Diagonal(h)

q1 = [0,
  6.3217364546846e-01,
  -1.6411963429825e-01,
  3.6495407984639e-02,
  -4.5494191548490e-03,
  0, 0, 0, 0]'
q2 = [-q1[2], 0,
  8.0515625504417e-01,
  -2.0755653563249e-01,
  3.4573926056780e-02,
  0, 0, 0, 0]'
q3 = [-q1[3], -q2[3], 0,
  7.9402676057785e-01,
  -1.6965680649860e-01,
  1.6666666666667e-02,
  0, 0, 0]'
q4 = [-q1[4], -q2[4], -q3[4], 0,
  7.5629896626333e-01,
  -3//20, 1//60, 0, 0]'
q5 = [-q1[5] -q2[5] -q3[5] -q4[5] 0 3//4 -3//20 1//60 0]
q6 = [0 0 -1//60 3//20 -3//4 0 3//4 -3//20 1//60]
q7 = [0 0 0 -1//60 3//20 -3//4 0 3//4 -3//20]
q8 = [0 0 0 0 -1//60 3//20 -3//4 0 3//4]
q9 = [0 0 0 0 0 -1//60 3//20 -3//4 0]
Q = vcat(q1, q2, q3, q4, q5, q6, q7, q8, q9)

D1 = H \ (Q - 1//2*e*e'); display(D1)

# order 8
h = [
  1.4523997892351e-01,
  7.6864793350174e-01,
  9.9116487068535e-01,
  9.9992473335107e-01,
  1.0002097054636e+00,
  9.9996591555866e-01,
  1, 1, 1, 1
  ]
e = [1//1, 0, 0, 0, 0, 0, 0, 0, 0, 0]
H = Diagonal(h)

q1 = [0,
  6.6697342753834e-01,
  -2.2919342278749e-01,
  7.4283116457276e-02,
  -1.2020661178873e-02,
  -4.2460029252999e-05,
  0, 0, 0, 0]'
q2 = [-q1[2], 0,
  8.8241196934163e-01,
  -2.6653314104602e-01,
  5.5302527504316e-02,
  -4.2079282615860e-03,
  0, 0, 0, 0]'
q3 = [-q1[3], -q2[3], 0,
  8.2904844081126e-01,
  -2.1156614214635e-01,
  3.9307676460659e-02,
  -3.5714285714286e-03,
  0, 0, 0]'
q4 = [-q1[4], -q2[4], -q3[4], 0,
  8.0305501223679e-01,
  -2.0078040553808e-01,
  3.8095238095238e-02,
  -3.5714285714286e-03,
  0, 0]'
q5 = [-q1[5], -q2[5], -q3[5], -q4[5], 0,
  8.0024692689207e-01,
  -1//5, 4//105, -1//280, 0]'
q6 = [-q1[6] -q2[6] -q3[6] -q4[6] -q5[6] 0 4//5 -1//5 4//105 -1//280]
q7 = [0 0 1//280 -4//105 1//5 -4//5 0 4//5 -1//5 4//105]
q8 = [0 0 0 1//280 -4//105 1//5 -4//5 0 4//5 -1//5]
q9 = [0 0 0 0 1//280 -4//105 1//5 -4//5 0 4//5]
q10 = [0 0 0 0 0 1//280 -4//105 1//5 -4//5 0]
Q = vcat(q1, q2, q3, q4, q5, q6, q7, q8, q9, q10)

D1 = H \ (Q - 1//2*e*e'); display(D1)


"""
Parallel Factor Analysis (PARAFAC) decomposition
"""
immutable PARAFAC2{T<:Number, N} <: TensorDecomposition{T, N}
    factors::NTuple{N, Matrix{T}}
    D::Vector{Matrix{T}}
    A::Matrix{T}
    props::Dict{Symbol, Any}

    function Base.call{T}(::Type{PARAFAC2},
                   X::Vector{Matrix{T}},
                   F::Matrix{Float64},
                   D::Vector{Matrix{Float64}},
                   A::Matrix{Float64})
        factors = (map(function (Xi, Di)
            U = svd(A_mul_Bt(F .* Di, Xi * A))
            U[3] * (U[1]'F)
        end, X, D)...)
        return new{T, length(factors)}(factors, D, A, Dict{Symbol,Any}())
    end

end

"""
PARAFAC2 model
"""
function parafac2{S<:Matrix}(X::Vector{S},
                             r::Integer;
                             tol::Float64=1e-5,
                             maxiter::Integer=100,
                             compute_error::Bool=false,
                             verbose::Bool=true)
    m = length(X)
    n = size(X[1], 2)
    for i in 2:m
        size(X[i], 2) == n || error("All input matrices should have the same number of rows.")
    end

    F = eye(r)
    D = Matrix{Float64}[ones(1, r) for _ in 1:m]
    A = eigs(sum(map(Xi -> Xi'Xi, X)), nev=r)[2]
    G = Matrix{Float64}[eye(r), eye(r), ones(r, r) * m]
    H = Matrix{Float64}[(size(X[i], 1) > size(X[i], 2) ? qr(X[i])[2]: X[i]) for i in 1:m]
    P = Array(Matrix{Float64}, m)

    niters = 0
    converged = false
    resid = vecnorm(vcat(X...))
    while !converged && niters < maxiter
        map!(function (Hi, Di)
            U = svd((F .* Di) * (Hi * A)')
            U[3]*U[1]'
        end, P, H, D)
        T = cat(3, [P[i]'H[i] for i in 1:m]...)

        B = vcat(D...)
        F = _row_unfold(T, 1) * khatrirao(B, A) / (G[3] .* G[2])
        At_mul_B!(G[1], F, F)
        A = _row_unfold(T, 2) * khatrirao(B, F) / (G[3] .* G[1])
        At_mul_B!(G[2], A, A)
        B = _row_unfold(T, 3) * khatrirao(A, F) / (G[2] .* G[1])
        At_mul_B!(G[3], B, B)

        D = Matrix{Float64}[B[i, :] for i in 1:m]

        resid_old = resid
        resid = sum(k -> sumabs2(H[k] - P[k] * A_mul_Bt(F .* D[k], A)), 1:m)
        converged = abs(resid - resid_old) < tol * resid_old

        niters += 1
    end

    verbose && _iter_status(converged, niters, maxiter)
    res = PARAFAC2(X, F, D, A)
    if compute_error
        _set_rel_residue(res, sqrt(resid) / vecnorm(vcat(X...)))
    end
    return res
end


using QuasiCopula, LinearAlgebra, GLM
using Random, Distributions, DataFrames, ToeplitzMatrices
using Test, BenchmarkTools
BLAS.set_num_threads(1)
Threads.nthreads()

p = 3    # number of fixed effects, including intercept

# true parameter values
Random.seed!(12345)
# try next
βtrue = rand(Uniform(-2, 2), p)
σ2true = [0.5]
ρtrue = [0.5]

function get_V(ρ, n)
  vec = zeros(n)
  vec[1] = 1.0
  for i in 2:n
      vec[i] = vec[i - 1] * ρ
  end
  V = ToeplitzMatrices.SymmetricToeplitz(vec)
  V
end

trueparams = [βtrue; ρtrue; σ2true] #hold true parameters

#simulation parameters
samplesize = 10000
ni = 25

st = time()
currentind = 1
d = Bernoulli()
link = LogitLink()
D = typeof(d)
Link = typeof(link)
T = Float64

gcs = Vector{GLMCopulaARObs{T, D, Link}}(undef, samplesize)

V = get_V(ρtrue[1], ni)

# true Gamma
Γ = σ2true[1] * V

# for reproducibility I will simulate all the design matrices here
Random.seed!(12345)
X_samplesize = [randn(ni, p - 1) for i in 1:samplesize]

for i in 1:samplesize
  X = [ones(ni) X_samplesize[1]]
  η = X * βtrue
  μ = exp.(η) ./ (1 .+ exp.(η))
  vecd = Vector{DiscreteUnivariateDistribution}(undef, ni)
  for i in 1:ni
      vecd[i] = Bernoulli(μ[i])
  end
  nonmixed_multivariate_dist = NonMixedMultivariateDistribution(vecd, Γ)
  # simuate single vector y
  y = Vector{Float64}(undef, ni)
  res = Vector{Float64}(undef, ni)
  rand(nonmixed_multivariate_dist, y, res)
  V = [ones(ni, ni)]
  gcs[i] = GLMCopulaARObs(y, X, d, link)
end

# form model
gcm = GLMCopulaARModel(gcs)
# precompile
println("precompiling Bernoulli AR fit")
gcm2 = deepcopy(gcm);
QuasiCopula.fit!(gcm2, IpoptSolver(print_level = 0, max_iter = 20));

fittime = @elapsed QuasiCopula.fit!(gcm)
@show fittime
@show gcm.β
@show gcm.σ2
@show gcm.ρ
@show gcm.∇β
@show gcm.∇σ2
@show gcm.∇ρ

@test logl(gcm) == loglikelihood!(gcm, false, false)
@show get_CI(gcm)

# mse and time under our model
mseβ, mseρ, mseσ2 = MSE(gcm, βtrue, ρtrue, σ2true)
@show mseβ
@show mseσ2
@show mseρ

@test mseβ < 0.01
@test mseσ2 < 1
@test mseρ < 0.01

println("checking memory allocation for Bernoulli AR")
logl_gradient_memory = @benchmark loglikelihood!($gcm.data[1], $gcm.β, $gcm.ρ[1], $gcm.σ2[1], true, false)
@test logl_gradient_memory.memory == 0.0

# note for multi-threading we will allocate some at the model level loglikelihood!(gcm, true, false)


#================================================================

Common matrix and vector operations used in multiple locations
throughout the module.

================================================================#

# Compute the norms of vectors and columns of matrices
col_norms(x::Union{AbstractVector,AbstractMatrix}) =
    map(norm, eachcol(x))

col_norms(x::Union{Vector,Matrix}) =
    map(BLAS.nrm2, eachcol(x))

col_norms(x::SparseVector) =
    [BLAS.nrm2(x.nzval)]

col_norms(x::SparseMatrixCSC{T}) where {T} = begin
    output = Vector{T}(undef, size(x,2))
    @inbounds for ii = 1:size(x,2)
        result = T(0)
        start_idx, end_idx = x.colptr[ii], x.colptr[ii+1]-1
        @simd for idx = start_idx:end_idx
            result += x.nzval[idx].^2
        end
        output[ii] = √result
    end
    return output
end


using Test

include("./Tools.jl")
@testset "Test Tools" begin
#evalpoly testing
@testset "Test evalpoly function" begin
coefs = [1.0,7.0,9.0]
x = 3.0
@test evalpoly(x,coefs) ≈ sum([x^(i-1)*coefs[i] for i = 1:length(coefs)])
end;

@testset "Test get index first greater function" begin
#getIndexFirstGreater testing
a = [-4.0,1.0,3.0,7.0]
@test getBoundingIndsSorted(2.0,a) == (2,3)
@test getBoundingIndsSorted(-5.0,a) == (1,-1)
@test getBoundingIndsSorted(8.0,a) == (length(a),-1)
end;
end;


let
    function f_lm(x)
      [x[1], 2.0 - x[2]]
    end
    function g_lm(x)
      [1.0 0.0; 0.0 -1.0]
    end

    initial_x = [100.0, 100.0]

    results = Optim.levenberg_marquardt(f_lm, g_lm, initial_x)
    @assert norm(Optim.minimizer(results) - [0.0, 2.0]) < 0.01

    function rosenbrock_res(x, r)
        r[1] = 10.0 * (x[2] - x[1]^2 )
        r[2] =  1.0 - x[1]
        return r
    end

    function rosenbrock_jac(x, j)
        j[1, 1] = -20.0 * x[1]
        j[1, 2] =  10.0
        j[2, 1] =  -1.0
        j[2, 2] =   0.0
        return j
    end

    r = zeros(2)
    j = zeros(2,2)

    frb(x) = rosenbrock_res(x, r)
    grb(x) = rosenbrock_jac(x, j)

    initial_xrb = [-1.2, 1.0]

    results = Optim.levenberg_marquardt(frb, grb, initial_xrb)

    @assert norm(Optim.minimizer(results) - [1.0, 1.0]) < 0.01

    # check estimate is within the bound PR #278
     result = Optim.levenberg_marquardt(frb, grb, [150.0, 150.0]; lower = [10.0, 10.0], upper = [200.0, 200.0])
     @test Optim.minimizer(result)[1] >= 10.0
     @test Optim.minimizer(result)[2] >= 10.0

    # tests for #178, taken from LsqFit.jl, but stripped
    let
        srand(12345)

        model(x, p) = p[1]*map(exp,-x.*p[2])

        xdata = linspace(0,10,20)
        ydata = model(xdata, [1.0 2.0]) + 0.01*randn(length(xdata))

        f_lsq = p -> p[1]*map(exp, -xdata.*p[2])-ydata
        g_lsq = Calculus.jacobian(f_lsq)
        results = Optim.levenberg_marquardt(f_lsq, g_lsq, [0.5, 0.5])

        @assert norm(Optim.minimizer(results) - [1.0, 2.0]) < 0.05
    end

    let
        srand(12345)

        model(x, p) = p[1]*map(exp, -x./p[2])+p[3]

        xdata = 1:100
        ydata = model(xdata, [10.0, 10.0, 10.0]) + 0.1*randn(length(xdata))

        f_lsq = p -> model(xdata,p)-ydata
        g_lsq = Calculus.jacobian(f_lsq)

        # tests for box constraints, PR #196
        @test_throws ArgumentError Optim.levenberg_marquardt(f_lsq, g_lsq, [15.0, 15.0, 15.0], lower=[5.0, 11.0])
        @test_throws ArgumentError Optim.levenberg_marquardt(f_lsq, g_lsq, [5.0, 5.0, 5.0], upper=[15.0, 9.0])
        @test_throws ArgumentError Optim.levenberg_marquardt(f_lsq, g_lsq, [15.0, 10.0, 15.0], lower=[5.0, 11.0, 5.0])
        @test_throws ArgumentError Optim.levenberg_marquardt(f_lsq, g_lsq, [5.0, 10.0, 5.0], upper=[15.0, 9.0, 15.0])

        lower=[5.0, 11.0, 5.0]
        results = Optim.levenberg_marquardt(f_lsq, g_lsq, [15.0, 15.0, 15.0], lower=lower)
        Optim.minimizer(results)
        @test Optim.converged(results)
        @test all(Optim.minimizer(results) .>= lower)

        upper=[15.0, 9.0, 15.0]
        results = Optim.levenberg_marquardt(f_lsq, g_lsq, [5.0, 5.0, 5.0], upper=upper)
        Optim.minimizer(results)
        @test Optim.converged(results)
        @test all(Optim.minimizer(results) .<= upper)

        # tests for PR #267
        Optim.levenberg_marquardt(f_lsq, g_lsq, [15.0, 15.0, 15.0], show_trace=true)
    end
end


import GRAYTR: lerp, quadratic

@testset "Utility functions" begin

@testset "Linear extrapolation/interpolation" begin
    @test lerp(0, 0, 1) == 0
    @test lerp(1, 0, 1) == 1
    @test lerp(0.0, 0, 1) ≈ 0.0
    @test lerp(1.0, 0, 1) ≈ 1.0

    for t = [-0.2, 0.15, 0.5, 1.25]
        @test lerp(t, 0.0, 1.0) ≈ t
        @test lerp(t, 0.0, 2.0) ≈ 2.0 * t
        @test lerp(t, -1.0, 1.0) ≈ 2.0 * t - 1.0
    end
end

@testset "Quadratic solver" begin
    @testset "No solution" begin
        solvable, t1, t2 = quadratic(2.0, 1.0, 1.0)
        @test !solvable
        @test t1 == t2 == Inf
    end

    @testset "Golden ratio" begin
        solvable, t1, t2 = quadratic(1.0, -1.0, -1.0)
        @test solvable
        @test t1 ≈ (1 - sqrt(5)) / 2
        @test t2 ≈ (1 + sqrt(5)) / 2
    end
end

end # outer testset


"""

Functions that wrap the coordinates

They modify the coordinates of the input vector.

```
function wrap!(x::AbstractVector{T}, sides::T, center::T) where T <: AbstractVector
```

Wrap to a given center of coordinates

"""
function wrap!(x::AbstractVector{T},
               sides::T,
               center::T) where T <: AbstractVector
  for i in eachindex(x)
    x[i] = wrapone(x[i],sides,center)
  end
  nothing
end

@inline function wrapone(x::T, sides::T, center::T) where T <: AbstractVector
  s = @. (x-center)%sides
  s = @. wrapx(s,sides) + center
  return s
end

@inline function wrapx(x,s)
  if x > s/2
    x = x - s
  elseif x < -s/2
    x = x + s
  end
  x
end

"""

```
wrap!(x::AbstractVector{T}, sides::T) where T <: AbstractVector
```

Wrap to origin.

"""
function wrap!(x::AbstractVector{T}, sides::T) where T <: AbstractVector
  for i in eachindex(x)
    x[i] = wrapone(x[i],sides)
  end
  nothing
end

@inline function wrapone(x::T, sides::T) where T <: AbstractVector
  s = @. x%sides
  s = @. wrapx(s,sides)
  return s
end


A = [
  1 2 3
  2 3 4
  3 4 5
]

B = [
  4 5 6
  5 6 7
  6 7 8
]

C = A + B

println(C)


### A Pluto.jl notebook ###
# v0.14.7

using Markdown
using InteractiveUtils

# This Pluto notebook uses @bind for interactivity. When running this notebook outside of Pluto, the following 'mock version' of @bind gives bound variables a default value (instead of an error).
macro bind(def, element)
    quote
        local el = $(esc(element))
        global $(esc(def)) = Core.applicable(Base.get, el) ? Base.get(el) : missing
        el
    end
end

# ╔═╡ 1a5c4360-c5ec-11eb-073d-8dc280b34894
begin
	using CSV
	using DataFrames
	using GLM
	using Plots
	using PlutoUI
end

# ╔═╡ abf556e0-810d-443c-87b1-d0f95f9c080d
data = CSV.read("D:\\Workspace\\Machine Learning\\Canada GDP Analysis\\Data\\canada_gdp.csv",DataFrame)

# ╔═╡ 02178664-f416-4505-932a-019ca57f17cb
scatter(data.year,data.GDP)

# ╔═╡ e6a9282e-b580-4fd4-bb16-e83a8e5e0105
reg = fit(LinearModel,@formula(GDP~year),data)

# ╔═╡ b55b5c14-5761-4297-8f56-c1a6fb86e67c
linear_regression = predict(reg)

# ╔═╡ 5c0c577a-3571-4d5f-90da-c5c94942673a
begin
	scatter(data.year,data.GDP)
	plot!(data.year,linear_regression)
end

# ╔═╡ e8c03cd5-12a6-45d3-80b4-20feae154e32
md"Enter a year to predict: $@bind y TextField()"

# ╔═╡ d383695f-5bb5-4b8b-8876-d856405ec8aa
pyear = parse(Int64,y)

# ╔═╡ fbb51392-e820-4317-9334-e0c722d9a423
pGDP = coef(reg)[2]*pyear+coef(reg)[1]

# ╔═╡ Cell order:
# ╠═1a5c4360-c5ec-11eb-073d-8dc280b34894
# ╠═abf556e0-810d-443c-87b1-d0f95f9c080d
# ╠═02178664-f416-4505-932a-019ca57f17cb
# ╠═e6a9282e-b580-4fd4-bb16-e83a8e5e0105
# ╠═b55b5c14-5761-4297-8f56-c1a6fb86e67c
# ╠═5c0c577a-3571-4d5f-90da-c5c94942673a
# ╠═e8c03cd5-12a6-45d3-80b4-20feae154e32
# ╠═d383695f-5bb5-4b8b-8876-d856405ec8aa
# ╠═fbb51392-e820-4317-9334-e0c722d9a423


"""
> Yura Malitsky. Projected reflected gradient methods for monotone variational inequalities.
> SIAM Journal on Optimization, 25(1):502–520, 2015.
"""
function ProjectedReflectedGradientVI(stp::AbstractStopping; rho0::Float64 = 0.5)
  xk = stp.current_state.x
  xkp = similar(xk)
  yk = copy(xk)
  rho = rho0
  Fx = similar(xk)

  OK = update_and_start!(stp)
  while !OK
    abcresidual!(stp.pb, xk, rho, yk, Fx)
    project!(stp.pb, Fx, xk) # possible failure here
    yk .= 2 .* xkp .- xk

    if norm(xk - xkp, Inf) < stp.meta.atol * rho
      stp.meta.optimal = true
    end
    OK = update_and_stop!(stp, x = xk)
  end

  return xk
end

export ProjectedReflectedGradientVI

function ProjectedReflectedGradientVI(
  model::AbstractVIModel,
  x0::AbstractVector;
  rho0::Float64 = 0.5,
  kwargs...,
)
  stp = GenericStopping(model, x0; kwargs...)
  return ProjectedReflectedGradientVI(stp, rho0 = rho0)
end

function abcresidual!(model, xk, rho, yk, Fx) # xk - rho * F(yk)
  residual!(model, yk, Fx)
  Fx .*= -rho
  Fx .+= xk
  return Fx
end


using StanSample

#StanBase.set_cmdstan_home!(homedir() * "/Projects/StanSupport/cmdstan_stanc3")

ProjDir = @__DIR__
cd(ProjDir)

bernoulli_model = "
  functions{
    #include model_specific_funcs.stan
    #include shared_funcs.stan // a comment
    //#include shared_funcs.stan // a comment
  }
  data {
    int<lower=1> N;
    int<lower=0,upper=1> y[N];
  }
  parameters {
    real<lower=0,upper=1> theta;
  }
  model {
    model_specific_function();
    theta ~ beta(my_function(),1);
    y ~ bernoulli(theta);
  }
";

tmpdir = joinpath(ProjDir, "tmp")
stanmodel = SampleModel("bernoulli", bernoulli_model)

observeddata = Dict("N" => 10, "y" => [0, 1, 0, 1, 0, 0, 0, 0, 0, 1])

(sample_file, log_file) = stan_sample(stanmodel, data=observeddata)

if sample_file !== Nothing
  # Convert to an MCMCChains.Chains object
  chns = read_samples(stanmodel)

  # Describe the MCMCChains using MCMCChains statistics
  cdf = describe(chns)
  display(cdf)

  # Show cmdstan summary in DataFrame format
  sdf = read_summary(stanmodel)
  display(sdf)
  println()

  # Retrieve mean value of theta from the summary
  sdf[:theta, :mean]
end


module MNIST_BATCH
using Flux
using Flux.Data.MNIST, Statistics
using Flux: onehotbatch, onecold, crossentropy,throttle
using Base.Iterators: repeated,partition

using CUDAnative
using CuArrays
CuArrays.allowscalar(false)

#=
Very important !!
ϵ is used to prevent loss NaN
=#
const ϵ = 1.0f-32

# Load training labels and images from Flux.Data.MNIST
@info("Loading data...")

train_imgs = MNIST.images()
train_labels = MNIST.labels()

# use 1nd GPU : default
CUDAnative.device!(0)
# use 2nd GPU
#CUDAnative.device!(1)

# Bundle images together with labels and group into minibatch
function make_minibatch(imgs,labels,batch_size)
  X = hcat(float.(reshape.(imgs,:))...) |> gpu
  Y = float.(onehotbatch(labels,0:9)) |> gpu

  data_set = [(X[:,i],Y[:,i]) for i in partition(1:length(labels),batch_size)]
  return data_set
end

@info("Making model...")
# Model
m = Chain(
  Dense(28^2,32,relu),
  Dense(32,10),
  softmax
) |> gpu
loss(x,y) = crossentropy(m(x) .+ ϵ, y)
accuracy(x,y) = mean(onecold(m(x)|>cpu) .== onecold(y|>cpu))

batch_size = 500
train_dataset = make_minibatch(train_imgs,train_labels,batch_size)

opt = ADAM()

@info("Training model...")

epochs = 200
# used for plots
accs = Array{Float32}(undef,0)

dataset_len = length(train_dataset)
for i in 1:epochs
  for (idx,dataset) in enumerate(train_dataset)
    Flux.train!(loss,params(m),[dataset],opt)
    acc = accuracy(dataset...)
    if idx == dataset_len
      @info("Epoch# $(i)/$(epochs) - loss: $(loss(dataset...)), accuracy: $(acc)")
      push!(accs,acc)
    end
  end
end

# Test Accuracy
tX = hcat(float.(reshape.(MNIST.images(:test),:))...) |> gpu
tY = float.(onehotbatch(MNIST.labels(:test),0:9)) |> gpu

println("Test loss:", loss(tX,tY))
println("Test accuracy:", accuracy(tX,tY))

end

using Plots;gr()
plot(MNIST_BATCH.accs)


##############################################################################################
module L96
##############################################################################################
# imports and exports
using Random, Distributions
using LinearAlgebra, SparseArrays
export dx_dt, jacobian, l96s_tay_2_step!, ρ, α

##############################################################################################
##############################################################################################
# Type union declarations for multiple dispatch
# and type aliases

# vectors and ensemble members of sample
VecA = Union{Vector{Float64}, SubArray{Float64, 1}}

# dictionary for model parameters
ParamDict = Union{Dict{String, Array{Float64}}, Dict{String, Vector{Float64}}}

##############################################################################################
##############################################################################################
# auxiliary function to return modular indices for the lorenz model

function mod_indx!(indx::Int64, dim::Int64)
    indx = mod(indx, dim)
    if indx==0
        indx = dim
    end
    return indx
end

##############################################################################################
# time derivative

function dx_dt(x::VecA, t::Float64, dx_params::ParamDict)
	"""Time derivative for Lorenz-96 model, x is a single model state of size state_dim."""

    # unpack the (only) derivative parameter for l96
    F = dx_params["F"][1]::Float64
    x_dim = length(x)
    dx = Vector{Float64}(undef, x_dim)

    for j in 1:x_dim
        # index j minus 2, modulo the system dimension
        j_m_2 = mod_indx!(j - 2, x_dim)

        # index j minus 1, modulo the system dimension
        j_m_1 = mod_indx!(j - 1, x_dim)

        # index j plus 1, modulo the system dimension
        j_p_1 = mod_indx!(j + 1, x_dim)

        dx[j] = (x[j_p_1] - x[j_m_2])*x[j_m_1] - x[j] + F
    end
    return dx
end

##############################################################################################
# vectorized time derivative for demonstration purposes only,
# note this is much less efficient in speed and in memory than the above version

function dx_dt_alt(x::VecA, t::Float64, dx_params::ParamDict)

    # unpack the (only) derivative parameter for l96
    F = dx_params["F"][1]::Float64
    x_dim = length(x)
    dx = Vector{Float64}(undef, x_dim)

    # shift minus and plus indices, array compute the derivative over arbitrary ensemble sizes
    x_m_2 = [x[end-1:end]; x[1:end-2]]
    x_m_1 = [x[end:end]; x[1:end-1]]
    x_p_1 = [x[2:end]; x[1:1]]

    # compute the vectorized derivative
    dx .= (x_p_1-x_m_2) .* x_m_1 - x .+ F
end

##############################################################################################
# linearized time derivative

function jacobian(x::Vector{Float64}, t::Float64, dx_params::ParamDict)
    """"Computes the Jacobian of Lorenz-96 about the state x.

    Note that this is designed to load entries in a zeros array and return a sparse array
    to make a compromise between memory and computational resources."""

    x_dim = length(x)
    dxF = zeros(x_dim, x_dim)

    # looping columns j of the jacobian, loading the standard matrix
    for j in 1:x_dim

        # index j minus 2, modulo the system dimension
        j_m_2 = mod_indx!(j - 2, x_dim)

        # index j minus 1, modulo the system dimension
        j_m_1 = mod_indx!(j - 1, x_dim)

        # index j plus 1, modulo the system dimension
        j_p_1 = mod_indx!(j + 1, x_dim)

        # index j plus 2, modulo the system dimension
        j_p_2 = mod_indx!(j + 2, x_dim)

        # load the jacobian entries in corresponding rows
        dxF[j_p_2, j] = -x[j_p_1]
        dxF[j_p_1, j] = x[j_p_2] - x[j_m_1]
        dxF[j, j] = -1.0
        dxF[j_m_1, j] = x[j_m_2]
    end
    return sparse(dxF)
end

##############################################################################################
# Auxiliary functions for the 2nd order Taylor-Stratonovich expansion below. These need
# to be computed once, only as a function of the order of truncation of the Fourier
# series, p for full timeseries.

function ρ(p::Int64)
    1.0/12.0 - 0.5 * π^(-2.0) * sum(1.0 ./ Vector{Float64}(1:p).^2.0)
end

function α(p::Int64)
    (π^2.0) / 180.0 - 0.5 * π^(-2.0) * sum(1.0 ./ Vector{Float64}(1:p).^4.0)
end

##############################################################################################
# 2nd order strong taylor SDE step
# This method is derived in
# Grudzien, C. et al.: On the numerical integration of the Lorenz-96 model,
# with scalar additive noise, for benchmark twin experiments,
# Geosci. Model Dev., 13, 1903–1924, https://doi.org/10.5194/gmd-13-1903-2020, 2020.
# This depends on rho and alpha as above
# NOTE: this Julia version still pending validation as in the above manuscript

function l96s_tay2_step!(x::Vector{Float64}, t::Float64, kwargs::Dict{String,Any})
    """One step of integration rule for l96 second order taylor rule

    The ρ and α are to be computed by the auxiliary functions, depending only on p,
    and supplied for all steps. This is the general formulation which includes,
    eg. dependence on the truncation of terms in the auxilliary function C with
    respect to the parameter p.  In general, truncation at p=1 is all that is
    necessary for order 2.0 convergence, and in this case C below is identically
    equal to zero.  This auxilliary function can be removed (and is removed) in other
    implementations for simplicity."""

    # Infer model and parameters
    sys_dim = length(x)
    dx_params = kwargs["dx_params"]::ParamDict
    h = kwargs["h"]::Float64
    diffusion = kwargs["diffusion"]::Float64
    p = kwargs["p"]::Int64
    ρ = kwargs["ρ"]::Float64
    α = kwargs["α"]::Float64

    # Compute the deterministic dxdt and the jacobian equations
    dx = dx_dt(x, 0.0, dx_params)
    Jac_x = jacobian(x, 0.0, dx_params)

    ## random variables
    # Vectors ξ, μ, ϕ are sys_dim X 1 vectors of iid standard normal variables,
    # ζ and η are sys_dim X p matrices of iid standard normal variables.
    # Functional relationships describe each variable W_j as the transformation of
    # ξ_j to be of variace given by the length of the time step h. Functions of random
    # Fourier coefficients a_i, b_i are given in terms μ / η and ϕ / ζ respectively.

    # draw standard normal samples
    rndm = rand(Normal(), sys_dim, 2*p + 3)
    ξ = rndm[:, 1]

    μ = rndm[:, 2]
    ϕ = rndm[:, 3]

    ζ = rndm[:, 4: p+3]
    η = rndm[:, p+4: end]

    ### define the auxiliary functions of random fourier coefficients, a and b

    # denominators for the a series
    denoms = repeat((1.0 ./ Vector{Float64}(1:p)), 1, sys_dim)

    # vector of sums defining a terms
    a = -2.0 * sqrt(h * ρ) * μ - sqrt(2.0*h) * sum(ζ' .* denoms, dims=1)' / π

    # denominators for the b series
    denoms = repeat((1.0 ./ Vector{Float64}(1:p).^2.0), 1, sys_dim)

    # vector of sums defining b terms
    b = sqrt(h * α) * ϕ + sqrt(h / (2.0 * π^2.0) ) * sum(η' .* denoms, dims=1)'

    # vector of first order Stratonovich integrals
    J_pdelta = (h / 2.0) * (sqrt(h) * ξ + a)

    ### auxiliary functions for higher order stratonovich integrals ###
    # C function is optional for higher precision but does not change order of convergence
    function C(l, j)
        if p == 1
            return 0.0
        end
        c = zeros(p, p)
        # define the coefficient as a sum of matrix entries where r and k do not agree
        indx = Set(1:p)

        for r in 1:p
            # vals are all values not equal to r
            vals = setdiff(indx, Set(r))
            for k in vals
                # and for column r, define all row entries below, with zeros on diagonal
                c[k, r] = (r / (r^2 - k^2)) * ((1.0 / k) * ζ[l, r] * ζ[j, k] + (1.0 / r) *
                                               η[l, r] * η[j, k])
            end
        end

        # return the sum of all values scaled by -1/(2pi^2)
        -0.5 * π^(-2.0) * sum(c)
    end

    function Ψ(l, j)
        # Ψ - generic function of the indicies l and j, define Ψ plus and Ψ minus index-wise
        h^2.0 * ξ[l] * ξ[j] / 3.0 + h * a[l] * a[j] / 2.0 +
        h^(1.5) * (ξ[l] * a[j] + ξ[j] * a[l]) / 4.0 -
        h^(1.5) * (ξ[l] * b[j] + ξ[j] * b[l]) / (2.0 * π) - h^2.0 * (C(l,j) + C(j,l))
    end

    # define the approximations of the second order Stratonovich integral
    Ψ_plus = Vector{Float64}(undef, sys_dim)
    Ψ_minus = Vector{Float64}(undef, sys_dim)
    for i in 1:sys_dim
        Ψ_plus[i] = Ψ(mod_indx!((i-1), sys_dim), mod_indx!((i+1), sys_dim))
        Ψ_minus[i] = Ψ(mod_indx!((i-2), sys_dim), mod_indx!((i-1), sys_dim))
    end

    # the final vectorized step forward is given as
    x .= collect(Iterators.flatten(
            x + dx * h + h^2.0 * 0.5 * Jac_x * dx +  # deterministic taylor step
            diffusion * sqrt(h) * ξ +                # stochastic euler step
            diffusion * Jac_x * J_pdelta +           # stochastic first order taylor step
            diffusion^2.0 * (Ψ_plus - Ψ_minus)       # stochastic second order taylor step
           ))
end

##############################################################################################
# end module

end


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

using Markdown
using InteractiveUtils

# This Pluto notebook uses @bind for interactivity. When running this notebook outside of Pluto, the following 'mock version' of @bind gives bound variables a default value (instead of an error).
macro bind(def, element)
    quote
        local el = $(esc(element))
        global $(esc(def)) = Core.applicable(Base.get, el) ? Base.get(el) : missing
        el
    end
end

# ╔═╡ 1dfd0480-6b0d-11eb-1f62-dbdc1a677571
begin
#	import Pkg;
#	Pkg.add("Plots")
#	Pkg.add("PlutoUI")
	using PlutoUI
	using Plots
	using LinearAlgebra
	using SparseArrays
end

# ╔═╡ e6375832-6b29-11eb-38b2-7582cac61e64
md"""

# **The strange dice** - _The central limit theorem_

**Bernoulli and Laplace** invented this interesting dice game with colored faces (red for odd numbers and green for even numbers of pips).
The side facing the gambler works as a bonus / malus factor.

In the case of a green side facing the gambler, the points get doubled, in case of a red face the points get subtracted by minus 1.

At the trinautic tournament the whole crew is rolling the strange dice and the points are summed up.

Look at the distribution below that shows the probabilities for the different number of points.

**Use the slider to add crew members to play and look how the average distribution of points (total points divided by the number of gamblers) changes**

"""

# ╔═╡ 86b1d820-6b2b-11eb-19cf-6390ab5fccc8
HTML("

")

# ╔═╡ 9f155e80-6be6-11eb-22c1-1b4f6e30ec72
md"Choose number of 👉 $(@bind iterations Slider(1:60, show_value = true)) crew members that roll the strange dice."

# ╔═╡ 7ecb15e0-6b26-11eb-1d2b-37dd4e07ef05
HTML("

")

# ╔═╡ 185ef260-6b0b-11eb-1fd8-b5c70ec616d7
p = [1/6 1/6 1/6 1/6 1/6 1/6]

# ╔═╡ 943dea30-6b0b-11eb-19f0-75227e40a492
A = [ 3/4 1/4; 1/2 1/2; 3/4 1/4; 1/4 3/4; 1/2 1/2; 1/4 3/4]

# ╔═╡ 036a9bb0-6b0c-11eb-0ece-d54174c4d7a0
probability_table = A.*p'

# ╔═╡ 3eb465e0-6b14-11eb-19a4-df75db84f12d
sum(probability_table)

# ╔═╡ 427d9fb0-6b10-11eb-043c-653dbc050130
begin
	B = zeros(13,2)
	B[(1:6)',2] = probability_table[:,2]
	B[:,1] = (0:12)'
	index = collect(1:6)*2
	B[index.+1,2] = B[index.+1,2] + probability_table[:,1]
end

# ╔═╡ 29969590-6b12-11eb-11d7-9d831132fe4f
begin
	plot(
    B[:,1], B[:,2],
    line = (1.0, 0.0, :bar),
    normalize = false,
    bins = 10,
	bar_width = 0.1,
    marker = (6, 0.5, :x),
    markerstrokewidth = 5.,
    color = [:steelblue],
    fill = 0.9,
    orientation = :v,
    title = "The strange dice distribution",
	ylabel = "probability mass function",
	xlabel = "Points",
	label = :none,
)
plot!( B[:,1], B[:,2],
    line = (1.0, 0, :path),
    normalize = false,
    bins = 10,
	bar_width = 0.2,
    marker = (4, 0.5, :o),
    markerstrokewidth = 5.,
    color = [:steelblue],
    fill = 0.9,
    orientation = :v,
    title = "The strange dice distribution",
	ylabel = "probability mass function",
	xlabel = "Points",
	label = :none,
)
end

# ╔═╡ 6df0d2c0-6b10-11eb-04d1-1b71b887b8bb
B,sum(B[:,2])

# ╔═╡ 462e27c0-6b14-11eb-2b2f-3bd578fb8712
dice_mean = sum(B[:,2].*B[:,1])

# ╔═╡ 691d5aa0-6b71-11eb-2f53-db4a2217c8e0
dice_std = sqrt(sum(B[:,2].* (B[:,1].-dice_mean).^2))

# ╔═╡ 18e26640-6be7-11eb-20b0-5542c5620d8d
iter = iterations-1

# ╔═╡ 7dbbfe90-6b20-11eb-39c8-c99fd55ccf05
sparse_vector = repeat(B[:,2]',iter*12+1)

# ╔═╡ 80f0ce50-6b21-11eb-2b31-6982188b8a41
sparse_vector[:,1]

# ╔═╡ e107cae0-6b1d-11eb-3ff5-b7861639d659
P = Array(spdiagm(0 => sparse_vector[:,1], 1 => sparse_vector[:,2], 2 => sparse_vector[:,3], 3 => sparse_vector[:,4], 4 => sparse_vector[:,5], 5 => sparse_vector[:,6], 6 => sparse_vector[:,7], 7 => sparse_vector[:,8], 8 => sparse_vector[:,9], 9 => sparse_vector[:,10], 10=> sparse_vector[:,11], 11 => sparse_vector[:,12], 12 => sparse_vector[:,13]))

# ╔═╡ 58fe1a00-6b22-11eb-207c-095c2b30a0e6
pp = append!(B[:,2], zeros((iter)*12))

# ╔═╡ b3279d30-6b22-11eb-0658-dbe663f09558
p_final = Array(pp'*P^(iter))

# ╔═╡ 6f072a20-6b23-11eb-0add-8b62821b708a
r = Array(LinRange(0,12,(iter-1)*12+13*2-1))

# ╔═╡ b7a9ad70-6b23-11eb-295b-4324ef257047
begin

	plot(
    [B[:,1], r, [dice_mean, dice_mean]], [B[:,2]./0.9,p_final[1,:]./((r[2]-r[1])*0.9), [0,1]],
    line = ([1 0.3 1], [0 0.0 3], [:bar :bar :line]),
	linestyle = [:auto :auto :dash],
    normalize = true,
    marker = (6, 0.8, [:none :none]),
    markerstrokewidth = 5.,
    color = [:steelblue :orangered :black],
    fill = [1. 1.],
	bar_width = [1*0.9 (r[2]-r[1])*0.9],
    orientation = :v,
    title = string("Average distribution of \nstrange dice with ", string(iter+1), " gamblers"),
	ylabel = "probability / width of bar",
	xlabel = "Points",
	usetex = true,
	label = [:none :none "Mean"],
	ylim = [0,1],
)
	plot!([dice_mean - dice_std, dice_mean + dice_std, 3, dice_mean+dice_std, dice_mean-dice_std], 0.5.*[1, 1,NaN, 1, 1],
		line = (2., 2., :path),
		color = :magenta,
		arrow = true,
		label = "Std",
		)
	plot!([dice_mean - dice_std/sqrt(iter), dice_mean + dice_std/sqrt(iter), 3, dice_mean+dice_std/sqrt(iter), dice_mean-dice_std/sqrt(iter)], 0.08*[1, 1,NaN, 1, 1]/sqrt(((r[2]-r[1]))),
		line = (2., 2., :path),
		color = :green,
		arrow = true,
		label = string("Std/sqrt(",string(iter+1), ")"),
		)

end

# ╔═╡ 89f8f100-6b24-11eb-3292-c15016c60a5b
p_final[1,:]

# ╔═╡ cec01d8e-6be2-11eb-1465-ad0f3ffcc2f5
function pretty(M::Matrix{T} where T<:String)
	max_length = maximum(length.(M))
	dv="
"
	HTML(dv*join([join("
".*M[i,:].*"
", " ") for i in 1:size(M,1)]
			, "
$dv")*"
")
end

# ╔═╡ e7e029f0-6b9c-11eb-179d-1d2fca9eb8af
[string.(rand("🧔👩👨👧🧑", 1, iterations));string.(rand("🎲", 1, iterations))] |> pretty

# ╔═╡ Cell order:
# ╟─1dfd0480-6b0d-11eb-1f62-dbdc1a677571
# ╟─e6375832-6b29-11eb-38b2-7582cac61e64
# ╟─86b1d820-6b2b-11eb-19cf-6390ab5fccc8
# ╟─29969590-6b12-11eb-11d7-9d831132fe4f
# ╟─e7e029f0-6b9c-11eb-179d-1d2fca9eb8af
# ╟─9f155e80-6be6-11eb-22c1-1b4f6e30ec72
# ╟─b7a9ad70-6b23-11eb-295b-4324ef257047
# ╟─7ecb15e0-6b26-11eb-1d2b-37dd4e07ef05
# ╟─185ef260-6b0b-11eb-1fd8-b5c70ec616d7
# ╟─943dea30-6b0b-11eb-19f0-75227e40a492
# ╟─036a9bb0-6b0c-11eb-0ece-d54174c4d7a0
# ╟─3eb465e0-6b14-11eb-19a4-df75db84f12d
# ╟─427d9fb0-6b10-11eb-043c-653dbc050130
# ╟─6df0d2c0-6b10-11eb-04d1-1b71b887b8bb
# ╟─691d5aa0-6b71-11eb-2f53-db4a2217c8e0
# ╟─462e27c0-6b14-11eb-2b2f-3bd578fb8712
# ╟─18e26640-6be7-11eb-20b0-5542c5620d8d
# ╟─7dbbfe90-6b20-11eb-39c8-c99fd55ccf05
# ╟─80f0ce50-6b21-11eb-2b31-6982188b8a41
# ╟─e107cae0-6b1d-11eb-3ff5-b7861639d659
# ╟─58fe1a00-6b22-11eb-207c-095c2b30a0e6
# ╟─b3279d30-6b22-11eb-0658-dbe663f09558
# ╟─6f072a20-6b23-11eb-0add-8b62821b708a
# ╟─89f8f100-6b24-11eb-3292-c15016c60a5b
# ╟─cec01d8e-6be2-11eb-1465-ad0f3ffcc2f5


### A Pluto.jl notebook ###
# v0.17.0

using Markdown
using InteractiveUtils

# ╔═╡ 39da52d0-3e5b-4dad-9de9-6a8b0337f124
begin
	# Run once if dependencies not installed
	using Pkg
	Pkg.add("Plots")
end

# ╔═╡ 4345a821-21fa-4a84-91a9-c3c539db2725
begin
	using Plots
end

# ╔═╡ 88843947-6008-41cc-b204-3f822e542ae3
md"""
# Chapter 1
## Section 2
"""

# ╔═╡ 98cdb6dc-3a40-11ec-1eb2-f7e085a3b6e6
function trapezoid(f, n::Integer, a::T, b::T)::T where T <: Real
	dx = (b - a) / n
	if a < b
		a, b = a + dx, b - dx
		sign = 1.
	else
		a, b = b + dx, a - dx
		sign = -1
	end
	integral = 0
	h = 2. * dx
	for x in a:h:b
		integral += 0.5 * dx * (f(x + dx) + 2 * f(x) + f(x - dx))
	end
	integral
end

# ╔═╡ e62b3b53-60c8-4450-978c-7558398fa085
trapezoid(sin, 256, 0., float(pi))

# ╔═╡ 7aa7729a-9c6c-457b-a271-8cd0775a65bf
function simpson(f, n::Integer, a::T, b::T)::T where T <: Real
	dx = (b - a) / n
	if a < b
		a, b = a + dx, b - dx
		sign = 1.
	else
		a, b = b + dx, a - dx
		sign = -1
	end
	integral = 0
	h = 2. * dx
	for x in a:h:b
		integral +=  (dx / 3) * (f(x + dx) + 4 * f(x) + f(x - dx))
	end
	integral
end

# ╔═╡ c1449e54-2bfa-4fce-a7b4-227451327b48
function bode(f, n::Integer, a::T, b::T)::T where T <: Real
	dx = (b - a) / n
	h = 4. * dx
	if a < b
		a, b = a + h / 2., b - h / 2.
		sign = 1.
	else
		a, b = b + h / 2., a - h / 2.
		sign = -1
	end
	integral = 0
	for x in a:h:b
		integral +=  (2 * dx / 45) * (
			7 * f(x + 2. * dx) + 32. * f(x + dx)
			+  12. * f(x) +
			7 * f(x - 2. * dx) + 32. * f(x - dx))
	end
	integral
end

# ╔═╡ 1dea0b53-fc17-43f3-b1df-495fff0529e9
begin
	# Exercice 1.2
	x = collect(2:8) # Weird result betwin 7 ~ 14
	func = [trapezoid, simpson, bode]
	y = Array{Float64}(undef, (length(func), length(x)))
	for (k, f) in enumerate(func)
		for i in 1:length(x)
			y[k, i] = abs(2. - f(sin, 1 << x[i], 0., float(pi)))
		end
	end
	plot(x, y', label=["trap" "simp" "bode"])
end

# ╔═╡ 76f93ada-aa5a-4b6b-abed-f006dbe6f89a
# Exercice 1.3
md"""

The function: $t^{-2/3} (1 - t) ^{-1/3)}$ have 2 singularity 0 and 1

We will need ton transform is integrale to be able to compute it

First we spliting the integral into two part

$$\int_0^1 t^{-2/3} (1 - t) ^{-1/3)} = \int_0^{0.5} t^{-2/3} (1 - t) ^{-1/3dt + \int_{0.5}^1 t^{-2/3} (1 - t) ^{-1/3}dt$$

Now we have isolated the singularities. Now let change variables

"""

# ╔═╡ 46f5c4bb-8f93-4c32-81ea-717a05232aec
begin

end

# ╔═╡ Cell order:
# ╟─88843947-6008-41cc-b204-3f822e542ae3
# ╠═39da52d0-3e5b-4dad-9de9-6a8b0337f124
# ╠═4345a821-21fa-4a84-91a9-c3c539db2725
# ╠═98cdb6dc-3a40-11ec-1eb2-f7e085a3b6e6
# ╠═e62b3b53-60c8-4450-978c-7558398fa085
# ╠═7aa7729a-9c6c-457b-a271-8cd0775a65bf
# ╠═c1449e54-2bfa-4fce-a7b4-227451327b48
# ╠═1dea0b53-fc17-43f3-b1df-495fff0529e9
# ╠═76f93ada-aa5a-4b6b-abed-f006dbe6f89a
# ╠═46f5c4bb-8f93-4c32-81ea-717a05232aec


#! What is the convention and or the most efficient way to store these multidim arrays: samples from column/row correspond to one distributions?
#! Here we choose row correspond to one distributions and number of row is the number of realization
#! It seems to be the convention
#TODO Only work for distributions with known fit_mle (cannot be product_distribution of mixture because of the typeof)
"""
Simply extend the `fit_mle` function to multivariate Product distributions.
"""
function fit_mle(g::Product, x::AbstractMatrix)
    S = size(x, 1) # row convention
    vec_g = g.v
    @argcheck S == length(vec_g)
    return product_distribution([fit_mle(typeof(vec_g[s]), y) for (s, y) in enumerate(eachrow(x))])
end

function fit_mle(g::Product, x::AbstractMatrix, γ::AbstractVector)
    S = size(x, 1) # row convention
    vec_g = g.v
    @argcheck S == length(vec_g)
    return product_distribution([fit_mle(typeof(vec_g[s]), y, γ) for (s, y) in enumerate(eachrow(x))])
end

"""
Now `fit_mle(Bernoulli(0.2), x)` is accepted in addition of `fit_mle(Bernoulli, x)` this allows compatibility with how `fit_mle(g::Product)` and `fit_mle(g::MixtureModel)` are written.
"""
function fit_mle(g::Distribution{Univariate,Discrete}, args...)
    fit_mle(typeof(g), args...)
end

function fit_mle(g::Distribution{Univariate,Continuous}, args...)
    fit_mle(typeof(g), args...)
end


abstract type AbstractCompactLattice{T,N_ARR,N_CDT} <: AbstractLattice{T,N_ARR,N_CDT} end

@doc """
A lattice of points, each an `N_ARR` coordinate, evenly distributed.

Can be constructed using default AbstractLattice constructor
"""
struct CompactLattice{T,N_ARR} <: AbstractCompactLattice{T,N_ARR,N_ARR}
    arr::Array{NTuple{N_ARR,T},N_ARR}
end
difference(lattice::AbstractCompactLattice, edge) = abs_difference(edge)

# TODO: Should be completely uniform in periodic case?
export simpson_weights
function simpson_weights(lattice::CompactLattice{T,1}) where T
    @assert all(size(lattice) .% 2 .== 1)
    w = ones(T, size(lattice)...)
    w[1] = 0.5
    w[end] = 0.5
    return w
end
function simpson_weights(lattice::CompactLattice{T,2}) where T
    # http://mathfaculty.fullerton.edu/mathews/n2003/SimpsonsRule2DMod.html
    @assert all(size(lattice) .% 2 .== 1)
    w = ones(T, size(lattice)...)
    w[2:2:end-1,:] .*= 4.0
    w[3:2:end-2,:] .*= 2.0
    w[:,2:2:end-1] .*= 4.0
    w[:,3:2:end-2] .*= 2.0
    return w
end
Base.getindex(lat::CompactLattice, dx) = getindex(lat.arr, dx)

# Use default lattice plotting


"""
    DiamondSquare

This type generates a neutral landscape using the diamond-squares
algorithm, which produces fractals with variable spatial autocorrelation.

https://en.wikipedia.org/wiki/Diamond-square_algorithm

The algorithm is named diamond-square because it is an iterative procedure of
"diamond" and "square" steps.

The degree of spatial autocorrelation is controlled by a parameter `H`,
which varies from 0.0 (low autocorrelation) to 1.0 (high autocorrelation) ---
note this is non-inclusive and H = 0 and H = 1 will not behave as expected.
The result of the diamond-square algorithm is a fractal with dimension D = 2 + H.

A similar algorithm, midpoint-displacement, functions almost
identically, except that in DiamondSquare, the square step interpolates
edge midpoints from the nearest two corners and the square's center, where as
midpoint-displacement only intepolates from the nearest corners (see `MidpointDisplacement`).

"""
struct DiamondSquare <: NeutralLandscapeMaker
    H::Float64
    function DiamondSquare(H::T) where {T <: Real}
        @assert 0 <=  H < 1
        new(H)
    end
end

"""
    DiamondSquare()

    Initializes DiamondSqaure with H = 0.5.
"""
DiamondSquare() = DiamondSquare(0.5)

"""
    MidpointDisplacement()

    Creates a midpoint-displacement algorithm object `MidpointDisplacement`. The degree of spatial autocorrelation is controlled by a parameter `H`,
    which varies from 0.0 (low autocorrelation) to 1.0 (high autocorrelation) --- note this is non-inclusive and H = 0 and H = 1 will not behavive as expected.

    A similar algorithm, diamond-square, functions almost
    identically, except that in diamond-square, the square step interpolates
    edge midpoints from the nearest two corners and the square's center, where as
    `MidpointDisplacement` only intepolates from the nearest corners (see `DiamondSquare`).
"""
struct MidpointDisplacement <: NeutralLandscapeMaker
    H::Float64
    function MidpointDisplacement(H::T) where {T <: Real}
        @assert 0 <= H < 1
        new(H)
    end
end
MidpointDisplacement() = MidpointDisplacement(0.5)

"""
    _landscape!(mat, alg::Union{DiamondSquare, MidpointDisplacement}; kw...)

    Check if `mat` is the right size.
    If mat is not the correct size (DiamondSquare can only run on a lattice of size NxN where N = (2^n)+1 for integer n),
    allocates the smallest lattice large enough to contain `mat` that can run DiamondSquare.
"""
function _landscape!(mat, alg::Union{DiamondSquare, MidpointDisplacement}; kw...) where {IT <: Integer}

    rightSize::Bool = _isPowerOfTwo(size(mat)[1]-1) && _isPowerOfTwo(size(mat)[2]-1)
    latticeSize::Int = size(mat)[1]

    dsMat = mat
    if !rightSize
        dim1, dim2 = size(mat)
        smallestContainingLattice::Int = 2^ceil(log2(max(dim1, dim2))) + 1
        @warn "$alg cannot be run on the input dimensions ($dim1 x $dim2),
                and will instead run on the next smallest valid size ($smallestContainingLattice x $smallestContainingLattice).
                This can slow performance as it involves additional memory allocation."
        dsMat = zeros(smallestContainingLattice, smallestContainingLattice)
    end
    _diamondsquare!(dsMat, alg)

    mat .= dsMat[1:size(mat)[1], 1:size(mat)[2]]
end

"""
    _diamondsquare!(mat, alg)

    Runs the diamond-square algorithm on a matrix `mat` of size
    `NxN`, where `N=(2^n)+1` for some integer `n`, i.e (N=5,9,17,33,65)

    Diamond-square is an iterative procedure, where the lattice is divided
    into subsquares in subsequent rounds. At each round, the subsquares shrink in size,
    as previously uninitialized values in the lattice are interpolated as a mean of nearby points plus random displacement.
    As the rounds increase, the magnitude of this displacement decreases. This creates spatioautocorrelation, which is controlled
    by a single parameter `H` which varies between `0` (no autocorrelation) and `1` (high autocorrelation)

"""
function _diamondsquare!(mat, alg)
    latticeSize = size(mat)[1]
    numberOfRounds::Int = log2(latticeSize-1)
    _initializeDiamondSquare!(mat, alg)

    for round in 0:(numberOfRounds-1)    # counting from 0 saves us a headache later
        subsquareSideLength::Int = 2^(numberOfRounds-(round))
        numberOfSubsquaresPerAxis::Int = ((latticeSize-1) / subsquareSideLength)-1

        for x in 0:numberOfSubsquaresPerAxis         # iterate over the subsquares within the lattice at this side length
            for y in 0:numberOfSubsquaresPerAxis
                subsquareCorners = _subsquareCornerCoordinates(x,y,subsquareSideLength)

                _diamond!(mat, alg, round, subsquareCorners)
                _square!(mat, alg, round, subsquareCorners)
            end
        end
    end
end

"""
    _initializeDiamondSquare!(mat, alg)

    Initialize's the `DiamondSquare` algorithm by displacing the four corners of the
    lattice using `displace`, scaled by the algorithm's autocorrelation `H`.
"""
function _initializeDiamondSquare!(mat, alg)
    latticeSize = size(mat)[1]
    corners = _subsquareCornerCoordinates(0,0, latticeSize-1)
    for mp in corners
        mat[mp...] = _displace(alg.H, 1)
    end
end

"""
    _subsquareCornerCoordinates(x::Int, y::Int, sideLength::Int)

    Returns the coordinates for the corners of the subsquare (x,y) given a side-length `sideLength`.
"""
function _subsquareCornerCoordinates(x::Int, y::Int, sideLength::Int)
    corners = [1 .+ sideLength.*i for i in [(x,y), (x+1, y), (x, y+1), (x+1, y+1)]]
end

"""
     _diamond!(mat, alg::DiamondSquare, round::Int, corners::AbstractVector{Tuple{Int, Int}})

     Runs the diamond step of the `DiamondSquare` algorithm on the square defined by
     `corners` on the matrix `mat`. The center of the square is interpolated from the
     four corners, and is displaced. The displacement is drawn according to `alg.H` and round using `displace`
"""
function _diamond!(mat, alg, round::Int, corners::AbstractVector{Tuple{Int, Int}})
    centerPt = _centerCoordinate(corners)
    mat[centerPt...] = _interpolate(mat, corners) + _displace(alg.H, round)
end

"""
    _square!(mat, alg::DiamondSquare, round::Int, corners::AbstractVector{Tuple{Int,Int}})

    Runs the square step of the `DiamondSquare` algorithm on the square defined
    by `corners` on the matrix `mat`. The midpoint of each edge of this square is interpolated
    by computing the mean value of the two corners on the edge and the center of the square, and the
    displacing it. The displacement is drawn according to `alg.H` and round using `displace`

"""
function _square!(mat, alg::DiamondSquare, round::Int, corners::AbstractVector{Tuple{Int, Int}})
    bottomLeft,bottomRight,topLeft,topRight = corners
    leftEdge, bottomEdge, topEdge, rightEdge = _edgeMidpointCoordinates(corners)
    centerPoint = _centerCoordinate(corners)

    mat[leftEdge...] = _interpolate(mat, [topLeft,bottomLeft,centerPoint]) + _displace(alg.H, round)
    mat[bottomEdge...] = _interpolate(mat, [bottomLeft,bottomRight,centerPoint]) + _displace(alg.H, round)
    mat[topEdge...] = _interpolate(mat, [topLeft,topRight,centerPoint]) + _displace(alg.H, round)
    mat[rightEdge...] = _interpolate(mat, [topRight,bottomRight,centerPoint]) + _displace(alg.H, round)
end

"""
    _square!(mat, alg::MidpointDisplacement, round::Int, corners::AbstractVector{Tuple{Int,Int}})

    Runs the square step of the `MidpointDisplacement` algorithm on the square defined
    by `corners` on the matrix `mat`. The midpoint of each edge of this square is interpolated
    by computing the mean value of the two corners on the edge and the center of the square, and the
    displacing it. The displacement is drawn according to `alg.H` and round using `displace`
"""
function _square!(mat, alg::MidpointDisplacement, round::Int, corners::AbstractVector{Tuple{Int, Int}})
    bottomLeft,bottomRight,topLeft,topRight = corners
    leftEdge, bottomEdge, topEdge, rightEdge = _edgeMidpointCoordinates(corners)
    mat[leftEdge...] = _interpolate(mat, [topLeft,bottomLeft]) + _displace(alg.H, round)
    mat[bottomEdge...] = _interpolate(mat, [bottomLeft,bottomRight]) + _displace(alg.H, round)
    mat[topEdge...] = _interpolate(mat, [topLeft,topRight]) + _displace(alg.H, round)
    mat[rightEdge...] = _interpolate(mat, [topRight,bottomRight]) + _displace(alg.H, round)
end

"""
    _interpolate(mat, points::AbstractVector{Tuple{Int,Int}})

    Computes the mean of a set of points, represented as a list of indecies to a matrix `mat`.
"""
function _interpolate(mat, points::AbstractVector{Tuple{Int,Int}})
    return mean(mat[pt...] for pt in points)
end

"""
    _displace(H::Float64, round::Int)

    `displace` produces a random value as a function of  `H`, which is the
    autocorrelation parameter used in `DiamondSquare` and must be between `0`
    and `1`, and `round` which describes the current tiling size for the
    DiamondSquare() algorithm.

    Random value are drawn from a Gaussian distribution using `Distribution.Normal`
    The standard deviation of this Gaussian, σ, is set to (1/2)^(round*H), which will
    move from 1.0 to 0 as `round` increases.

"""
function _displace(H::Float64, round::Int)
    σ = 0.5^(round*H)
    return(rand(Normal(0, σ)))
end

"""
    _centerCoordinate(corners::AbstractVector{Tuple{Int,Int}})

    Returns the center coordinate for a square defined by `corners` for the
    `DiamondSquare` algorithm.
"""
function _centerCoordinate(corners::AbstractVector{Tuple{Int,Int}})
    bottomLeft,bottomRight,topLeft,topRight = corners
    centerX::Int = (_xcoord(bottomLeft)+_xcoord(bottomRight)) ÷ 2
    centerY::Int = (_ycoord(topRight)+_ycoord(bottomRight)) ÷ 2
    return (centerX, centerY)
end

"""
    _xcoord(pt::Tuple{Int,Int})

    Returns the x-coordinate from a lattice coordinate `pt`.
"""
 _xcoord(pt::Tuple{Int,Int}) = pt[1]

 """
     _ycoord(pt::Tuple{Int,Int})

     Returns the y-coordinate from a lattice coordinate `pt`.
 """
 _ycoord(pt::Tuple{Int,Int}) = pt[2]

"""
    _edgeMidpointCoordinates(corners::AbstractVector{Tuple{Int,Int}})

    Returns an array of midpoints for a square defined by `corners` for the `DiamondSquare` algorithm.
"""
function _edgeMidpointCoordinates(corners::AbstractVector{Tuple{Int,Int}})
    # bottom left, bottom right, top left, top right
    bottomLeft,bottomRight,topLeft,topRight = corners

    leftEdgeMidpoint::Tuple{Int,Int} = (_xcoord(bottomLeft), (_ycoord(bottomLeft)+_ycoord(topLeft))÷2 )
    bottomEdgeMidpoint::Tuple{Int,Int} = ( (_xcoord(bottomLeft)+ _xcoord(bottomRight))÷2, _ycoord(bottomLeft) )
    topEdgeMidpoint::Tuple{Int,Int} = ( (_xcoord(topLeft)+_xcoord(topRight))÷2, _ycoord(topLeft))
    rightEdgeMidpoint::Tuple{Int,Int} = ( _xcoord(bottomRight), (_ycoord(bottomRight)+_ycoord(topRight))÷2)

    edgeMidpoints = [leftEdgeMidpoint, bottomEdgeMidpoint, topEdgeMidpoint, rightEdgeMidpoint]
    return edgeMidpoints
end

"""
     _isPowerOfTwo(x::IT) where {IT <: Integer}

     Determines if `x`, an integer, can be expressed as `2^n`, where `n` is also an integer.
"""
function _isPowerOfTwo(x::IT) where {IT <: Integer}
    return x & (x-1) == 0
end


println("# Test Distributions")

# Multinomial-Dirichlet
println(" * Dirichlet - Multinomial")

d = DirichletMultinomial(5, 1.0)

println(" * - add data")

add_data!(d, [2, 3, 3])

@test d.n == 3
@test d.counts[2] == 1
@test d.counts[3] == 2

println(" * - add sparse data")

add_data!(d, sparse([1,0,1,0,0]))

@test d.n == 4
@test d.counts[1] == 1
@test d.counts[3] == 3

println(" * - remove data")

remove_data!(d, [3])

@test d.n == 3
@test d.counts[3] == 2

println(" * - remove sparse data")

remove_data!(d, sparse([1,0,0,0,0]))

@test d.n == 2
@test d.counts[1] == 0

println(" * - log pdf")
@test logpred(d, [3])[1] == logpred(d, sparse([0,0,1,0,0]))[1]

# student-t Distribution
println(" * Generalized student-t Distribution")

df = 1
mean = 0
sigma = 1

d = Distributions.TDist(df)
@test_approx_eq_eps BNP.tlogpdf(0, df, mean, sigma) Distributions.logpdf(d, 0) 1e-4

df = 1
mean = 1
sigma = 1

d = Distributions.TDist(df)
@test_approx_eq_eps BNP.tlogpdf(0, df, mean, sigma) Distributions.logpdf(d, 0 - mean) 1e-4

df = 1
mean = 0
sigma = 10

d = Distributions.TDist(df)

# Normal-Gamma
println(" * Gamma-Normal")

d = GammaNormal()

(" * - add data")

x = 1.0

add_data!(d, x)
@test d.n == 1

add_data!(d, vec([1.0, 2.0, 1.5, 1.0]))
@test d.n == 5

println(" * - remove data")

remove_data!(d, 2.0)
@test d.n == 4

println(" * - log pdf")

x2 = [2.0]
@test logpred(d, x)[1] > logpred(d, x2)[1]

# Normal-Gamma
println(" * GaussianDiagonal")

d = GaussianDiagonal{NormalNormal}(NormalNormal[NormalNormal() for i in 1:2])

println(" * - add data")
add_data!(d, rand(2))
@test d.n == 1

println("# FINISHED")


# Part 1
layout = Array{Int8}(hcat(map(line -> 'L' .== collect(line), readlines("day11_input.txt"))...)')

function singlestep(current)
    updated = copy(current)
    region = CartesianIndices(current)
    ifirst, ilast = first(region), last(region)
    iunit = oneunit(ifirst)
    for i in region
        neighbours = [current[j] for j in max(ifirst, i - iunit):min(ilast, i + iunit) if j != i]
        if (current[i] == 1 && count(neighbours .== 2) == 0) updated[i] = 2 end
        if (current[i] == 2 && count(neighbours .== 2) >= 4) updated[i] = 1 end
    end
    updated
end

function steadystate(func, current)
    previous = zeros(size(current))
    while current != previous
        previous, current = current, func(current)
    end
    current
end

println("Part 1: occupied count = $(count(steadystate(singlestep, layout) .== 2))")

# Part 2
function singlestep2(current)
    updated = copy(current)
    region = CartesianIndices(current)
    iunit = oneunit(first(region))
    for i in region
        neighbours = []
        for direction in -iunit:iunit
            if !iszero(direction)
                j = i + direction
                while j in region
                    if current[j] != 0
                        push!(neighbours, current[j])
                        break
                    end
                    j += direction
                end
            end
        end
        if (current[i] == 1 && count(neighbours .== 2) == 0) updated[i] = 2 end
        if (current[i] == 2 && count(neighbours .== 2) >= 5) updated[i] = 1 end
    end
    updated
end

println("Part 2: occupied count = $(count(steadystate(singlestep2, layout) .== 2)) ")


"""
    PELT(segment_cost, n[, β = log(n)])

Runs the PELT algorithm using cost function `segment_cost` for a time series of length `n` and a penalty `β` and returns the position of changepoints and optimal segmentation cost.

See also: [`@PELT`](@ref), [`@segment_cost`](@ref), [`CROPS`](@ref)

# Returns
* `(CP::Vector{Int}, cost::Float64)`:
  * `CP::Vector{Int}`: Vector of indices of detected changepoints
  * `cost::Float64`: Cost of optimal segmentation

# Example
```julia-repl
n, λ = 1000, 100
μ, σ = Normal(0.0, 10.0), 1.0
sample, cps = @changepoint_sampler n λ Normal(μ, σ)  # Sample changepoints
seg_cost = NormalMeanChange(sample, σ)               # Create segment cost function
pelt_cps, pelt_cost = PELT(seg_cost, n)              # Run PELT
```

# References
Killick, R., Fearnhead, P. and Eckley, I.A. (2012) Optimal detection of changepoints with a linear computational cost, JASA 107(500), 1590-1598
"""
function PELT( segment_cost::Function , n::Int; pen::Float64 = log(n) )

    # F[t] is optimal cost of segmentation upto time t
    F = Array{Float64}(undef, n+1)
    F[1] = -pen
    F[2] = 0

    # last chpt prior to time t
    chpts =  Array{Int64}(undef, n)
    chpts[1] = 0

    # vector of candidate chpts at t
    R = Int64[0]

    for t in 2:n
        cpt_cands = R
        seg_costs = Array{Float64}(undef, length(cpt_cands))
        for i in 1:length(cpt_cands)
            seg_costs[i] = segment_cost(cpt_cands[i], t)
        end

        F[t+1] , tau = findmin( F[cpt_cands .+ 1] .+ seg_costs .+ pen )
        chpts[t] = cpt_cands[tau]

        # pruning step
        ineq_prune = (F[cpt_cands .+ 1] .+ seg_costs) .< F[t+1]
        R = push!(cpt_cands[ineq_prune], t - 1)
    end

    # get changepoints
    CP = Array{Int64}(undef,0)
    last = chpts[n]
    push!(CP,last)
    while last > 0
        last = chpts[last]
        push!(CP,last)
    end
    sort!(CP)

    return CP , F[n+1]

end


# Create example data (continent population in million people)
continents = ["Africa", "America", "Asia", "Europe", "Oceania"]
population_1960 = [285, 425, 1687, 606, 16]
population_2010 = [1044, 944, 4170, 735, 36]
population_matrix = [population_1960 population_2010]
# Plot with respect to 500 millions of people
barplot(continents, population_matrix, baseline=500)
legend("1960", "2010") # add legend


using Random

struct CyclicGroup <: GroupsCore.Group
    order::UInt
end

struct CyclicGroupElement <: GroupsCore.GroupElement
    residual::UInt
    parent::CyclicGroup
    CyclicGroupElement(n::Integer, C::CyclicGroup) = new(n % C.order, C)
end

Base.one(C::CyclicGroup) = CyclicGroupElement(0, C)

Base.eltype(::Type{CyclicGroup}) = CyclicGroupElement
Base.iterate(C::CyclicGroup) = one(C), 1
Base.iterate(C::CyclicGroup, state) =
    (state < C.order ? (CyclicGroupElement(state, C), state + 1) : nothing)
Base.IteratorSize(::Type{CyclicGroup}) = Base.HasLength()

GroupsCore.order(::Type{T}, C::CyclicGroup) where {T<:Integer} = T(C.order)
GroupsCore.gens(C::CyclicGroup) = [CyclicGroupElement(1, C)]

function Base.rand(
    rng::Random.AbstractRNG,
    rs::Random.SamplerTrivial{<:CyclicGroup},
)
    C = rs[]
    return CyclicGroupElement(rand(0:C.order-1), C)
end

GroupsCore.parent(c::CyclicGroupElement) = c.parent
GroupsCore.parent_type(::Type{CyclicGroupElement}) = CyclicGroup
Base.:(==)(g::CyclicGroupElement, h::CyclicGroupElement) =
    parent(g) === parent(h) && g.residual == h.residual

# since CyclicGroupElement is isbits, there is no need to define for julia>=1.3.0
if VERSION < v"1.3.0"
    Base.deepcopy_internal(g::CyclicGroupElement, ::IdDict) =
        CyclicGroupElement(deepcopy(g.residual), parent(g))
end

Base.inv(g::CyclicGroupElement) =
    (C = parent(g); CyclicGroupElement(order(UInt, C) - g.residual, C))

function Base.:(*)(g::CyclicGroupElement, h::CyclicGroupElement)
    @assert parent(g) === parent(h)
    C = parent(g)
    return CyclicGroupElement(g.residual + h.residual, C)
end

################## Implementing Group Interface Done!

#=
### Possible performance modifications:

NOTE: Since CyclicGroupElement is immutable there is no need to implement in-place mutable arithmetic.

Base.isone(g::CyclicGroupElement) = iszero(g.residual)

function GroupsCore.order(::Type{I}, g::CyclicGroupElement) where {I<:Integer}
    isone(g) && return one(I)
    o = order(I, parent(g))
    return div(o, gcd(I(g.residual), o))
end

Base.hash(g::CyclicGroupElement, h::UInt) = hash(g.residual, hash(parent(g), h))

=#

### end of Group[Element] methods

# Some eye-candy if you please
Base.show(io::IO, C::CyclicGroup) =
    print(io, "Group of residues modulo $(order(Int, C))")
Base.show(io::IO, c::CyclicGroupElement) =
    print(io, Int(c.residual), " (mod ", order(Int, parent(c)), ")")