export Boussinesq

"""
    Boussinesq(param;kwargs)

Define an object of type `AbstractModel` in view of solving the initial-value problem for
`abcd`-Boussinesq models (with `b=d` and `c=0`).
See [Bona, Chen, and Saut](https://doi.org/10.1007/s00332-002-0466-4)

# Argument
`param` is of type `NamedTuple` and must contain
- dimensionless parameters `ϵ` (nonlinearity) and `μ` (dispersion);
- numerical parameters to construct the mesh of collocation points, if `mesh` is not provided as a keyword argument.

## Optional keyword arguments
- two parameters `a` (default is `-1/3`) and `b` (default is `+1/3`) which determine the model solved. You need `a+2*b=1/3` for validity as a long wave model (without surface tension).
- `mesh`: the mesh of collocation points. By default, `mesh = Mesh(param)`;
- `ktol`: tolerance of the low-pass Krasny filter (default is `0`, i.e. no filtering);
- `dealias`: dealiasing with Orlicz rule `1-dealias/(dealias+2)` (default is `0`, i.e. no dealiasing);
- `label`: a label for future references (default is `"Boussinesq"`);

# Return values
Generate necessary ingredients for solving an initial-value problem via `solve!`:
1. a function `Boussinesq.f!` to be called in explicit time-integration solvers;
2. a function `Boussinesq.mapto` which from `(η,v)` of type `InitialData` provides the raw data matrix on which computations are to be executed;
3. a function `Boussinesq.mapfro` which from such data matrix returns the Tuple of real vectors `(η,v,x)`, where
    - `η` is the values of surface deformation at collocation points `x`;
    - `v` is the derivative of the trace of the velocity potential at `x`.

"""
mutable struct Boussinesq <: AbstractModel

	label 	:: String
	f!		:: Function
	mapto	:: Function
	mapfro	:: Function
	info	:: String

    function Boussinesq(param::NamedTuple;
						mesh = Mesh(param),
						a=-1//3,b=1//3,
						dealias=0,ktol=0,
						label="Boussinesq"
						)

		m=WhithamBoussinesq(param;Boussinesq=true,
							mesh=mesh,
							a=a,b=b,
							dealias=dealias,ktol=ktol,
							label=label)

		new(m.label, m.f!, m.mapto, m.mapfro, m.info)
    end
end


using Robotlib, SymPy
Tworld2base = I4
# """
# `τ = gravity{P}(q::VecOrMat{P}, rm, m, dh::DH, Tworld2base = I4)`
# calculates the gravity torques given a joint angles `q`, (distances to center of mass × mass) `rm` and masses `m`. The DH-parameters for the robot must also be provided
# """
function gravity(q::VecOrMat{P}, rm, m, dh::DH, Tworld2base = I4) where P
    Tn  = dh2Tn(dh,q)
    gravity(q, rm, m, Tn, Tworld2base)
end

"""
    gravity{P}(q, rm, m, Tn, Tworld2base)

`q` is the joint coordinates, `rm` are distances to center of masses times masses, `m` are the masses and `Tn` are the transformations between consecutive joints.
"""
function gravity(q::VecOrMat{P}, rm, m, Tn::AbstractArray, Tworld2base = I4) where P
    n   = size(Tn,3)-1
    # Find elements that are ≈ 0 and make them zero, smiplifies a lot!
    if P == SymPy.Sym
        Tnz = Bool[round(Float64(evalf(subs(Tn[i,j,k],(q[1],1),(q[2],1),(q[3],1),(q[4],1),(q[5],1),(q[6],1),(q[7],1)))), digits=10) == 0 for i=1:4, j=1:4, k=1:n+1]
        Tn[Tnz] = 0
        # Replace the numerical value for π/2 by the SymPi constant PI, this allows for nice simplifications
        Tn = Sym[subs(Tn[i,j,k],("1.5707963267948966",SymPy.PI/2))  for i=1:4, j=1:4, k=1:n+1]
    end

    Ti = zeros(typeof(q[1]),4,4,n+1)
    Tn = cat(3,Tworld2base,Tn)
    Ti[:,:,1] = Tn[:,:,1]

    for i = 2:n
        Ti[:,:,i] = Ti[:,:,i-1]*Tn[:,:,i] # Cumulative forward kinematics transformations
    end

    gv         = [0, 0, -9.82]
    τ          = Array(P,3,n+1) # Extend by one to run everything in loop
    τ[:,end]   = 0
    τhat       = Array(P,n)
    force      = Array(P,3,n) # All forces are given in base frame

    for i = n:-1:1
        Ri     = Tn[1:3,1:3,i+1] # Rotation from i to i+1
        gi     = Ti[1:3,1:3,i]'gv # Gravity vector in i
        τ[:,i] = skew(rm[:,i]) * gi + Ri*τ[:,i+1] # Torque in i is gravity torque + rotated torque from i+1
        τi     = τ[:,i]
        for k = (i+1):n # I can not get it to work with accumulated forces which would be much faster
            rik = (trinv(Ti[:,:,i])*Ti[1:4,4,k])[1:3] # Vector from i to k, given in i
            τi += skew(rik)*Ti[1:3,1:3,i]'force[:,k] # Torque around i, given in i
        end

        τhat[i]     = τi[3] # The motor torque is the local torque around z
        force[:,i]  = gv*m[i]
    end
    return -τhat # Minus sign to match sensor torque
end

function create_gravmatrix(dh::DH)

    n_joints = size(dh.dhpar,1)

    rm       = Sym[symbols("rm$i$j",real=true) for i = 1:3, j=1:n_joints]
    m        = Sym[symbols("m$j",real=true) for j=1:n_joints]
    q        = Sym[symbols("q$j",real=true) for j=1:n_joints]
    tau      = Sym[symbols("tau$j",real=true) for j=1:n_joints]

    tauhat = gravity(q, rm, m, dh, Tbase);
    w      = [rm[:];m[:]];

    AA     = Array(Sym,n_joints, 4n_joints);
    AAz    = falses(size(AA))

    eq     = tauhat;

    println("Finding non-zero coefficients")
    for i = 1:n_joints
        for j = 1:size(AA,2)
            AA[i,j] = (coeff(expand(eq[i]),w[j])) # This is a nice place to put simplify
            AAz[i,j] = abs(evalf(subs(AA[i,j],(q[1],1),(q[2],1),(q[3],1),(q[4],1),(q[5],1),(q[6],1),(q[7],1)))) < 1e-5
        end
        print(i)
    end
    println("")

    # non_identifiable    = (sum(AA,1) .== 0)[:]
    # # non_identifiable = all(AAz,1)[:]
    # w                   = w[!non_identifiable];
    # AA                  = AA[:,!non_identifiable];
    n_joints,n_params   = size(AA)

    println("Trying to find expressions common in all terms")
    change,res  = cse(AA[:]);
    res         = reshape(res,size(AA));

    println("Printing results to file")
    fid = open("gravityFridaLS.jl","w");
    println(fid,"@fastmath function gravityFridaLS(q)")
    for i = 1:n_joints
        println(fid,"q$i = q[$i]")
    end
    for i = 0:length(change)-1
        sei = change[i+1][2]
        println(fid,"x$i = $sei")
    end

    println(fid,"A = zeros($n_joints,$n_params)")

    for j = 1:n_joints
        for p = 1:n_params
            if res[j,p] != 0 #&& !AAz[j,p]
                println(fid,"A[$j,$p] = $(res[j,p])")
            end
        end
    end
    println(fid,"return A")
    println(fid,"end")
    close(fid)
    println("Done")
end
# dh       = DHYuMi()
# baseAnglesLeft = [-0.63 , 0.95 , -0.18]
# Rbase          = rpy2R(baseAnglesLeft,"xyz")
# Tbase          = I4
# Tbase[1:3,1:3] = Rbase

"""
    create_gravmodel(filename, r,m, dh, Tbase=I4)
Given known masses and vectors to center of masses, this function creates a new function to calculate the gravity torque.
If only `r*m` is known, call the function with `m=1` and `r=rm`
`r ∈ ℜ(3 × n_joints)`
`m ∈ ℜ(n_joints)`

The created function is written to a file with and is called like `τ = gravity(q)`
"""
function create_gravmodel(filename, r::AbstractMatrix,m, dh::DH, Tbase=I4)
    n_joints       = size(dh.dhpar,1)
    rm             = (r.*m)'
    q              = Sym[symbols("q$j",real=true) for j=1:n_joints]
    tauhat         = gravity(q, rm, m, dh, Tbase);
    n_joints       = size(r,1)

    println("Trying to find expressions common in all terms")
    change,res  = cse(tauhat);

    println("Printing results to file")
    fid = open(filename*".jl","w");
    println(fid,"@fastmath function gravity(q)")
    for i = 1:n_joints
        println(fid,"q$i = q[$i]")
    end
    for i = 0:length(change)-1
        sei = change[i+1][2]
        println(fid,"x$i = $sei")
    end

    println(fid,"tau = zeros($n_joints)")

    for j = 1:n_joints
            if res[j] != 0 #&& !AAz[j,p]
                println(fid,"tau[$j] = $(res[j])")
            end
    end
    println(fid,"return tau")
    println(fid,"end")
    close(fid)
    println("Done")
end


# This script demonstrates how to use MERAKit.jl to find an approximation to the ground
# state of the critical Ising model.
using Random
using LinearAlgebra
using TensorKit
using MERAKit
# DemoTools holds some utility functions needed by this script, for creating Hamiltonians
# and storing and reading MERAs to/from disk. It exports nothing, so all calls need to be
# qualified as DemoTools.whatever.
include("demo_tools.jl")
using .DemoTools

function main()
    # Fix the random seed for the sake of reproducibility. MERA algorithms sometimes get
    # stuck in local minima, so you should usually try a few different seeds for the same
    # optimisation. This one is going to do well though.
    Random.seed!(1)

    # Create the local Hamiltonian term `h` for the critical Ising model. Use Z2 symmetric
    # tensors, and block two sites together, so that the local state space dimension is 4.
    # The blocking simply obviates the need for a trivial, unitary layer at the bottom of
    # the MERA.
    h = DemoTools.ising_hamiltonian(; symmetry = :Z2, block_size = 2)
    exact_energy = -4/pi
    V_physical = space(h, 1)

    # Choose the type of MERA to use.
    meratype = TernaryMERA
    layers = 3
    # It's typically beneficial to gradually increase the bond dimension during the
    # optimisation. We do this here using the following steps: Each element of
    # `V_virtual_steps` is a vector space for the virtual indices of the MERA, and we run
    # the optimisation for each of them, using as the starting point the result of the
    # previous step.
    # Note that for Z2 symmetric tensors we have to specify how the bond dimension is
    # divided between the symmetry sectors. It is a priori not obvious what the optimal
    # choice is, and although here it happens to be such that both sectors get the same
    # dimension, this isn't always the case. In general, finding the optimal proportions
    # requires manually trying different options.
    # Note also that we could use different vector spaces for the different indices in the
    # MERA, both the ones between layers and the ones within layers (between isometries and
    # disentanglers, making the disentanglers isometric instead of unitary). This would
    # probably be better for accuracy/speed, but using the same `V_virtual` everywhere works
    # well too, so for simplicity we stick to that.
    V_virtual_steps = (
        Z2Space(0 => 2, 1 => 2),
        Z2Space(0 => 3, 1 => 3),
        Z2Space(0 => 4, 1 => 4),
    )

    # Specify the optimisation method. We use here Riemannian LBFGS as described in
    # https://arxiv.org/abs/2007.03638
    method = :lbfgs
    verbosity = 2
    gradient_delta = 1e-7
    # In practice the above threshold for the gradient norm, which determines convergence,
    # will not be reached. Instead, we set the maximum number of iterations to be done. We
    # specify this separately for the different bond dimensions in `V_virtual_steps`.
    # Choosing these numbers is more art than science, but many different choices would work
    # roughly equally well.
    maxiter_steps = (1000, 500, 500)

    # Once the optimisation is done, the resulting MERA will be written to a file at this
    # path.
    path = "./demo_result.jld2"

    # To keep track of how the optimisation is progressing, we will every once in a while
    # print out the energy error, the entropies of the density matrices at different scales
    # in the MERA, and the scaling dimensions. For this we use the `finalize!` keyword
    # argument of `minimize_expectation`, which can be a function that is run at the end of
    # every iteration. `finalize!` could be used to do plenty of things, including modifying
    # the current MERA `m` during the optimisation, but we simply use it to print things out
    # every `checkpoint_frequency` iterations. See the documentation of MERAKit.jl for more
    # details on `finalize!`.
    checkpoint_frequency = 100
    function finalize!(m, energy, g, repnum)
        repnum % checkpoint_frequency != 0 && (return m, energy, g)
        rhoees = densitymatrix_entropies(m)
        scaldims = scalingdimensions(m)
        @info("Energy error: $(energy - exact_energy)")
        @info("Density matrix entropies: $rhoees")
        @info("Scaling dimensions: $scaldims")
        return m, energy, g
    end

    # The main loop, that for every vector space in `V_virtual_steps` optimises the MERA.
    # For the first one, we use a randomly initialised MERA as the starting point, for the
    # latter ones we expand the bond dimension by padding tensors with zeros.
    local m
    for i in 1:length(V_virtual_steps)
        V_virtual = V_virtual_steps[i]
        if i == 1
            @info("Creating a random MERA with bond dimension $(dim(V_virtual)).")
            # The lowest index is of the physical bond dimension, the others are virtual.
            Vs = (V_physical, ntuple(x -> V_virtual, layers-1)...)
            m = random_MERA(meratype, Float64, Vs)
        else
            @info("Expanding to bond dimension $(dim(V_virtual)).")
            for j in 2:layers
                m = expand_bonddim(m, j-1, V_virtual)
                m = expand_internal_bonddim(m, j, V_virtual)
            end
        end
        # Many more optimisation parameters can be specified, see the documentation on
        # `minimize_expectation`.
        pars = (
            gradient_delta = gradient_delta,
            maxiter = maxiter_steps[i],
            method = method,
            verbosity = verbosity
        )
        m = minimize_expectation(m, h, pars; finalize! = finalize!)
    end

    # Write the resulting MERA to disc. It can then be read using
    # `DemoTools.load_mera(path)`.
    DemoTools.store_mera(path, m)
    return m
