## TODO Update Refactor

using Plots
using PortfolioOpt
using PortfolioOpt.TestUtils

include("./examples/literature_examples/example_utils.jl")

############ Read Prices #############
Prices = get_test_data()
numD, numA = size(Prices) # A: Assets    D: Days

############# Calculating returns #####################
returns_series = percentchange(Prices)
returns = values(returns_series)
# risk free asset
rf = fill(3.2e-6, numD) # readcsv(".\\rf.csv")

############ Efficient frontier ###################
t = 100
k_back = 60
Σ, r̄ = mean_variance(returns[(t - k_back):(t - 1), :])

range = -0.01:0.0005:0.025;
len = size(range, 1);
E = zeros(len);
σ = zeros(len);
E_quad = zeros(len);
σ_quad = zeros(len);
E_quad_Rf = zeros(len);
σ_quad_Rf = zeros(len);

x_sharpe = max_sharpe(Σ, r̄, rf[t])
E_sharpe = r̄'x_sharpe
σ_sharpe = sqrt(x_sharpe'Σ * x_sharpe)
iter = 0
for R in range
    global iter += 1
    # analytical
    x = mean_variance_noRf_analytical(Σ, r̄, R)
    # quadratic optimization no risk free asset
    model, w = base_model(numA)
    po_minvar_limitmean_noRf!(model, w, Σ, r̄, R)
    x_quad, obj, r = compute_solution(model, w)
    # quadratic optimization with risk free asset
    model, w = base_model(numA)
    po_minvar_limitmean_Rf!(model, w, Σ, r̄, R, rf[t], 1)
    x_quad_Rf, obj, r = compute_solution(model, w)

    E[iter] = sum(x'r̄)
    σ[iter] = sqrt((x'Σ * x)[1])
    E_quad[iter] = r[1]
    σ_quad[iter] = sqrt(x_quad'Σ * x_quad)
    E_quad_Rf[iter] = R
    σ_quad_Rf[iter] = sqrt(x_quad_Rf'Σ * x_quad_Rf)
end

# plot frontier
plt = plot(
    σ_quad,
    E_quad;
    title="Efficient Frontier Mean-Variance",
    xlabel="σ",
    ylabel="r",
    label="Quadratic",
    legend=:outertopright,
);
plot!(plt, [0; σ_sharpe], [rf[t]; E_sharpe]; label="Sharpe");
plot!(plt, σ, E; label="Analytical");
plot!(plt, σ_quad_Rf, E_quad_Rf; label="Quadratic with RF");
plt


# Extended Kalman filter functions

"""
    predict(filter::ExtendedKalmanFilter, b0::GaussianBelief, u::AbstractVector)

Uses Extended Kalman filter to run prediction step on gaussian belief b0,
given control vector u.
"""
function predict(filter::ExtendedKalmanFilter, b0::GaussianBelief,
                u::AbstractVector{<:Number})

    # Motion update
    μp = predict(filter.d, b0.μ, u)
    F = jacobian(filter.d, b0.μ, u)

    Σp = F * b0.Σ * F' + filter.d.W
    return GaussianBelief(μp, Σp)
end

"""
    measure(filter::ExtendedKalmanFilter, bp::GaussianBelief, y::AbstractVector;
        u::AbstractVector = [false])

Uses Extended Kalman filter to run measurement update on predicted gaussian
belief bp, given measurement vector y. If u is specified and filter.o.D has
been declared, then matrix D will be factored into the y predictions.
"""
function measure(filter::ExtendedKalmanFilter, bp::GaussianBelief, y::AbstractVector{a};
                u::AbstractVector{b} = [false]) where {a<:Number, b<:Number}

    # Measurement update
    yp = measure(filter.o, bp.μ, u)
    H = jacobian(filter.o, bp.μ, u)

    # Kalman Gain
    K = bp.Σ * H' * inv(H * bp.Σ * H' + filter.o.V)

    # Measurement update
    μn = bp.μ + K * (y - yp)
    Σn = (I - K * H) * bp.Σ
    return GaussianBelief(μn, Σn)
end


function gaussconditional(μ::Array{Float64,1}, Σ::Array{Float64,2}, a_idx::Array{Int64,1}, b_idx::Array{Int64,1})

    # make sure the given indices do not overlap
    @assert(isempty(intersect(a_idx, b_idx)))

    Σ_aa = Σ[a_idx, a_idx]

    Σ_ab = Σ[a_idx, b_idx]

    Σ_ba = Σ[b_idx, a_idx]

    Σ_bb = Σ[b_idx, b_idx]

    μ_a  = @view μ[a_idx]

    μ_b  = @view μ[a_idx]

    # L = cholesky(Σ_bb).L
    #
    # Q = L \ Σ_ba
    #
    # μ_a_given_b(x) = μ_a + Σ_ab*(L' \ (L \ (x .- μ_b)))
    #
    # Σ_a_given_b    = PDMat(Σ_aa - Q'*Q)

    μ_a_given_b(x) = μ_a + Σ_ab*(Σ_bb \ (x .- μ_b))

     Σ_a_given_b    = PDMat(Σ_aa - Σ_ab*(Σ_bb\Σ_ba))

    return μ_a_given_b, Σ_a_given_b

end

function MvNormalConditional(g::AbstractMvNormal, a_idx::Array{Int64,1}, b_idx::Array{Int64,1})

    μcond, Σcond = gaussconditional(g.μ, g.Σ.mat, a_idx, b_idx)

    x -> MvNormal(μcond(x), Σcond)

end

function MvNormalPartition(g::AbstractMvNormal, keep_idx::Array{Int64,1})

    μnew = g.μ[keep_idx]

    Σnew = PDMat(g.Σ.mat[keep_idx, keep_idx])

    MvNormal(μnew, Σnew)

end

function testMvNormalConditional()

    # let's do a comparison on an example calculated in two different ways

end


# Load Julia packages (libraries) needed.

using StatisticalRethinking

ProjDir = @__DIR__

# Simulate uncorrelated data

m_5_1 = "
data {
 int < lower = 1 > N; // Sample size
 vector[N] D; // Outcome
 vector[N] A; // Predictor
}

parameters {
 real a; // Intercept
 real bA; // Slope (regression coefficients)
 real < lower = 0 > sigma;    // Error SD
}

model {
  vector[N] mu;               // mu is a vector
  a ~ normal(0, 0.2);         //Priors
  bA ~ normal(0, 0.5);
  sigma ~ exponential(1);
  mu = a + bA * A;
  D ~ normal(mu , sigma);     // Likelihood
}
";

# Define the SampleModel and set the output format to :mcmcchains.

m5_1s = SampleModel("m5.1", m_5_1);

# Input data for cmdstan

ad_data = Dict(
  "N" => size(df, 1),
  "D" => df[!, :D],
  "A" => df[!, :A]
);

# Sample using StanSample

rc = stan_sample(m5_1s, data=ad_data);

if success(rc)

  # Describe the draws

  dfa1 = read_samples(m5_1s; output_format=:dataframe)
  Particles(dfa1)

end


"""
random(a, b [, T = Float64])

Returns a random number between `a` and `b`. The result is converted to the type `T`.
"""
random(a::Real, b::Real, ::Type{T} = Float64) where {T <: Real} = convert(T, (b - a) * rand(Float64) + a)

"""
random(len [, T = Float64])

Returns a random number between `0` and `len`. The result is converted to the type `T`.
"""
random(len::Real, ::Type{T} = Float64) where {T} = random(0, len, T)

"""
        point(v::Vector)

Converts a [`Vector`](@ref) to a [`Point`](@ref) object.
"""
point(v::Vector{<:Real}) = Luxor.Point(v.x, v.y)

"""
        mapr(s::Real, intervals::Matrix)

Return the value `s` linearly mapped from the intervals in `intervals`.

# Examples

```julia
julia> mapr(1, [0 1; 0 10])
10.0

julia> mapr(0.5, [0 1; 0 10])
5.0

julia> mapr(-2, [-5 5; 0 5])
1.5
```
"""
function mapr(s::Real, intervals::Matrix{<:Real})
    a₁, a₂ = intervals[1, :]
    b₁, b₂ = intervals[2, :]
    return b₁ + (s - a₁) * (b₂ - b₁) / (a₂ - a₁)
end

"""
        s ⟶ [a₁ a₂; b₁ b₂]

Return the value `s` linearly mapped from the `a` interval to the `b` interval.

# Examples

```julia
julia> 1 ⟶ [0 1; 0 10]
10.0

julia> 0.5 ⟶ [0 1; 0 10]
5.0

julia> -2 ⟶ [-5 5; 0 5]
1.5
```
"""
⟶(s::Real, intervals::Matrix{<:Real}) = mapr(s, intervals)


#
# Factors Iteration
#
# Iterating over indices, potential, etc.

Base.start(ϕ::Factor) = 1
Base.next(ϕ::Factor, i) = (sub2at(ϕ, ind2sub(ϕ, i)...), i+1)
Base.done(ϕ::Factor, i) = (i == length(ϕ))

"""
    pattern(ϕ::Factor, [dim])

Return the index of the states of `dim`, given its current position in ϕ.
"""
pattern(ϕ::Factor, dim::Symbol) = pattern(ϕ, [dim])

function pattern(ϕ::Factor, dims::Vector{Symbol})
    isempty(dims) && return Int[]

    _check_dims_valid(dims, ϕ)

    inds = indexin(dims, ϕ)
    len = length(ϕ)
    p = Array{Int}(len, length(dims))

    for i in 1:len
        t = ind2sub(ϕ, i)
        p[i, :] = [t[inds]...]
    end

    return p
end

function pattern(ϕ::Factor)
    sz = size(ϕ)

    len = length(ϕ)
    p = Array{Int}(len, ndims(ϕ))

    for i in 1:len
        p[i, :] = [ind2sub(sz, i)...]
    end

    return p
end

"""
    pattern_states(ϕ::Factor, [dim])

Returns the states of `dim`, given its current position in ϕ, for all possible states.
"""
pattern_states(ϕ::Factor, d::Symbol) = getdim(ϕ, d)[pattern(ϕ, d)]

function pattern_states(ϕ::Factor, dims::Vector{Symbol})
    isempty(dims) && return []

    pat = pattern(ϕ, dims)
    dims = getdim(ϕ, dims)

    return hcat([d[pat[:, i]] for (i, d) in enumerate(dims)]...)
end

function pattern_states(ϕ::Factor)
    pat = pattern(ϕ)

    return hcat([d[pat[:, i]] for (i, d) in enumerate(ϕ.dimensions)]...)
end


using Test
using OMEinsum
using OMEinsum: IndexSize, get_size_dict

@testset "index size" begin
    is = IndexSize(('i', 'j'), (2, 3))
    @test is['i'] == 2
    is = is + is
    @test is == IndexSize(('i', 'j', 'i', 'j'), (2, 3, 2, 3))
    @test is['i'] == 2
    a, b, c = randn(2,3), randn(3,4), randn(2,4)
    xs = (a, b, c)
    ixs = (('i', 'j'), ('j', 'k'), ('i', 'k'))
    is = get_size_dict(ixs, xs)
    @test is == IndexSize(('i', 'j', 'j', 'k', 'i', 'k'), (2,3,3,4,2,4))
    @test is['i'] == 2
    @test is['j'] == 3
    @test is['k'] == 4

    @test IndexSize('i'=>2) == IndexSize(('i',), (2,))
    noindex = IndexSize((), ())
    @test noindex isa IndexSize
    @test noindex + is == is
end

@testset "unspecified index sizes" begin
    v = randn(5)
    a = randn(6,6)
    xs = (v, a)
    ixs = (('j',), ('k', 'k'))
    @test einsum(ein"j,kk->j", xs, get_size_dict(ixs, xs)) ≈ ein"j,kk->j"(xs...)
    res = ein"j,kk->iij"(xs...; size_info=IndexSize('i'=>9))
    @test size(res) == (9,9,5)
end


using ParallelAccelerator

@acc function mttkrp(I::Int, J::Int,
                     K::Int, L::Int,
                     B::Array{Float64, 3}, C::Matrix{Float64},
                     D::Matrix{Float64})

    A = Matrix{Float64}(I, J)
    for i = 1:I
        for j = 1:J
            s = 0.0
            for k = 1:K
                t = 0.0
                for l = 1:L
                    t = t + B[i, k, l] * D[l, j]
                end
                s = s + t * C[k, j]
            end
            A[i, j] = s
        end
    end
    return A
end

I, J, K, L = (200, 200, 200, 200)

B = randn(I, K, L)
C = randn(K, J)
D = randn(L, J)

res1 = mttkrp(I, J, K, L, B, C, D)

@inbounds function mttkrp_o(I::Int, J::Int,
                     K::Int, L::Int,
                     B::Array{Float64, 3}, C::Matrix{Float64},
                     D::Matrix{Float64})

    tic()
    A = Matrix{Float64}(I, J)
    for i = 1:I
        for j = 1:J
            s = 0.0
            for k = 1:K
                t = 0.0
                for l = 1:L
                    t = t + B[i, k, l] * D[l, j]
                end
                s = s + t * C[k, j]
            end
            A[i, j] = s
        end
    end
    println("Julia ", toq())
    return A
 end

res2 = mttkrp_o(I, J, K, L, B, C, D)

for i = 1:I
    exit = false
    for j = 1:J
        if res2[i, j] != res1[i, j]
            println("different $(res2[i, j]) $(res1[i,j]) $(res1[j, i])")
            exit = true
            break
        end
    end
    if exit
        break
    end
end

function compare(A, B)
    for i in 1:I
        for j in 1:J
            if A[i, j] != B[i, j]
                println("different $((i, j)) : $(A[i, j]) != $(B[i, j])")
                return false
            end
        end
    end
    return true
end

println(compare(res1, res2))


swapa(a, i, j) = begin a[i], a[j] = a[j], a[i] end

function pd!(a, first, last)
    while (c = 2 * first - 1) < last
        if c < last && a[c] < a[c + 1]
            c += 1
        end
        if a[first] < a[c]
            swapa(a, c, first)
            first = c
        else
            break
        end
    end
end

hfy!(a, n) = (f = div(n, 2); while f >= 1 pd!(a, f, n); f -= 1 end)

heapsort!(a) = (n = length(a); hfy!(a, n); l = n; while l > 1 swapa(a, 1, l); l -= 1; pd!(a, 1, l) end; a)

a = shuffle(collect(1:12))
println("Unsorted: $a")
println("Heap sorted: ", heapsort!(a))


# 1D Grid examples
# ===============
using ExtendableGrids

# ## Interval from vector
#
function interval_from_vector()
    X=collect(0:0.05:1)
    grid=simplexgrid(X)
end
# ![](interval_from_vector.svg)

#
# ##Interval with local refinement
#
function interval_localref()

    XLeft=geomspace(0.0,0.5,0.1, 0.01)
    XRight=geomspace(0.5,1.0,0.01,0.1)
    X=glue(XLeft, XRight)
    grid=simplexgrid(X)
end
# ![](interval_localref.svg)

#
# ## Interval with  multiple regions
#
function interval_multiregion()

    X=collect(0:0.05:1)
    grid=simplexgrid(X)
    cellmask!(grid,[0.0],[0.5],3)
    bfacemask!(grid,[0.5], [0.5],4)
    grid
end
# ![](interval_multiregion.svg)
#
# ## Multiple regions and subgrid
#
function interval_subgrid()
    X=collect(0:0.01:1)
    grid=simplexgrid(X)
    bfacemask!(grid,[0.5],[0.5],3)
    cellmask!(grid,[0.0],[0.25],2)
    cellmask!(grid,[0.20],[0.5],3)
    subgrid(grid,[2,3])
end
# ![](interval_subgrid.svg)


export graph_degopt, get_topo_order_degopt

"""
    (graph,crefs)=graph_degopt(k;T=ComplexF64,input=:A)
    (graph,crefs)=graph_degopt(x,z;input=:A)
    (graph,crefs)=graph_degopt(d::Degopt;input=:A)

Corresponds to the (for a fixed number of multiplications) degree-optimal
polynomial

    B1=A
    B2=(x *I+x *A)(x *I+x *A)
    B3=(x *I+x *A+x*B2)(x *I+x *A+x *B2)
     ..

and

    Out=z*I+z*A1+z*B2+z*B3+z*B4...

The `x`-values are given in the argument `x`, which is a
`Vector{Tuple{Vector,Vector}}`, containing the elements of each sum. The
`z`-vector contains the elements to form the output, and `input` determines the
name of the matrix A above. If the parameter `k` is supplied instead of the
coefficients, all coeffs will be set to one. The general recursion is given in
(9) in the paper referenced below.

**Reference**

1. P. Bader, S. Blanes, and F. Casas, "Computing the matrix exponential
   with an optimized Taylor polynomial approximation", Mathematics, 7(12), 2019.
   DOI: [10.3390/math7121174](https://doi.org/10.3390/math7121174)
"""
function graph_degopt(x, z; input = :A)
    T = promote_type(eltype(eltype(eltype(x))), eltype(z))
    (graph, crefs) = graph_degopt_B(x, T, input = input)

    k = size(x, 1)

    # Add the polynomial in the Bk coeffs

    z_nodes = [:I; input]
    for s = 2:k+1
        push!(z_nodes, Symbol("B$s"))
    end
    key = Symbol("T2k$(k+3)")
    crefs_new = add_sum!(graph, key, z, z_nodes, Symbol("T2k"))
    append!(crefs, crefs_new)

    # Set the output
    empty!(graph.outputs)
    push!(graph.outputs, key)

    return (graph, crefs)
end
function graph_degopt(k; T = ComplexF64, input = :A)
    x = Vector{Tuple{Vector{T},Vector{T}}}(undef, k)
    for i = 1:k
        x[i] = (ones(T, i + 1), ones(T, i + 1))
    end

    z = ones(T, k + 2)

    return graph_degopt(x, z; input = :A)
end
function graph_degopt(degopt::Degopt; input = :A)
    return graph_degopt(degopt.x, degopt.y, input = input)
end

# Normally x::Vector{Tuple{Vector{Number},Vector{Number}}}
# Containing the coefficients in the recursion
function graph_degopt_B(x, T; input = :A)
    k = size(x, 1)

    graph = Compgraph(T)
    useful_syms = [:I, input]

    cref = []
    for s = 2:k+1
        base0 = "B$s"
        base_a = "Ba$s"
        base_b = "Bb$s"

        # First poly
        c_a = x[s-1][1]
        crefs_a = add_sum!(
            graph,
            Symbol(base_a),
            c_a,
            useful_syms,
            Symbol("$(base_a)_"),
        )

        append!(cref, crefs_a)

        # Second poly
        c_b = x[s-1][2]
        crefs_b = add_sum!(
            graph,
            Symbol(base_b),
            c_b,
            useful_syms,
            Symbol("$(base_b)_"),
        )
        append!(cref, crefs_b)

        # Multiply them together
        add_mult!(graph, Symbol(base0), Symbol(base_a), Symbol(base_b))
        push!(useful_syms, Symbol(base0))
        empty!(graph.outputs)
        push!(graph.outputs, Symbol(base0))
    end

    return (graph, cref)
end

"""
    order=get_topo_order_degopt(k)

A special implementation of [`get_topo_order`](@ref) for degree-optimal polynomials
generated with [`graph_degopt`](@ref). The natural order of computation is to compute
row by row.

See also [`get_degopt_crefs`](@ref).
"""
function get_topo_order_degopt(k)
    (x, z) = get_degopt_crefs(k)
    computation_order = Vector{Symbol}(undef, 0)
    for i = 1:k
        for n = 1:2
            for j = 2:(i+1)
                push!(computation_order, x[i][n][j][1])
            end
        end
        push!(computation_order, Symbol("B$(i+1)"))
    end
    for i = 2:(2+k)
        push!(computation_order, z[i][1])
    end
    return computation_order
end


# testing nearal ODE

using DiffEqFlux, OrdinaryDiffEq, Flux, Optim, Plots

u0 = Float32[2.; 0.]
datasize = 30
tspan = (0.0f0,1.5f0)

function trueODEfunc(du,u,p,t)
    true_A = [-0.1 2.0; -2.0 -0.1]
    du .= ((u.^3)'true_A)'
end
t = range(tspan[1],tspan[2],length=datasize)
prob = ODEProblem(trueODEfunc,u0,tspan)
ode_data = Array(solve(prob,Tsit5(),saveat=t))

# dudt2 = FastChain((x,p) -> x.^3,
#             FastDense(2,50,tanh),
#             FastDense(50,2))
# n_ode = NeuralODE(dudt2,tspan,Tsit5(),saveat=t)

dudt2 = Chain(x -> x.^3,
             Dense(2,50,tanh),
             Dense(50,2))
n_ode = NeuralODE(dudt2,tspan,Tsit5(),saveat=t)

function predict_n_ode(p)
  n_ode(u0,p)
end

function loss_n_ode(p)
    pred = predict_n_ode(p)
    loss = sum(abs2,ode_data .- pred)
    loss,pred
end

loss_n_ode(n_ode.p) # n_ode.p stores the initial parameters of the neural ODE

cb = function (p,l,pred;doplot=false) #callback function to observe training
  display(l)
  # plot current prediction against data
  if doplot
    pl = scatter(t,ode_data[1,:],label="data")
    scatter!(pl,t,pred[1,:],label="prediction")
    display(plot(pl))
  end
  return false
end

# Display the ODE with the initial parameter values.
cb(n_ode.p,loss_n_ode(n_ode.p)..., doplot=true)

res1 = DiffEqFlux.sciml_train(loss_n_ode, n_ode.p, ADAM(0.05), cb = cb, maxiters = 300)
cb(res1.minimizer,loss_n_ode(res1.minimizer)...;doplot=true)
res2 = DiffEqFlux.sciml_train(loss_n_ode, res1.minimizer, LBFGS(), cb = cb)
cb(res2.minimizer,loss_n_ode(res2.minimizer)...;doplot=true)

## Try GPU

# u0 = Float32[2.; 0.] |> gpu
# dudt = Chain(Dense(2,50,tanh),Dense(50,2)) |> gpu
#
# # p,re = DiffEqFlux.destructure(dudt)
# p,re = Flux.destructure(dudt)
#
# dudt_(u,p,t) = re(p)(u)
# prob = ODEProblem(ODEfunc, u0,tspan, p)
#
# # Runs on a GPU
# sol = solve(prob,Tsit5(),saveat=0.1)
#
# n_ode = NeuralODE(gpu(dudt2),tspan,Tsit5(),saveat=0.1)

#


n = 8 # rand(1:10)
@test matrixdepot("frank", n) == matrixdepot("frank", Float64, n, 0)

A = zeros(Float64, n, n)
for i = 1:n, j = 1:n
    A[i,j] = i<=j ? n + 1 -j : j == i - 1 ? n - j : zero(Float64)
end
@test matrixdepot("frank", n) ≈ A

println("'frank' passed test...")


if get(ENV, "CMBLENSING_TEST_CUDA", false) == "1"
    using Adapt, CUDA
    CUDA.allowscalar(false)
    maybegpu(x) = adapt(CuArray,x)
    storage = CuArray
    @info "Running tests on $(repr("text/plain", CUDA.device()))"
else
    maybegpu(x) = x
    storage = Array
    @info "Running tests on CPU"
end

##

using CMBLensing
using CMBLensing: @SMatrix, @SVector, AbstractCℓs, basis, Basis,
    LinearInterpolation, Measurement, RK4Solver, ±, typealias, BatchedReal

##

using FileIO
using FFTW
using FiniteDifferences
using LinearAlgebra
using Random
using Random: default_rng
using Requires
using Serialization
using SparseArrays
using Test
using Zygote

##

try
    push!(LOAD_PATH, "@v#.#") # assumes you have CirculantCov in your global environment
    using CirculantCov
catch
    @warn """CirculantCov.jl not found, not testing EquiRect fields.
    Run `pkg> add https://github.com/EthanAnderes/CirculantCov.jl` to add this package.
    """
end

##

macro test_real_gradient(f, x, tol=:(rtol=1e-3))
    esc(:(@test real(gradient($f,$x)[1]) ≈ central_fdm(5,1)($f,$x) $tol))
end

Nsides     = [(8,8), (4,8), (8,4)]
Nsides_big = [(128,128), (64,128), (128,64)]

Random.seed!(1)

has_batched_fft = (FFTW.fftw_provider != "mkl") || (storage != Array)

##

@testset "CMBLensing" begin

##

# basic printing sanity checks, which were super annoying to get right...
# see also: https://discourse.julialang.org/t/dispatching-on-the-result-of-unwrap-unionall-seems-weird/25677

@testset "Printing" begin

    # concrete types:
    for f in [maybegpu(FlatMap(rand(4,4))), maybegpu(FlatQUMap(rand(4,4),rand(4,4)))]
        @test occursin("pixel",sprint(show, MIME("text/plain"), f))
    end

end

##

@testset "Flat" begin

    Ny,Nx = Nside = first(Nsides)

    @testset "Constructors" begin

        @testset "batch=$D" for D in [(), (3,)]

            Ix = maybegpu(rand(Ny,Nx,D...))
            Il = maybegpu(rand(Ny÷2+1,Nx,D...))

            @testset "$(basis(F))" for (F,ks,args,kwargs) in [
                (FlatMap,        (:Ix,),        (Ix,),      ()),
                (FlatFourier,    (:Il,),        (Il,),      (;Ny,)),
                (FlatQUMap,      (:Qx,:Qx),     (Ix,Ix),    ()),
                (FlatQUFourier,  (:Ql,:Ql),     (Il,Il),    (;Ny,)),
                (FlatEBMap,      (:Ex,:Bx),     (Ix,Ix),    ()),
                (FlatEBFourier,  (:El,:Bl),     (Il,Il),    (;Ny,)),
                (FlatIQUMap,     (:Ix,:Qx,:Qx), (Ix,Ix,Ix), ()),
                (FlatIQUFourier, (:Il,:Ql,:Ql), (Il,Il,Il), (;Ny,)),
                (FlatIEBMap,     (:Ix,:Ex,:Bx), (Ix,Ix,Ix), ()),
                (FlatIEBFourier, (:Il,:El,:Bl), (Il,Il,Il), (;Ny,)),
            ]
                local f
                @test (f = F(args...; kwargs...)) isa F
                @test @inferred(F(getproperty.(Ref(f),ks)..., f.metadata)) == f
                @test (io=IOBuffer(); serialize(io,f); seekstart(io); deserialize(io) == f)
                @test (save(".test_field.jld2", "f", f); load(".test_field.jld2", "f") == f)
                @test_throws ErrorException F(args..., ProjLambert(Nx=Nx+1, Ny=Ny+1))

            end

        end

        rm(".test_field.jld2", force=true)

    end

    @testset "Basis conversions" begin
        @testset "$(typealias(Bin)) → $(typealias(Bout))" for (f,Bin,Bout) in [
            (f,Bin,Bout)
            for (f,Bs) in [
                (FlatMap(rand(Nside...)), (Map,Fourier)),
                (FlatQUMap(rand(Nside...),rand(Nside...)), (QUMap,QUFourier,EBMap,EBFourier))
            ]
            for Bin in Bs
            for Bout in Bs
        ]

            @test basis(@inferred(Bout(Bin(f)))) == Bout
            @test Bin(Bout(Bin(f))) ≈ f

        end
    end

end

##

@testset "Algebra" begin

    @testset "Nside = $Nside" for Nside in Nsides

        fs = ((B0,f0),(B2,f2),(Bt,ft)) = [
            (Fourier,    maybegpu(FlatMap(rand(Nside...)))),
            (EBFourier,  maybegpu(FlatQUMap(rand(Nside...),rand(Nside...)))),
            (Fourier,    maybegpu(FieldTuple(FlatMap(rand(Nside...)),FlatMap(rand(Nside...)))))
            # inference currently broken for this case:
            # (IEBFourier, maybegpu(FlatIQUMap(rand(Nside...),rand(Nside...),rand(Nside...)))),
        ]
        # MKL doesnt seem to support batched FFTs, not that theyre really useful on CPU
        has_batched_fft && append!(fs, [
            (Fourier,    maybegpu(FlatMap(rand(Nside...,2)))),
            (EBFourier,  maybegpu(FlatQUMap(rand(Nside...,2),rand(Nside...,2)))),
        ])

        @testset "f :: $(typeof(f)) " for (B,f) in fs

            local Ðf, Ðv, g, H

            @test similar(f) isa typeof(f)
            @test zero(f) isa typeof(f)
            @test similar(f,Float32) isa Field
            @test eltype(similar(f,Float32)) == Float32

            # used in lots of tests
            @test f ≈ f
            @test !(f ≈ 2f)

            # broadcasting
            @test (@inferred f .+ f) isa typeof(f)
            @test (@inferred f .+ Float32.(f)) isa typeof(f)

            # promotion
            @test (@inferred f + B(f)) isa typeof(f)
            @test (@inferred f + B(Float32.(f))) isa typeof(f)

            # Diagonal broadcasting
            @test (@inferred Diagonal(f) .* Diagonal(f) .* Diagonal(f)) isa typeof(Diagonal(f))

            # inverses
            @test (@inferred pinv(Diagonal(f))) isa Diagonal{<:Any,<:typeof(f)}
            @test_throws Exception inv(Diagonal(0*f))

            # trace
            @test all(tr(Diagonal(f)' * Diagonal(f)) ≈ f'f)
            @test all(tr(Diagonal(f) * Diagonal(f)') ≈ f'f)
            @test_broken all(tr(f*f') ≈ f'f) # broken by (intentionally) removing OuterProdOp

            # Field dot products
            D = Diagonal(f)
            @test (@inferred f' * f) isa Real
            @test (@inferred f' * B(f)) isa Real
            @test (@inferred f' * D * f) isa Real
            @test sum(f, dims=:) ≈ sum(f[:])

            # Explicit vs. lazy DiagOp algebra
            @test (Diagonal(Ð(f)) + Diagonal(Ð(f))) isa DiagOp{<:Field{basis(Ð(f))}}
            @test (Diagonal(Ł(f)) + Diagonal(Ð(f))) isa LazyBinaryOp
            @test (Diagonal(Ł(f)) + Diagonal(Ł(f))) isa DiagOp{<:Field{basis(Ł(f))}}

            if !(f isa FieldTuple)

                # gradients
                @test (Ðf = @inferred ∇[1]*f) isa Field
                @test all(∇[1]'*f ≈ -∇[1]*f)
                @test all(-∇[1]'*f ≈ ∇[1]*f)
                @test (@inferred mul!(Ðf,∇[1],Ð(f))) isa Field
                @test (Ðv = @inferred ∇*f) isa FieldVector
                @test (@inferred mul!(Ðv,∇,Ð(f))) isa FieldVector
                @test ((g,H) = map(Ł, (@inferred gradhess(f)))) isa NamedTuple{<:Any, <:Tuple{FieldVector, FieldMatrix}}

                # FieldVector dot product
                @test (@inferred Diagonal.(g)' * g) isa typeof(g[1])
                @test (@inferred mul!(similar(g[1]), Diagonal.(g)', g)) isa typeof(g[1])

                # FieldMatrix-FieldVector product
                @test (@inferred Diagonal.(H) * g) isa FieldVector
                @test (@inferred Diagonal.(H) * Diagonal.(g)) isa FieldOrOpVector
                @test (@inferred mul!(Diagonal.(similar.(g)), Diagonal.(H), Diagonal.(g))) isa FieldOrOpVector

            end

        end

        # # tuple adjoints
        # f0b = identity.(batch(f0, batch_length(f)))
        # v = similar.(@SVector[f0b, f0b])
        # @test (@inferred mul!(f0b, spin_adjoint(f), f)) isa Field{<:Any,S0}
        # @test (@inferred mul!(v, spin_adjoint(f), @SVector[f,f])) isa FieldVector{<:Field{<:Any,S0}}

        # mixed-spin
        @test (@inferred f0 .* f2) isa typeof(f2)

        # matrix type promotion
        @test (@inferred FlatMap(rand(Float64,2,2)) .+ FlatMap(view(rand(Float32,2,2),:,:))) isa FlatMap{<:Any,Float64,Matrix{Float64}}

    end

end

##

@testset "Log/Trace" begin

    @testset "logdet(Diagonal(::Map))" begin
        @test logdet(Diagonal(FlatMap([1 -2; 3 -4])))                                           ≈  log(24)
        @test logdet(Diagonal(FlatQUMap([1 -2; 3 -4], [1 -2; 3 -4])))                           ≈ 2log(24)
        @test logdet(Diagonal(FlatIQUMap([1 -2; 3 -4], [1 -2; 3 -4], [1 -2; 3 -4])))            ≈ 3log(24)
        @test logdet(Diagonal(FieldTuple(FlatMap([1 -2; 3 -4]), FlatMap([1 -2; 3 -4]))))        ≈ 2log(24)
        @test all(logdet(Diagonal(FlatMap(cat([1 -2; 3 -4],[1 -2; 3 -4],dims=3))))::BatchedReal ≈  log(24))
    end

    @testset "logdet(Diagonal(::Fourier)) Nside=$Nside" for Nside in Nsides_big
        x = rand(Nside...)
        @test logdet(Diagonal(Fourier(FlatMap(x))))                                                   ≈ real( logdet(Diagonal(fft(x)[:])))
        @test logdet(Diagonal(QUFourier(FlatQUMap(x,x))))                                             ≈ real(2logdet(Diagonal(fft(x)[:])))
        @test logdet(Diagonal(IQUFourier(FlatIQUMap(x,x,x))))                                         ≈ real(3logdet(Diagonal(fft(x)[:])))
        @test logdet(Diagonal(FieldTuple(Fourier(FlatMap(x)), Fourier(FlatMap(x)))))                  ≈ real(2logdet(Diagonal(fft(x)[:])))
        has_batched_fft && @test all(logdet(Diagonal(Fourier(FlatMap(cat(x,x,dims=3)))))::BatchedReal ≈ real( logdet(Diagonal(fft(x)[:]))))
    end

    @testset "tr(Diagonal(::Map))" begin
        @test tr(Diagonal(FlatMap([1 -2; 3 -4])))                                           ≈  -2
        @test tr(Diagonal(FlatQUMap([1 -2; 3 -4], [1 -2; 3 -4])))                           ≈  -4
        @test tr(Diagonal(FlatIQUMap([1 -2; 3 -4], [1 -2; 3 -4], [1 -2; 3 -4])))            ≈  -6
        @test tr(Diagonal(FieldTuple(FlatMap([1 -2; 3 -4]), FlatMap([1 -2; 3 -4]))))        ≈  -4
        @test all(tr(Diagonal(FlatMap(cat([1 -2; 3 -4],[1 -2; 3 -4],dims=3))))::BatchedReal ≈  -2)
    end

    @testset "tr(Diagonal(::Fourier)) Nside=$Nside" for Nside in Nsides_big
        x = rand(Nside...)
        @test tr(Diagonal(Fourier(FlatMap(x))))                                                   ≈  tr(Diagonal(fft(x)[:]))
        @test tr(Diagonal(QUFourier(FlatQUMap(x,x))))                                             ≈ 2tr(Diagonal(fft(x)[:]))
        @test tr(Diagonal(IQUFourier(FlatIQUMap(x,x,x))))                                         ≈ 3tr(Diagonal(fft(x)[:]))
        @test tr(Diagonal(FieldTuple(Fourier(FlatMap(x)), Fourier(FlatMap(x)))))                  ≈ 2tr(Diagonal(fft(x)[:]))
        has_batched_fft && @test all(tr(Diagonal(Fourier(FlatMap(cat(x,x,dims=3)))))::BatchedReal ≈ real(tr(Diagonal(fft(x)[:]))))
    end

end

##

@testset "FlatS2" begin
    @testset "Nside = $Nside" for Nside in Nsides
        C = maybegpu(Diagonal(EBFourier(FlatEBMap(rand(Nside...), rand(Nside...)))))
        f = maybegpu(FlatQUMap(rand(Nside...), rand(Nside...)))
        @test C*f ≈ FlatQUFourier(C[:QQ]*f[:Q]+C[:QU]*f[:U], C[:UU]*f[:U]+C[:UQ]*f[:Q])
    end
end

##

@testset "BatchedReal" begin

    @testset "Nside = $Nside" for Nside in Nsides

        r  = 1.
        rb = batch([1.,2])

        @testset "f :: $(typeof(f))" for (f,fb) in [
            (maybegpu(FlatMap(rand(Nside...))),                  maybegpu(FlatMap(rand(Nside...,2)))),
            (maybegpu(FlatQUMap(rand(Nside...),rand(Nside...))), maybegpu(FlatQUMap(rand(Nside...,2),rand(Nside...,2))))
        ]

            @test @inferred(r * f)  == f
            @test @inferred(r * fb) == fb
            @test unbatch(@inferred(rb * f)) == [f, 2f]
            @test unbatch(@inferred(rb * fb)) == [batch_index(fb,1), 2batch_index(fb,2)]

        end

    end

end

##

@testset "Misc" begin

    @testset "Nside = $Nside" for Nside in Nsides

        f = maybegpu(FlatMap(rand(Nside...)))

        @test                  (MidPass(100,200) .* Diagonal(Fourier(f))) isa Diagonal
        @test_throws Exception  MidPass(100,200) .* Diagonal(        f)

    end

end

##

@testset "Cℓs" begin

    @test InterpolatedCℓs(1:100, rand(100))      isa AbstractCℓs{Float64}
    @test InterpolatedCℓs(1:100, rand(100) .± 1) isa AbstractCℓs{Measurement{Float64}}
    @test (InterpolatedCℓs(1:100, 1:100) * ℓ²)[50] == 50^3

end

##

@testset "ParamDependentOp" begin

    D = Diagonal(maybegpu(FlatMap(rand(4,4))))

    @test ParamDependentOp((;x=1, y=1)->x*y*D)() ≈ D
    @test ParamDependentOp((;x=1, y=1)->x*y*D)(z=2) ≈ D
    @test ParamDependentOp((;x=1, y=1)->x*y*D)(x=2) ≈ 2D
    @test ParamDependentOp((;x=1, y=1)->x*y*D)((x=2,y=2)) ≈ 4D # tuple calling form
    @test_throws MethodError ParamDependentOp((;x=1, y=1)->x*y*D)(2) # only Tuple unnamed arg is OK

end

##

@testset "Chains" begin

    chains = CMBLensing.wrap_chains([
        [Dict(:i=>1, :b=>2), Dict(:i=>2       ), Dict(:i=>3, :b=>2)],
        [Dict(:i=>1, :b=>3), Dict(:i=>2, :b=>3), Dict(:i=>3, :b=>3)],
    ])

    # basic
    @test chains[1, 1, :i] == 1
    @test chains[:, 1, :i] == [1,1]
    @test chains[:, :, :i] == [[1, 2, 3], [1, 2, 3]]

    # slices
    @test chains[1, 1:2, :i] == [1, 2]
    @test chains[:, 1:2, :i] == [[1,2], [1,2]]

    # implied : in first dims
    @test chains[:i] == [[1,2,3],[1,2,3]]
    @test chains[1,:i] == [1,2,3]

    # missing
    @test all(chains[1,:b] .=== [2, missing, 2])

end;

##

@testset "Zygote" begin

    @testset "Nside = $Nside" for Nside in Nsides

        @testset "$FMap" for (FMap, FFourier, Npol) in [
            (FlatMap,   FlatFourier,   1),
            (FlatQUMap, FlatQUFourier, 2)
        ]

            Ny,Nx = Nside
            Ixs = collect(maybegpu(rand(Nside...)) for i=1:Npol)
            Ils = rfft.(Ixs)
            f,g,h = @repeated(maybegpu(FMap(Ixs...)),3)
            v = @SVector[f,f]
            D = Diagonal(f)
            A = 2

            @testset "Fields" begin

                @testset "sum" begin
                    @test ((δ = gradient(f -> sum(Map(f)),     Map(f))[1]); basis(δ)==basis(Map(f))     && δ ≈ one(Map(f)))
                    @test ((δ = gradient(f -> sum(Map(f)), Fourier(f))[1]); basis(δ)==basis(Fourier(f)) && δ ≈ Fourier(one(Map(f))))
                end

                @testset "B1=$B1, B2=$B2, B3=$B3" for B1=[Map,Fourier], B2=[Map,Fourier], B3=[Map,Fourier]

                    @test gradient(f -> dot(B1(f),B2(f)), B3(f))[1] ≈ 2f
                    @test gradient(f -> norm(B1(f)), B3(f))[1] ≈ f/norm(f)
                    @test gradient(f -> B1(f)' * B2(g), B3(f))[1] ≈ g
                    @test gradient(f -> sum(Diagonal(Map(f)) * B2(g)), B3(f))[1] ≈ g
                    @test gradient(f -> sum(Diagonal(Map(∇[1]*f)) * B2(g)), B3(f))[1] ≈ ∇[1]'*g
                    @test gradient(f -> B1(f)'*(D\B2(f)), B3(f))[1] ≈ D\f + D'\f
                    @test gradient(f -> (B1(f)'/D)*B2(f), B3(f))[1] ≈ D\f + D'\f
                    @test gradient(f -> B1(f)'*(D*B2(f)), B3(f))[1] ≈ D*f + D'*f
                    @test gradient(f -> (B1(f)'*D)*B2(f), B3(f))[1] ≈ D*f + D'*f
                    if B2==Map
                        @test gradient(f -> B1(f)'*Diagonal(B2(f))*f, B3(f))[1] ≈ @. 3*$B2(f)^2
                    else
                        @test_broken gradient(f -> B1(f)'*Diagonal(B2(f))*f, B3(f))[1] ≈ @. 3*$B2(f)^2
                    end

                end

                @testset "Broadcasting" begin
                    @test gradient(f -> sum(@. f*f + 2*f + 1), f)[1] ≈ 2*f+2
                    @test gradient(f -> sum(@. f^2 + 2*f + 1), f)[1] ≈ 2*f+2
                end

            end

            @testset "FieldVectors" begin

                @test gradient(f -> Map(∇[1]*f)' *     Map(v[1]) + Map(∇[2]*f)' *     Map(v[2]), f)[1] ≈ ∇' * v
                @test gradient(f -> Map(∇[1]*f)' * Fourier(v[1]) + Map(∇[2]*f)' * Fourier(v[2]), f)[1] ≈ ∇' * v
                @test gradient(f -> sum(Diagonal(Map(∇[1]*f)) * v[1] + Diagonal(Map(∇[2]*f)) * v[2]), f)[1] ≈ ∇' * v

                @test gradient(f -> @SVector[f,f]' * Map.(@SVector[g,g]), f)[1] ≈ 2g
                @test gradient(f -> @SVector[f,f]' * Fourier.(@SVector[g,g]), f)[1] ≈ 2g

                @test gradient(f -> sum(sum(@SVector[f,f])),                            f)[1] ≈ 2*one(f)
                @test gradient(f -> sum(sum(@SVector[f,f]      .+ @SVector[f,f])),      f)[1] ≈ 4*one(f)
                @test gradient(f -> sum(sum(@SMatrix[f f; f f] .+ @SMatrix[f f; f f])), f)[1] ≈ 8*one(f)

                # these were broken by (intentionally) removing OuterProdOp .they
                # seem like a fairly unusual, but keeping them here as broken for now...
                @test_broken gradient(f -> sum(Diagonal.(Map.(∇*f))' * Fourier.(v)), f)[1] ≈ ∇' * v
                @test_broken gradient(f -> sum(Diagonal.(Map.(∇*f))' * Map.(v)), f)[1] ≈ ∇' * v
                @test_broken gradient(f -> sum(sum(Diagonal.(@SMatrix[f f; f f]) * @SVector[f,f])), f)[1] ≈ 8*f

            end

            if f isa FlatS0

                @testset "logdet" begin
                    @test gradient(x->logdet(x*Diagonal(Map(f))),     1)[1] ≈ size(Map(f))[1]
                    @test gradient(x->logdet(x*Diagonal(Fourier(f))), 1)[1] ≈ size(Map(f))[1]
                    L = ParamDependentOp((;x=1)->x*Diagonal(Fourier(f)))
                    @test gradient(x->logdet(L(x=x)), 1)[1] ≈ size(Map(f))[1]
                end

                @test gradient(x -> norm(x*Fourier(f)), 1)[1] ≈ norm(f)
                @test gradient(x -> norm(x*Map(f)), 1)[1] ≈ norm(f)

                L₀ = Diagonal(Map(f))
                @test gradient(x -> norm((x*L₀)*f), 1)[1] ≈ norm(L₀*f)
                L₀ = Diagonal(Fourier(f))
                @test gradient(x -> norm((x*L₀)*f), 1)[1] ≈ norm(L₀*f)

                @test gradient(x -> norm((x*Diagonal(Map(f)))*f), 1)[1] ≈ norm(Diagonal(Map(f))*f)
                @test gradient(x -> norm((x*Diagonal(Fourier(f)))*f), 1)[1] ≈ norm(Diagonal(Fourier(f))*f)

            end

            @testset "Array Bailout" begin

                @testset "Fourier" begin

                    @test_real_gradient(A -> logdet(Diagonal(FFourier((A.*Ils)...; Ny))), A)
                    @test_real_gradient(A -> logdet(Diagonal(A*FFourier(Ils...; Ny))), A)
                    @test_real_gradient(A -> logdet(A*Diagonal(FFourier(Ils...; Ny))), A)

                    @test_real_gradient(A -> f' * (Diagonal(FFourier((A.*Ils)...; Ny))) * f, A)
                    @test_real_gradient(A -> f' * (Diagonal(A*FFourier(Ils...; Ny))) * f, A)
                    @test_real_gradient(A -> f' * (A*Diagonal(FFourier(Ils...; Ny))) * f, A)

                    @test_real_gradient(A -> norm(FFourier((A.*Ils)...; Ny)), A)
                    @test_real_gradient(A -> norm(A*FFourier(Ils...; Ny)), A)

                end

                @testset "Map" begin

                    @test_real_gradient(A -> logdet(Diagonal(FMap((A.*Ixs)...))), A)
                    @test_real_gradient(A -> logdet(Diagonal(A*FMap(Ixs...))), A)
                    @test_real_gradient(A -> logdet(A*Diagonal(FMap(Ixs...))), A)

                    @test_real_gradient(A -> f' * (Diagonal(FMap((A.*Ixs)...))) * f, A)
                    @test_real_gradient(A -> f' * (Diagonal(A*FMap(Ixs...))) * f, A)
                    @test_real_gradient(A -> f' * (A*Diagonal(FMap(Ixs...))) * f, A)

                    @test_real_gradient(A -> norm(FMap((A.*Ixs)...)), A)
                    @test_real_gradient(A -> norm(A*FMap(Ixs...)), A)

                end

            end

        end

    end

    @testset "LinearInterpolation" begin
        @test gradient(x->LinearInterpolation([1,2,3],[1,2,3])(x), 2)[1] == 1
        @test gradient(x->LinearInterpolation([1,2,3],[1,x,3])(2), 2)[1] == 1
        @test gradient(x->LinearInterpolation([1,x,3],[1,2,3])(2), 2)[1] == -1
    end

end

##

@testset "Lensing" begin

    local f, ϕ, Lϕ
    Cℓ = camb().unlensed_total

    @testset "$L" for (L,atol) in [(BilinearLens,300), (LenseFlow,0.1)]

        @testset "Nside = ($Ny,$Nx)" for (Ny,Nx) in Nsides_big

            @testset "T :: $T" for T in (Float32, Float64)

                proj = ProjLambert(;Ny,Nx,T,storage)

                Cϕ = maybegpu(Cℓ_to_Cov(:I, proj, Cℓ.ϕϕ))
                @test (ϕ = @inferred simulate(Cϕ)) isa FlatS0

                ## S0
                Cf = maybegpu(Cℓ_to_Cov(:I, proj, Cℓ.TT))
                @test (f = @inferred simulate(Cf)) isa FlatS0
                @test (Lϕ = cache(LenseFlow(ϕ),f)) isa CachedLenseFlow
                @test (@inferred Lϕ*f) isa FlatS0
                # adjoints
                f,g = simulate(Cf),simulate(Cf)
                @test f' * (Lϕ * g) ≈ (f' * Lϕ) * g
                # gradients
                δf, δϕ = simulate(Cf), simulate(Cϕ)
                @test_real_gradient(α -> norm(L(ϕ+α*δϕ)*(f+α*δf)), 0, atol=atol)

                # S2 lensing
                Cf = maybegpu(Cℓ_to_Cov(:P, proj, Cℓ.EE, Cℓ.BB))
                @test (f = @inferred simulate(Cf)) isa FlatS2
                @test (Lϕ = cache(LenseFlow(ϕ),f)) isa CachedLenseFlow
                @test (@inferred Lϕ*f) isa FlatS2
                # adjoints
                f,g = simulate(Cf),simulate(Cf)
                @test f' * (Lϕ * g) ≈ (f' * Lϕ) * g
                # gradients
                δf, δϕ = simulate(Cf), simulate(Cϕ)
                @test_real_gradient(α -> norm(L(ϕ+α*δϕ)*(f+α*δf)), 0, atol=atol)

            end

        end

    end

end

##

@testset "Posterior" begin

    Cℓ = camb()
    L = LenseFlow{RK4Solver{7}}
    T = Float64

    @testset "Nside = $Nside" for Nside in Nsides_big

        @testset "pol = $pol" for pol in (:I,:P)

            @unpack f,f̃,ϕ,ds,ds₀ = load_sim(
                Cℓ       = Cℓ,
                θpix     = 3,
                Nside    = Nside,
                T        = T,
                beamFWHM = 3,
                pol      = pol,
                storage  = storage,
                rng      = default_rng(),
                pixel_mask_kwargs = (edge_padding_deg=1,)
            )
            @unpack Cf,Cϕ = ds₀
            f°,ϕ° = mix(f,ϕ,ds)

            @test lnP(0,f,ϕ,ds) ≈ lnP(1,    f̃,  ϕ , ds) rtol=1e-4
            @test lnP(0,f,ϕ,ds) ≈ lnP(:mix, f°, ϕ°, ds) rtol=1e-4

            δf,δϕ = simulate(Cf, rng=default_rng()), simulate(Cϕ, rng=default_rng())

            @test_real_gradient(α->lnP(0,    f +α*δf, ϕ +α*δϕ, ds), 0, atol=0.5)
            @test_real_gradient(α->lnP(1,    f̃ +α*δf, ϕ +α*δϕ, ds), 0, atol=130)
            @test_real_gradient(α->lnP(:mix, f°+α*δf, ϕ°+α*δϕ, ds), 0, atol=0.5)

        end

    end

end

##

@require CirculantCov="edf8e0bb-e88b-4581-a03e-dda99a63c493" begin

@testset "EquiRect" begin

    θspan  = (π/2 .- deg2rad.((-40,-70)))
    φspan  = deg2rad.((-60, 60))
    φspan′ = deg2rad.((-50, 50))
    Cℓ     = camb()
    rtol   = 1e-5

    @testset "T = $T" for T in (Float32, Float64)

        @testset "Nside = $Nside" for Nside in [(32,64)]

            Ny, Nx = Nside

            # non-periodic
            proj′ = ProjEquiRect(;Ny, Nx, T, θspan, φspan=φspan′)
            @test (proj′.Ny == length(proj′.θ) == Ny) && (proj′.Nx == length(proj′.φ) == Nx)

            # Make a linear list `θedges[1], θ[1], θedges[2], θ[2], ..., θedges[n], θ[n], θedges[n+1]`
            # and test that it is strictly increasing.
            ∂θedges = vcat(vcat(proj′.θedges[1:end-1]', proj′.θ')[:],  proj′.θedges[end])
            @test all(diff(∂θedges) .> 0)

            proj′Δφpix   = rem2pi(proj′.φ[2]-proj′.φ[1],   RoundDown)
            proj′Δφspan  = rem2pi(proj′.φ[end]-proj′.φ[1], RoundDown) + proj′Δφpix
            inputΔφspan = φspan′ |> x->rem2pi(x[end]-x[1], RoundDown)
            @test proj′Δφspan ≈ inputΔφspan

            # with integer fraction of 2π (and Float64)
            proj = ProjEquiRect(;Ny, Nx, T, θspan, φspan)
            P = typeof(proj)
            f0   = EquiRectMap(randn(T, Nside...), proj)
            f2   = EquiRectQUMap(randn(T, Nside...), randn(T, Nside...), proj)

            @test f0 isa EquiRectMap
            @test f2 isa EquiRectQUMap

            # transform
            @test Map(AzFourier(f0)) ≈ f0
            @test QUMap(QUAzFourier(f2)) ≈ f2
            @test real(eltype(Map(AzFourier(f0)))) == T
            @test real(eltype(QUMap(QUAzFourier(f2)))) == T

            # transform (testing equality independent of dot)
            @test AzFourier(f0)[:Ix]   ≈ f0[:Ix]
            @test QUAzFourier(f2)[:Px] ≈ f2[:Px]
            @test Map(f0)[:Il]   ≈ f0[:Il]
            @test QUMap(f2)[:Pl] ≈ f2[:Pl]

            # dot product independent of basis
            @test dot(f0,f0) ≈ dot(AzFourier(f0), AzFourier(f0))
            @test dot(f2,f2) ≈ dot(QUAzFourier(f2), QUAzFourier(f2))

            # creating block-diagonal covariance operators
            Cf0 = Cℓ_to_Cov(:I, f0.proj, Cℓ.total.TT)
            Cf2 = Cℓ_to_Cov(:P, f2.proj, Cℓ.total.EE, Cℓ.total.BB)
            Bf0 = CMBLensing.Cℓ_to_Beam(:I, f0.proj, Cℓ.total.TT)
            Bf2 = CMBLensing.Cℓ_to_Beam(:P, f2.proj, Cℓ.total.TT)
            @test Cf0 isa BlockDiagEquiRect{AzFourier,T}
            @test Cf2 isa BlockDiagEquiRect{QUAzFourier,T}
            @test Bf0 isa BlockDiagEquiRect{AzFourier,T}
            @test Bf2 isa BlockDiagEquiRect{QUAzFourier,T}

            # sqrt
            @test (sqrt(Cf0) * sqrt(Cf0) * f0) ≈ (Cf0 * f0) rtol=rtol
            @test (sqrt(Cf2) * sqrt(Cf2) * f2) ≈ (Cf2 * f2) rtol=rtol

            # simulation
            @test real(eltype(simulate(Cf0) :: EquiRectS0)) == T
            @test real(eltype(simulate(Cf2) :: EquiRectS2)) == T

            # pinv
            @test (pinv(Cf0) * Cf0 * f0) ≈ f0 rtol=rtol
            @test (pinv(Cf2) * Cf2 * f2) ≈ f2 rtol=rtol

            @test (Cf0 \ Cf0 * f0) ≈ f0 rtol=rtol
            @test (Cf2 \ Cf2 * f2) ≈ f2 rtol=rtol

            @test (Cf0 / Cf0 * f0) ≈ f0 rtol=rtol
            @test (Cf2 / Cf2 * f2) ≈ f2 rtol=rtol

            # some operator algebra on ops
            @test (Cf0 + Cf0) * f0 ≈ Cf0 * (2 * f0) ≈ (2 * Cf0) * f0
            @test (Cf2 + Cf2) * f2 ≈ Cf2 * (2 * f2) ≈ (2 * Cf2) * f2

            # logdet
            @test logdet(Cf0) ≈ logabsdet(Cf0)[1]
            @test logdet(Cf2) ≈ logabsdet(Cf2)[1]

            # adjoint
            g0 = simulate(Cf0)
            g2 = simulate(Cf2)
            @test f0' * (Cf0 * g0) ≈ (f0' * Cf0) * g0 rtol=rtol
            @test f2' * (Cf2 * g2) ≈ (f2' * Cf2) * g2 rtol=rtol

            # Test for correct Fourier symmetry in monopole and nyquist f2
            let f2kk = f2[:Pl], f2xx = f2[:Px]
                v = f2kk[1:end÷2,1]
                w = f2kk[end÷2+1:end,1]
                @test v ≈ conj.(w)
                if iseven(size(f2xx,2))
                    v′ = f2kk[1:end÷2,end]
                    w′ = f2kk[end÷2+1:end,end]
                    @test v′ ≈ conj.(w′)
                end
            end

            # gradients w.r.t. fields
            @test gradient(f -> dot(f,f), f0)[1] ≈ 2*f0
            @test gradient(f -> dot(f,f), f2)[1] ≈ 2*f2
            @test gradient(f0 -> f0' * (pinv(Cf0) * f0), f0)[1] ≈ (2 * pinv(Cf0) * f0)
            @test gradient(f2 -> f2' * (pinv(Cf2) * f2), f2)[1] ≈ (2 * pinv(Cf2) * f2)

            # gradients through entries of BlockDiagEquiRect
            @test gradient(α -> logabsdet(α * Cf0)[1], 1)[1] ≈ size(Cf0,1)
            @test gradient(α -> logabsdet(α * Cf2)[1], 1)[1] ≈ size(Cf2,1)
            @test_broken gradient(α -> f0' * (pinv(α * Cf0) * f0), 1)[1] isa Real
            @test_broken gradient(α -> f2' * (pinv(α * Cf2) * f2), 1)[1] isa Real

        end

    end

end

end

##

end


#=
    Results should be as sampling from a uniform distribution in [0,1).
    Written to do that rather well, rather than to do that very fast.
=#

function rand1(::Type{ArbFloat{P}}) where {P}
    i = rand(Int128)
    while i ==  typemin(Int128) || abs(i) >= (typemax(Int128)-1)
       i = rand(Int128)
    end
    bf = BigFloat(abs(i)) / BigFloat(typemax(Int128))
    return ArbFloat{P}(bf)
end

function rand(::Type{ArbFloat{P}}) where {P}
    n = cld(P,128)
    rs = zeros(ArbFloat{P}, n)
    for i in 1:n
        rs[i] = rand1(ArbFloat{P})
    end
    r = reduce(*, one(ArbFloat{P}), rs)
    if n > 1
       if n == 2
          r = sqrt(r)
       else
          r = r^(1/n)
       end
    end
    return midpoint(r)
end

function rand(::Type{ArbFloat{P}}, N::Int) where {P}
    n = max(1, N)
    rs = zeros(ArbFloat{P}, n)
    for i in 1:n
        rs[i] = rand(ArbFloat{P})
    end
    return rs
end

function randn(::Type{ArbFloat{P}}) where {P}
    n = cld(P,64)
    rs = zeros(ArbFloat{P}, n)
    for i in 1:n
        rs[i] = randn(Float64)
    end
    r = reduce(*, one(ArbFloat{P}), rs)
    return n==1 ? r : (n==2 ? sqrt(r) : r^(1/n))
end

function randn(::Type{ArbFloat{P}}, N::Int) where {P}
    n = max(1,N)
    rs = zeros(ArbFloat{P}, n)
    for i in 1:n
        rs[i] = randn(ArbFloat{P})
     end
     return rs
end


### A Pluto.jl notebook ###
# v0.12.10

using Markdown
using InteractiveUtils

# ╔═╡ e077ab16-22e5-11eb-15ed-5f0654bf7328
using Pkg, DrWatson

# ╔═╡ e75a9a06-22e5-11eb-2e36-8d4f62b830ed
begin
	@quickactivate "StatisticsWithJuliaPlutoNotebooks"
	using DataFrames, CSV, Statistics, Dates, Plots, Measures
end

# ╔═╡ ca5fbdc8-22e5-11eb-0a60-cb0a127a3ca5
md"## Listing4.19"

# ╔═╡ fa68607e-22e5-11eb-0558-c9a4d9f77426
begin
	data = CSV.read(datadir("temperatures.csv"))
	brisbane = data.Brisbane
	goldcoast = data.GoldCoast

	diff = brisbane - goldcoast
	dates = [Date(
				Year(data.Year[i]),
				Month(data.Month[i]),
				Day(data.Day[i])
			) for i in 1:nrow(data)]

	fortnightRange = 250:263
	brisFortnight = brisbane[fortnightRange]
	goldFortnight = goldcoast[fortnightRange]
end;

# ╔═╡ 47673b4e-250b-11eb-16f2-2f9e1f103b04
describe(data)

# ╔═╡ 22fd2570-250b-11eb-02af-73597d381ffb
begin
	default(xlabel="Time", ylabel="Temperature")
	default(label=["Brisbane" "Gold Coast"])

	p1 = plot(dates, [brisbane goldcoast],
			c=[:blue :red])
	p2 = plot(dates[fortnightRange], [brisFortnight goldFortnight],
			c=[:blue :red], m=(:dot, 5, Plots.stroke(1)))
	p3 = plot(dates, diff,
			c=:black, ylabel="Temperature Difference",legend=false)
	p4 = histogram(diff, bins=-4:0.5:6,
			ylims=(0,140), legend = false,
			xlabel="Temperature Difference", ylabel="Frequency")
	plot(p1,p2,p3,p4, size = (800,500), margin = 5mm)
end

# ╔═╡ 475ff888-22e6-11eb-2354-09f8f40a8e12
md"## End of listing4.19"

# ╔═╡ Cell order:
# ╟─ca5fbdc8-22e5-11eb-0a60-cb0a127a3ca5
# ╠═e077ab16-22e5-11eb-15ed-5f0654bf7328
# ╠═e75a9a06-22e5-11eb-2e36-8d4f62b830ed
# ╠═fa68607e-22e5-11eb-0558-c9a4d9f77426
# ╠═47673b4e-250b-11eb-16f2-2f9e1f103b04
# ╠═22fd2570-250b-11eb-02af-73597d381ffb
# ╟─475ff888-22e6-11eb-2354-09f8f40a8e12


function leaky_fixed_rnn(activation, alpha, W, W_in, x, y)
    return (1-alpha).*x + alpha*activation.((W*x)+(W_in*y))
end

function double_leaky_fixed_rnn(alpha, f_activation, s_activation, f_lambda, s_lambda, W, W_in, x, y)
    return (1-alpha).*x + f_lambda*f_activation.((W*x)+(W_in*y)) + s_lambda*s_activation.((W*x)+(W_in*y))
end


using MathProgBase, GLPKMathProgInterface

include("misc.jl")

function leastAbsolutes(X, y)
  (n, d) = size(X)
  X = [ones(n, 1) X]
  d += 1
  y = y[:,1]
  c = [ones(n); zeros(d)]
  A = [eye(n) -X; eye(n) X]
  b = [-y; y]

  display(size(y))
  display(size(c))
  display(size(A))
  display(size(b))

  l = [zeros(n); ones(d) * -Inf]

  sol = linprog(c, A, '>', b, l, Inf, GLPKSolverLP())
  w = sol.sol[n+1:n+d]
  if sol.status == :Optimal
      println("Optimal objective value is $(sol.objval)")
      println("Optimal solution vector is: $(w)")
  else
      println("Error: solution status $(sol.status)")
  end
  predict(Xtilde) = [ones(size(Xtilde,1),1) Xtilde]*w

  # Return model
  return LinearModel(predict, w)
end


# Function that updates partial force and energy vectors, used for parallel computation

function compute_uf_partials_QUAD!(it :: Int64, f :: Forces,
                                   i :: Int64, j :: Int64,
                                   xj :: Float64, yj :: Float64, r :: Float64,
                                   input :: MDInput)

   f.upartial[it] = f.upartial[it] + input.eps*(input.sig-r)^2

   drdxi = -xj/r
   drdyi = -yj/r

   fx = 2*input.eps*(input.sig-r)*drdxi
   fy = 2*input.eps*(input.sig-r)*drdyi

   f.fpartial[it][i,1] = f.fpartial[it][i,1] + fx
   f.fpartial[it][i,2] = f.fpartial[it][i,2] + fy
   f.fpartial[it][j,1] = f.fpartial[it][j,1] - fx
   f.fpartial[it][j,2] = f.fpartial[it][j,2] - fy

end


function MALA(U_fun::Function, grad_U::Function, h::Real, dim::Integer, n_iter::Integer, Xi_init = zeros(dim))

  Xi = zeros(dim, n_iter + 1)
  Xi[:,1] = Xi_init
  Z = randn(dim, n_iter)
  U = rand(n_iter)
  q(x, y) = exp(-1/(2*h) * sumabs2(y - x + (h/2)*grad_U(x)))
  acceptance_count = 0
  for i=1:n_iter
    Xi_current = Xi[:,i]
    Xi_proposed = Xi_current - (h/2) * grad_U(Xi_current) + sqrt(h) * Z[:,i];
    accept_prob = min(1, exp(U_fun(Xi_current)-U_fun(Xi_proposed)) * q(Xi_proposed, Xi_current) /q(Xi_current,Xi_proposed))
    if U[i] <= accept_prob
      Xi[:,i+1] = Xi_proposed
      acceptance_count += 1
    else
      Xi[:,i+1] = Xi_current
    end
  end
  println("MALA: Fraction of accepted proposals: ", acceptance_count / n_iter)
  return Xi
end

function consensusMALA(U_fun::Vector{Function}, grad_U::Vector{Function}, h::Real, dim::Integer, n_samples::Integer, Xi_init::Vector{Float64} = zeros(dim))
  n_batches = length(U_fun)
  samples = zeros(dim, n_samples, n_batches)
  W = Vector{Array{Float64,2}}(n_batches)  # weights
  cumulative_samples = zeros(dim, n_samples)
  cumulative_weights = zeros(dim,dim)
  for i=1:n_batches
    samples[:,:,i] = MALA(U_fun[i], grad_U[i], h, dim, n_samples, Xi_init)[:,2:end]
    # println(size(samples[:,:,i]))
    # print(cov(samples[:,:,i],2)); println()
    W[i] = inv(cov(samples[:,:,i],2))
    cumulative_samples += W[i] * samples[:,:,i]
    cumulative_weights += W[i]
  end
  println(cumulative_weights)
  consensus_samples = cumulative_weights \ cumulative_samples
end


function calc_shapes(ncell::Int64, sol::Array{Float64})

    E = zeros(ncell+2); B = zeros(ncell+2); F = zeros(ncell+2); S = zeros(ncell+2)
    dfde = zeros(ncell+2); del = zeros(ncell+2)

    for i = 1:ncell+2
        del[i] = sol[1,i]
        E[i] = sol[2,i]/del[i] - 1.
        F[i] = 4.8274*E[i]^4 - 5.9816*E[i]^3 + 4.0274*E[i]^2 + 0.23247*E[i] + 0.15174
        if E[i] < -0.0616
            B[i] = -225.86*E[i]^3 - 3016.6*E[i]^2 - 208.68*E[i] - 17.915
        elseif E[i] > -0.0395
            B[i] = 131.9*E[i]^3 - 167.32*E[i]^2 + 76.642*E[i] - 11.068
        else
            B[i] = 0.5*(-225.86*E[i]^3 - 3016.6*E[i]^2 - 208.68*E[i] - 17.915
                        + 131.9*E[i]^3 - 167.32*E[i]^2 + 76.642*E[i] - 11.068)
        end
        if E[i] < -0.0582
            S[i] = 451.55*E[i]^3 + 2010*E[i]^2 + 138.96*E[i] + 11.296
        elseif E[i] > -0.042
            S[i] = -96.739*E[i]^3 + 117.74*E[i]^2 - 46.432*E[i] + 6.8074
        else
            S[i] = 0.5*(451.55*E[i]^3 + 2010*E[i]^2 + 138.96*E[i] + 11.296
                        - 96.739*E[i]^3 + 117.74*E[i]^2 - 46.432*E[i] + 6.8074)
        end
        dfde[i] = 4*4.8274*E[i]^3 - 3*5.9816*E[i]^2 + 2*4.0274*E[i] + 0.23247
    end

    return E, F, B, S, dfde, del
end

function calc_eigen(ncell::Int64, E::Array{Float64}, F::Array{Float64},
                    dfde::Array{Float64}, ue::Array{Float64})

    lamb1 = zeros(ncell+2); lamb2 = zeros(ncell+2)

    for i = 1:ncell+2
        a_q = 1.
        b_q = -ue[i] * (dfde[i] - 1)
        c_q = ue[i] * ue[i] * (E[i]*dfde[i] - F[i])
        lamb1[i] = (-b_q + sqrt(b_q*b_q - 4*a_q*c_q))/(2*a_q)
        lamb2[i] = (-b_q  -sqrt(b_q*b_q - 4*a_q*c_q))/(2*a_q)

        #Always have lamb1 > lamb2
        if lamb2[i]>lamb1[i]
            temp  = lamb2[i]
            lamb2[i] = lamb1[i]
            lamb1[i] = temp
        end
    end

    return lamb1, lamb2
end

function calc_dt(ncell::Int64, cfl::Float64, dx::Float64, lamb1::Array{Float64},
                 lamb2::Array{Float64})

    dt = 10000
    for i = 2:ncell+1
        dti = cfl*(dx/(abs(lamb1[i]) + abs(lamb2[i])))
        if dti < dt
            dt = dti
        end
    end

    return dt
end

function calc_fluxes(ncell::Int64, E::Array{Float64}, F::Array{Float64},
                     del::Array{Float64}, ue::Array{Float64}, lamb1::Array{Float64}, lamb2::Array{Float64})

    Apos = zeros(2,2); Aneg = zeros(2,2)
    flux = zeros(2,2,ncell+2)

    for i = 1:ncell+2
        if lamb1[i] >= 0. && lamb2[i] >= 0.
            flux[1,1,i] = ue[i]*E[i]*del[i]
            flux[1,2,i] = ue[i]*F[i]*del[i]
        elseif lamb1[i] <= 0. && lamb2[i] <= 0.
            flux[1,:,i] = 0.
        else
            Apos[1,1] = ue[i]*lamb1[i]/(lamb1[i] - lamb2[i])*(-1. - lamb2[i]/ue[i])
            Apos[1,2] = ue[i]*lamb1[i]/(lamb1[i] - lamb2[i])
            Apos[2,1] = -ue[i]*lamb1[i]/(lamb1[i] - lamb2[i])*(1. + lamb1[i]/ue[i])*(1. + lamb2[i]/ue[i])
            Apos[2,2] = ue[i]*lamb1[i]/(lamb1[i] - lamb2[i])*(1. + lamb1[i]/ue[i])
            flux[1,1,i] = Apos[1,1]*del[i] + Apos[1,2]*(E[i] + 1.)*del[i]
            flux[1,2,i] = Apos[2,1]*del[i] + Apos[2,2]*(E[i] + 1.)*del[i]
        end
    end

    for i = 1:ncell+2
        if lamb1[i] >= 0. && lamb2[i] >= 0.
            flux[2,1,i] = 0
            flux[2,2,i] = 0
        elseif lamb1[i] < 0. && lamb2[i] < 0.
            flux[2,1,i]= ue[i]*E[i]*del[i]
            flux[2,2,i]= ue[i]*F[i]*del[i]
        else
            Aneg[1,1] = ue[i]*lamb2[i]/(lamb1[i]-lamb2[i])*(1. + lamb1[i]/ue[i])
            Aneg[1,2] = -ue[i]*lamb2[i]/(lamb1[i]-lamb2[i])
            Aneg[2,1] = ue[i]*lamb2[i]/(lamb1[i]-lamb2[i])*(1. + lamb1[i]/ue[i])*(1. + lamb2[i]/ue[i])
            Aneg[2,2] = ue[i]*lamb2[i]/(lamb1[i]-lamb2[i])*(-1. - lamb2[i]/ue[i])
            flux[2,1,i] = Aneg[1,1]*del[i] + Aneg[1,2]*(E[i]+1)*del[i]
            flux[2,2,i] = Aneg[2,1]*del[i] + Aneg[2,2]*(E[i]+1)*del[i]
        end
    end

    return flux
end

function matuCyl(tTot::Float64, ncell::Int64 = 200, cfl::Float64 = 0.8, nstage::Int64 = 2)

    dx = pi/ncell

    #Timestep determined dynamically
    dt = 0

    uinf = 1.; c = 1.; re = 1000.

    x = zeros(ncell+2); ue = zeros(ncell+2); uex = zeros(ncell+2); uet = zeros(ncell+2)
    E = zeros(ncell+2); B = zeros(ncell+2); F = zeros(ncell+2); S = zeros(ncell+2); dfde = zeros(ncell+2);
    del = zeros(ncell+2);

    sol = zeros(2,ncell+2); solt = zeros(2,ncell+2)
    lamb1 = zeros(ncell+2); lamb2 = zeros(ncell+2)
    Jsep = zeros(2,ncell+2)
    Csep = zeros(ncell+2)

    flux = zeros(2,2,ncell+2)
    cf = zeros(ncell+2)

    #Set sources (cell centred)
    for i = 2:ncell+1
        x[i] = (i-1.5)*dx
        ue[i] = 2*sin(x[i])
        uex[i] = 2*cos(x[i])
    end
    x[1] = 2*x[2] - x[3]
    x[ncell+2] = 2*x[ncell+1] - x[ncell]
    uet[:] = zeros(ncell+2)
    ue[1] = 0.
    ue[ncell+2] = 0.
    uex[1] = cos(0.)
    uex[ncell+2] = cos(pi)

    #Set initial conditions
    E_init = 0.4142
    for i = 1:ncell+2
        E[i] = E_init
        B[i] = 131.9*E_init^3 - 167.32*E_init^2 + 76.642*E_init - 11.068
        del[i] = sqrt.(B[i]*0.005)
        F[i] = 4.8274*E_init^4 - 5.9816*E_init^3 + 4.0274*E_init^2 + 0.23247*E_init + 0.15174
    end

    t = 0.005

    sol[1,:] = del[:]
    sol[2,:] = del[:].*(E[:] .+ 1)

    rhs = zeros(2, ncell+2)

    istep = 0
    while t < tTot
        istep += 1

        #Update ghost cells
        sol[:,1] = 2*sol[:,2] - sol[:,3]
        sol[:,ncell+2] = 2*sol[:,ncell+1] - sol[:,ncell]

        #Calculate derived quantities
        E, F, B, S, dfde, del = calc_shapes(ncell, sol)

        #Compute eigenvalues
        lamb1, lamb2 = calc_eigen(ncell, E, F, dfde, ue)

        #Compute timestep
        dt = calc_dt(ncell, cfl, dx, lamb1, lamb2)

        flux = calc_fluxes(ncell, E, F, del, ue, lamb1, lamb2)

        rhs[:,:] = zeros(2, ncell+2)

        #The ghost rhs are never used - dont matter
        for i = 2:ncell+1
            rhs[1,i] = B[i]/(2*del[i]) - del[i]*uet[i]/ue[i] - (E[i]+1)*del[i]*uex[i]
            rhs[2,i] = S[i]/del[i] - 2*E[i]*del[i]*uet[i]/ue[i] - 2*F[i]*del[i]*uex[i]
        end

        for i = 2:ncell
            solt[1,i] = sol[1,i] - (dt/dx)*(flux[1,1,i] - flux[1,1,i-1] + flux[2,1,i+1]
                                            -flux[2,1,i]) + dt*rhs[1,i]
            solt[2,i] = sol[2,i] - (dt/dx)*(flux[1,2,i] - flux[1,2,i-1] + flux[2,2,i+1]
                                            -flux[2,2,i]) + dt*rhs[2,i]
        end

        i = ncell+1
        solt[1,i] = sol[1,i] - (dt/dx)*(flux[1,1,i] - flux[1,1,i-1]) + dt*rhs[1,i]
        solt[2,i] = sol[2,i] - (dt/dx)*(flux[1,2,i] - flux[1,2,i-1]) + dt*rhs[2,i]

        #Update ghost cells
        solt[:,1] = 2*solt[:,2] - solt[:,3]
        solt[:,ncell+2] = 2*solt[:,ncell+1] - solt[:,ncell]

        #Calculate derived quantities
        E, F, B, S, dfde, del = calc_shapes(ncell, solt)

        #Compute eigenvalues
        lamb1, lamb2 = calc_eigen(ncell, E, F, dfde, ue)

        fluxt = calc_fluxes(ncell, E, F, del, ue, lamb1, lamb2)

        rhs[:,:] = zeros(2, ncell+2)
        #The ghost rhs are never used - dont matter
        for i = 2:ncell+1
            rhs[1,i] = B[i]/(2*del[i]) - del[i]*uet[i]/ue[i] - (E[i]+1)*del[i]*uex[i]
            rhs[2,i] = S[i]/del[i] - 2*E[i]*del[i]*uet[i]/ue[i] - 2*F[i]*del[i]*uex[i]
        end

        i = 2
        sol[1,i] = 0.5*(sol[1,i] + solt[1,i]) - (0.5*dt/dx)*(flux[1,1,i] - flux[1,1,i-1]
                                                             - flux[2,1,i] + 2*flux[2,1,i+1] - flux[2,1,i+2] + fluxt[1,1,i] - fluxt[1,1,i-1]
                                                             + fluxt[2,1,i+1] - fluxt[2,1,i]) + 0.5*dt*rhs[1,i]
        sol[2,i] = 0.5*(sol[2,i] + solt[2,i]) - (0.5*dt/dx)*(flux[1,2,i] - flux[1,2,i-1]
                                                             - flux[2,2,i] + 2*flux[2,2,i+1]-flux[2,2,i+2]+fluxt[1,2,i] - fluxt[1,2,i-1]
                                                             + fluxt[2,2,i+1] - fluxt[2,2,i])+0.5*dt*rhs[2,i]
        for i = 3:ncell
            sol[1,i] = 0.5*(sol[1,i] + solt[1,i]) - (0.5*dt/dx)*(flux[1,1,i] - 2*flux[1,1,i-1] + flux[1,1,i-2]
                                                                 - flux[2,1,i] + 2*flux[2,1,i+1] - flux[2,1,i+2] + fluxt[1,1,i] - fluxt[1,1,i-1]
                                                                 + fluxt[2,1,i+1] - fluxt[2,1,i]) + 0.5*dt*rhs[1,i]
            sol[2,i] = 0.5*(sol[2,i] + solt[2,i]) - (0.5*dt/dx)*(flux[1,2,i] - 2*flux[1,2,i-1] + flux[1,2,i-2]
                                                                 - flux[2,2,i] + 2*flux[2,2,i+1] - flux[2,2,i+2] + fluxt[1,2,i] - fluxt[1,2,i-1]
                                                                 + fluxt[2,2,i+1] - fluxt[2,2,i]) + 0.5*dt*rhs[2,i]
        end
        i = ncell+1
        sol[1,i] = 0.5*(sol[1,i] + solt[1,i]) - (0.5*dt/dx)*(flux[1,1,i] - 2*flux[1,1,i-1] + flux[1,1,i-2]
                                                             +fluxt[1,1,i] - fluxt[1,1,i-1]) + 0.5*dt*rhs[1,i]
        sol[2,i] = 0.5*(sol[2,i] + solt[2,i]) - (0.5*dt/dx)*(flux[1,2,i] - 2*flux[1,2,i-1] + flux[1,2,i-2]
                                                             +fluxt[1,2,i] - fluxt[1,2,i-1]) + 0.5*dt*rhs[2,i]

        t = t + dt
        if t > tTot
            break
        end

        for i = 2:ncell+1
            cf[i] = B[i]*ue[i]*c/(del[i]*uinf*re)
            Csep[i] = (del[i+1] - del[i])/(del[i] - del[i-1])
            if Csep[i] > 10.
                println("singularity (separation) detected at t=$t, x=$(x[i])")
            end
        end
        #println("$istep \t", "$dt \t", "$t \t", "$(maximum(Csep)) \t", "$(minimum(Csep)) \n")
        println("$t \t")
        display(plot(x/pi,[del, E] , layout=(2,1), legend = false))
        sleep(0.1)

    end

mat = zeros(0, 8)
mat = mat'
for i = 2:ncell+1
    mat = hcat(mat,[x[i], ue[i], uex[i], uet[i], del[i], E[i], cf[i], Csep[i]])
end
mat = mat'

return mat
end

# f = open("t$istep", "w")
# write(f, ["#x \t", "#ue \t", "#uex \t", "#uet \t", "#del \t", "#E \t", "#cf \t", "#Csep \n"])
# writedlm(f, mat)
# close(f)


"""
    log_hypergeometric(a, b, x)

Compute the logarithm of the hypergeometric function
``{}_{2}F_{1}(-a, -b; -(a + b - 1); -exp(x))`` where `a` and `b` are positive
integers, and `a >= b`.
"""
function log_hypergeometric(
  a::F,
  b::F,
  x::F
)::F where {
  F <: AbstractFloat
}
  if b > 1
    # define H_{b} = y_{b} / y_{b - 1}, i. e. y_{b} = H_{b} * y_{b - 1}
    # y_{0} = 1, y_{1} = 1 + (a * z) / (a + 1 - k)
    tmp = log1pexp(x)

    # v = [log(H), log(y), compensation_error]
    v = [tmp, tmp, F(0)]

    n = F(1)
    while n < b
      n += 1
      kbn_2F1_iteration!(v, n, a, x, F(1))
    end

    # finally add the compensation error
    v[2] + v[3]
  elseif b == 1
    # 2F1(-a, -1; -a, -z) = 1 + z
    log1pexp(x)
  else
    0
  end
end

"""
    log_hypergeometric_joint(a, b, x)

Compute the logarithm of the three hypergeometric functions
``{}_{2}F_{1}(-a, -b; -(a + b - 1); -exp(x))``,
``{}_{2}F_{1}(-(a - 1), -(b - 1); -(a + b); -exp(x))``, and
``{}_{2}F_{1}(-(a - 2), -(b - 2); -(a + b + 1); -exp(x))`` at the same time.
`a` and `b` are positive integers with `a >= b`. We assume that `a` and `b` are
both greater than 1.
"""
function log_hypergeometric_joint(
  a::F,
  b::F,
  x::F
)::Tuple{F, F, F} where {
  F <: AbstractFloat
}
  # define H_{b} = y_{b} / y_{b - 1}, i. e. y_{b} = H_{b} * y_{b - 1}
  # y_{0} = 1, y_{1} = 1 + (a * z) / (a + 1 - k)

  # k = 1
  # 2F1(-a, -b; -(a + b - 1); -z)
  tmp = log1pexp(x)
  v1 = [tmp, tmp, F(0)]

  # k = 0
  # 2F1(-(a - 1), -(b - 1); -(a + b); -z)
  tmp = log1pexp(log((a - 1) / a) + x)
  v2 = [tmp, tmp, F(0)]

  if b > 2
    # k = -1
    # 2F1(-(a - 2), -(b - 2); -(a + b + 1); -z)
    tmp = log1pexp(log((a - 2) / a) + x)
    v3 = [tmp, tmp, F(0)]

    n = F(1)
    while n < b - 2
      n += 1
      kbn_2F1_iteration!(v1, n, a, x, F(1))
      kbn_2F1_iteration!(v2, n, a - 1, x, F(0))
      kbn_2F1_iteration!(v3, n, a - 2, x, F(-1))
    end

    # We still need few iterations
    kbn_2F1_iteration!(v1, b - 1, a, x, F(1))
    kbn_2F1_iteration!(v2, b - 1, a - 1, x, F(0))

    kbn_2F1_iteration!(v1, b, a, x, F(1))

    # add compensation errors and return the final values
    v1[2] + v1[3], v2[2] + v2[3], v3[2] + v3[3]
  else
    # log(2F1(-a, -2; -a; -z)) = v1[2] + H
    # log(2F1(-(a - 1), -1; -(a + 2); -z)) = v2[2]
    # log(2F1(-(a - 2), 0; -(a + 3); -z)) = 0
    kbn_2F1_iteration!(v1, F(2), a, x, F(1))
    v1[2], v2[2], 0
  end
end

"""
    kbn_2F1_iteration!(v, n, a, x, k)

Basic Kahan iteration operations for computing
``{}_{2}F_{1}(-a, -b; -(a + b - k); -exp(x))``.
"""
@inline function kbn_2F1_iteration!(
  v::Vector{F},
  n::F,
  a::F,
  x::F,
  k::F
) where {
  F <: AbstractFloat
}
  # new value for log(y_{b} / y_{b - 1})
  v[1] = log1pexp(x + log((a - n + 1) / (a + n - k)) +
                  log1pexp(log((n - 1) * (n - 1 - k) /
                               ((a - n + 1) * (a + n - k - 1))) - v[1]))

  # running sum including error
  s = v[2] + v[1]

  # Neumaier modification to Kahan summation
  v[3] += if abs(v[2]) >= abs(v[1])
    (v[1] - s) + v[2]
  else
    (v[2] - s) + v[1]
  end

  v[2] = s

  nothing
end


# An implementation of the Lanczos version of the conjugate gradient method
# for a family of shifted systems of the form (A + αI) x = b.
#
# The implementation follows
# A. Frommer and P. Maass, Fast CG-Based Methods for Tikhonov-Phillips
# Regularization, SIAM Journal on Scientific Computing, 20(5),
# pp. 1831-1850, 1999.
#
# Dominique Orban, <dominique.orban@gerad.ca>
# Princeton, NJ, March 2015.

export cg_lanczos, cg_lanczos_shift_seq

"""The Lanczos version of the conjugate gradient method to solve the
symmetric linear system

  Ax = b

The method does _not_ abort if A is not definite.
"""
function cg_lanczos(A :: AbstractLinearOperator, b :: AbstractVector{T};
                    atol :: Float64=1.0e-8, rtol :: Float64=1.0e-6, itmax :: Int=0,
                    check_curvature :: Bool=false, verbose :: Bool=false) where T <: Number

  n = size(b, 1);
  (size(A, 1) == n & size(A, 2) == n) || error("Inconsistent problem size");
  verbose && @printf("CG Lanczos: system of %d equations in %d variables\n", n, n);

  # Initial state.
  x = zeros(T, n);
  β = @knrm2(n, b)
  β == 0 && return x, LanczosStats(true, [0.0], false, 0.0, 0.0, "x = 0 is a zero-residual solution")
  v = b / β;
  v_prev = copy(v);
  p = copy(b);
  iter = 0;
  itmax == 0 && (itmax = 2 * n);

  # Initialize some constants used in recursions below.
  σ = β;
  ω = 0;
  γ = 1;
  Anorm2 = 0.0;
  β_prev = 0.0;

  # Define stopping tolerance.
  rNorm = σ;
  rNorms = [rNorm;];
  ε = atol + rtol * rNorm;
  verbose && @printf("%5d  %8.1e\n", iter, rNorm);

  indefinite = false;
  solved = rNorm <= ε;
  tired = iter >= itmax;
  status = "unknown";

  # Main loop.
  while ! (solved || tired || (check_curvature & indefinite))
    # Form next Lanczos vector.
    v_next = A * v;
    δ = @kdot(n, v, v_next);  # BLAS.dot(n, v, 1, v_next, 1) doesn't seem to pay off.

    # Check curvature. Exit fast if requested.
    # It is possible to show that σⱼ² (δⱼ - ωⱼ₋₁ / γⱼ₋₁) = pⱼᵀ A pⱼ.
    γ = 1 / (δ - ω / γ);
    indefinite |= (γ <= 0.0);
    (check_curvature & indefinite) && continue;

    @kaxpy!(n, -δ, v, v_next)  # Faster than v_next = Av - δ * v;
    if iter > 0
      @kaxpy!(n, -β, v_prev, v_next)  # Faster than v_next = v_next - β * v_prev;
      @. v_prev .= v
    end
    β = @knrm2(n, v_next)
    @. v .= v_next / β
    Anorm2 += β_prev^2 + β^2 + δ^2;  # Use ‖T‖ as increasing approximation of ‖A‖.
    β_prev = β;

    # Compute next CG iterate.
    @kaxpy!(n, γ, p, x)  # Faster than x = x + γ * p;

    ω = β * γ;
    σ = -ω * σ;
    ω = ω * ω;
    @kscal!(n, ω, p)
    @kaxpy!(n, σ, v, p)  # Faster than p = σ * v + ω * p;
    rNorm = abs(σ);
    push!(rNorms, rNorm);
    iter = iter + 1;
    verbose && @printf("%5d  %8.1e\n", iter, rNorm);
    solved = rNorm <= ε;
    tired = iter >= itmax;
  end

  status = tired ? "maximum number of iterations exceeded" : (check_curvature & indefinite) ? "negative curvature" : "solution good enough given atol and rtol"
  stats = LanczosStats(solved, rNorms, indefinite, sqrt(Anorm2), 0.0, status);  # TODO: Estimate Acond.
  return (x, stats);
end

"""The Lanczos version of the conjugate gradient method to solve a family
of shifted systems

  (A + αI) x = b  (α = α₁, α₂, ...)

The method does _not_ abort if A + αI is not definite.
"""
function cg_lanczos_shift_seq(A :: AbstractLinearOperator, b :: AbstractVector{Tb}, shifts :: AbstractVector{Ts};
                              atol :: Float64=1.0e-8, rtol :: Float64=1.0e-6, itmax :: Int=0,
                              check_curvature :: Bool=false, verbose :: Bool=false) where {Tb, Ts <: Number}

  n = size(b, 1);
  (size(A, 1) == n & size(A, 2) == n) || error("Inconsistent problem size");

  nshifts = size(shifts, 1);
  verbose && @printf("CG Lanczos: system of %d equations in %d variables with %d shifts\n", n, n, nshifts);

  # Initial state.
  ## Distribute x similarly to shifts.
  x = zeros(Tb, n, nshifts);
  β = @knrm2(n, b)
  β == 0 && return x, LanczosStats(true, [0.0], false, 0.0, 0.0, "x = 0 is a zero-residual solution")
  v = b / β;
  v_prev = copy(v);

  # Initialize each p to b.
  p = b * ones(nshifts)';

  # Initialize some constants used in recursions below.
  σ = β * ones(nshifts);
  δhat = zeros(nshifts);
  ω = zeros(Tb, nshifts);
  γ = ones(Tb, nshifts);

  # Define stopping tolerance.
  rNorms = β * ones(Tb, nshifts);
  rNorms_history = [rNorms;];
  ε = atol + rtol * β;

  # Keep track of shifted systems that have converged.
  converged = rNorms .<= ε;
  iter = 0;
  itmax == 0 && (itmax = 2 * n);

  # Keep track of shifted systems with negative curvature if required.
  indefinite = falses(nshifts);

  # Build format strings for printing.
  if verbose
    fmt = "%5d" * repeat("  %8.1e", nshifts) * "\n";
    # precompile printf for our particular format
    local_printf(data...) = Core.eval(Main, :(@printf($fmt, $(data)...)))
    local_printf(iter, rNorms...)
  end

  solved = all(converged);
  tired = iter >= itmax;
  status = "unknown";

  # Main loop.
  while ! (solved || tired)
    # Form next Lanczos vector.
    v_next = A * v;
    δ = @kdot(n, v, v_next)
    @kaxpy!(n, -δ, v, v_next)  # Faster than v_next = Av - δ * v;
    if iter > 0
      @kaxpy!(n, -β, v_prev, v_next)  # Faster than v_next = v_next - β * v_prev;
      v_prev = v;
    end
    β = @knrm2(n, v_next)
    v = v_next / β;

    # Check curvature: v'(A + sᵢI)v = v'Av + sᵢ ‖v‖² = δ + sᵢ because ‖v‖ = 1.
    # It is possible to show that σⱼ² (δⱼ + sᵢ - ωⱼ₋₁ / γⱼ₋₁) = pⱼᵀ (A + sᵢ I) pⱼ.
    for i = 1 : nshifts
      δhat[i] = δ + shifts[i]
      γ[i] = 1 ./ (δhat[i] - ω[i] ./ γ[i])
    end
    indefinite .|= (γ .<= 0.0);

    # Compute next CG iterate for each shifted system that has not yet converged.
    # Stop iterating on indefinite problems if requested.
    not_cv = check_curvature ? findall(.! (converged .| indefinite)) : findall(.! converged);

    # Loop is a bit faster than the vectorized version.
    for i in not_cv
      x[:, i] += γ[i] * p[:, i];  # Strangely, this is faster than a loop.

      ω[i] = β * γ[i];
      σ[i] *= -ω[i];
      ω[i] *= ω[i];
      p[:, i] = v * σ[i] + p[:, i] * ω[i];  # Faster than loop.

      # Update list of systems that have converged.
      rNorms[i] = abs(σ[i]);
      converged[i] = rNorms[i] <= ε;
    end

    length(not_cv) > 0 && append!(rNorms_history, rNorms);

    # Is there a better way than to update this array twice per iteration?
    not_cv = check_curvature ? findall(.! (converged .| indefinite)) : findall(.! converged);
    iter = iter + 1;
    verbose && local_printf(iter, rNorms...)

    solved = length(not_cv) == 0;
    tired = iter >= itmax;
  end

  status = tired ? "maximum number of iterations exceeded" : "solution good enough given atol and rtol"
  stats = LanczosStats(solved, reshape(rNorms_history, nshifts, round(Int, sum(size(rNorms_history))/nshifts))', indefinite, 0.0, 0.0, status);  # TODO: Estimate Anorm and Acond.
  return (x, stats);
end


"""
    graph_small_world(n::Integer, m::Integer, p::Real)

Small-world network model of Watts and Strogatz (1998).

# Arguments
* `n` — number of nodes,
* `k` — number of neighbors for each node, should be even,
* `p` — probability of random rewiring of each edge.

# References
Watts, D. J., & Strogatz, S. H. (1998).
Collective dynamics of “small-world” networks. Nature, 393(6684), 440–442.
"""
function graph_small_world(n::Integer, k::Integer, p::Real; kwargs...)
    @assert(n > 5,         "number of nodes (n) should be > 5")
    @assert(0 < k < n / 2, "number of neighbors (k) is not in a valid range")
    @assert(k % 2 == 0,    "number of neighbors (k) should be even")
    @assert(0 < p < 1,     "rewiring probability (p) is not in (0, 1) interval")
    g = Graph(n; kwargs...)
    for i = 1:n, offset = 1:Int(k / 2)
        j = (i + offset - 1) % n + 1
        if rand() < p || hasedge(g, i, j)
            while true
                j = rand(1:n)
                if j != i && !hasedge(g, i, j)
                    break
                end
            end
        end
        addedge!(g, i, j)
    end
    return g
end

@deprecate(graph_smallworld, graph_small_world)


#=
Generates an array of draws from a discrete random variable with a
specified vector of probabilities.

@author : Spencer Lyon <spencer.lyon@nyu.edu>

@date: 2014-07-10

References
----------

http://quant-econ.net/jl/finite_markov.html?highlight=discrete_rv

TODO: as of 07/10/2014 it is not possible to define the property
      interface we see in the python version. Once issue 1974 from
      the main Julia repository is resolved, we can revisit this and
      implement that feature.
=#

"""
Generates an array of draws from a discrete random variable with
vector of probabilities given by q.

##### Fields

- `q::AbstractVector`: A vector of non-negative probabilities that sum to 1
- `Q::AbstractVector`: The cumulative sum of q
"""
type DiscreteRV{TV1<:AbstractVector, TV2<:AbstractVector}
    q::TV1
    Q::TV2
    function (::Type{DiscreteRV{TV1,TV2}}){TV1,TV2}(q, Q)
        abs(Q[end] - 1) > 1e-10 && error("q should sum to 1")
        new{TV1,TV2}(q, Q)
    end
end

function DiscreteRV{TV<:AbstractVector}(q::TV)
    Q = cumsum(q)
    DiscreteRV{TV,typeof(Q)}(q, Q)
end

"""
Make a single draw from the discrete distribution

##### Arguments

- `d::DiscreteRV`: The `DiscreteRV` type represetning the distribution

##### Returns

- `out::Int`: One draw from the discrete distribution
"""
draw(d::DiscreteRV) = searchsortedfirst(d.Q, rand())

"""
Make multiple draws from the discrete distribution represented by a
`DiscreteRV` instance

##### Arguments

- `d::DiscreteRV`: The `DiscreteRV` type representing the distribution
- `k::Int`:

##### Returns

- `out::Vector{Int}`: `k` draws from `d`
"""
draw(d::DiscreteRV, k::Int) = Int[draw(d) for i=1:k]


# This file was generated, do not modify it.

using RDatasets
using DataFrames

boston = dataset("MASS", "Boston");

typeof(boston)

names(boston)

first(boston, 4)

boston.Crim[1:5]

boston[3, 5]

boston[1:5, [:Crim, :Zn]]

boston[1:5, 1:2]

b1 = select(boston, [:Crim, :Zn, :Indus])
first(b1, 2)

b2 = select(boston, Not(:NOx))
first(b2, 2)

select!(b1, Not(:Crim))
first(b1, 2)

using StatsBase
describe(boston, :min, :max, :mean, :median, :std)

foo(x) = sum(abs.(x)) / length(x)
d = describe(boston, :mean, :median, :foo => foo)
first(d, 3)

using Statistics

mat = convert(Matrix, boston)
mat[1:3, 1:3]

boston.Crim_x_Zn = boston.Crim .* boston.Zn;

mao = dataset("gap", "mao")
describe(mao, :nmissing)

std(mao.Age)

std(skipmissing(mao.Age))

iris = dataset("datasets", "iris")
first(iris, 3)

unique(iris.Species)

gdf = groupby(iris, :Species);

subdf_setosa = gdf[1]
describe(subdf_setosa, :min, :mean, :max)

df = DataFrame(a=1:3, b=4:6)
combine(df, :a => sum, nrow)

foo(v) = v[1:2]
combine(df, :a => maximum, :b => foo)

bar(v) = v[end-1:end]
combine(df, :a => foo, :b => bar)

combine(groupby(iris, :Species), :PetalLength => mean)

gdf = groupby(iris, :Species)
combine(gdf, :PetalLength => mean => :MPL, :PetalLength => std => :SPL)

combine(gdf, names(iris, Not(:Species)) .=> std)

names(iris, Not(:Species))


# %% Sensititivy to trajectory length
using DrWatson
@quickactivate :FractalDimension # uses DynamicalSystems, PyPlot
include(srcdir("style.jl"))

N = Int(1e5)
Ds = range(4; step = 2, length = 6)

qH = 1
qC = 2
data = :lorenz96_chaotic

# Calculate values for H, C
eHs, eCs, Hs, Cs = [Vector{Float64}[] for i in 1:4]
for i in 1:length(Ds)
    # These are the parameters that change between the different elements of the plot
    D = Ds[i]

    # Here we simply pack all parameters into a dictionary
    # (other parameters are (probably) globals)
    params = @strdict N D qH qC data

    # This is the main call that calculates everything
    output, s = produce_or_load(
        datadir("main"), params, make_C_H;
        prefix = string(data), suffix = "jld2", force = true,
        ignores = ["data"],
    )
    @unpack eH, eC, H, C = output
    push!(eHs, eH); push!(Hs, H); push!(eCs, eC); push!(Cs, C)
end

# Do the actual plot
labels = ["\$D=$(D)\$" for D in Ds]
legendtitle = "impact of dimensionality \$D\$ (Lorenz96 system)"

fig, axs = mainplot(
    Hs, Cs, es, labels, legendtitle;
    qH, qC, tol = 0.25,
    # dimension_fit = FractalDimension.logarithmic_corrected_fit_lsqfit
)

# wsave(plotsdir("paper", "dimensionality"), fig)


function f(x,θ)
  α,β,γ,δ = θ
  x1,x2=x
   dx1 = α*x1 - β*x1*x2
   dx2 = δ*x1*x2 - γ*x2
  [dx1,dx2]
end

function solve(f,x0::AbstractVector,θ,dt,N)
  X = hcat([zero(x0) for i=1:N]...)
  X[:,1]=x0
  for t=1:N-1
     X[:,t+1]=X[:,t]+dt*f(X[:,t],θ)
  end
  X
end

θ0 = [0.1,0.2,0.3,0.2]
x0 = [1.0,1.0]
dt = 0.1
N  = 1000
X=solve(f,x0,θ0,dt,N)

using Plots

p=plot(X[1,:],xlabel="t",label="x",color=:blue)
p=plot!(X[2,:],xlabel="t",label="y",color=:red)

K=100
X0 = [x0 .+ 0.1*randn(2) for k=1:K]
θ1 =[[θ0[1]+0.01randn();θ0[2:end]] for k=1:K]
Xens=[X=solve(f,X0[i],θ1[i],dt,N) for i=1:K]

p=plot(X[1,:],xlabel="t",label="x",color=:blue)
p=plot!(X[2,:],xlabel="t",label="y",color=:red)
for i=1:K
  p=plot!(Xens[i][1,:],label="",color=:blue,alpha=0.2)
  p=plot!(Xens[i][2,:],label="",color=:red,alpha=0.2)
end
savefig("LV_MC_param.svg")

ind = 100:100:1000
scatter!(p,ind,[X[1,ind]],errorbar=[0.2],label="measurements")
savefig("LV_MC_param_data.svg")

Ds=zeros(100)
Ind=[]
for i=1:100
 Ds[i] = sum((X[1,ind]-Xens[i][1,ind]).^2)
 if Ds[i]<1
  push!(Ind,i)
 end
end

p=plot(X[1,:],xlabel="t",label="x",color=:blue)
p=plot!(X[2,:],xlabel="t",label="y",color=:red)
for i=Ind
  p=plot!(Xens[i][1,:],label="",color=:blue,alpha=0.2)
  p=plot!(Xens[i][2,:],label="",color=:red,alpha=0.2)
end
scatter!(p,ind,[X[1,ind]],errorbar=[0.2],label="measurements")
savefig("LV_MC_param_assim.svg")

MXe=mean(Xens)
SXe=std(Xens)
plot(MXe[1,1:30:end],label="x",color=:blue,errorbar=SXe[1,1:30:end])
plot!(MXe[2,1:30:end],label="y",color=:red,errorbar=SXe[2,1:30:end])

using LinearAlgebra
function solve_res(f,x0::AbstractVector,sqΣ0, θ,dt,N,Nr)
   n = length(x0)
   n2 = 2*length(x0)
   Qp = sqrt(n)*[I(n) -I(n)]

  X = hcat([zero(x0) for i=1:N]...)
  S = hcat([zero(x0) for i=1:N]...)
  X[:,1]=x0
  Xp = x0 .+ sqΣ0*Qp
  sqΣ = sqΣ0
  Σ = sqΣ* sqΣ'
  S[:,1]= diag(Σ)
  for t=1:N-1
    if rem(t,Nr)==0
      Xp .= X[:,t] .+ sqΣ * Qp
    end
    for i=1:n2 # all quadrature points
      Xp[:,i].=Xp[:,i] + [dt*f(Xp[1:2,i],[Xp[3,i];θ[2:end]]);0]
    end
    mXp=mean(Xp,dims=2)
    X[:,t+1]=mXp
    Σ=Matrix((Xp.-mXp)*(Xp.-mXp)'/n2)
    S2=cov(Xp,dims=2)
    @show Σ
    @show S2
    S[:,t+1]=sqrt.(diag(Σ))
    # @show Σ

    sqΣ = cholesky(Σ).L

  end
  X,S
end

QXr,QSr=solve_res(f,[1.0,1.0,0.1],diagm([0.1,0.1,0.01]),θ0,0.1,1000,100)
plot(QXr[1,1:30:end],label="x",color=:blue,errorbar=QSr[1,1:30:end])
plot!(QXr[2,1:30:end],label="y",color=:red,errorbar=QSr[2,1:30:end])

function solve(f,x0::AbstractVector,Σ0, θ,dt,N)
   n = length(x0)
   n2 = 2*length(x0)
   Qp = sqrt(n)*[I(n) -I(n)]

  X = hcat([zero(x0) for i=1:N]...)
  S = hcat([zero(x0) for i=1:N]...)
  X[:,1]=x0
  Σ=Hermitian(Σ0)
  sqΣ = cholesky(Σ).L
  Xp = x0 .+ sqΣ*Qp
  S[:,1]= diag(Σ)
  for t=1:N-1
    # if rem(t,Nr)==0
    #   Xp .= X[:,t] .+ sqΣ * Qp
    # end
    for i=1:n2 # all quadrature points
      Xp[:,i].=Xp[:,i] + [dt*f(Xp[1:2,i],[Xp[3,i];θ[2:end]]);0]
    end
    mXp=mean(Xp,dims=2)
    X[:,t+1]=mXp
    Σ=cov(Xp,dims=2)
    S[:,t+1]=sqrt.(diag(Σ))
    # @show Σ
  end
  X,S,Xp
end

## Extension to arbitrary

QX,QS=solve(f,[1.0,1.0,0.1],diagm([0.1,0.1,0.01].^2),θ0,0.1,1000)
plot(QX[1,1:30:end],label="x",color=:blue,errorbar=QS[1,1:30:end])
plot!(QX[2,1:30:end],label="y",color=:red,errorbar=QS[2,1:30:end])

savefig("LV_Quadrics.svg")

function filter(f,x0::AbstractVector,Σ0, θ,dt,Ne,Y,σy)
  XX=[]
  SS=[]
  Σ=Σ0
  x=x0
  for t=1:length(Y)
    @show x
    @show diag(Σ)
    Xt,St,Xp=solve(f,x,Σ,θ,dt,Ne) # prediction
    if false
      @show Xp
      Yp = Xp[1,:] # measure only the first variable
      @show Yp
      mYp = mean(Yp)
      mXp = mean(Xp,dims=2)
      SYp = cov(Yp)+σy
      @show SYp
      C = mean([(Xp[:,i].-mXp)*(Yp[i].-mYp) for i=1:6])
      @show C
      G = C*inv(SYp)
      @show G
      x = vec(Xt[:,end] + G*(Y[t]-mYp))
      Σ = Matrix((Xp.-mXp)*(Xp.-mXp)'/n2) - G*SYp*G'

      # Σ = cov(Xp,dims=2) - G*SYp*G'
    else
        x = vec(mean(Xp,dims=2))
        Σ = Matrix((Xp.-x)*(Xp.-x)'/6.0)

    end

    push!(XX,Xt)
    push!(SS,St)
  end
  XX,SS
end

Y = X[1,100:100:end]
Xh,Sh=filter(f,[1.0,1.0,0.1],diagm([0.1,0.1,0.01].^2),θ0,0.1,100,Y[1:end],1e8)
XF=hcat(Xh...)
SF=hcat(Sh...)

step=1
plot([1:step:size(XF,2)],XF[1,1:step:end],label="x",color=:blue,errorbar=SF[1,1:step:end])
plot!([1:step:size(XF,2)],XF[2,1:step:end],label="y",color=:red,errorbar=SF[2,1:step:end])
scatter!([100:100:size(XF,2)],Y,label="measured",marker=:xcross)


function create_grid(Np, Npoin_cg, Ne, x, DFloat)

    coord_cg=zeros(DFloat,Npoin_cg)
    # intma_cg=zeros(Int64,Np,Ne)
    # periodicity_cg=zeros(Int64,Npoin_cg)

    #Set some constants
    #Always assume that the domain is divided in equispaced intervals
    xmin=-1
    xmax=+1
    dx=(xmax-xmin)/Ne

    #Generate Grid Points
    # I=1
    # coord_cg[1]=xmin
    # for e=1:Ne
    #     x0=xmin + (e-1)*dx
    #     for i=2:Np
    #         I+=1
    #         coord_cg[I]= x0 + 0.5 * ( x[i] + 1 ) * dx
    #     end
    # end

    grid_points = copy(coord_cg)
    nelements = Ne
    npoints = Np
    ξ = copy(x)
    x_min = xmin

    grid_points[1] = x_min
    for i in 1:nelements
      # calculate left element edge location
      # and index offset
      i_left = (i - 1) * (npoints - 1)
      x_left = (i - 1) * dx + x_min

      # loop of number of interpolation points
      # within each element starting at the second
      # point to avoid duplication
      for j in 2:npoints
        println((i, i_left, j))
        # calculate interpolation point location
        # relative to left element edge
        x_relative = 0.5 * (ξ[j] + 1) * dx
        grid_points[i_left + j] = x_left + x_relative
      end
    end

    coord_cg = copy(grid_points)

    # #Generate INTMA
    # I=1
    # for e=1:Ne
    #     intma_cg[1,e]=I
    #     for i=2:Np
    #         I+=1
    #         intma_cg[i,e]=I
    #     end
    # end

    # #Generate Periodicity_CG array
    # for I=1:Npoin_cg
    #     periodicity_cg[I]=I
    # end
    # periodicity_cg[Npoin_cg]=1

    return coord_cg#, intma_cg, periodicity_cg)
end

# Np = poly_order + 1 (dof)
# Npoin_cg = num_elements * poly_order + 1
# ne = num_elements
# x = interpolation points on reference element (-1,1) there should be P+1

# |. . . , . . . , . . .| = Ne * (P + 1) - (Ne-1) = Ne * P + Ne - Ne + 1 = Ne * P + 1

"""
  QuadratureStyle
"""
abstract type AbstractQuadratureStyle end

Base.@kwdef struct LegendreGaussLobatto{𝒯} <: AbstractQuadratureStyle
    polynomial_order::𝒯
end

function Base.getproperty(quadrature_style::AbstractQuadratureStyle, name::Symbol)
    if name == :npoints
        return quadrature_style.polynomial_order + 1
    elseif name == :interpolation_points
        return create_interpolation_points(quadrature_style)
    elseif name == :quadrature_weights
        return create_quadrature_weights(quadrature_style)
    else
        return getfield(quadrature_style, name)
    end
end

function Base.propertynames(::AbstractQuadratureStyle)
    return (:polynomial_order, :npoints, :interpolation_points, :quadrature_weights)
end

function create_interpolation_points(quadrature_style::LegendreGaussLobatto)
    polynomial_order = quadrature_style.polynomial_order
    npoints          = quadrature_style.npoints

    xgl=zeros(DFloat,npoints)
    p= polynomial_order #Order of the Polynomials
    ph=floor(typeof(npoints), (p+1)/2 )

    for i=1:ph
        x=cos( (2*i-1)*pi/(2*p+1) )
        for k=1:20
            (L0,L0_1,L0_2)=legendre_poly(p,x); #Compute Nth order Derivatives of Legendre Polys
            #Get new Newton Iteration
            dx=-(1-x^2)*L0_1/(-2*x*L0_1 + (1-x^2)*L0_2)
            x=x+dx
            if (abs(dx) < 1.0e-20)
                break
            end
        end
        xgl[p+2-i]=x
    end

    #Check for Zero Root
    if (p+1 != 2*ph)
        x=0
        (L0,L0_1,L0_2)=legendre_poly(p,x)
        xgl[ph+1]=x
    end

    #Find remainder of roots via symmetry
    for i=1:ph
        xgl[i]=-xgl[p+2-i]
    end

    return xgl
end

function create_quadrature_weights(quadrature_style::LegendreGaussLobatto)
    polynomial_order = quadrature_style.polynomial_order
    npoints          = quadrature_style.npoints

    wgl=zeros(DFloat,npoints)
    p=polynomial_order; #Order of the Polynomials
    ph=floor(typeof(npoints), (p+1)/2 )

    for i=1:ph
        x=cos( (2*i-1)*pi/(2*p+1) )
        for k=1:20
            (L0,L0_1,L0_2)=legendre_poly(p,x); #Compute Nth order Derivatives of Legendre Polys
            #Get new Newton Iteration
            dx=-(1-x^2)*L0_1/(-2*x*L0_1 + (1-x^2)*L0_2)
            x=x+dx
            if (abs(dx) < 1.0e-20)
                break
            end
        end
        wgl[p+2-i]=2/(p*(p+1)*L0^2)
    end

    #Check for Zero Root
    if (p+1 != 2*ph)
        x=0
        (L0,L0_1,L0_2)=legendre_poly(p,x)
        wgl[ph+1]=2/(p*(p+1)*L0^2)
    end

    #Find remainder of roots via symmetry
    for i=1:ph
        wgl[i]=+wgl[p+2-i]
    end

    return wgl
end

function legendre_poly(p,x)
    L1, L1_1, L1_2 = 0,0,0
    L0, L0_1, L0_2 = 1,0,0

    for i=1:p
        L2=L1;L2_1=L1_1;L2_2=L1_2
        L1=L0;L1_1=L0_1;L1_2=L0_2
        a=(2*i-1)/i
        b=(i-1)/i
        L0=a*x*L1 - b*L2
        L0_1=a*(L1 + x*L1_1) - b*L2_1
        L0_2=a*(2*L1_1 + x*L1_2) - b*L2_2
    end
    return (L0,L0_1,L0_2)
end

"""
  Domain
"""
abstract type AbstractDomain end

Base.@kwdef struct IntervalDomain{𝒯} <: AbstractDomain
    min::𝒯
    max::𝒯

    function IntervalDomain{𝒯}(min, max)
        @assert min < max
        new{𝒯}(min, max)
    end
end

Base.length(domain::IntervalDomain) = domain.max - domain.min

Base.size(domain::IntervalDomain) = length(domain)

"""
  Topology
  "Informally, a topology tells how elements of a set of points relate spatially to each other." (wikipedia.org)
"""
abstract type AbstractTopology end

Base.@kwdef struct ElementTopology{𝒯,𝒰} <: AbstractTopology
    domain::𝒯
    nelements::𝒰
end

function Base.getproperty(topology::ElementTopology, name::Symbol)
    if name == :nfaces
        return topology.nelements + 1
    elseif name == :interface_locations
        return create_interface_locations(topology.domain, topology.nfaces)
    else
        return getfield(topology, name)
    end
end

function Base.propertynames(::ElementTopology)
    return (:domain, :nelements, :nfaces, :interface_locations)
end

function create_interface_locations(domain::IntervalDomain, nfaces::Int)
    return Array(range(domain.min, domain.max, length = nfaces))
end

"""
  ContinuousStyle
"""
abstract type AbstractContinuityStyle end

struct ContinuousStyle <: AbstractContinuityStyle end

"""
  Grid
"""
abstract type AbstractGrid end

Base.@kwdef struct SpectralElementGrid{𝒯,𝒰,𝒱} <: AbstractGrid
    topology::𝒯
    quadrature_style::𝒰
    continuity_style::𝒱
end

function Base.getproperty(grid::SpectralElementGrid, name::Symbol)
    if name == :points
        return create_grid_points(grid.topology, grid.quadrature_style, grid.continuity_style)
    else
        return getfield(grid, name)
    end
end

function Base.propertynames(::SpectralElementGrid)
    return (:topology, :quadrature_style, :continuity_style, :points)
end

function create_grid_points(topology::𝒯, quadrature_style::𝒰, ::ContinuousStyle) where {𝒯,𝒰}
    # Unpack relevant parameters
    nelements = topology.nelements
    x_faces   = topology.interface_locations

    npoints   = quadrature_style.npoints
    ξ_points  = quadrature_style.interpolation_points

    grid_points[1] = x_min
    for i in 1:nelements
      # calculate left element edge location
      # and index offset
      dx     = x_faces[i+1] - x_faces[i]
      i_left = (i - 1) * (npoints - 1)
      x_left = (i - 1) * dx + x_min

      # loop of number of interpolation points
      # within each element starting at the second
      # point to avoid duplication
      for j in 2:npoints
        println((i, i_left, j))
        # calculate interpolation point location
        # relative to left element edge
        x_relative = 0.5 * (ξ_points[j] + 1) * dx
        grid_points[i_left + j] = x_left + x_relative
      end
    end

    return grid_points
end


import MitosisStochasticDiffEq as MSDE
using StochasticDiffEq
using Mitosis
using LinearAlgebra
using SparseArrays
using DiffEqNoiseProcess
using Test, Random

"""
forwardsample(f, g, p, s, W, x) using the Euler-Maruyama scheme
on a time-grid s with associated noise values W
"""
function forwardsample(f, g, p, s, Ws, x)
    xs = typeof(x)[]
    for i in eachindex(s)[1:end-1]
        dt = s[i+1] - s[i]
        push!(xs, x)
        x = x + f(x, p, s[i])*dt + g(x, p, s[i])*(Ws[i+1]-Ws[i])
    end
    push!(xs, x)

    return xs
end

@testset "sampling tests" begin
  # define SDE function
  f(u,p,t) = p[1]*u + p[2] - 1.5*sin.(u*2pi)
  g(u,p,t) = p[3] .- 0.2*(1 .-sin.(u))

  # set estimate of model parameters or true model parameters
  p = [-0.1,0.2,0.9]

  # time range
  tstart = 0.0
  tend = 1.0
  dt = 0.02
  trange = tstart:dt:tend

  # intial condition
  u0 = 1.1

  kernel = MSDE.SDEKernel(f,g,trange,p)
  # sample using MSDE and EM default
  sol, solend = MSDE.sample(kernel, u0)

  kernel = MSDE.SDEKernel(f,g,collect(trange),p)
  uend, (ts, u, noise) = MSDE.sample(kernel, u0, save_noise=true)

  @test isapprox(u, forwardsample(f,g,p,ts,noise,u0), atol=1e-12)
end

@testset "multivariate sampling tests" begin
  seed = 12345
  Random.seed!(seed)
  d = 2
  u0 = randn(2)
  θlin = (randn(d,d), randn(d), Diagonal([0.1, 0.1]))

  Σ(θ) = Diagonal(θ[2]) # just to generate the noise_rate_prototype

  f(u,p,t) = p[1]*u + p[2]
  f!(du,u,p,t) = (du .= p[1]*u + p[2])
  gvec(u,p,t) = diag(p[3])
  function gvec!(du,u,p,t)
    du[1] = p[3][1,1]
    du[2] = p[3][2,2]
  end
  g(u,p,t) = p[3]
  # Make `g` write the sparse matrix values
  function g!(du,u,p,t)
    du[1,1] = p[3][1,1]
    du[2,2] = p[3][2,2]
  end

  function gstepvec!(dx, _, u, p, t, dw, _)
    dx .+= diag(p[3]).*dw
  end

  function gstep!(dx, _, u, p, t, dw, _)
    dx .+= p[3]*dw
  end

  # Define a sparse matrix by making a dense matrix and setting some values as not zero
  A = zeros(2,2)
  A[1,1] = 1
  A[2,2] = 1
  A = sparse(A)

  # time range
  tstart = 0.0
  tend = 1.0
  dt = 0.02
  trange = tstart:dt:tend

  # define kernels
  k1 = MSDE.SDEKernel(f,gvec,trange,θlin)
  k2 = MSDE.SDEKernel(f,g,trange,θlin,Σ(θlin))
  k3 = MSDE.SDEKernel(f!,g!,trange,θlin,A)
  k4 = MSDE.SDEKernel(Mitosis.AffineMap(θlin[1], θlin[2]), Mitosis.ConstantMap(θlin[3]), trange, θlin, Σ(θlin))
  k5 = MSDE.SDEKernel!(f!,gvec!,gstepvec!,trange,θlin; ws = copy(u0))
  k6 = MSDE.SDEKernel!(f!,g!,gstep!,trange,θlin,A; ws = copy(A))

  @testset "StochasticDiffEq EM() solver" begin
    uend1, (ts1, u1, noise1) = MSDE.sample(k1, u0, EM(false), save_noise=true)
    NG = NoiseGrid(ts1, noise1)
    Z = pCN(NG, 1.0)
    uend2, (ts2, u2, noise2) = MSDE.sample(k2, u0, EM(false), Z, save_noise=true)
    Z = pCN(NG, 1.0)
    uend3, (ts3, u3, noise3) = MSDE.sample(k3, u0, EM(false), Z)
    Z = pCN(NG, 1.0)
    uend4, (ts4, u4, noise4) = MSDE.sample(k4, u0, EM(false), Z)
    Z = pCN(NG, 1.0)
    uend5, (ts5, u5, noise5) = MSDE.sample(k5, u0, EM(false), Z)
    Z = pCN(NG, 1.0)
    uend6, (ts6, u6, noise6) = MSDE.sample(k6, u0, EM(false), Z)

    #@show solend1
    @test isapprox(u1, u2, atol=1e-12)
    @test isapprox(uend1, uend2, atol=1e-12)
    @test isapprox(u1, u3, atol=1e-12)
    @test isapprox(uend1, uend3, atol=1e-12)
    @test isapprox(u1, u4, atol=1e-12)
    @test isapprox(uend1, uend4, atol=1e-12)
    @test isapprox(u1, u5, atol=1e-12)
    @test isapprox(uend1, uend5, atol=1e-12)
    @test isapprox(u1, u6, atol=1e-12)
    @test isapprox(uend1, uend6, atol=1e-12)
  end

  @testset "internal solver" begin
    @testset "without passing a noise" begin
      Random.seed!(seed)
      uend1, (ts1, u1, noise1) = MSDE.sample(k1, u0, MSDE.EulerMaruyama!(), save=true)
      Random.seed!(seed)
      uend2, (ts2, u2, noise2) = MSDE.sample(k2, u0, MSDE.EulerMaruyama!(), save=true)
      Random.seed!(seed)
      # inplace must be written out manually
      @test_broken uend3, (ts3, u3, noise3) = MSDE.sample(k3, u0, MSDE.EulerMaruyama!(), save=true)
      Random.seed!(seed)
      uend4, (ts4, u4, noise4) = MSDE.sample(k4, u0, MSDE.EulerMaruyama!(), save=true)
      Random.seed!(seed)
      uend5, (ts5, u5, noise5) = MSDE.sample(k5, u0, MSDE.EulerMaruyama!(), save=true)
      Random.seed!(seed)
      uend6, (ts6, u6, noise6) = MSDE.sample(k6, u0, MSDE.EulerMaruyama!(), save=true)

      @test length(ts1) == length(trange)
      @test ts1[end] == trange[end]
      @test uend1 == uend2
      @test_broken uend1 == uend3
      @test uend1 == uend4
      @test uend1 == uend5
      @test uend1 == uend6

      Random.seed!(seed)
      uend7, (ts7, u7, noise7) = MSDE.sample(k1, u0, MSDE.EulerMaruyama!(), save=false)
      @test uend1 == uend7
      @test u7 === nothing
    end

    @testset "passing a noise grid" begin
      # pass noise process and compare with EM()
      Ws = cumsum([[zero(u0)];[sqrt(trange[i+1]-ti)*randn(size(u0))
              for (i,ti) in enumerate(trange[1:end-1])]])
      NG = NoiseGrid(trange,Ws)

      uendEM, (tsEM, uEM, noiseEM) = MSDE.sample(k1, u0, EM(false), NG)
      uend1, (ts1, u1, noise1) = MSDE.sample(k1, u0, MSDE.EulerMaruyama!(), NG)
      uend2, (ts2, u2, noise2) = MSDE.sample(k2, u0, MSDE.EulerMaruyama!(), NG)
      @test_broken uend3, (ts3, u3, noise3) = MSDE.sample(k3, u0, MSDE.EulerMaruyama!(), NG)
      uend4, (ts4, u4, noise4) = MSDE.sample(k4, u0, MSDE.EulerMaruyama!(), NG)
      uend5, (ts5, u5, noise5) = MSDE.sample(k5, u0, MSDE.EulerMaruyama!(), NG)
      uend6, (ts6, u6, noise6) = MSDE.sample(k6, u0, MSDE.EulerMaruyama!(), NG)

      @test u1 ≈ uEM rtol=1e-12
      @test uendEM ≈ uend1 rtol=1e-12
      @test uendEM ≈ uend2 rtol=1e-12
      @test_broken uendEM ≈ uend3 rtol=1e-12
      @test uendEM ≈ uend4 rtol=1e-12
      @test uendEM ≈ uend5 rtol=1e-12
      @test uendEM ≈ uend6 rtol=1e-12
    end

    @testset "passing the noise values" begin
      # pass noise process and compare with EM()
      Ws = cumsum([[zero(u0)];[sqrt(trange[i+1]-ti)*randn(size(u0))
              for (i,ti) in enumerate(trange[1:end-1])]])
      NG = NoiseGrid(trange,Ws)

      uendEM, (tsEM, uEM, noiseEM)  = MSDE.sample(k1, u0, EM(false), NG)
      uend1, (ts1, u1, noise1) = MSDE.sample(k1, u0, MSDE.EulerMaruyama!(), Ws)
      uend2, (ts2, u2, noise2) = MSDE.sample(k2, u0, MSDE.EulerMaruyama!(), Ws)
      @test_broken uend3, (ts3, u3, noise3) = MSDE.sample(k3, u0, MSDE.EulerMaruyama!(), Ws)
      uend4, (ts4, u4, noise4) = MSDE.sample(k4, u0, MSDE.EulerMaruyama!(), Ws)
      uend5, (ts5, u5, noise5) = MSDE.sample(k5, u0, MSDE.EulerMaruyama!(), Ws)
      uend6, (ts6, u6, noise6) = MSDE.sample(k6, u0, MSDE.EulerMaruyama!(), Ws)

      @test u1 ≈ uEM rtol=1e-12
      @test uendEM ≈ uend1 rtol=1e-12
      @test uendEM ≈ uend2 rtol=1e-12
      @test_broken uendEM ≈ uend3 rtol=1e-12
      @test uendEM ≈ uend4 rtol=1e-12
      @test uendEM ≈ uend5 rtol=1e-12
      @test uendEM ≈ uend6 rtol=1e-12
    end

    @testset "custom P" begin
      # checks that defining and passing P manually works

      struct customP{θType}
        θ::θType
      end

      function MSDE.tangent!(du, u, dz, P::customP)
        du[3] .= (P.θ[1]*u[3]+P.θ[2])*dz[2] + P.θ[3]*dz[3]

        (dz[1], dz[2], du[3])
      end

      function MSDE.exponential_map!(u, du, P::customP)
        x = u[3]
        @. x += du[3]
        (u[1] + du[1], u[2] + du[2], x)
      end

      # pass noise process and compare with EM()
      Ws = cumsum([[zero(u0)];[sqrt(trange[i+1]-ti)*randn(size(u0))
              for (i,ti) in enumerate(trange[1:end-1])]])
      NG = NoiseGrid(trange,Ws)

      uendEM, (tsEM, uEM, noiseEM) = MSDE.sample(k1, u0, EM(false), NG)
      uend1, (ts1, u1, noise1) = MSDE.sample(k1, u0, MSDE.EulerMaruyama!(), Ws, P=customP(θlin))
      uend2, (ts2, u2, noise2) = MSDE.sample(k2, u0, MSDE.EulerMaruyama!(), Ws, P=customP(θlin))
      uend3, (ts3, u3, noise3) = MSDE.sample(k3, u0, MSDE.EulerMaruyama!(), Ws, P=customP(θlin))
      uend4, (ts4, u4, noise4)  = MSDE.sample(k4, u0, MSDE.EulerMaruyama!(), Ws)
      uend5, (ts5, u5, noise5) = MSDE.sample(k5, u0, MSDE.EulerMaruyama!(), Ws)
      uend6, (ts6, u6, noise6) = MSDE.sample(k6, u0, MSDE.EulerMaruyama!(), Ws)

      @test u1 ≈ uEM rtol=1e-12
      @test uendEM ≈ uend1 rtol=1e-12
      @test uendEM ≈ uend2 rtol=1e-12
      @test uendEM ≈ uend3 rtol=1e-12
      @test uendEM ≈ uend4 rtol=1e-12
      @test uendEM ≈ uend5 rtol=1e-12
      @test uendEM ≈ uend6 rtol=1e-12
    end

  end
end


# Algorithm 2.2
# Fixed point iteration, solves g(p) = p

function bisect_approx(f, p0, tolerance, maxIter)
	for i ∈ 1:maxIter
		p = f(p0)
		if abs(p - p0) < tolerance
			return p
		end
		p0 = p
	end
	throw(ArgumentError("Method failed"))
end

println(bisect_approx(x -> ((10/(4 + x))^(1/2)), 1.5, 0.1, 4))