end

# Run the actual script. The whole thing is in a function mainly to avoid polluting the
# global namespace.
main()


### A Pluto.jl notebook ###
# v0.12.11

using Markdown
using InteractiveUtils

# ╔═╡ 2ebffa2a-feac-11ea-27d0-3b7cf030672a
using Pkg, DrWatson

# ╔═╡ 2ec03170-feac-11ea-3cd8-e9c1d16cd34c
begin
  #@quickactivate "StatisticalRethinkingStan"
  using StanSample
  using StatisticalRethinking
end

# ╔═╡ 37e0328a-feab-11ea-031f-df82e9c87154
md"## Clip-06-13-15s.jl"

# ╔═╡ 2ec0c25c-feac-11ea-2439-07e617727eb5
begin
  N = 100
  df = DataFrame(
    :h0 => rand(Normal(10,2 ), N),
    :treatment => vcat(zeros(Int, Int(N/2)), ones(Int, Int(N/2)))
  );
  df.fungus = [rand(Binomial(1, 0.5 - 0.4 * df[i, :treatment]), 1)[1] for i in 1:N]
  df.h1 = [df[i, :h0] + rand(Normal(5 - 3 * df[i, :fungus]), 1)[1] for i in 1:N]
  scale!(df, [:h0, :treatment, :fungus, :h1])
end;

# ╔═╡ 2ed0d892-feac-11ea-1ad3-5dfeacd40f38
Text(precis(df, io=String))

# ╔═╡ 2ed17b7e-feac-11ea-2d49-3ba0f6783ba4
sim_p = DataFrame(:sim_p => rand(LogNormal(0, 0.25), 10000));

# ╔═╡ 2edbc4b2-feac-11ea-0977-37cb5da52497
Text(precis(sim_p; io=String))

# ╔═╡ 2edc81fe-feac-11ea-3f3c-2ff7b2c13458
stan6_6 = "
data {
  int <lower=1> N;
  vector[N] h0;
  vector[N] h1;
}
parameters{
  real<lower=0> p;
  real<lower=0> sigma;
}
model {
  vector[N] mu;
  p ~ lognormal(0, 0.25);
  sigma ~ exponential(1);
  mu = h0 * p;
  h1 ~ normal(mu, sigma);
}
";

# ╔═╡ 2ee7bd6a-feac-11ea-1120-776954c32bd8
begin
	m6_6s = SampleModel("m6.6s", stan6_6)
	m6_6_data = Dict(:N => nrow(df), :h0 => df.h0, :h1 => df.h1)
	rc6_6s = stan_sample(m6_6s; data=m6_6_data)
	success(rc6_6s) && (post6_6s_df = read_samples(m6_6s, :dataframe))
end;

# ╔═╡ 2ef5aa9c-feac-11ea-1da4-67356630f549
if success(rc6_6s)
	part6_6s = Particles(post6_6s_df)
end

# ╔═╡ 2f01856c-feac-11ea-0b6b-537d35689cfc
success(rc6_6s) && (Text(precis(post6_6s_df; io=String)))

# ╔═╡ 2f093c3a-feac-11ea-2325-37124b7e6bdf
md"## End of clip-06-13-15s.jl"

# ╔═╡ Cell order:
# ╠═37e0328a-feab-11ea-031f-df82e9c87154
# ╠═2ebffa2a-feac-11ea-27d0-3b7cf030672a
# ╠═2ec03170-feac-11ea-3cd8-e9c1d16cd34c
# ╠═2ec0c25c-feac-11ea-2439-07e617727eb5
# ╠═2ed0d892-feac-11ea-1ad3-5dfeacd40f38
# ╠═2ed17b7e-feac-11ea-2d49-3ba0f6783ba4
# ╠═2edbc4b2-feac-11ea-0977-37cb5da52497
# ╠═2edc81fe-feac-11ea-3f3c-2ff7b2c13458
# ╠═2ee7bd6a-feac-11ea-1120-776954c32bd8
# ╠═2ef5aa9c-feac-11ea-1da4-67356630f549
# ╠═2f01856c-feac-11ea-0b6b-537d35689cfc
# ╟─2f093c3a-feac-11ea-2325-37124b7e6bdf


module EnsembleAnalysis

using DiffEqBase, StaticArrays, Statistics, RecursiveArrayTools

# Getters
get_timestep(sim,i) = (getindex(sol,i) for sol in sim)
get_timepoint(sim,t) = (sol(t) for sol in sim)
function componentwise_vectors_timestep(sim,i)
  arr = [get_timestep(sim,i)...]
  if typeof(arr[1]) <: AbstractArray
   return vecarr_to_vectors(VectorOfArray(arr))
  else
   return arr
  end
 end
 function componentwise_vectors_timepoint(sim,t)
   arr = [get_timepoint(sim,t)...]
   if typeof(arr[1]) <: AbstractArray
    return vecarr_to_vectors(VectorOfArray(arr))
   else
    return arr
   end
  end

# Timestep statistics
timestep_mean(sim,i) = componentwise_mean(get_timestep(sim,i))
timestep_mean(sim,::Colon) = timeseries_steps_mean(sim)
function timestep_median(sim,i)
  arr = componentwise_vectors_timestep(sim,i)
  if typeof(first(arr)) <: AbstractArray
    return reshape([median(x) for x in arr],size(sim[1][i])...)
  else
    return median(arr)
  end
end
timestep_median(sim,::Colon) = timeseries_steps_median(sim)
function timestep_quantile(sim,q,i)
  arr = componentwise_vectors_timestep(sim,i)
  if typeof(first(arr)) <: AbstractArray
    return reshape([quantile(x,q) for x in arr],size(sim[1][i])...)
  else
    return quantile(arr,q)
  end
end
timestep_quantile(sim,q,::Colon) = timeseries_steps_quantile(sim,q)
timestep_meanvar(sim,i) = componentwise_meanvar(get_timestep(sim,i))
timestep_meanvar(sim,::Colon) = timeseries_steps_meanvar(sim)
timestep_meancov(sim,i,j) = componentwise_meancov(get_timestep(sim,i),get_timestep(sim,j))
timestep_meancov(sim,::Colon,::Colon) = timeseries_steps_meancov(sim)
timestep_meancor(sim,i,j) = componentwise_meancor(get_timestep(sim,i),get_timestep(sim,j))
timestep_meancor(sim,::Colon,::Colon) = timeseries_steps_meancor(sim)
timestep_weighted_meancov(sim,W,i,j) = componentwise_weighted_meancov(get_timestep(sim,i),get_timestep(sim,j),W)
timestep_weighted_meancov(sim,W,::Colon,::Colon) = timeseries_steps_weighted_meancov(sim,W)

function timeseries_steps_mean(sim)
  DiffEqArray([timestep_mean(sim,i) for i in 1:length(sim[1])],sim[1].t)
end
function timeseries_steps_median(sim)
  DiffEqArray([timestep_median(sim,i) for i in 1:length(sim[1])],sim[1].t)
end
function timeseries_steps_quantile(sim,q)
  DiffEqArray([timestep_quantile(sim,q,i) for i in 1:length(sim[1])],sim[1].t)
end
function timeseries_steps_meanvar(sim)
  m,v = timestep_meanvar(sim,1)
  means = [m]
  vars = [v]
  for i in 2:length(sim[1])
    m,v = timestep_meanvar(sim,i)
    push!(means,m)
    push!(vars,v)
  end
  DiffEqArray(means,sim[1].t),DiffEqArray(vars,sim[1].t)
end
function timeseries_steps_meancov(sim)
  reshape([timestep_meancov(sim,i,j) for i in 1:length(sim[1]) for j in 1:length(sim[1])],length(sim[1]),length(sim[1]))
end
function timeseries_steps_meancor(sim)
  reshape([timestep_meancor(sim,i,j) for i in 1:length(sim[1]) for j in 1:length(sim[1])],length(sim[1]),length(sim[1]))
end
function timeseries_steps_weighted_meancov(sim,W)
  reshape([timestep_meancov(sim,W,i,j) for i in 1:length(sim[1]) for j in 1:length(sim[1])],length(sim[1]),length(sim[1]))
end

timepoint_mean(sim,t) = componentwise_mean(get_timepoint(sim,t))
function timepoint_median(sim,t)
  arr = componentwise_vectors_timepoint(sim,t)
  if typeof(first(arr)) <: AbstractArray
    return reshape([median(x) for x in arr],size(sim[1][1])...)
  else
    return median(arr)
  end
end
function timepoint_quantile(sim,q,t)
  arr = componentwise_vectors_timepoint(sim,t)
  if typeof(first(arr)) <: AbstractArray
    return reshape([quantile(x,q) for x in arr],size(sim[1][1])...)
  else
    return quantile(arr,q)
  end
end
timepoint_meanvar(sim,t) = componentwise_meanvar(get_timepoint(sim,t))
timepoint_meancov(sim,t1,t2) = componentwise_meancov(get_timepoint(sim,t1),get_timepoint(sim,t2))
timepoint_meancor(sim,t1,t2) = componentwise_meancor(get_timepoint(sim,t1),get_timepoint(sim,t2))
timepoint_weighted_meancov(sim,W,t1,t2) = componentwise_weighted_meancov(get_timepoint(sim,t1),get_timepoint(sim,t2),W)

function DiffEqBase.EnsembleSummary(sim::DiffEqBase.AbstractEnsembleSolution{T,N},
                         t=sim[1].t;quantiles=[0.05,0.95]) where {T,N}
  if typeof(sim[1]) <: DiffEqBase.DESolution
    m,v = timeseries_point_meanvar(sim,t)
    qlow = timeseries_point_quantile(sim,quantiles[1],t)
    qhigh = timeseries_point_quantile(sim,quantiles[2],t)
  else
    m,v = timeseries_steps_meanvar(sim)
    qlow = timeseries_steps_quantile(sim,quantiles[1])
    qhigh = timeseries_steps_quantile(sim,quantiles[2])
  end

  trajectories = length(sim)
  DiffEqBase.EnsembleSummary{T,N,typeof(t),typeof(m),typeof(v),typeof(qlow),typeof(qhigh)}(t,m,v,qlow,qhigh,trajectories,sim.elapsedTime,sim.converged)
end

function timeseries_point_mean(sim,ts)
  DiffEqArray([timepoint_mean(sim,t) for t in ts],ts)
end
function timeseries_point_median(sim,ts)
  DiffEqArray([timepoint_median(sim,t) for t in ts],ts)
end
function timeseries_point_quantile(sim,q,ts)
  DiffEqArray([timepoint_quantile(sim,q,t) for t in ts],ts)
end
function timeseries_point_meanvar(sim,ts)
  m,v = timepoint_meanvar(sim,first(ts))
  means = [m]
  vars = [v]
  for t in Iterators.drop(ts,1)
    m,v = timepoint_meanvar(sim,t)
    push!(means,m)
    push!(vars,v)
  end
  DiffEqArray(means,ts),DiffEqArray(vars,ts)
end
function timeseries_point_meancov(sim,ts1,ts2)
  reshape([timepoint_meancov(sim,t1,t2) for t1 in ts1 for t2 in ts2],length(ts1),length(ts2))
end
function timeseries_point_meancor(sim,ts1,ts2)
  reshape([timepoint_meancor(sim,t1,t2) for t1 in ts1 for t2 in ts2],length(ts1),length(ts2))
end
function timeseries_point_weighted_meancov(sim,W,ts1,ts2)
  reshape([timepoint_meancov(sim,W,t1,t2) for t1 in ts1 for t2 in ts2],length(ts1),length(ts2))
end

function componentwise_mean(A)
  x0 = first(A)
  n = 0
  mean = zero(x0)
  for x in A
    n += 1
    if typeof(x0) <: AbstractArray && !(typeof(x0) <: SArray)
      mean .+= x
    else
      mean += x
    end
  end
  if typeof(x0) <: AbstractArray && !(typeof(x0) <: SArray)
    mean ./= n
  else
    mean /= n
  end
  mean
end

# Welford algorithm
# https://en.wikipedia.org/wiki/Algorithms_for_calculating_variance
function componentwise_meanvar(A;bessel=true)
  x0 = first(A)
  n = 0
  mean = zero(x0)
  M2 = zero(x0)
  delta = zero(x0)
  delta2 = zero(x0)
  for x in A
    n += 1
    if typeof(x0) <: AbstractArray && !(typeof(x0) <: SArray)
      delta .= x .- mean
      mean .+= delta./n
      delta2 .= x .- mean
      M2 .+= delta.*delta2
    else
      delta = x .- mean
      mean += delta./n
      delta2 = x .- mean
      M2 += delta.*delta2
    end
  end
  if n < 2
    return NaN
  else
    if bessel
      if typeof(x0) <: AbstractArray && !(typeof(x0) <: SArray)
        M2 .= M2 ./ (n .- 1)
      else
        M2 = M2 ./ (n .- 1)
      end
    else
      if typeof(x0) <: AbstractArray && !(typeof(x0) <: SArray)
        M2 .= M2 ./ n
      else
        M2 = M2 ./ n
      end
    end
    return mean,M2
  end
end

function componentwise_meancov(A,B;bessel=true)
  x0 = first(A)
  y0 = first(B)
  n = 0
  meanx = zero(x0)
  meany = zero(y0)
  C = zero(x0)
  dx = zero(x0)
  for (x,y) in zip(A,B)
    n += 1
    if typeof(x0) <: AbstractArray && !(typeof(x0) <: SArray)
      dx .= x .- meanx
      meanx .+= dx./n
      meany .+= (y.-meany)./n
      C .+= dx .* (y .- meany)
    else
      dx = x .- meanx
      meanx += dx./n
      meany += (y.-meany)./n
      C += dx .* (y .- meany)
    end
  end
  if n < 2
    return NaN
  else
    if bessel
      if typeof(x0) <: AbstractArray && !(typeof(x0) <: SArray)
        C .= C ./ (n .- 1)
      else
        C = C ./ (n .- 1)
      end
    else
      if typeof(x0) <: AbstractArray && !(typeof(x0) <: SArray)
        C .= C ./ n
      else
        C = C ./ n
      end
    end
    return meanx,meany,C
  end
end

function componentwise_meancor(A,B;bessel=true)
  mx,my,cov = componentwise_meancov(A,B;bessel=bessel)
  mx,vx = componentwise_meanvar(A;bessel=bessel)
  my,vy = componentwise_meanvar(B;bessel=bessel)
  if typeof(vx) <: AbstractArray
    vx .= sqrt.(vx)
    vy .= sqrt.(vy)
  else
    vx = sqrt.(vx)
    vy = sqrt.(vy)
  end
  mx,my,cov./(vx.*vy)
end

function componentwise_weighted_meancov(A,B,W;weight_type=:reliability)
  x0 = first(A)
  y0 = first(B)
  w0 = first(W)
  n = 0
  meanx = zero(x0)
  meany = zero(y0)
  wsum = zero(w0)
  wsum2 = zero(w0)
  C = zero(x0)
  dx = zero(x0)
  for (x,y,w) in zip(A,B,W)
    n += 1
    if typeof(x0) <: AbstractArray && !(typeof(x0) <: SArray)
      wsum .+= w
      wsum2 .+= w.*w
      dx .= x .- meanx
      meanx .+= (w ./ wsum) .* dx
      meany .+= (w ./ wsum) .* (y .- meany)
      C .+= w .* dx .* (y .- meany)
    else
      wsum += w
      wsum2 += w.*w
      dx = x .- meanx
      meanx += (w ./ wsum) .* dx
      meany += (w ./ wsum) .* (y .- meany)
      C += w .* dx .* (y .- meany)
    end
  end
  if n < 2
    return NaN
  else
    if weight_type == :population
      if typeof(x0) <: AbstractArray && !(typeof(x0) <: SArray)
        C .= C ./ wsum
      else
        C = C ./ wsum
      end
    elseif weight_type == :reliability
      if typeof(x0) <: AbstractArray && !(typeof(x0) <: SArray)
        C .= C ./ (wsum .- wsum2 ./ wsum)
      else
        C = C ./ (wsum .- wsum2 ./ wsum)
      end
    elseif weight_type == :frequency
      if typeof(x0) <: AbstractArray && !(typeof(x0) <: SArray)
        C .= C ./ (wsum .- 1)
      else
        C = C ./ (wsum .- 1)
      end
    else
      error("The weight_type which was chosen is not allowed.")
    end
    return meanx,meany,C
  end
end

export get_timestep, get_timepoint, apply_timestep, apply_timepoint,
       componentwise_vectors_timestep, componentwise_vectors_timepoint

export componentwise_mean, componentwise_meanvar

export timestep_mean, timestep_median, timestep_quantile, timestep_meanvar,
       timestep_meancov, timestep_meancor, timestep_weighted_meancov

export timeseries_steps_mean, timeseries_steps_median, timeseries_steps_quantile,
       timeseries_steps_meanvar, timeseries_steps_meancov,
       timeseries_steps_meancor, timeseries_steps_weighted_meancov

export timepoint_mean, timepoint_median, timepoint_quantile,
       timepoint_meanvar, timepoint_meancov,
       timepoint_meancor, timepoint_weighted_meancov

export timeseries_point_mean, timeseries_point_median, timeseries_point_quantile,
       timeseries_point_meanvar, timeseries_point_meancov,
       timeseries_point_meancor, timeseries_point_weighted_meancov

end


"""
    NormInfinityBridge{T}

The `NormInfinityCone` is representable with LP constraints, since
``t \\ge \\max_i \\lvert x_i \\rvert`` if and only if
``t \\ge x_i`` and ``t \\ge -x_i`` for all ``i``.
"""
struct NormInfinityBridge{T,F,G} <:
       SetMapBridge{T,MOI.Nonnegatives,MOI.NormInfinityCone,F,G}
    constraint::MOI.ConstraintIndex{F,MOI.Nonnegatives}
end

function concrete_bridge_type(
    ::Type{<:NormInfinityBridge{T}},
    G::Type{<:MOI.AbstractVectorFunction},
    ::Type{MOI.NormInfinityCone},
) where {T}
    F = MOI.Utilities.promote_operation(+, T, G, G)
    return NormInfinityBridge{T,F,G}
end

function MOI.Bridges.map_set(
    ::Type{<:NormInfinityBridge},
    set::MOI.NormInfinityCone,
)
    return MOI.Nonnegatives(2 * (MOI.dimension(set) - 1))
end

function MOI.Bridges.inverse_map_set(
    ::Type{<:NormInfinityBridge},
    set::MOI.Nonnegatives,
)
    return MOI.NormInfinityCone(div(MOI.dimension(set), 2) + 1)
end

function MOI.Bridges.map_function(
    ::Type{<:NormInfinityBridge{T}},
    func,
) where {T}
    scalars = MOI.Utilities.eachscalar(func)
    t = scalars[1]
    lb = scalars[2:end]
    ub = MOI.Utilities.operate(-, T, lb)
    f_new = MOI.Utilities.operate(vcat, T, ub, lb)
    for i in 1:(2*(length(scalars)-1))
        MOI.Utilities.operate_output_index!(+, T, i, f_new, t)
    end
    return f_new
end

function MOI.Bridges.inverse_map_function(
    ::Type{<:NormInfinityBridge{T}},
    func,
) where {T}
    scalars = MOI.Utilities.eachscalar(func)
    t = MOI.Utilities.operate!(/, T, sum(scalars), T(length(scalars)))
    d = div(length(scalars), 2)
    x = MOI.Utilities.operate!(
        /,
        T,
        MOI.Utilities.operate!(-, T, scalars[(d+1):end], scalars[1:d]),
        T(2),
    )
    return MOI.Utilities.operate(vcat, T, t, x)
end

# Given a_i is dual on t - x_i >= 0 and b_i is dual on t + x_i >= 0,
# the dual on (t, x) in NormInfinityCone is (u, v) in NormOneCone, where
# v_i = -a_i + b_i and u = sum(a) + sum(b).
function MOI.Bridges.adjoint_map_function(::Type{<:NormInfinityBridge}, func)
    scalars = MOI.Utilities.eachscalar(func)
    t = sum(scalars)
    d = div(length(scalars), 2)
    x = (scalars[(d+1):end] - scalars[1:d])
    return vcat(t, x)
end

function MOI.Bridges.inverse_adjoint_map_function(
    ::Type{<:NormInfinityBridge{T}},
    func::AbstractVector{T},
) where {T}
    # This is used by `MOI.ConstraintDualStart`.
    # The result should belong to `MOI.Nonnegatives` and the sum of the elements
    # should be `t`.
    # Only one of dual of the upper and lower bound should be active so one of
    # the two duals is zero. We know which one since they should be nonnegative.
    # Then if `t = sum abs(x_i)`, we will indeed have only one of them being
    # zero.
    t = func[1]
    y = func[2:end]
    lb = [x > 0 ? x : zero(x) for x in y]
    ub = [x < 0 ? -x : zero(x) for x in y]
    x = [ub; lb]
    return x .+ (t - sum(x)) / length(x)
end

"""
    NormOneBridge{T}

The `NormOneCone` is representable with LP constraints, since
``t \\ge \\sum_i \\lvert x_i \\rvert`` if and only if there exists a vector y
such that
``t \\ge \\sum_i y_i`` and ``y_i \\ge x_i``, ``y_i \\ge -x_i`` for all ``i``.
"""
struct NormOneBridge{T,F,G} <: AbstractBridge
    y::Vector{MOI.VariableIndex}
    nn_index::MOI.ConstraintIndex{F,MOI.Nonnegatives}
end

function bridge_constraint(
    ::Type{NormOneBridge{T,F,G}},
    model::MOI.ModelLike,
    f::G,
    s::MOI.NormOneCone,
) where {T,F,G}
    f_scalars = MOI.Utilities.eachscalar(f)
    d = MOI.dimension(s)
    y = MOI.add_variables(model, d - 1)
    rhs = MOI.ScalarAffineFunction(MOI.ScalarAffineTerm.(one(T), y), zero(T))
    ge = MOI.Utilities.operate(-, T, f_scalars[1], rhs)
    lb = f_scalars[2:d]
    ub = MOI.Utilities.operate(-, T, lb)
    lb = MOI.Utilities.operate!(+, T, lb, MOI.VectorOfVariables(y))
    ub = MOI.Utilities.operate!(+, T, ub, MOI.VectorOfVariables(y))
    f_new = MOI.Utilities.operate(vcat, T, ge, ub, lb)
    nn_index = MOI.add_constraint(model, f_new, MOI.Nonnegatives(2d - 1))
    return NormOneBridge{T,F,G}(y, nn_index)
end

function MOI.supports_constraint(
    ::Type{NormOneBridge{T}},
    ::Type{<:MOI.AbstractVectorFunction},
    ::Type{MOI.NormOneCone},
) where {T}
    return true
end

function MOI.Bridges.added_constrained_variable_types(::Type{<:NormOneBridge})
    return Tuple{Type}[]
end

function MOI.Bridges.added_constraint_types(
    ::Type{<:NormOneBridge{T,F}},
) where {T,F}
    return Tuple{Type,Type}[(F, MOI.Nonnegatives)]
end

function concrete_bridge_type(
    ::Type{<:NormOneBridge{T}},
    G::Type{<:MOI.AbstractVectorFunction},
    ::Type{MOI.NormOneCone},
) where {T}
    S = MOI.Utilities.scalar_type(G)
    F = MOI.Utilities.promote_operation(
        vcat,
        T,
        MOI.Utilities.promote_operation(+, T, S, S),
        MOI.Utilities.promote_operation(-, T, S, S),
    )
    return NormOneBridge{T,F,G}
end

MOI.get(b::NormOneBridge, ::MOI.NumberOfVariables)::Int64 = length(b.y)

MOI.get(b::NormOneBridge, ::MOI.ListOfVariableIndices) = copy(b.y)

function MOI.get(
    ::NormOneBridge{T,F},
    ::MOI.NumberOfConstraints{F,MOI.Nonnegatives},
)::Int64 where {T,F}
    return 1
end

function MOI.get(
    b::NormOneBridge{T,F},
    ::MOI.ListOfConstraintIndices{F,MOI.Nonnegatives},
) where {T,F}
    return [b.nn_index]
end

function MOI.delete(model::MOI.ModelLike, c::NormOneBridge)
    MOI.delete(model, c.nn_index)
    MOI.delete(model, c.y)
    return
end

function MOI.get(
    model::MOI.ModelLike,
    ::MOI.ConstraintFunction,
    c::NormOneBridge{T,F,G},
) where {T,F,G}
    nn_func = MOI.Utilities.eachscalar(
        MOI.get(model, MOI.ConstraintFunction(), c.nn_index),
    )
    t = MOI.Utilities.operate!(/, T, nn_func[1] + sum(nn_func), T(2))
    d = div(length(nn_func) - 1, 2)
    x = MOI.Utilities.operate!(
        /,
        T,
        MOI.Utilities.operate!(-, T, nn_func[(d+2):end], nn_func[2:(d+1)]),
        T(2),
    )
    return MOI.Utilities.convert_approx(
        G,
        MOI.Utilities.remove_variable(
            MOI.Utilities.operate(vcat, T, t, x),
            c.y,
        ),
    )
end

function MOI.get(model::MOI.ModelLike, ::MOI.ConstraintSet, c::NormOneBridge)
    dim = div(
        MOI.dimension(MOI.get(model, MOI.ConstraintSet(), c.nn_index)) + 1,
        2,
    )
    return MOI.NormOneCone(dim)
end

function MOI.supports(
    model::MOI.ModelLike,
    attr::Union{MOI.ConstraintPrimalStart,MOI.ConstraintDualStart},
    ::Type{NormOneBridge{T,F,G}},
) where {T,F,G}
    return MOI.supports(model, attr, MOI.ConstraintIndex{F,MOI.Nonnegatives})
end

function MOI.set(
    model::MOI.ModelLike,
    attr::MOI.ConstraintPrimalStart,
    bridge::NormOneBridge{T},
    value,
) where {T}
    x_value = value[1 .+ (1:length(bridge.y))]
    y_value = abs.(x_value)
    for i in eachindex(bridge.y)
        MOI.set(model, MOI.VariablePrimalStart(), bridge.y[i], y_value[i])
    end
    nn_value = vcat(
        value[1] - reduce(+, y_value, init = zero(T)),
        y_value - x_value,
        y_value + x_value,
    )
    MOI.set(model, attr, bridge.nn_index, nn_value)
    return
end

function MOI.get(
    model::MOI.ModelLike,
    attr::Union{MOI.ConstraintPrimal,MOI.ConstraintPrimalStart},
    bridge::NormOneBridge,
)
    nn_primal = MOI.get(model, attr, bridge.nn_index)
    t = (nn_primal[1] + sum(nn_primal)) / 2
    d = length(bridge.y)
    x = (nn_primal[(d+2):end] - nn_primal[2:(d+1)]) / 2
    return vcat(t, x)
end

# Given a_i is dual on y_i - x_i >= 0 and b_i is dual on y_i + x_i >= 0 and c is
# dual on t - sum(y) >= 0, the dual on (t, x) in NormOneCone is
# (u, v) in NormInfinityCone, where v_i = -a_i + b_i and u = c.
function MOI.get(
    model::MOI.ModelLike,
    attr::Union{MOI.ConstraintDual,MOI.ConstraintDualStart},
    bridge::NormOneBridge,
)
    nn_dual = MOI.get(model, attr, bridge.nn_index)
    d = length(bridge.y)
    x = nn_dual[(d+2):end] - nn_dual[2:(d+1)]
    return vcat(nn_dual[1], x)
end

# value[1 + i] = nn_dual[1 + d + i] - nn_dual[1 + i]
# and `nn_dual` is nonnegative. By complementarity slackness, only one of each
# `nn_dual` can be nonzero (except if `x = 0`) so we can set
# depending on the sense of `value[1 + i]`.
function MOI.set(
    model::MOI.ModelLike,
    ::MOI.ConstraintDualStart,
    bridge::NormOneBridge,
    value,
)
    d = length(bridge.y)
    nn_dual = zeros(eltype(value), 2d + 1)
    nn_dual[1] = value[1]
    for i in eachindex(bridge.y)
        v = value[1+i]
        if v < 0
            nn_dual[1+i] = -v
        else
            nn_dual[1+d+i] = v
        end
    end
    MOI.set(model, MOI.ConstraintDualStart(), bridge.nn_index, nn_dual)
    return
end


struct CDCapacitanceReport{U1,U2,U3,U4,U5,U6,U7,U8,U9,U10}
    Δt::U1
    ΔV::U2
    area::U3
    I::U4
    C::U5
    C_specific::U6
    E::U7
    E_specific::U8
    P::U9
    P_specific::U10
end

function capacitance(df, I, fixed_ΔV)
    t, V = df[!, Symbol("Time (s)")], df[!, Symbol("Potential (V)")]
    idxs = findall(v->v > 0, V)
    t, V = t[idxs], V[idxs]

    ∫Vdt = integrate(t, V) * u"V*s"
    Δt = (t[end] - t[begin]) * u"s"
    ΔV = !isnothing(fixed_ΔV) ? fixed_ΔV : -(reverse(extrema(V))...)

    return Δt, ΔV, ∫Vdt, 2*I*∫Vdt/ΔV^2
end

function CDCapacitanceReport(datafile, df, folder, setup)
    I = parse(Float64, filevalue(datafile)) * uparse(datafile.legend_units)
    porosity = parse(Int, foldervalue(datafile))
    @unpack a, A, fixed_ΔV = setup

    Δt, ΔV, area, C = capacitance(df, I, fixed_ΔV)
    C_specific = specific_capacitance(C, porosity, folder, a, A)
    E = energy(C, ΔV)
    E_specific = energy(C_specific, ΔV)
    P = power(E, Δt)
    P_specific = power(E_specific, Δt)

    CDCapacitanceReport(Δt, ΔV, area, I, C, C_specific, E, E_specific, P, P_specific)
end

function add_aux!(df, cr::CDCapacitanceReport)
    push!(df[!, :Area], ustrip(u"V*s", cr.area))
    push!(df[!, :Δt], ustrip(u"s", cr.Δt))
    push!(df[!, :ΔV], ustrip(u"V", cr.ΔV))
    push!(df[!, :I], ustrip(u"mA", cr.I))

    return nothing
end

function cd_report_units()
    join(["V s", "s", "V", "mA", replace_unicode("μF"), "F/g", "W*h", "Wh/kg", "W", "W/kg"], ',')
end

function cd_result()
    DataFrame(
        :Area=>Float64[],
        :Δt=>Float64[],
        :ΔV=>Float64[],
        :I=>Float64[],
        :C=>Float64[],
        :C_specific=>Float64[],
        :E=>Float64[],
        :E_specific=>Float64[],
        :P=>Float64[],
        :P_specific=>Float64[])
end


function hole_correction_pair_correlation(x1::AbstractVector{T},s1::Specie{T}, x2::AbstractVector{T},s2::Specie{T}) where T
    overlapping = norm(x1 - x2) > outer_radius(s1) * s1.exclusion_distance + outer_radius(s2) * s2.exclusion_distance

    return  overlapping ? one(T) : zero(T)
end

function smooth_pair_corr_distance(pair_corr_distance::Function, a12::T; smoothing::T = T(0), max_distance::T = T(20*a12),
        polynomial_order::Int = 15, mesh_size::Int = 10*polynomial_order + 4
    ) where T

    if smoothing > T(1.0)
        @warn "smoothing should be in [0,1], setting smoothing = 1"
        smoothing = T(1.0)
    elseif smoothing < T(0.0)
        @warn "smoothing should be in [0,1], setting smoothing = 0"
        smoothing < T(0.0)
    end

    zs = collect(LinRange(zero(T), max_distance, mesh_size))

    inds = findall(abs.(zs .- a12) .< a12*smoothing/T(2) )
    deleteat!(zs,inds)
    data = pair_corr_distance.(zs)

    P = Legendre()
    ls = 0:polynomial_order

    Pmat = P[2zs ./ max_distance .- T(1.0), ls .+ 1];

    # projector_mat = inv(transpose(Pmat) * (Pmat)) * transpose(Pmat);
    # pls = projector_mat * data
    # Pmat * pls ~ data
    pls = Pmat \ data

    return function (z)
        Ps = P[2z / max_distance - T(1.0), ls .+ 1]
        return sum(Ps .* pls)
    end
end

# function constant_number_density_fun(material::Material)
#
#     @warn "the constant_number_density doesn't currently verify if the whole particle is inside the shape"
#
#     prob = function (x1::AbstractVector, s1::Specie)
#         if x1 ∈ sh
#             number_density(s1)
#         else
#             zero(typeof(V))
#         end
#     end
#
#     return prob
# end
"""
    gls_pair_radial_fun(pair_corr_distance::Function, a12::T; polynomial_order::Int, mesh_size::Int)

Return a function ``gls_fun``. For any radial distances ``r_1`` and ``r_2`` we have ``gls = gls_fun(r1,r2)``  such that ``g(r_1,r_2,\\cos \\theta_{12}) = \\sum_{\\ell_1 =0} \\frac{2\\ell_1 + 1}{4\\pi} gls[\\ells+1] P_{\\ell_1}(\\cos \\theta_{12}) ``, where ``g(r_1,r_2,\\cos \\theta_{12})`` is the radially symmetric pair-correlation, so it depends only on the radial distances ``r_1`` and ``r_2``, and the angle between two position vectors ``\\theta_{12}``.

The function `gls_fun` is calculated from the function `pair_corr_distance`, where `pair_corr_distance(sqrt(r1^2 + r2^2 - 2r1 * r2 * cos(θ12)))` gives the pair correlation.
"""
function gls_pair_radial_fun(pair_corr_distance::Function, a12::T;
            polynomial_order::Int = 15,
            mesh_size::Int = 10*polynomial_order + 4,
            sigma_approximation = true
        ) where T

    P = Legendre()
    ls = 0:polynomial_order

    if sigma_approximation
        sigmas = [one(T); sin.(pi .* ls[2:end] ./ (polynomial_order+1)) ./ (pi .* ls[2:end] ./ (polynomial_order+1))]
    else
        sigmas = ones(T,polynomial_order+1)
    end
    S = diagm(sigmas)

    us = LinRange(-1.0, 1.0, mesh_size)
    Pmat = P[us, ls .+ 1];
    Pmat = [Pmat[i] * T(2i[2] - 1) / (4pi) for i in CartesianIndices(Pmat)];

    # Pmat = S * Pmat;

    # projector_mat =  S * inv(transpose(Pmat) * (Pmat)) * transpose(Pmat);

    return function (r1,r2)
        if (r1 + r2 < a12)
            return zeros(typeof(a12), polynomial_order + 1)
        else
            data = pair_corr_distance.(sqrt.(r1^2 .+ r2^2 .- 2r1 .* r2 .* us))
            # pls = projector_mat * data
            # data ~ Pmat * pls
            pls = Pmat \ data

            return S * pls
        end
    end
end

"""
    pair_radial_fun(pair_corr::Function, a12::T; polynomial_order::Int, mesh_size::Int)

Return a function `pair_radial` such that `pair_radial(r1,r2, cos(θ12))` gives the pair correlation particles at the radial distances r1 and r2, with and angle of `θ12` between them.

The function `pair_radial` is calculated from the function `pair_corr_distance`, where `pair_corr_distance(sqrt(r1^2 + r2^2 - 2r1 * r2 * cos(θ12)))` gives the pair correlation.
"""
function pair_radial_fun(pair_corr_distance::Function, a12::T; polynomial_order::Int = 15, kws...) where T

    gls_fun = gls_pair_radial_fun(pair_corr_distance, a12; polynomial_order = polynomial_order, kws...)
    P = Legendre()

    return function (r1,r2,u)
        Pus = P[u, 1:(polynomial_order + 1)] .* (2 .* (0:polynomial_order) .+ 1) ./ (4pi)

        return sum(Pus .* gls_fun(r1,r2))
    end
end


""" Construct variational implicit surface given unordered set of points in R3 """

# struct surface{M,N,T}
#     centers::MVector{M,SVector{N,T}}
#     distances::MVector{M,T}
#     coeff::SVector{M,T}
#     A::MMatrix{P,P,T}
#     b::MVector{P,T}
# end

""" Duchon's interpolant """
function rbf(x, y)
    return norm(x-y)^3
end

# """ Form interpolation matrix A """
# function interpolation_matrix!(A::MMatrix{P,P,T}, centers::SVector{M,SVector{N,T}}) where {M,N,P,T}
#     for i = 1:M
#         A[i,P-2:P] = centers[i]
#         A[P-2:P,i] = centers[i]
#         for j = 1:M
#             A[i,j] = rbf(centers[i],centers[j])
#         end
#     end
#     A[M+1,1:M] .= 1
#     A[1:M,M+1] .= 1
#     A[M+1:P,M+1:P] .= 0
# end

""" Form interpolation matrix A """
function interpolation_matrix!(A::MMatrix{P,P,T}, centers::SVector{M,SVector{N,T}}) where {M,N,P,T}
    @inbounds for i = 1:M
        # Upper right block
        A[i,P-2] = centers[i][1]
        A[i,P-1] = centers[i][2]
        A[i,P] = centers[i][3]
        # Lower left block
        A[P-2,i] = centers[i][1]
        A[P-1,i] = centers[i][2]
        A[P,i] = centers[i][3]
        # Upper left block
        @inbounds for j = 1:M
            A[i,j] = rbf(centers[i],centers[j])
        end
        # Column of ones
        A[M+1,i] = 1
        #  Row of ones
        A[i,M+1] = 1
    end
    # Bottom right 4x4 block of zeros
    @inbounds for i = 1:4
        @inbounds for j = 1:4
            A[M+i,M+j] = 0
        end
    end
    return nothing
end

""" Form b vector corresponding to signed distance values """
function fit_vector!(b::MVector{M,T}, distances::SVector{N,T}) where {M,N,T}
    @inbounds for i = 1:N
        b[i] = distances[i]
    end
    @inbounds for i = 1:4
        b[N+i] = 0
    end
    return nothing
end

""" Evaluate implicit surface value given point x (single evaluation) """
function eval_surface(coeff::SVector{M,T}, centers::SVector{N,SVector{P,T}}, x) where {M,N,P,T}
    # f = c[M-3] + c[M-2]*x[1] + c[M-1]*x[2] + c[M]*x[3]
    f = coeff[M-3]
    f += dot(SVector{3,T}(coeff[M-2], coeff[M-1], coeff[M]), x)
    @inbounds for i = 1:N
        f += coeff[i]*rbf(x, centers[i])
    end
    return f
end

""" Main call """
function implicit_surface(verts::Vector{gt.Point{3,T}},
                          points::SVector{N,SVector{3,T}},
                          signed_distances::SVector{N,T}) where {N,T}
    centers, distances = construct_centers(verts, points, signed_distances)
    coeff = construct_surface(centers, distances)
    return coeff, centers
end

""" Construct implicit surface """
function construct_surface(centers::SVector{M,SVector{N,T}}, distances::SVector{M,T}) where {M,N,T}
    A = MMatrix{M+4,M+4,T}(undef)
    b = MVector{M+4,T}(undef)
    interpolation_matrix!(A, centers)
    fit_vector!(b, distances)
    coeff = A \ b
    return SVector{M+4,Float64}(coeff)
end

""" Construct centers of implicit surface """
function construct_centers(verts::Vector{gt.Point{3,T}},
                           points::SVector{N,SVector{3,T}},
                           point_distances::SVector{N,T}) where {N,T}
    M = length(verts)
    centers = MVector{M+N,SVector{3,Float64}}(undef)
    distances = MVector{M+N,Float64}(undef)
    @inbounds for i = 1:M
        centers[i] = SVector{3,Float64}(verts[i])
        distances[i] = 0
    end
    @inbounds for i = 1:N
        centers[M+i] = points[i]
        distances[M+i] = point_distances[i]
    end
    return SVector(centers), SVector(distances)
end

# """ Create initial rbf centers from geometry, where N is the number of environment points """
# function init_centers(verts::Vector{GeometryTypes.Point{3,T}}, N::Int) where {T}
#     M = length(verts)
#     centers = MVector{M+N,SVector{3,Float64}}(undef)
#     distances = MVector{M+N,Float64}(undef)
#     for i = 1:M
#         centers[i] = SVector{3,Float64}(verts[i])
#         distances[i] = 0
#     end
#     return centers, distances
# end

# """ Add points and distances to centers """
# function add_centers!(centers::MVector{M,SVector{3,T}}, distances::MVector{M,T}, points::SVector{N,SVector{3,T}}, point_distances::SVector{N,T}) where {M,N,T}
#
#     add_centers = MVector{M+N, SVector{3,Float64}}(undef)
#     add_dist = MVector{M+N, Float64}(zeros(M+N))
#     for i = 1:N
#         centers[i] = SVector{3,Float64}(verts[i])
#     end
#     for i = 1:M
#         centers[N+i] = points[i]
#         dist[N+i] = distances[i]
#     end
#     return SVector(centers), SVector(dist)
# end

# Functions for surface fitting
""" Get ∇f(x,y,z) and H(x,y,z) at given eval point (x,y,z) """
function surface_gradient!(∇f,
                           coeff::SVector{M,T},
                           centers::SVector{N,SVector{P,T}},
                           point) where {M,N,P,T}
    ForwardDiff.gradient!(∇f, x -> eval_surface(coeff, centers, x), point)
    # return ∇f
    # H = ForwardDiff.hessian(x -> eval_surface(coeff, centers, x), point)
    # return ∇f, H
    # result = ForwardDiff.gradient!(result, x -> eval_surface(coeff, centers, x), point)
    # result = ForwardDiff.hessian!(result, x -> eval_surface(coeff, centers, x), point)
    return nothing
end

""" First order newton method to project grid points onto zero level set """
function fzero(coeff, centers, grid, ∇f; max_iter=20, atol=1e-6, xtol=1e-3)
    n_grid = length(grid)
    zero_points = MVector{n_grid, SVector{3,Float64}}(undef)
    cnt = 1
    for i = 1:n_grid
        iter = 1
        x = grid[i]
        while true
            y = eval_surface(coeff, centers, x)
            if iter == max_iter
                break
            elseif dot(y, y) < atol
                zero_points[cnt] = SVector{3,Float64}(x)
                cnt += 1
                break
            else
                surface_gradient!(∇f, coeff, centers, x)
                Δx = y * (∇f / dot(∇f, ∇f))
                println(Δx)
                if dot(Δx, Δx) < xtol
                    break
                else
                    x -= Δx
                end
                iter += 1
            end
        end
    end
    return zero_points, cnt-1
end

# Meshing functions
""" Return points and faces of the isosurface of the implicit surface """
function make_geom(coeff::SVector{M,T}, centers::SVector{N, SVector{P,T}}, origin::SVector{3,T}, dims::SVector{3,T}, pts::Int) where {M,N,P,T}
    function f_eval(x)
        f = coeff[N-3] + coeff[N-2]*x[1] + coeff[N-1]*x[2] + coeff[N]*x[3]
        for i = 1:N
            f += coeff[i]*rbf(x, centers[i])
        end
        return f
    end
    points, faces = isosurface(f_eval, MarchingTetrahedra(), origin=origin, widths=dims, samples=(pts,pts,pts))
    return points, faces
end

""" Construct mesh compatible with MeshCat
Args (output of isosurface):
    -points: Array of vertices
    -faces: Vertex connectivity graph
"""
function make_mesh(points, faces)
    npoints = size(points,1)
    nfaces = size(faces,1)
    meshpoints = Vector{Point3f0}(undef,npoints)
    for i = 1:npoints
        meshpoints[i] = Point3f0(points[i][1],points[i][2],points[i][3])
    end
    meshfaces = Vector{GLTriangleFace}(undef,nfaces)
    for i = 1:nfaces
        meshfaces[i] = TriangleFace(faces[i][1],faces[i][2],faces[i][3])
    end
    return Mesh(meshpoints, meshfaces)
end

""" Make implicit surface mesh with sdf """
function make_mesh_sdf(coeff::SVector{M,T}, centers::SVector{N, SVector{P,T}}, origin::SVector{3,T}, dims::SVector{3,T}, resolution::T) where {M,N,P,T}
    function f_eval(x)
        f = coeff[M-3]
        f += SVector{3,T}(coeff[M-2], coeff[M-1], coeff[M])'*x
        for i = 1:N
            f += coeff[i]*rbf(x, centers[i])
        end
        return f
    end
    sdf = gt.SignedDistanceField(f_eval, gt.HyperRectangle(gt.Vec(origin), gt.Vec(dims)), resolution)
    mesh = gt.HomogenousMesh(sdf, MarchingTetrahedra())
    return mesh
end


macro brenttrace()
    quote
        if tracing
            dt = Dict()
            if extended_trace
                dt["x_minimum"] = x_minimum
                dt["x_lower"] = x_lower
                dt["x_upper"] = x_upper
            end
            update!(tr,
                    iteration,
                    f_minimum,
                    NaN,
                    dt,
                    store_trace,
                    show_trace,
                    show_every,
                    callback)
        end
    end
end

immutable Brent <: Optimizer end

function optimize{T <: AbstractFloat}(
        f::Function, x_lower::T, x_upper::T,
        mo::Brent;
        rel_tol::T = sqrt(eps(T)),
        abs_tol::T = eps(T),
        iterations::Integer = 1_000,
        store_trace::Bool = false,
        show_trace::Bool = false,
        callback = nothing,
        show_every = 1,
        extended_trace::Bool = false)

    if !(x_lower < x_upper)
        error("x_lower must be less than x_upper")
    end

    # Save for later
    initial_lower = x_lower
    initial_upper = x_upper

    const golden_ratio::T = 0.5 * (3.0 - sqrt(5.0))

    x_minimum = x_lower + golden_ratio*(x_upper-x_lower)
    f_minimum = f(x_minimum)
    f_calls = 1 # Number of calls to f

    step = zero(T)
    step_old = zero(T)

    x_minimum_old = x_minimum
    x_minimum_old_old = x_minimum

    f_minimum_old = f_minimum
    f_minimum_old_old = f_minimum

    iteration = 0
    converged = false

    # Trace the history of states visited
    tr = OptimizationTrace{typeof(mo)}()
    tracing = store_trace || show_trace || extended_trace || callback != nothing
    @brenttrace

    while iteration < iterations

        p = zero(T)
        q = zero(T)

        tolx = rel_tol * abs(x_minimum) + abs_tol

        x_midpoint = (x_upper+x_lower)/2

        if abs(x_minimum - x_midpoint) <= 2*tolx - (x_upper-x_lower)/2
            converged = true
            break
        end

        iteration += 1

        if abs(step_old) > tolx
            # Compute parabola interpolation
            # x_minimum + p/q is the optimum of the parabola
            # Also, q is guaranteed to be positive

            r = (x_minimum - x_minimum_old) * (f_minimum - f_minimum_old_old)
            q = (x_minimum - x_minimum_old_old) * (f_minimum - f_minimum_old)
            p = (x_minimum - x_minimum_old_old) * q - (x_minimum - x_minimum_old) * r
            q = 2(q - r)

            if q > 0
                p = -p
            else
                q = -q
            end
        end

        if abs(p) < abs(q*step_old/2) && p < q*(x_upper-x_minimum) && p < q*(x_minimum-x_lower)
            step_old = step
            step = p/q

            # The function must not be evaluated too close to x_upper or x_lower
            x_temp = x_minimum + step
            if ((x_temp - x_lower) < 2*tolx || (x_upper - x_temp) < 2*tolx)
                step = (x_minimum < x_midpoint) ? tolx : -tolx
            end
        else
            step_old = (x_minimum < x_midpoint) ? x_upper - x_minimum : x_lower - x_minimum
            step = golden_ratio * step_old
        end

        # The function must not be evaluated too close to x_minimum
        if abs(step) >= tolx
            x_new = x_minimum + step
        else
            x_new = x_minimum + ((step > 0) ? tolx : -tolx)
        end

        f_new = f(x_new)
        f_calls += 1

        if f_new <= f_minimum
            if x_new < x_minimum
                x_upper = x_minimum
            else
                x_lower = x_minimum
            end
            x_minimum_old_old = x_minimum_old
            f_minimum_old_old = f_minimum_old
            x_minimum_old = x_minimum
            f_minimum_old = f_minimum
            x_minimum = x_new
            f_minimum = f_new
        else
            if x_new < x_minimum
                x_lower = x_new
            else
                x_upper = x_new
            end
            if f_new <= f_minimum_old || x_minimum_old == x_minimum
                x_minimum_old_old = x_minimum_old
                f_minimum_old_old = f_minimum_old
                x_minimum_old = x_new
                f_minimum_old = f_new
            elseif f_new <= f_minimum_old_old || x_minimum_old_old == x_minimum || x_minimum_old_old == x_minimum_old
                x_minimum_old_old = x_new
                f_minimum_old_old = f_new
            end
        end

        @brenttrace
    end

    return UnivariateOptimizationResults("Brent's Method",
                                         initial_lower,
                                         initial_upper,
                                         x_minimum,
                                         Float64(f_minimum),
                                         iteration,
                                         iteration == iterations,
                                         converged,
                                         rel_tol,
                                         abs_tol,
                                         tr,
                                         f_calls)
end


function findpeaks(activation; filterlength=10, minh=0, threshold=-Inf, minw=0, maxw=Inf, doplot=false, kwargs...)
    activationf = filtfilt(ones(filterlength),[filterlength],activation)
    doplot && plot!(activationf, lab="Filtered Activation", kwargs...)

    peaks = allpeaks(activationf)
    doplot &&  scatter!(peaks, activationf[peaks], lab="All peaks", kwargs...)

    peaks = peaks[activationf[peaks] .>= minh] # Remove too small peaks
    doplot &&  scatter!(peaks, activationf[peaks]*1.1, lab="Below minh removed", kwargs...)

    peaks = remove_threshold(peaks,activationf,threshold)
    peaks = remove_width(peaks,activationf,minw,maxw)
    doplot &&  scatter!(peaks, activationf[peaks]*1.2, lab="Too wide or narrow removed" , m=(10,:xcross), kwargs...)
    return peaks
end

function remove_threshold(peaks,activation,threshold)
    base = max.(activation[peaks .- 1],activation[peaks .+ 1])
    peaks = peaks[activation[peaks].-base .>= threshold]
end

function remove_width(peaks,activation,minw,maxw)
    if minw <= 0 && maxw == Inf
        return peaks # Nothing needs to be done
    end
    i = 1
    while i < length(peaks)
        speaks = sortperm(activation[peaks], rev=true) # sorted peak locations (in peaks vector)

        if speaks[i] > 1
            wl = abs(peaks[speaks[i]] - peaks[speaks[i]-1]) # highest peak location - location to the left
            if wl < minw || wl > maxw
                deleteat!(peaks,speaks[i]-1)
                continue
            end
        end
        if speaks[i] < length(peaks)
            wr = abs(peaks[speaks[i]] - peaks[speaks[i]+1])
            if wr < minw || wr > maxw
                deleteat!(peaks,speaks[i]+1)
                continue
            end
        end
        i += 1
    end
    if !isempty(peaks)
        @assert minimum(diff(peaks)) >= minw
    end
    peaks
end

function allpeaks(activation)
    da = diff(activation)
    peaks = findall((diff(da) .< 0) .& (da[1:end-1] .> 0) .& (da[2:end] .< 0)) .+ 1
end


###
#   An implementation of Transitional Markov Chain Monte Carlo in Julia
#
#            Institute for Risk and Uncertainty, Uni of Liverpool
#
#                       Authors: Ander Gray, Adolphus Lye
#
#                       Email: Ander.Gray@liverpool.ac.uk,
#                              Adolphus.Lye@liverpool.ac.uk
#
#
#   This Transitional MCMC algorithm is inspirted by OpenCOSSAN:
#                   https://github.com/cossan-working-group/OpenCossan
#
#
#   Algorithm originally proposed by:
#           J. Ching, and Y. Chen (2007). Transitional Markov Chain Monte Carlo method
#           for Bayesian model updating, Model class selection, and Model averaging.
#           Journal of Engineering Mechanics, 133(7), 816-832.
#           doi:10.1061/(asce)0733-9399(2007)133:7(816)
#
###

function tmcmc(log_fD_T, fT, sample_fT, Nsamples, burnin=20, thin=3, beta2=0.01)

    j1 = 0;                     # Iteration number
    βj = 0;                     # Tempering parameter
    θ_j = sample_fT(Nsamples);  # Samples of prior
    Lp_j = zeros(Nsamples, 1);   # Log liklihood of first iteration

    Log_ev = 0                  # Log Evidence

    Ndims = size(θ_j, 2)         # Number of dimensions (input)

    while βj < 1

        j1 = j1 + 1
        println()
        println("Beginning iteration $j1")

        ###
        # Parallel evaluation of the likelihood
        ###
        print("Computing likelihood with $(nworkers()) workers....")
        # Lp_j = log_fD_T(θ_j')
        ins = [θ_j[i,:] for i in 1:size(θ_j, 1)]
        Lp_j = pmap(log_fD_T, ins)
        Lp_j = reduce(vcat, Lp_j)
        # Lp_j = pamp(log_fD_T,θ_j')
        println("Done!")

        ###
        #   Computing new βj
        #   Uses bisection method
        ###
        println("Computing Bj")

        low_β = βj; hi_β = 2; Lp_adjust = maximum(Lp_j);
        x1 = (hi_β + low_β) / 2;

        while (hi_β - low_β) / ((hi_β + low_β) / 2) > 1e-6
            x1 = (hi_β + low_β) / 2;
            wj_test = exp.((x1 .- βj ) .* (Lp_j .- Lp_adjust));
            cov_w   = std(wj_test) / mean(wj_test);
            if cov_w > 1; hi_β = x1; else; low_β = x1; end
        end

        βj1 = min(1, x1)
        println("B_$(j1) = $(βj1)")

        ###
        #   Computation of normalised weights
        ###
        println("Computing weights")

        w_j = exp.((βj1 - βj) .* (Lp_j .- Lp_adjust))       # Nominal weights from likilhood and βjs

        Log_ev = log(mean(w_j)) + (βj1 - βj) * Lp_adjust + Log_ev   # Log evidence in current iteration

        # Normalised weights
        wn_j = w_j ./ sum(w_j);

        Th_wm = θ_j .* wn_j                 # Weighted mean of samples

        ###
        #   Calculation of COV matrix of proposal
        ###
        SIGMA_j = zeros(Ndims, Ndims)

        for l = 1:Nsamples
            SIGMA_j = SIGMA_j + beta2 .* wn_j[l] .* (θ_j[l,:]' .- Th_wm)' * (θ_j[l,:]' .- Th_wm)
        end

        # Ensure that cov is symetric
        SIGMA_j = (SIGMA_j' + SIGMA_j) / 2

        prop = mu -> proprnd(mu, SIGMA_j, fT)           # Anonymous function for proposal

        target = x -> log_fD_T(x) .* βj1 .+ log.(fT(x)) # Anonymous function for transitional distribution

        # Weighted resampling of θj (indecies with replacement)
        randIndex = sample(1:Nsamples, Weights(wn_j), Nsamples, replace=true)

        ins = [θ_j[randIndex[i], :] for i = 1:Nsamples]

        print("Markov chains with $(nworkers()) workers...")
        # pmap is a parallel map
        θ_j1 = pmap(x -> runChains(target, prop, x, burnin, thin), ins)

        θ_j1 = reduce(vcat, θ_j1)

        println("Done!")

        # println("Mean α = $(meanα)")

        βj = βj1
        θ_j = θ_j1
    end
    return θ_j, Log_ev
end

function proprnd(mu, covMat, prior)

    samp = rand(MvNormal(mu, covMat), 1)
    while iszero(prior(samp))
        samp = rand(MvNormal(mu, covMat),1)
    end
    return samp[:]
end

function runChains(target, prop, θ_js, burnin, thin)
    samps, α  = MHsampleSimple(target, prop, θ_js, 1, burnin, thin)
    return samps
end

function runChains2(target, prop, θ_js, burnin, thin)
    Nsamples = size(θ_js, 1)
    Ndims    = size(θ_js, 2)
    θ_j1 = zeros(Nsamples, Ndims)
    α = zeros(Nsamples, 1)
    for i = 1:Nsamples
        samps, αs = MHsampleSimple(target, prop, θ_j[i], 1, burnin, thin)
        θ_j1[i,:] = samps
        α[i] = αs
    end

    return θ_j1, α

end


struct Word{T}
    alphabet::Vector{T}
    letters::Vector{Int}

    function Word(a::AbstractVector{T}, l::AbstractVector{<:Integer}) where {T}
        all(i -> 1 <= i <= length(a), l) ||
            throw(ArgumentError("Invalid word over alphabet $a: $w"))
        return new{T}(a, l)
    end
end
Base.length(w::Word) = length(w.letters)
function Base.show(io::IO, w::Word)
    if isone(w)
        print(io, "(id)")
    else
        join(io, w.alphabet[w.letters], "·")
    end
end

Base.:(==)(w::Word, v::Word) = w.alphabet == v.alphabet && w.letters == v.letters
Base.hash(w::Word, h::UInt) = hash(w.alphabet, hash(w.letters, hash(Word, h)))

Base.one(w::Word) = Word(w.alphabet, Int[])
Base.isone(w::Word) = length(w) == 0

function Base.:*(w::Word, z::Word)
    @assert w.alphabet == z.alphabet
    return Word(w.alphabet, [w.letters; z.letters])
end

function StarAlgebras.star(w::Word)
    newletters = similar(w.letters)
    A = w.alphabet

    # star(:a) = :b
    # star(:b) = :a
    # star(:c) = :c

    for (i, l) in enumerate(Iterators.reverse(w.letters))
        k = if A[l] === :a
            2
        elseif A[l] === :b
            1
        else
            l
        end
        newletters[i] = k
    end
    return Word(w.alphabet, newletters)
end

function words(alphabet; radius)

    words = [Word(alphabet, Int[])] # word identity

    for r in 1:radius
        append!(
            words,
            [
                Word(alphabet, collect(w)) for
                w in Iterators.product(fill(1:length(alphabet), r)...)
            ],
        )
    end

    return words
end


## Convert orignal Coconet training data into Julia format
using BSON
using PyCall
using Coconet

#===============================================================================
    SECTION 1: Python to Julia
===============================================================================#

path = "datasets"
name = "Tonnetz_Bach.bson"

let np = pyimport("numpy")
    # Load data as PyObject
    data = np.load(
        joinpath(path, "Jsb16thSeparated.npz"),
        allow_pickle = true,
        encoding = "latin1",
    )

    # Convert to Julia vectors
    test_data = [[timestep for timestep in song] for song in get(data, :test)]
    train_data = [[timestep for timestep in song] for song in get(data, :train)]
    valid_data = [[timestep for timestep in song] for song in get(data, :valid)]

    # Convert to arrays
    test_data = [hcat(song...) for song in test_data]
    train_data = [hcat(song...) for song in train_data]
    valid_data = [hcat(song...) for song in valid_data]

    # Transpose and re-convert to arrays
    test_data = [hcat(song') for song in test_data]
    train_data = [hcat(song') for song in train_data]
    valid_data = [hcat(song') for song in valid_data]

    # Save as dict
    bson(
        joinpath(path, name),
        test = test_data,
        train = train_data,
        valid = valid_data,
    )
end
#===============================================================================
    SECTION 2: Notes to tonnetz vectors
===============================================================================#

# Load in as dict
bach = BSON.load(joinpath(path, name))

notes = []
dataset = []
# for each component of the dataset
for test in [bach[v] for v in [:test, :train, :valid]]
    set = []

    # for each song
    for i = 1:length(test)
        song = zeros(Float32, 3, size(test[i])...)

        # for each instrument
        for instrument = 1:size(test[i])[2]
            # for each timestep
            for timestep = 1:size(test[i])[1]
                pitch = test[i][timestep, instrument]
                x, y, z = 0, 0, 0

                if !isnan(pitch)
                    z = 1
                    x, y = Coconet.pitch_to_tonnetz(Integer(pitch))
                end

                song[:, timestep, instrument] = [x,y,z]
                push!(notes, song[:, timestep, instrument])
            end # for each timestep
        end # for each instrument

        push!(set, song)
    end #for each song

    push!(dataset, set)
end # for each set

# Normalization
μ = sum(notes) / length(notes)
σ = sum([(note - μ).^2 for note in notes]) / length(notes)
for i in 1:length(dataset)
    for j in 1:length(dataset[i])
        for k in 1:size(dataset[i][j], 2)
            for l in 1:size(dataset[i][j], 3)
                dataset[i][j][:, k, l] = (dataset[i][j][:, k, l] - μ) ./ σ
            end
        end
    end
end

#===============================================================================
    SECTION 3: Save it
===============================================================================#
bson(
    joinpath(path, name),
    test = dataset[1],
    train = dataset[2],
    valid = dataset[3],
    μ = μ,
    σ = σ,
)


@testset "Grid diameter identify" begin
    l = 12
    p = 66

    img = zeros(1, 500)
    for i = 20:length(img)
        if (mod(i, p) < l)
            img[i] = 1
        end
    end

    img1 = img + rand(1:500, 1, 500) / 2500

    img2 = zeros(1, 500)
    for i = 20:length(img)
        img2[i] = (img[i-4] + img[i-3] + img[i-2] + img[i-1] + img[i]) / 5
    end

    img3 = img2 + rand(1:500, 1, 500) / 2500

    # sharp grids
    stript, space = grid_resolution(img)
    @test stript ≈ l atol = 1
    @test space + stript ≈ p atol = 1

    # sharp grids with random
    stript, space = grid_resolution(img1)
    @test stript ≈ l atol = 1
    @test space + stript ≈ p atol = 1

    # moving average grids window 5
    stript, space = grid_resolution(img2)
    @test stript ≈ l atol = 1
    @test space + stript ≈ p atol = 1

    # moving average grids window 5 with random
    stript, space = grid_resolution(img3)
    @test stript ≈ l atol = 1
    @test space + stript ≈ p atol = 1
end


import ..UncertainDatasets.AbstractUncertainValueDataset
import ..UncertainDatasets.UncertainDataset
import ..UncertainValues.UncertainScalarKDE

abstract type SamplingConstraint end

"""
    NoConstraint

A (non)constraint indicating that the distribution furnishing an uncertain value
should be sampled its entire support.
"""
struct NoConstraint <: SamplingConstraint end

#########################################################################
# Sampling constraints for regular data (i.e. not age/depth/time index...,
# but the values associated to those indices).
#########################################################################

abstract type ValueSamplingConstraint <: SamplingConstraint end

"""
    TruncateLowerQuantile(lower_quantile::Float64)

A constraint indicating that the distribution furnishing an uncertain value
should be truncated below at some quantile.
"""
struct TruncateLowerQuantile <: ValueSamplingConstraint
    lower_quantile::Float64
end

"""
    TruncateUpperQuantile(upper_quantile::Float64)

A constraint indicating that the distribution furnishing an uncertain value
should be truncated above at some quantile.
"""
struct TruncateUpperQuantile <: ValueSamplingConstraint
    upper_quantile::Float64
end

"""
    TruncateQuantiles(lower_quantile::Float64, upper_quantile::Float64)

A constraint indicating that the distribution furnishing an uncertain value
should be truncated at some quantile quantile
`(lower_quantile, upper_quantile)`.
"""
struct TruncateQuantiles{T1<:Real, T2<:Real} <: ValueSamplingConstraint
    lower_quantile::T1
    upper_quantile::T2

    function TruncateQuantiles(lower_quantile::T1, upper_quantile::T2) where {T1, T2}
        err_msg = "Need 0 <= lower_quantile < upper_quantile <= 1"

        if !(lower_quantile < upper_quantile && 0.0 <= lower_quantile < upper_quantile <= 1.0)
            throw(DomainError(err_msg * " (got lo = $lower_quantile, hi = $upper_quantile)"))
        else
            new{T1, T2}(lower_quantile, upper_quantile)
        end
    end
end

"""
    TruncateStd(nσ::Number)

A constraint indicating that the distribution furnishing an uncertain value
should be truncated at the mean ± `nσ` (`n` standard deviations).

## Notes

- Beware when you apply the `TruncateStd` constraint to a (usually a numeric)
    population with a small value range. With `nσ` small, you might end up with
    a population mean *between* the actual values, so that the range
    `[mean(pop) - nσ*std(pop), mean(pop) + nσ*std(pop)]` returns `nothing`.
"""
struct TruncateStd{T<:Number} <: ValueSamplingConstraint
    nσ::T

    function TruncateStd(nσ::T) where T

        if nσ <= 0
            err_str = "TruncateStd must be initialised with nσ strictly positive"
            throw(DomainError(err_str *  " (got nσ = $nσ)"))
        else
            new{T}(nσ)
        end
    end
end

"""
    TruncateMinimum(min::Number)

A constraint indicating that the distribution furnishing an uncertain value
should be truncated below at some specified minimum value.
"""
struct TruncateMinimum{T<:Number} <: ValueSamplingConstraint
    min::T
end

"""
    TruncateMaximum(max::Number)

A constraint indicating that the distribution furnishing an uncertain value
should be truncated above at some specified maximum value.
"""
struct TruncateMaximum{T<:Number} <: ValueSamplingConstraint
    max::T
end

"""
    TruncateRange(min::Number, max::Number)

A constraint indicating that the distribution furnishing an uncertain value
should be truncated at some range `[min, max]`.
"""
struct TruncateRange{T1, T2} <: ValueSamplingConstraint
    min::T1
    max::T2

    function TruncateRange(min::T1, max::T2) where {T1, T2}
        if min < max
            return new{T1, T2}(min, max)
        else
            err_msg = "Cannot create TruncateRange instance. Need min < max"
            throw(DomainError(err_msg * " (got min = $min, max = $max)"))
        end
    end
end

export
SamplingConstraint,
NoConstraint,

ValueSamplingConstraint,
TruncateLowerQuantile,
TruncateUpperQuantile,
TruncateQuantiles,
TruncateMinimum,
TruncateMaximum,
TruncateRange,
TruncateStd

#########################################################################
# Sampling constraints for sample indices (time index, age, depth, etc...)
# Often, these need to be sampled to obey some physical criteria (i.e.,
# observations are from physical samples lying above each other, so the order
# of the observations cannot be mixed).
#########################################################################
include("ordered_sequences/ordered_sequence_algorithms.jl")

"""
    IndexSamplingConstraint

An abstract type for sampling constraints valid only for indices.
"""
abstract type IndexSamplingConstraint <: SamplingConstraint end

#########################################################################
# Sequential sampling constraints
#########################################################################

abstract type SequentialSamplingConstraint end

# Add the ordered sampling scheme to the seqential sampling constraints,
# because that's all they affect. Defaults to `StartToEnd`.
(::Type{SSC})() where SSC<:SequentialSamplingConstraint = SSC(StartToEnd())

"""
    StrictlyIncreasing

A sampling constraint indicating element-wise sampling of the uncertain values in a dataset,
such that the values of the draw are strictly increasing in magnitude.

Typically used when there are known, physical constraints on the measurements.
For example, geochemical measurements of sediments at different depths of a sediment core
are taken at physically separate depths in the core. Thus, the order of the indices cannot
be flipped, and must be strictly decreasing/increasing.
"""
struct StrictlyIncreasing{OA<:OrderedSamplingAlgorithm} <: SequentialSamplingConstraint
    ordered_sampling_algorithm::OA
end

"""
    StrictlyDecreasing

A sampling constraint indicating element-wise sampling of the uncertain values in a dataset,
such that the values of the draw are strictly decreasing in magnitude.

Typically used when there are known, physical constraints on the measurements.
For example, geochemical measurements of sediments at different depths of a sediment core
are taken at physically separate depths in the core. Thus, the order of the indices cannot
be flipped, and must be strictly decreasing/increasing.
"""
struct StrictlyDecreasing{OA<:OrderedSamplingAlgorithm} <: SequentialSamplingConstraint
    ordered_sampling_algorithm::OA
end

export
IndexSamplingConstraint,
SequentialSamplingConstraint,
StrictlyIncreasing,
StrictlyDecreasing


using ArrayInterface, Test
using Base: setindex
import ArrayInterface: has_sparsestruct, findstructralnz, fast_scalar_indexing
@test ArrayInterface.ismutable(rand(3))

using StaticArrays
@test ArrayInterface.ismutable(@SVector [1,2,3]) == false
@test ArrayInterface.ismutable(@MVector [1,2,3]) == true

using LinearAlgebra, SparseArrays

D=Diagonal([1,2,3,4])
@test has_sparsestruct(D)
rowind,colind=findstructralnz(D)
@test [D[rowind[i],colind[i]] for i in 1:length(rowind)]==[1,2,3,4]
@test length(rowind)==4
@test length(rowind)==length(colind)

Bu = Bidiagonal([1,2,3,4], [7,8,9], :U)
@test has_sparsestruct(Bu)
rowind,colind=findstructralnz(Bu)
@test [Bu[rowind[i],colind[i]] for i in 1:length(rowind)]==[1,7,2,8,3,9,4]
Bl = Bidiagonal([1,2,3,4], [7,8,9], :L)
@test has_sparsestruct(Bl)
rowind,colind=findstructralnz(Bl)
@test [Bl[rowind[i],colind[i]] for i in 1:length(rowind)]==[1,7,2,8,3,9,4]

Tri=Tridiagonal([1,2,3],[1,2,3,4],[4,5,6])
@test has_sparsestruct(Tri)
rowind,colind=findstructralnz(Tri)
@test [Tri[rowind[i],colind[i]] for i in 1:length(rowind)]==[1,2,3,4,4,5,6,1,2,3]
STri=SymTridiagonal([1,2,3,4],[5,6,7])
@test has_sparsestruct(STri)
rowind,colind=findstructralnz(STri)
@test [STri[rowind[i],colind[i]] for i in 1:length(rowind)]==[1,2,3,4,5,6,7,5,6,7]

Sp=sparse([1,2,3],[1,2,3],[1,2,3])
@test has_sparsestruct(Sp)
rowind,colind=findstructralnz(Sp)
@test [Tri[rowind[i],colind[i]] for i in 1:length(rowind)]==[1,2,3]

@test ArrayInterface.ismutable(spzeros(1, 1))
@test ArrayInterface.ismutable(spzeros(1))

@test !fast_scalar_indexing(qr(rand(10, 10)).Q)
@test !fast_scalar_indexing(qr(rand(10, 10), Val(true)).Q)
@test !fast_scalar_indexing(lq(rand(10, 10)).Q)

using BandedMatrices

B=BandedMatrix(Ones(5,5), (-1,2))
B[band(1)].=[1,2,3,4]
B[band(2)].=[5,6,7]
@test has_sparsestruct(B)
rowind,colind=findstructralnz(B)
@test [B[rowind[i],colind[i]] for i in 1:length(rowind)]==[5,6,7,1,2,3,4]
B=BandedMatrix(Ones(4,6), (-1,2))
B[band(1)].=[1,2,3,4]
B[band(2)].=[5,6,7,8]
rowind,colind=findstructralnz(B)
@test [B[rowind[i],colind[i]] for i in 1:length(rowind)]==[5,6,7,8,1,2,3,4]

using BlockBandedMatrices
BB=BlockBandedMatrix(Ones(10,10),([1,2,3,4],[4,3,2,1]),(1,0))
BB[Block(1,1)].=[1 2 3 4]
BB[Block(2,1)].=[5 6 7 8;9 10 11 12]
rowind,colind=findstructralnz(BB)
@test [BB[rowind[i],colind[i]] for i in 1:length(rowind)]==
    [1,5,9,2,6,10,3,7,11,4,8,12,
    1,1,1,1,1,1,1,1,1,1,1,1,1,1,
    1,1,1,1,1,1,1,1,1,1,1,1,1,1,
    1,1,1,1,1]

dense=collect(Ones(8,8))
for i in 1:8
    dense[:,i].=[1,2,3,4,5,6,7,8]
end
BBB=BandedBlockBandedMatrix(dense, ([4, 4] ,[4, 4]), (1, 1), (1, 1))
rowind,colind=findstructralnz(BBB)
@test [BBB[rowind[i],colind[i]] for i in 1:length(rowind)]==
    [1,2,3,1,2,3,4,2,3,4,5,6,7,5,6,7,8,6,7,8,
     1,2,3,1,2,3,4,2,3,4,5,6,7,5,6,7,8,6,7,8]

@testset "setindex" begin
    @testset "$(typeof(x))" for x in [
        zeros(3),
        falses(3),
        spzeros(3),
    ]
        y = setindex(x, true, 1)
        @test iszero(x)  # x is not mutated
        @test y[1] == true
        @test iszero(x[CartesianIndices(size(x)) .== [CartesianIndex(1)]])

        y2 = setindex(x, one.(x), :)
        @test iszero(x)
        @test all(isone, y2)
    end

    @testset "$(typeof(x))" for x in [
        zeros(3, 3),
        falses(3, 3),
        spzeros(3, 3),
    ]
        y = setindex(x, true, 1, 1)
        @test iszero(x)  # x is not mutated
        @test y[1, 1] == true
        @test iszero(x[CartesianIndices(size(x)) .== [CartesianIndex(1, 1)]])

        y2 = setindex(x, one.(x), :, :)
        @test iszero(x)
        @test all(isone, y2)
    end

    @testset "$(typeof(x))" for x in [
        zeros(3, 3, 3),
        falses(3, 3, 3),
    ]
        y = setindex(x, true, 1, 1, 1)
        @test iszero(x)  # x is not mutated
        @test y[1, 1, 1] == true
        @test iszero(x[CartesianIndices(size(x)) .== [CartesianIndex(1, 1, 1)]])

        y2 = setindex(x, one.(x), :, :, :)
        @test iszero(x)
        @test all(isone, y2)
    end
end


using Flux
using Flux: @epochs
using Flux: params
using Flux: throttle
using Flux: glorot_uniform
using Flux: @functor
using Plots
using Distributions
using Zygote
using Flux.Optimise: update!
using ForwardDiff
using ReverseDiff
using FiniteDiff
using ZygoteRules

```
Workflow:

# 1. Define network architecture:
    # 1. define 1 LSTM-like layer
    # 2. define 1 dense layer with activation for the start
    # 3. define 1 dense layer without activation for the end
    # 4. connect them all together -> model
# 2. Define the finite element solution function
# 3. Define sampling function
# 4. Define the loss function
# 5. Train the model

```

# define a DGM layer
mutable struct LSTMLayer
    Ug :: Array{Float64}
    Uz :: Array{Float64}
    Ur :: Array{Float64}
    Uh :: Array{Float64}
    Wg :: Array{Float64}
    Wz :: Array{Float64}
    Wr :: Array{Float64}
    Wh :: Array{Float64}
    bg :: Array{Float64}
    bz :: Array{Float64}
    br :: Array{Float64}
    bh :: Array{Float64}
    act1
    act2
end

LSTMLayer(input_dim :: Integer, output_dim :: Integer, activation_1, activation_2) = LSTMLayer(glorot_uniform(output_dim, input_dim), glorot_uniform(output_dim, input_dim),
    glorot_uniform(output_dim, input_dim), glorot_uniform(output_dim, input_dim), glorot_uniform(output_dim, output_dim), glorot_uniform(output_dim, output_dim),
    glorot_uniform(output_dim, output_dim), glorot_uniform(output_dim, output_dim), zeros(output_dim), zeros(output_dim), zeros(output_dim), zeros(output_dim),
    activation_1, activation_2)

function (a :: LSTMLayer)(S,x)

    # assign all opbject variables
    Ug, Uz, Ur, Uh, Wg, Wz, Wr, Wh, bg, bz, br, bh, act1, act2 = a.Ug, a.Uz, a.Ur, a.Uh, a.Wg, a.Wz, a.Wr, a.Wh, a.bg, a.bz, a.br, a.bh, a.act1, a.act2

    # define all operations
    G = act1.(Ug*x + Wg*S .+ bg)
    Z = act1.(Uz*x + Wz*S .+ bz)
    R = act1.(Ur*x + Wr*S .+ br)
    H = act2.(Uh*x + Wh*(S.*R) .+ bh)
    S_new = (1 .- G).*H + Z.*S

    return S_new
end

# define dense layer with activation
mutable struct Dense_start
    W1 :: Array{Float64}
    b1 :: Array{Float64}
    act
end

Dense_start(input_dim::Integer, output_dim::Integer, act) = Dense_start(glorot_uniform(output_dim, input_dim), zeros(output_dim), act)

function (a :: Dense_start)(x)

    # assign all opbject variables
    W1, b1, act = a.W1, a.b1, a.act

    # define all operations
    S_1 = act.(W1*x .+ b1)

    return S_1
end

# define dense layer without activation
mutable struct Dense_end
    W :: Array{Float64}
    b :: Array{Float64}
end

Dense_end(input_dim::Integer, output_dim::Integer) = Dense_end(glorot_uniform(output_dim, input_dim), zeros(output_dim))

function (a :: Dense_end)(x)

    # assign all opbject variables
    W, b = a.W, a.b

    # define all operations
    y = W*x .+ b

    return y
end

# connect them together into one model
mutable struct DGM
    D1 :: Dense_start
    LSTMLayers :: Array{LSTMLayer}
    D2 :: Dense_end
end

DGM(layer_width, n_layers, spatial_dim) = DGM(Dense_start(spatial_dim + 1, layer_width, tanh),
[LSTMLayer(spatial_dim + 1, layer_width, tanh, tanh) for i in 1:n_layers], Dense_end(layer_width, 1))

function (a :: DGM)(t, x)

    # assign object variables
    D1, LSTMLayers, D2 = a.D1, a.LSTMLayers, a.D2

    #concatenate time and space inputs
    X = vcat(t, x)

    # define all operations
    S = D1(X)

    for lstm_layer in LSTMLayers
        S = lstm_layer(S, X)
    end

    y = D2(S)

    return y
end

# convey trainable parameters to Flux
@functor Dense_start
@functor LSTMLayer
@functor Dense_end
@functor DGM

# check if Flux recognizes parameters
m1 = DGM(50, 3, 1)
params(m1)

# define analytical solution function
function european_put(S, K, r, sigma, t)

   d1 = (log.(S./K) .+ (r + sigma^2/2)*(t_term-t))/(sigma*sqrt(t_term-t))
   d2 = d1 .- (sigma*sqrt(t_term-t))
   put_price = -S.*cdf.(Normal(0,1), -d1) .+ K*exp(-r * (t_term-t))*cdf.(Normal(0,1), -d2)

   return put_price
end

# define the fintie element solution for American put
function finite_element_solution()
    dt = t_term/n_steps
    t = collect(0:dt:t_term + dt)
    dx = sigma*sqrt(3*dt)
    alpha = 0.5*(sigma^2)*dt/(dx^2)
    beta = (r - 0.5*sigma^2)*dt/(2*dx)

    # diagnostics
    # println("r: $(r)")
    # println("sigma: $(sigma)")
    #
    # println("dt: $(dt)")
    # # println("t: $(t)")
    # println("dx: $(dx)")
    # println("alpha: $(alpha)")
    # println("beta: $(beta)")

    # log space grid
    x_max = ceil(10*sigma*sqrt(t_term)/dx)*dx
    x_min = -x_max
    x = collect(x_min:dx:x_max + dx)
    Ndx = length(x) - 1
    x_int = collect(2:Ndx)

    # diagnostics
    # println("x_max: $(x_max)")
    # println("x_min: $(x_min)")
    # # println("x: $(x)")
    # println("len(x): $(length(x))")
    # println("Ndx: $(Ndx)")
    # # println("x_int: $(x_int)")
    # println("len(x_int): $(length(x_int))")

    # put option price grid
    h = ones(Ndx + 1, n_steps + 1)
    h[:, end] = max.(K .- S0*exp.(x), 0)
    hc = ones(Ndx + 1)
    S = S0.*exp.(x)

    # # println("h: $(x)")
    # println("len(h): $(length(h))")
    # # println("hc: $(x)")
    # println("len(hc): $(length(hc))")
    # # println("S: $(x)")
    # println("len(S): $(length(S))")

    # calculate optimal exercise boundary
    exer_bd = ones(n_steps + 1)
    exer_bd[end] = K

    # back-step in time
    for n in collect(n_steps + 1:-1:2)
        # calculate continuation values
        hc[x_int] =  h[x_int, n] .- r*dt*h[x_int, n] .+ beta.*(h[x_int .+ 1, n] .- h[x_int .- 1, n]) + alpha*(h[x_int .+ 1, n] -2*h[x_int, n] + h[x_int .- 1,n])
        hc[1] = 2*hc[2] - hc[3]
        hc[Ndx + 1] = 2*hc[Ndx] - hc[Ndx]

        # compare with intrinsic values
        exer_val = K .- S
        h[:, n-1] = max.(hc, exer_val)

        # calculate optimal exercise boundaries
        checkIdx = (hc > exer_val)*1
        idx = argmax(checkIdx) - 1
        if max(checkIdx) > 0
            exer_bd[n-1] = S[idx]
        end
    end
    putPrice = h
    return t, S, putPrice, exer_bd
end

function sampling_function(space_internal, space_terminal)

    # within-domain sample
    t_internal = rand(Uniform(t_initial, t_term), space_internal)
    S_internal = rand(Uniform(S_low, S_high*S_multiplier), space_internal)

    # terminal sample
    t_terminal = t_term .*ones(space_terminal)
    S_terminal = rand(Uniform(S_low, S_high*S_multiplier), space_terminal)

    return t_internal', S_internal', t_terminal', S_terminal'
end

# define the loss function
function loss_function(N)

    t_internal, S_internal, t_terminal, S_terminal = sampling_function(nSim_interior, nSim_terminal)

    # differential operator loss
    ϵ = 0.01
    model_output_internal = model(t_internal, S_internal)
    ∂g∂x = (model(t_internal, S_internal .+ ϵ) - model_output_internal)./ϵ
    ∂g∂t = (model(t_internal .+ ϵ, S_internal) - model_output_internal)./ϵ
    ∂g∂xx = (model(t_internal, S_internal .+ 2*ϵ) - 2*model(t_internal, S_internal .+ ϵ) + model_output_internal)./(ϵ^2)

    operator_loss_vec = ∂g∂t + r.*S_internal.*∂g∂x + (0.5*(sigma^2)).*(S_internal.^2).*∂g∂xx - r.*model_output_internal

    payoff = relu.(K .- S_internal)
    value = model(t_internal, S_internal)
    L1 = mean((operator_loss_vec.*(value-payoff)).^2)

    temp = relu.(operator_loss_vec)
    L2 = mean(temp.^2)

    V_ineq = relu.(-(payoff - value))
    L3 = mean(V_ineq.^2)

    target_payoff = relu.(K .- S_terminal)
    fitted_payoff = model(t_terminal, S_terminal)

    L4 = mean((fitted_payoff - target_payoff).^2)

    return L1 + L2 + L3 + L4
end

# initialize all problem parameters
r = 0.05           # Interest rate
sigma = 0.5       # Volatility
K = 50             # Strike
t_term = 1              # Terminal time
S0 = 50           # Initial price

# Solution parameters (domain on which to solve PDE)
t_initial = 0 + 1e-10    # time lower bound
S_low = 0.0 + 1e-10  # spot price lower bound
S_high = 2*K         # spot price upper bound

# neural network parameters
n_steps = 10000
num_layers = 3
nodes_per_layer = 50
learning_rate = 0.001

# Training parameters
sampling_stages  = 100   # number of times to resample new time-space domain points
steps_per_sample = 10    # number of SGD steps to take before re-sampling

# Sampling parameters
nSim_interior = 1000
nSim_terminal = 100
S_multiplier  = 1.5   # multiplier for oversampling i.e. draw S from [S_low, S_high * S_multiplier]

# define model
model = DGM(nodes_per_layer, num_layers, 1)

S_model = reshape(S_values, 1, length(S_values))
t_model = 0.6667.*ones(1, length(S_values))

function train_model()
    # set optimizer as ADAM
    evalcb = function ()
    @show(loss_function())
    end
    opt = ADAM(learning_rate)
    dataset = [(1) for i in 1:10]
    Flux.train!(loss_function, params(model), zip(dataset), opt, cb = throttle(evalcb, 25))
end

@epochs 1000 train_model()

params(model)

# compare with analytical solution
compare_times = [0, 0.3333, 0.6667, 1]
n_plots = 100
S_values = [S_low + i*(S_high - S_low)/n_plots for i in 0:n_plots]
S_model = reshape(S_values, 1, length(S_values))

t, S, price, exer_bd = finite_element_solution()
S_idx = [i for i in 1:length(S) if(S[i] < S_high)]
t_idx = [1, 3334, 6668, n_steps]

for (i, t) in enumerate(compare_times)
    # analytical_soln(S, K, r, sigma, t)
    true_values = price[S_idx, t_idx[i]]
    t_model = t.*ones(1, length(S_values))
    fitted_values = model(t_model, S_model)
    fitted_values = reshape(fitted_values, length(fitted_values))
    display(plot!(S_values, fitted_values, label = "Fitted values"))
    display(plot!(S[S_idx], true_values, label = "True values"))
end


"""
    cascademodel(S::Int64, Co::Float64)

Return matrix of the type `UnipartiteNetwork` randomly assembled according to
the cascade model for a given nuber of `S` and connectivity `Co`.

See also: `nichemodel`, `mpnmodel`, `nestedhierarchymodel`

#### References

Cohen, J.E., Newman, C.M., 1985. A stochastic theory of community food webs I.
Models and aggregated data. Proceedings of the Royal Society of London. Series
B. Biological Sciences 224, 421–448. https://doi.org/10.1098/rspb.1985.0042
"""
function cascademodel(S::Int64, Co::Float64)

    # Evaluate input
    maxconnectance = ((S^2-S)/2)/(S*S)
    Co >= maxconnectance && throw(ArgumentError("Connectance for $(S) species cannot be larger than $(maxconnectance) "))
    Co <= 0 && throw(ArgumentError("Connectance C must be positive"))
    S <= 0 && throw(ArgumentError("Number of species L must be positive"))

    # Initiate matrix
    A = UnipartiteNetwork(zeros(Bool, (S, S)))

    # For each species randomly asscribe rank e
    e = sort(rand(S), rev=false)

    # Probability for linking two species
    p = 2*Co*S/(S - 1)

    for consumer in  1:S

        # Rank for a consumer
        rank = e[consumer]

        # Get a set of resources smaller than consumer
        potentialresources = findall(x -> x > rank, e)  # indices
        #resourcesranks = η[potentialresources]         # ranks
        #R = length(potentialresources)                 # length

        # Each gets a link with a probability p
        for resource in potentialresources

            # Take the resources and for each of them check a random number
            # if it is smaller than probability of linking two speceis create
            # a link in the A matrix

            # Random number for a potential resource
            rand() < p && (A[consumer, resource] = true)

        end

    end

    return A

end

"""
    cascademodel(N::T) where {T <: UnipartiteNetwork}

Applied to a `UnipartiteNetwork` return its randomized version.

#### References

Cohen, J.E., Newman, C.M., 1985. A stochastic theory of community food webs I.
Models and aggregated data. Proceedings of the Royal Society of London. Series
B. Biological Sciences 224, 421–448. https://doi.org/10.1098/rspb.1985.0042
"""
function cascademodel(N::T) where {T <: UnipartiteNetwork}

    return cascademodel(richness(N), connectance(N))

end

"""
    cascademodel(S::Int64, L::Int64)

Number of links can be specified instead of connectance.

#### References

Cohen, J.E., Newman, C.M., 1985. A stochastic theory of community food webs I.
Models and aggregated data. Proceedings of the Royal Society of London. Series
B. Biological Sciences 224, 421–448. https://doi.org/10.1098/rspb.1985.0042
"""
function cascademodel(S::Int64, L::Int64)
    Co = (L/(S*S))
    return cascademodel(S, Co)

end

"""

    cascademodel(parameters::Tuple)

Parameters tuple can also be provided in the form (S::Int64, Co::Float64)
or (S::Int64, L::Int64).

#### References

Cohen, J.E., Newman, C.M., 1985. A stochastic theory of community food webs I.
Models and aggregated data. Proceedings of the Royal Society of London. Series
B. Biological Sciences 224, 421–448. https://doi.org/10.1098/rspb.1985.0042
"""
function cascademodel(parameters::Tuple)
    return cascademodel(parameters[1], parameters[2])
end


@testset "Accuracy test for Laffine on NonhomogeneousAbsorbing" begin
    regular_grids = [0:4, -1:6]
    irregular_grids = [[-1.0; 0.0; 2.0; 5.0; 9.0], [0.2; 0.3; 0.5; 0.9; 1.6; 2.5]]
    for x̄ in [regular_grids; irregular_grids],
        (L_bc_generator, L_generator, L_affine_generator) in [(L₁₋bc, L₁₋, L₁₋affine), (L₁₊bc, L₁₊, L₁₊affine), (L₂bc, L₂, L₂affine)],
        (S_lower, S_upper) in [(0.0, 1.0), (1.0, 0.0), (1.5, -1.5), (1.2, 1.3), (1.0, 1.0)]
        # Setup
        # RHS function
        f(x) = x^2

        # grid setup
        # x̄ = 0:4
        x = interiornodes(x̄)
        M = length(x)
        Δ_0p = x̄[2] - x̄[1]
        Δ_Mp = x̄[end] - x̄[end-1]

        # boundary condition
        bc = (NonhomogeneousAbsorbing(S_lower), NonhomogeneousAbsorbing(S_upper))

        # corresponding boundary condition matrix
        B = transpose([[1; zeros(M+1)] [zeros(M+1); 1]])
        b = [S_lower; S_upper]

        # operator on the interior, with bc applied
        L_bc = L_bc_generator(x̄, bc)
        # operator on the exterior
        L = L_generator(x̄)

        # solution for interior v from L_bc
        v_bc = L_bc \ (f.(x) + L_affine_generator(x̄, f.(x), bc))
        # solution for interior v from L_bc, using Laffine
        v_bc_Laffine = L_bc \ (f.(x) + Laffine(L, f.(x), bc))
        # solution for extended v from L by stacking up boundary condition matrices
        v̄ = [L; B] \ [f.(x); b]

        # test if extrapolated solutions are identical
        @test extrapolatetoboundary(v_bc, x̄, bc) ≈ v̄
        @test extrapolatetoboundary(v_bc_Laffine, x̄, bc) ≈ v̄

    end
end

@testset "Accuracy test for Laffine on Absorbing on interior nodes" begin
    regular_grids = [0:4, -1:6]
    irregular_grids = [[-1.0; 0.0; 2.0; 5.0; 9.0], [0.2; 0.3; 0.5; 0.9; 1.6; 2.5]]
    for x̄ in [regular_grids; irregular_grids],
        (L_bc_generator, L_generator, L_affine_generator) in [(L₁₋bc, L₁₋, L₁₋affine), (L₁₊bc, L₁₊, L₁₊affine), (L₂bc, L₂, L₂affine)],
        loc in 0:3,
        S_upper in [1.0; 0.0; -1.5; 1.3; 1.0]
        # Setup
        # RHS function
        f(x) = x^2

        # grid setup
        # x̄ = 0:4
        x = interiornodes(x̄)
        M = length(x)
        Δ_0p = x̄[2] - x̄[1]
        Δ_Mp = x̄[end] - x̄[end-1]

        # boundary condition
        bc = (Absorbing(loc), NonhomogeneousAbsorbing(S_upper))

        # corresponding boundary condition matrix
        B = transpose([[1; zeros(M+1)] [zeros(M+1); 1]])
        b = [0; S_upper]

        # operator on the interior, with bc applied
        L_bc = L_bc_generator(x̄, bc)
        # operator on the exterior
        L = L_generator(x̄)
        # apply interior boundary conditions (if needed)
        RHS = f.(x)
        if (bc[1].loc > 0)
            L[1:bc[1].loc,:] = zeros(bc[1].loc,M+2)
            L[1:bc[1].loc,2:(1+bc[1].loc)] = Diagonal(ones(bc[1].loc))
            RHS[1:bc[1].loc] .= zero(eltype(L))
        end

        # solution for interior v from L_bc
        v_bc = L_bc \ (f.(x) + L_affine_generator(x̄, f.(x), bc))
        # solution for interior v from L_bc, using Laffine
        v_bc_Laffine = L_bc \ (f.(x) + Laffine(L, f.(x), bc))
        # solution for extended v from L by stacking up boundary condition matrices
        v̄ = [L; B] \ [RHS; b]

        # test if extrapolated solutions are identical
        @test extrapolatetoboundary(v_bc, x̄, bc) ≈ v̄
        @test extrapolatetoboundary(v_bc_Laffine, x̄, bc) ≈ v̄

    end
end


"""
    exponents = get_exponents(variables, dmax::Int)

Compute the set of all exponents in `variables` variables, of degree up to
`dmax`.
"""
function compute_exponents(variables::Set{Variable}, dmax::Int; compute_conj=false)
    cur_order = Set{Exponent}([Exponent()])
    result = deepcopy(cur_order)
    prev_order = Set{Exponent}()
    for i=1:dmax
        prev_order = deepcopy(cur_order)
        cur_order = Set{Exponent}()
        for var in variables
            if compute_conj
                union!(cur_order, Set([product(conj(var), elt) for elt in prev_order]))
            else
                union!(cur_order, Set([product(Exponent(var), elt) for elt in prev_order]))
            end
        end
        union!(result, cur_order)
    end
    return result
end

"""
    ishomo = is_homogeneous(p, kind)

    Check wether `p`, of kind `:Real` or `:Complex` is homogeneous or not.
"""
function is_homogeneous(p::Polynomial, kind::Symbol)
    ishomo = true
    for (expo, λ) in p
        ishomo = ishomo && is_homogeneous(expo, kind)
    end
    return ishomo
end

"""
    ishomo = is_homogeneous(expo, kind)

    Check wether `expo`, of kind `:Real` or `:Complex` is homogeneous or not.
"""
function is_homogeneous(expo::Exponent, kind::Symbol)
    explsum, conjsum = get_sumdegs(expo)
    if kind == :Real
        (conjsum != 0) && error(LOGGER, "is_homogeneous(): Exponent $expo has conjugated variables for a real hierarchy.")
        return explsum % 2 == 0
    elseif kind == :Complex
        return explsum == conjsum
    else
        error(LOGGER, "is_homogeneous(expo, kind): kind should be either :Real or :Complex ($kind here).")
    end
end

"""
    explsum, conjsum = get_sumdegs(expo)

    Compute `|α|`, `|β|` the sum of the real variables exponents and conjugated variables exponents.
"""
function get_sumdegs(expo::Exponent)
    explsum = conjsum = 0

    state1 = start(expo)
    while !done(expo, state1)
        (i1, state1) = next(expo, state1)
        explsum += last(i1).explvar
        conjsum += last(i1).conjvar
    end

    return explsum, conjsum
end

"""
    cstrtype = get_cstrtype(cstr::Constraint)

    Return a cstraint type among `:eq`, `:ineqhi`, `:ineqlo`, `:ineqdouble`.
"""
function get_cstrtype(cstr::Constraint)
    if cstr.lb == cstr.ub && isfinite(cstr.ub)
        return :eq
    elseif (cstr.lb == -Inf-im*Inf) && (cstr.ub != Inf+im*Inf)
        return :ineqhi
    elseif (cstr.lb != -Inf-im*Inf) && (cstr.ub == Inf+im*Inf)
        return :ineqlo
    elseif (cstr.lb != -Inf-im*Inf) && (cstr.ub != Inf+im*Inf)
        return :ineqdouble
    else
        error(LOGGER, "get_cstrtype(): unknown constraint type.\nConstraint is $cstr")
    end
end

function compute_degree(expo::Exponent)
    expldeg = conjdeg = 0
    for (var, deg) in expo.expo
        expldeg = max(expldeg, deg.explvar)
        conjdeg = max(conjdeg, deg.conjvar)
    end
    return Degree(expldeg, conjdeg)
end

function compute_degree(p::Polynomial)
    expldeg = conjdeg = 0
    for (expo, λ) in p
        for (var, deg) in expo
            expldeg = max(expldeg, deg.explvar)
            conjdeg = max(conjdeg, deg.conjvar)
        end
    end
    return Degree(expldeg, conjdeg)
end

function update_degree!(expo::Exponent)
    updeg = compute_degree(expo)
    expo.degree.explvar = updeg.explvar
    expo.degree.conjvar = updeg.conjvar
end

function update_degree!(p::Polynomial)
    updeg = compute_degree(p)
    p.degree.explvar = updeg.explvar
    p.degree.conjvar = updeg.conjvar
end