# scalar
set = Setup(
    "scalar", # matter
    "advection", # case
    "1d0f0v", # space
    "gks", # flux
    "", # collision: for scalar conservation laws there are none
    1, # species
    1, # interpolation order
    "vanleer", # limiter
    "period", # boundary
    0.5, # cfl
    1.0, # simulation time
)
pSpace = PSpace1D(0.0, 1.0, 100, 1)
vSpace = nothing
property = Scalar(1.0, 1e-6)
w0 = 1.0
prim0 = conserve_prim(w0, property.a)
ib = IB(x -> sin(2π * x), property)
ks = SolverSet(set, pSpace, vSpace, property, ib)

ctr, face = init_fvm(ks, ks.ps)
for i in eachindex(ctr)
    ctr[i].w = sin(2π * ks.pSpace.x[i])
end

dt = 1e-3
reconstruct!(ks, ctr)
evolve!(ks, ctr, face, dt)
update!(ks, ctr, face, dt, 0.0)


using BayesianLinearRegression
using Test
using LinearAlgebra

const BLR = BayesianLinearRegression

function runtests()
    X = randn(100,30)
    β = randn(30)
    y = X * β .+ randn(100)

    m = BayesianLinReg(X,y)
    fit!(m)

    @testset "log-determinant of Hessian" begin
        @test BLR.logdetH(m) ≈ logdet(hessian(m))
    end

    @testset "Inverse Hessian" begin
        @test m.Hinv * hessian(m) ≈ I
    end

    @testset "Sum of Squared Residuals" begin
        @test BLR.ssr(m) ≈ BLR.normSquared(y - predict(m,X; uncertainty=false))
    end

    # This comes from the equating γ = γ between the α and β updates
    @testset "N invariant" begin
        α = priorPrecision(m)
        β = noisePrecision(m)
        w² = BLR.normSquared(m.weights)
        ssr = BLR.ssr(m)
        @test α * w² + β * ssr ≈ m.N
    end
end

runtests()


"""
math : keeps track of a few mathy functions.
"""

module math

using StatsBase

function mylogsumexp(b)
	B = maximum(b, dims=2)
	return log.(sum(exp.(broadcast(-, b, B)), dims=2)) + B

end

function abs_tol(older, newer, thresh=0.0002)
	if abs(older - newer) < thresh
		return true
	else
		return false
	end
end

function rel_tol(older, newer, thresh=0.0004)
	fun = abs(1 - (older / newer))
	if fun < thresh
		return true
	else
		return false
	end
end

function softmax(vec, num_samples::Int=1, temperature::Float64=1.0)
	vec = vec ./ temperature
	max_value = maximum(vec)
	log_partition = log.(sum(exp.(vec .- max_value))) .+ max_value

	prob_vec = exp.(vec .- log_partition)

	wv = Weights(prob_vec)
	result = sample(1:length(prob_vec), wv, num_samples)
	return result
end

end


#############################################################################
# add_subtract.jl
# Handles unary negation, addition and subtraction of variables, constants
# and expressions.
# All expressions and atoms are subtpyes of AbstractExpr.
# Please read expressions.jl first.
#############################################################################

export +, -
export sign, curvature, monotonicity, evaluate

### Unary Negation

struct NegateAtom <: AbstractExpr
    head::Symbol
    id_hash::UInt64
    children::Tuple{AbstractExpr}
    size::Tuple{Int, Int}

    function NegateAtom(x::AbstractExpr)
        children = (x,)
        return new(:-, hash(children), children, x.size)
    end
end

function sign(x::NegateAtom)
    return -sign(x.children[1])
end

function monotonicity(x::NegateAtom)
    return (Nonincreasing(),)
end

function curvature(x::NegateAtom)
    return ConstVexity()
end

function evaluate(x::NegateAtom)
    return -evaluate(x.children[1])
end

-(x::AbstractExpr) = NegateAtom(x)

function conic_form!(x::NegateAtom, unique_conic_forms::UniqueConicForms)
    if !has_conic_form(unique_conic_forms, x)
        objective = conic_form!(x.children[1], unique_conic_forms)
        objective = -objective
        cache_conic_form!(unique_conic_forms, x, objective)
    end
    return get_conic_form(unique_conic_forms, x)
end

### Addition
struct AdditionAtom <: AbstractExpr
    head::Symbol
    id_hash::UInt64
    children::Array{AbstractExpr, 1}
    size::Tuple{Int, Int}

    function AdditionAtom(x::AbstractExpr, y::AbstractExpr)
        # find the size of the expression = max of size of x and size of y
        if x.size == y.size || y.size == (1, 1)
            sz = x.size
        elseif x.size == (1, 1)
            sz = y.size
        else
            error("Cannot add expressions of sizes $(x.size) and $(y.size)")
        end
        # see if we're forming a sum of more than two terms and condense them
        children = AbstractExpr[]
        if isa(x, AdditionAtom)
            append!(children, x.children)
        else
            push!(children, x)
        end
        if isa(y, AdditionAtom)
            append!(children, y.children)
        else
            push!(children, y)
        end
        return new(:+, hash(children), children, sz)
    end
end

function sign(x::AdditionAtom)
    return sum(Sign[sign(child) for child in x.children])
    # Creating an array of type Sign and adding all the sign of xhildren of x so if anyone is complex the resultant sign would be complex.
end

function monotonicity(x::AdditionAtom)
    return Monotonicity[Nondecreasing() for child in x.children]
end

function curvature(x::AdditionAtom)
    return ConstVexity()
end

function evaluate(x::AdditionAtom)
    # broadcast is used in reduction instead of using sum directly to support addition
    # between scalars and arrays
    return mapreduce(evaluate, (a, b) -> a .+ b, x.children)
end

function conic_form!(x::AdditionAtom, unique_conic_forms::UniqueConicForms)
    if !has_conic_form(unique_conic_forms, x)
        objective = ConicObj()
        for child in x.children
            child_objective = conic_form!(child, unique_conic_forms)
            if x.size != child.size
                child_objective = promote_size(child_objective, length(x))
            end
            objective += child_objective
        end
        cache_conic_form!(unique_conic_forms, x, objective)
    end
    return get_conic_form(unique_conic_forms, x)
end

+(x::AbstractExpr, y::AbstractExpr) = AdditionAtom(x, y)
+(x::Value, y::AbstractExpr) = AdditionAtom(Constant(x), y)
+(x::AbstractExpr, y::Value) = AdditionAtom(x, Constant(y))
-(x::AbstractExpr, y::AbstractExpr) = x + (-y)
-(x::Value, y::AbstractExpr) = Constant(x) + (-y)
-(x::AbstractExpr, y::Value) = x + Constant(-y)


module Functions

using ..TopOpt: dim, whichdevice, CPU, GPU, TopOpt, PENALTY_BEFORE_INTERPOLATION
using ..TopOptProblems, ..TrussTopOptProblems
using ..TopOptProblems: initialize_K, getdh
using ..FEA, ..CheqFilters
using ..Utilities, ForwardDiff, LinearAlgebra, Requires
using ..TrussTopOptProblems: getA, compute_local_axes

using Parameters: @unpack
using TimerOutputs, Ferrite, StaticArrays, StatsFuns
using SparseArrays, Statistics, ChainRulesCore, Zygote
using Nonconvex: Nonconvex
using Flux

export Volume,
    Compliance,
    Displacement,
    MeanCompliance,
    BlockCompliance,
    AbstractFunction,
    getfevals,
    getmaxfevals,
    maxedfevals,
    MicroVonMisesStress,
    MacroVonMisesStress,
    project,
    generate_scenarios,
    hutch_rand!,
    hadamard!,
    TrussStress,
    AssembleK,
    TrussElementKσ,
    ElementK,
    apply_boundary_with_zerodiag!,
    apply_boundary_with_meandiag!,
    NeuralNetwork,
    PredictFunction,
    TrainFunction

const to = TimerOutput()

abstract type AbstractFunction{T} <: Nonconvex.NonconvexCore.AbstractFunction end

include("function_utils.jl")
include("compliance.jl")
include("displacement.jl")
include("volume.jl")
include("stress.jl")
include("trace.jl")
include("mean_compliance.jl")
include("block_compliance.jl")

# buckling
include("apply_boundary.jl")
include("assemble_K.jl")
include("element_ksigma.jl")
include("element_k.jl")

# TODO no rrules yet
include("truss_stress.jl")

include("neural.jl")

end


import Optim: retract!, project_tangent!

# * Sphere

metric_norm(x, g::UniformScaling) = norm(x)*√(g.λ)
metric_norm(x, g) = √(dot(x,g,x))

metric_normalize!(x, ::UniformScaling{Bool}) = normalize!(x)
function metric_normalize!(x, g)
    x ./= metric_norm(x,g)
end

"""Spherical manifold {|x| = 1}."""
struct MetricSphere{M} <: Manifold
    "Metric used for inner products, i.e. `⟨a|b⟩_g = gₘₙaᵐbⁿ"
    g::M
end
MetricSphere() = MetricSphere(I)
Optim.retract!(S::MetricSphere, x) = metric_normalize!(x, S.g)
Optim.project_tangent!(S::MetricSphere,g,x) = (g .-= real(dot(x, S.g, g)).*x)

# * Stiefel

matfun(f::Function, ee) = ee.vectors*Diagonal(f.(ee.values))*ee.vectors'

# For orthogonal bases, e.g. FE-DVR
löwdin_transform(g::UniformScaling{Bool}) = g,g
# For quasi-orthogonal bases, e.g. finite-differences
löwdin_transform(g::UniformScaling) = inv(√(g.λ))*I,√(g.λ)*I
löwdin_transform(g::Diagonal) = inv(√(g)),√(g)
# For non-orthogonal bases, e.g. B-splines
function löwdin_transform(g::AbstractMatrix)
    ee = eigen(g)
    matfun(λ -> inv(√(λ)), ee), matfun(λ -> √(λ), ee)
end
function löwdin_transform(g::BandedMatrix)
    m,n = size(g)
    m == n || throw(DimensionMismatch("Löwdin transform only valid for square matrices"))
    l,u = bandwidths(g)
    if l == u == 0
        löwdin_transform(Diagonal(g))
    else
        löwdin_transform(Symmetric(g))
    end
end

löwdin_transform!(v, ::UniformScaling{Bool}) = v
löwdin_transform!(v, S⁻ᴴ) = (v .= S⁻ᴴ*v)

abstract type MetricStiefel <: Stiefel end
struct MetricStiefel_CholQR{L} <: MetricStiefel
    "Löwdin transform to orthogonal basis"
    S⁻ᴴ::L
    "Löwdin transform to non-orthogonal basis"
    Sᴴ::L
end
struct MetricStiefel_SVD{L} <: MetricStiefel
    "Löwdin transform to orthogonal basis"
    S⁻ᴴ::L
    "Löwdin transform to non-orthogonal basis"
    Sᴴ::L
end
function MetricStiefel(g, retraction=:SVD)
    S⁻ᴴ,Sᴴ = löwdin_transform(g)
    if retraction == :CholQR
        MetricStiefel_CholQR(S⁻ᴴ,Sᴴ)
    elseif retraction == :SVD
        MetricStiefel_SVD(S⁻ᴴ,Sᴴ)
    else
        throw(ArgumentError("Unknown retraction $retraction"))
    end
end
function retract!(St::MetricStiefel_SVD, X)
    löwdin_transform!(X, St.Sᴴ)
    U,S,V = svd(X)
    X .= U*V'
    löwdin_transform!(X, St.S⁻ᴴ)
end
function retract!(St::MetricStiefel_CholQR, X)
    löwdin_transform!(X, St.Sᴴ)
    overlap = X'X
    X .= X/cholesky(overlap).L
    löwdin_transform!(X, St.S⁻ᴴ)
end
# For functions depending only on the subspace spanned by X, we always
# have G = A*X for some A, and so X'G = G'X, and Stiefel == Grassmann
# Edelman et al. have G .-= X*G'X (2.53), corresponding to a different
# metric ("canonical metric"). We follow Absil et al. here and use the
# metric inherited from Nxn matrices.
function project_tangent!(St::MetricStiefel, G, X)
    löwdin_transform!(X, St.Sᴴ)
    löwdin_transform!(G, St.Sᴴ)
    XG = X'G
    G .-= X*((XG .+ XG')./2)
    löwdin_transform!(X, St.S⁻ᴴ)
    löwdin_transform!(G, St.S⁻ᴴ)
end

# * Sphere & Stiefel combo

struct SphereStiefelCombo{Sp,St,M,F} <: Manifold
    "Indices of columns that only need normalization"
    normal_indidices::Vector{Int}
    "Manifold used to normalize columns"
    sphere::Sp
    "Sets of indices of columns that should be mutually orthogonalized"
    ortho_indidices::Vector{Vector{Int}}
    "Manifold used to orthogonalize sets of columns"
    stiefel::St
    "Temporary storage for orthgonalization"
    ortho_X::M
    "Temporary storage for orthgonalization"
    ortho_G::M
    f::F
end

function copytotmp!(dst, src, ortho_indices)
    for (i,j) in enumerate(ortho_indices)
        copyto!(view(dst, :, i), view(src, :, j))
    end
end

function copyfromtmp!(dst, src, ortho_indices)
    for (i,j) in enumerate(ortho_indices)
        copyto!(view(dst, :, j), view(src, :, i))
    end
end

function retract!(ssc::SphereStiefelCombo, X)
    retract!(ssc.sphere, view(X, :, ssc.normal_indidices))
    for o in ssc.ortho_indidices
        copytotmp!(ssc.ortho_X, X, o)
        retract!(ssc.stiefel, view(ssc.ortho_X, :, 1:length(o)))
        copyfromtmp!(X, ssc.ortho_X, o)
    end
    X
end

function project_tangent!(ssc::SphereStiefelCombo, G, X)
    project_tangent!(ssc.sphere, view(G, :, ssc.normal_indidices), view(X, :, ssc.normal_indidices))
    for o in ssc.ortho_indidices
        sel = 1:length(o)
        copytotmp!(ssc.ortho_G, G, o)
        copytotmp!(ssc.ortho_X, X, o)
        project_tangent!(ssc.stiefel, view(ssc.ortho_G, :, sel), view(ssc.ortho_X, :, sel))
        copyfromtmp!(G, ssc.ortho_G, o)
        copyfromtmp!(X, ssc.ortho_X, o)
    end
    G
end

# * Manifold setup

function setup_manifold(f::FockProblem)
    fock = f.fock
    P = orbitals(fock.quantum_system)
    m,n = size(P)

    normal_indices = sort(reduce(vcat, filter(s -> length(s)==1, fock.symmetries), init=Int[]))
    ortho_indices = filter(s -> length(s)≠1, fock.symmetries)
    n_ortho = length.(ortho_indices)

    length(normal_indices) + sum(n_ortho) == n ||
        throw(DimensionMismatch("Number of orbitals does not agree with amount of symmetry specified"))

    sphere = Optim.PowerManifold(MetricSphere(fock.S), (m,), (length(normal_indices),))
    stiefel = MetricStiefel(fock.S)

    nn = isempty(n_ortho) ? 0 : maximum(n_ortho)

    T = eltype(P)
    ortho_X = Matrix{T}(undef, m, nn)
    ortho_G = Matrix{T}(undef, m, nn)

    SphereStiefelCombo(normal_indices, sphere,
                       ortho_indices, stiefel,
                       ortho_X, ortho_G, f)
end


using LogNormals
using Test, Distributions, StatsBase, Missings, MissingStrategies, Random
using OffsetArrays, RecursiveArrayTools, LinearAlgebra
using DistributionVectors # cormatri_for_acf
using Unitful

@testset "sem_cor" begin
    acf0 = [1,0.4,0.1]
    Sigma = cormatrix_for_acf(100, acf0);
    dmn = MvNormal(ones(100), Symmetric(Sigma));
    Random.seed!(1234)
    a = rand(dmn);
    am = allowmissing(a); am[3:4] .= missing;
    az = copy(a); az[3:4] .= 0.0;
    amu = am .* u"m"; # unitful missing vector of meters
    @testset "count_forlags" begin
        pred(x_i, x_iplusk) = ismissing(x_i) || ismissing(x_iplusk)
        x = [1,2,missing,missing,5,6]
        lags = 0:5
        nmiss = @inferred count_forlags(pred, x, lags)
        @test nmiss == [2,3,4,2,0,0]
        x = [1,2,missing,missing,5]
        lags = 0:4
        nmiss = @inferred count_forlags(pred, x, lags)
        @test nmiss == [2,3,3,1,0]
        nmissu = @inferred count_forlags(pred, x.*u"m", lags)
        @test nmissu == nmiss
        # test for custom indices
        nmiss2 = count_forlags(pred, OffsetArray(x,-1), lags)
        @test nmiss2 == nmiss
        # views still work and larger lags return 0
        nmiss3 = @inferred count_forlags(pred, view(x,1:2:5), lags)
        @test nmiss3 == [1,2,0,0,0]
    end;
    @testset "autocorr with missings" begin
        # for comparability center all around nomissing mean
        meana = mean(a)
        lags = 0:8
        acfe1 = @inferred autocor(a, lags) # StatsBase still works
        acfe2 = @inferred autocor(a, lags, PassMissing()) # can supply Missing clause
        @test acfe1 == acfe2
        # get only missing if not explicitly requesting SkipMissing or ExactMissing
        ismissing(@inferred(Vector{Float64},autocor(am, lags, PassMissing(); demean=false)))
        ismissing(@inferred(Vector{Float64},autocor(am, lags)))
        ismissing(@inferred(Vector{Float64},autocor(am, lags; demean=false)))
        acfe = @inferred autocor(a.-meana, lags; demean=false);
        acfemf = @inferred autocor(am.-meana, lags, SkipMissing(); demean=false);
        acfem = @inferred(autocor(am.-meana, lags, ExactMissing(); demean=false))
        acfes = @inferred autocor((a.-meana)[4:end], lags; demean=false) # different nobs
        azdemean = copy(a.-meana); azdemean[ismissing.(am)] .= 0.0
        acfez = @inferred autocor(azdemean, lags; demean=false);
        hcat(acfe, acfem, acfez, acfemf)[1:4,:]
        ## with SkipMissing() rather than ExactMissing(): missing same as vector with zeros
        @test acfemf == acfez
        @test acfem[1] == acfez[1] == 1
        # with excactmissing: correct var(x) by two missings (100-2)
        # lag1: three missings
        @test acfem[2] ≈ acfez[2]*98/97
        # lag2: four missings
        @test acfem[3] ≈ acfez[3]*98/96
        # lag3: three missings
        @test acfem[4] ≈ acfez[4]*98/97
        # lag4: two missings as in z
        @test acfem[5:end] ≈ acfez[5:end]
        # TODO after StatsBase.autocor can deal with unitful
        #acfemu = @inferred autocor((am.-meana).*u"m", lags; demean=false);
        #
        # variant without lags
        ismissing(@inferred autocor(am))
        ismissing(@inferred autocor(am, PassMissing()))
        @test acfemf ≈ (@inferred autocor(am.-meana, SkipMissing(); demean=false))[1:9]
        @test acfem ≈ (@inferred(autocor(am.-meana, ExactMissing(); demean=false)))[1:9]
        #SimpleTraits.istrait(IsEltypeSuperOfMissing{typeof(am.-meana)})
        #
        # with demean
        autocor(hcat(a,a), lags) # original variant
        tmp = @inferred autocor(am, lags, SkipMissing());
        tmp2 = @inferred autocor(hcat(am,am), lags, SkipMissing());
        @test tmp2[:,1] == tmp2[:,2] == tmp
        tmp = @inferred autocor(am, lags, ExactMissing());
        tmp2 = @inferred autocor(hcat(am,am), lags, ExactMissing());
        @test tmp2[:,1] == tmp2[:,2] == tmp
    end;
    @testset "autocor_effective" begin
        effa = @inferred autocor_effective(a, acf0)
        @test effa == acf0
        effa = @inferred autocor_effective(a, [acf0..., -0.001])
        @test effa == acf0
    end;
    @testset "effective_n_cor" begin
        neff = @inferred effective_n_cor(a, acf0)
        @test neff < length(a)
        @test neff ≈ 50.30 atol=0.01 # regression test
        # call with MissingStrategy but nonmissing type
        neffs = @inferred effective_n_cor(a, acf0, ExactMissing())
        @test neffs == neff
        #
        @test ismissing(@inferred(Missing,effective_n_cor(am, acf0, PassMissing())))
        @test ismissing(@inferred(Missing,effective_n_cor(am, acf0)))
        neffm = @inferred effective_n_cor(am, acf0, ExactMissing())
        @test neffm < neff
        @test neffm ≈ 49.61 atol=0.01 # regression test
        # n smaller than length(acf):
        neff2 = @inferred effective_n_cor(a[1:2], [acf0..., 0.05, 0.05])
        @test neff2 ≈ 1.43 atol=0.01 # regression test
        # no correlation estimate for lag 3
        neff3 = @inferred effective_n_cor(am[1:4], [acf0..., 0.05, 0.05], ExactMissing())
        @test neff3 ≈ 1.43 atol=0.01 # regression test
        #
        # without specifying acf
        ismissing(@inferred(Float64,effective_n_cor(am)))
        neff = @inferred effective_n_cor(am, ExactMissing())
        @test neff == effective_n_cor(am, autocor(am, ExactMissing()), ExactMissing())
    end;
    @testset "var_cor" begin
        va1 = @inferred var_cor(a, acf0, ExactMissing())
        va2 = @inferred var_cor(a, acf0, PassMissing())
        va = @inferred var_cor(a, acf0)
        # slight differences between var(x) and var(skipmissing(x))
        @test va ≈ va1 ≈ va2
        # need to specify handling of missing otherwise get missing back
        @test ismissing(@inferred(Float64,var_cor(am, acf0)))
        vam = @inferred var_cor(am, acf0, ExactMissing())
        #@code_warntype var_cor(am, acf0, ms=ExactMissing(); neff=nothing)
        @test isnan(var_cor(am[[3]],acf0, ExactMissing()))
        @test isnan(var_cor(am[2:3],acf0, ExactMissing())) # only one non-missing
        vamf = @inferred var_cor(am, acf0, SkipMissing())
        # larger uncertainty with missings
        @test vam > va
        # not correcting neff for missings overestimates neff
        # and hence underestimates uncertainty
        @test vamf < vam
        # calling without acf
        @test @inferred var_cor(a) > 0
        @test ismissing(@inferred(var_cor(am)))
        @test ismissing(@inferred(var_cor(am, PassMissing())))
        @test @inferred(var_cor(am, ExactMissing())) > 0
    end;
    @testset "semcor with same effective acf" begin
        se_a = @inferred sem_cor(a, acf0)
        ismissing(@inferred(Float64, sem_cor(am, acf0, PassMissing())))
        ismissing(@inferred(Float64, sem_cor(am, acf0)))
        se_am = @inferred sem_cor(am, acf0, ExactMissing())
        @test isnan(sem_cor(am[[3]],acf0, ExactMissing()))
        se_az = @inferred sem_cor(az, acf0)
        ar = filter(x -> !ismissing(x), am)
        se_ar = @inferred Missing sem_cor(ar, acf0)
        [se_a, se_am, se_az, se_ar]
        # the few missing do not lead to large deviations
        @test se_am ≈ se_a atol=0.02
        # skipping missing values has larger stderror, because shorter series
        @test se_ar > se_a
        # caring for missing pairs in correlation also has larger stderror
        # because some of larger relative error in sum_normals
        @test se_am > se_a
        # the SkipMissing() case may lead to lower estimates
        # of unceertainty than without missings?
    end;
    @testset "semcor calling without acf" begin
        @test @inferred sem_cor(a, ExactMissing()) ≈ sem_cor(a, autocor_effective(a))
        @test @inferred sem_cor(a, SkipMissing()) ≈ sem_cor(a, autocor_effective(a))
        @test @inferred sem_cor(a, PassMissing()) == sem_cor(a, autocor_effective(a))
        @test @inferred sem_cor(am, ExactMissing()) == sem_cor(am,
            autocor_effective(am, ExactMissing()), ExactMissing())
        @test @inferred sem_cor(am, SkipMissing()) == sem_cor(am,
            autocor_effective(am, SkipMissing()), SkipMissing())
    end;
    @testset "semcor with many random missings" begin
        b = allowmissing(a);
        Random.seed!(1234)
        b[sample(axes(a,1),60, replace=false)] .= missing;
        se_a = @inferred sem_cor(a, acf0)
        se_b = @inferred Missing sem_cor(b, acf0, ExactMissing())
        se_be = @inferred Missing sem_cor(b, ExactMissing())
        # uncertainty increases due to missings (fewer obs)
        @test se_b > se_a
        #@test se_be > se_a
        # regression tests with random seed 1234
        @test se_b ≈ 0.18 atol = 0.01
    end;
   @testset "semcor short series" begin
        @test isnan(@inferred(var_cor([1.0])))
        @test isnan(@inferred(sem_cor([1.0])))
        @test @inferred sem_cor(1.0:2) == 0.5
        @test @inferred sem_cor(1.0:3) == sqrt(var(1:3)/3)
        @test isnan(@inferred var_cor(fill(1.0,4))) # autocorrelation NaN
        #@test @inferred var_cor(fill(1.0,4)) == 0.0
        @test isnan(@inferred sem_cor(fill(1.0,4))) # autocorrelation NaN
        #@test @inferred sem_cor(fill(1.0,4)) == 0.0
        @test ismissing(@inferred(Float64, sem_cor([1.0,2,missing])))
        @test ismissing(@inferred(Float64, var_cor([1.0,2,missing])))
        @test @inferred var_cor([1.0,2,missing], ExactMissing()) == 0.5
        @test @inferred sem_cor([1.0,2,missing], ExactMissing()) == sqrt(0.5/2)
        @test isnan(@inferred sem_cor([1.0,missing,missing], ExactMissing()))
   end;
end;

@testset "autocorr with many missings" begin
    # ovrriding a from sd testset
    nsum = 1_000
    acf0 = [1,0.4,0.1]
    Sigma = cormatrix_for_acf(nsum, acf0);
    dmn = MvNormal(ones(nsum), Symmetric(Sigma));
    a = rand(dmn);
    probgap = 0.4
    clustersize = 4
    # need to average over some clusters to get reliable estimate
    nboot = 30
    resboot = map(1:nboot) do iboot
        igap0 = sample(0:(nsum-clustersize),trunc(Int,nsum*probgap/clustersize), replace=true);
        b = allowmissing(a);
        # outer product see https://stackoverflow.com/a/44592419
        #igap0.+Base.OneTo(clustersize)'
        b[igap0.+Base.OneTo(clustersize)'] .= missing;
        ae = autocor(b, 1:2, ExactMissing())
        af = autocor(b, 1:2, SkipMissing())
        (ae, af)
    end;
    aeb, afb = VectorOfArray.(vectuptotupvec(resboot));
    ae = mean(aeb; dims=2)
    af = mean(afb; dims=2)
    aorig = autocor(a, 1:2)
    @test ae ≈ aorig atol=0.02
    @test all((aorig .- af)./aorig .> 0.05) # >5% low bias
end;


# v0.6

function insertionsort!(A::Array{T}) where T <: Number
    for i in 1:length(A)-1
        value = A[i+1]
        j = i
        while j > 0 && A[j] > value
            A[j+1] = A[j]
            j -= 1
        end
        A[j+1] = value
    end
    return A
end

x = randn(5)
@show x insertionsort!(x)


# ---
# jupyter:
#   jupytext:
#     text_representation:
#       extension: .jl
#       format_name: light
#       format_version: '1.5'
#       jupytext_version: 1.4.2
#   kernelspec:
#     display_name: Julia 1.4.0
#     language: julia
#     name: julia-1.4
# ---

print("Interact Sample")

using Interact
using Plots
using Rotations

# +
function f(a)
    xs = -pi:0.01:pi
    ys = @. sin(xs - a)
    plot(xs, ys)
end

a_slider=slider(-3:0.1:3,value=0.2, label="a_slider")
p=map(f, observe(a_slider))
map(display,[a_slider, p])
# +
p=ones(3)/sqrt(3)
scatter3d(
    [p[1]],[p[2]],[p[3]],
    xlabel="x",ylabel="y",zlabel="z",
    color=:blue
)

function plot_rotated_pts(x_angle,y_angle,z_angle)
    xyz_rad = deg2rad.([x_angle,y_angle,z_angle])
    r=RotXYZ(xyz_rad...)
    q=r*p
    scatter3d!([q[1]],[q[2]],[q[3]],color=:green,legend=false)
    xlims!((-1.0,1.0)); ylims!((-1.0,1.0)); zlims!((-1.0,1.0))
end

x_axis=slider(-180:180, label="x_axis")
y_axis=slider(-180:180, label="y_axis")
z_axis=slider(-180:180, label="z_axis")

angle_plot=map(
    (x_deg,y_deg,z_deg)->plot_rotated_pts(x_deg,y_deg,z_deg),
    map(observe, [x_axis,y_axis,z_axis])...,
)
    #observe(x_axis),observe(y_axis),observe(z_axis)

map(display,[x_axis, y_axis, z_axis, angle_plot]);


# ---
# title: 645. Set Mismatch
# id: problem645
# author: Indigo
# date: 2021-07-04
# difficulty: Easy
# categories: Hash Table, Math
# link: <https://leetcode.com/problems/set-mismatch/description/>
# hidden: true
# ---
#
# The set `S` originally contains numbers from 1 to `n`. But unfortunately, due
# to the data error, one of the numbers in the set got duplicated to **another**
# number in the set, which results in repetition of one number and loss of
# another number.
#
# Given an array `nums` representing the data status of this set after the
# error. Your task is to firstly find the number occurs twice and then find the
# number that is missing. Return them in the form of an array.
#
# **Example 1:**
#
#
#
#     Input: nums = [1,2,2,4]
#     Output: [2,3]
#
#
# **Note:**
#
#   1. The given array size will in the range [2, 10000].
#   2. The given array's numbers won't have any order.
#
#
## @lc code=start
using LeetCode

function find_error_nums(nums::Vector{Int})
    cnt = fill(0x0, length(nums))
    for num in nums
        cnt[num] += 1
    end
    return findfirst(==(0x2), cnt), findfirst(==(0x0), cnt)
end
## @lc code=end


if lowercase(get(ENV, "CI", "false")) == "true"
    include("install_dependencies.jl")
end

using Simplices
using Simplices: SimplexSplitting

using Test
using LinearAlgebra
using Printf

println("Running tests for Simplices.jl ....")

include("even_sampling.jl")
include("test_delaunay.jl")
include("test_intersectionvol.jl")
include("tensordecomposition.jl")

println("Tests finished ....")


workspace()
using TimeWindowStats
using Base.Test, Dates
reload("TimeWindowStats")
dt1 = DateTime(2015,6,4,7,1,1)
dt2 = dt1 + Second(2)
dt3 = dt1 + Second(4)
dt4 = dt1 + Second(5)
dt5 = dt1 + Second(8)

c = Sum(5000,4)
update!(c,1., dt1)
@test c.head == 1
@test c.tail == 1
@test c.len == 1
@test c.v == 1.
@test typeof(c.xs) == Vector{Float64}
@test typeof(c.v) == Float64

update!(c,2., dt2)
@test c.head == 2
@test c.tail == 1
@test c.len == 2
@test c.v == 3.

update!(c,3., dt3)
@test c.head == 3
@test c.tail == 1
@test c.len == 3
@test c.v == 6.

update!(c,4., dt4)
@test c.head == 4
@test c.tail == 1
@test c.len == 4
@test c.size == 5
@test c.v == 10.

update!(c,5., dt4)
@test c.head == 5
@test c.tail == 1
@test c.len == 5
@test c.size == 6
@test c.v == 15.

c.head
c.xs

update!(c,6., dt5)
@test c.head == 6
@test c.tail == 3
@test c.len == 4
@test c.size == 6
@test c.v == 18.

update!(c,7., dt5)
@test c.head == 1
@test c.tail == 3
@test c.len == 5
@test c.size == 6
@test c.v == 25.

update!(c,8., dt5)
@test c.head == 6
@test c.tail == 1
@test c.len == 6
@test c.size == 7
@test c.v == 33.

update!(c,9., dt5)
@test c.head == 7
@test c.tail == 1
@test c.len == 7
@test c.size == 8
@test c.v == 42.

c = Sum(5000)
update!(c,1, dt1)
@test c.v == 1
@test typeof(c.xs) == Vector{Int64}
@test typeof(c.v) == Int64

c = Pct(5000)
update!(c,1,1, dt1)
@test c.v == 1.
update!(c,0,1, dt2)
@test c.v == .5

c = Mean(5000)
update!(c,1, dt1)
@test c.v == 1.
update!(c,2, dt2)
@test c.v == 1.5
update!(c,3, dt3)
@test c.v == 2.
update!(c,4, dt4)
@test c.v == 2.5
update!(c,5, dt5)
@test c.v == 4.


const pyspatial = PyNULL()
scipy = false
function __init__()
    copy!(pyspatial,  pyimport("scipy.spatial"))
    # is scipy available?
    global scipy = pyspatial == PyNULL() ? false : true
    return
end
#-------------------------------------------------------------------------------
function edgestoloop(el::Array{Int64, 2})
    elcol = reshape(el, length(el), 1)
    listep, I, J = uniquerows(elcol)
    J = reshape(J, size(el))
    ptoe = zeros(Int64, length(I), 2)
    for k = 1:size(el, 1)
        for l = 1:2
            if ptoe[J[k, l], 1] == 0
                ptoe[J[k, l], 1] = k
            else
                ptoe[J[k, l], 2] = k
            end
        end
    end
    if minimum(ptoe) == 0
        error("Pb I in edgestoloop")
    end
    loop = zeros(Int64, length(I)); loop[1:2] = J[1, :]; ec = 1 # courant edge
    for k = 3:length(I)
        ec = ec == ptoe[loop[k - 1], 1] ? ptoe[loop[k - 1], 2] : ptoe[loop[k - 1], 1]
        #ec = setdiff(ptoe[loop[k - 1], :], ec) # actualize courant edge
        # setdiff(J[ec, :], loop[k - 1])[1]
        loop[k] = J[ec, 1] == loop[k - 1] ? J[ec, 2] : J[ec, 1]
    end
    if minimum(loop) == 0
        error("Pb II in edgestoloop")
    end
    return elcol[I[loop]]
end
#-------------------------------------------------------------------------------
"""
    Tri, pyhull = convexhull(p::Array{Float64, 2},
                             qhull_option::String = "QJ Pp")
    convexhull()
Call spatial scipy library to compute convehull with qhull.

WARNING : there is perhaps a problem of orientation of triangles!
"""
function convexhull(p::Array{Float64, 2}, qhull_option::String = "QJ Pp")
    pyhull = pyspatial.ConvexHull(p, qhull_options = qhull_option)
    Tri = convert(Array{Int64, 2}, pyhull.simplices) .+ 1
    return Tri, pyhull
end
#-------------------------------------------------------------------------------
"""
    Au, Iu, Ju = uniquerows(A::AbstractArray{T}, preci::Int64 = 6)
Matlab like unique_rows function

N.B.
* use preci = 0 for integers
* maximum(abs.((A[Iu, :])[Ju, :] .- A)) == 0
"""
function uniquerows(A::AbstractArray{T}, preci::Int64 = 0) where T
    Ar = if preci > 0 # use preci = 0 for integers
        round.(Int64, A * 10 ^ preci)
    else
        A
    end
    Ju = groupslices(Ar, dims = 1)
    p = sortperm(Ju)
    Iu, c = [Ju[p[1]],], 1
    Jun = similar(Ju)
    Jun[p[1]] = c
    for k = 2:length(Ju)
        if Ju[p[k]] > Ju[p[k - 1]]
            push!(Iu, Ju[p[k]])
            c += 1
        end
        Jun[p[k]] = c
    end
    #@test maximum(abs.((Ar[Iu, :])[Jun, :] .- Ar)) == 0
    return A[Iu, :], Iu, Jun
end


module day6

readInput() = split.(readlines(joinpath(@__DIR__, "input.txt")), ')')

mutable struct Node
    name::Symbol
    parent::Union{Nothing, Node}
    children::Vector{Node}
end
Base.show(io::IO, n::Node) = print(io, n.name)

function makeTree(input)
    d = Dict{Symbol, Node}()
    for line in input
        leftSym = Symbol(first(line))
        left = get!(d, leftSym, Node(leftSym, nothing, []))
        rightSym = Symbol(last(line))
        right = get!(d, rightSym, Node(rightSym, left, []))
        right.parent = left
        push!(left.children, right)
    end
    d
end

function part1(input)
    d = makeTree(input)
    nbOrbits = 0
    for n in values(d)
        p = n
        while p.parent !== nothing
            p = p.parent
            nbOrbits += 1
        end
    end
    nbOrbits
end

function part2(input)
    d = makeTree(input)

    YOU = d[:YOU]
    SAN = d[:SAN]

    # List all parents from "YOU" with their distance
    parentNodes = Dict{Node, Int}()
    p, i = YOU, 0
    while p.parent !== nothing
        push!(parentNodes, p.parent => i)
        p = p.parent
        i += 1
    end

    # Iterate each parent from "SAN", counting the distance,
    # while they're not in the list of "YOU"'s parents
    p, i = SAN, 0
    while !(p.parent in keys(parentNodes))
        p = p.parent
        i += 1
    end

    return i + parentNodes[p.parent]
end

end


@testset "Test rescale" begin
	for _ in 1:100
		vec0 = sort(randn(rand(2:500)));
		vec2 = rand(1:100)*randn(2) |> sort
		vec1 = rescale(vec0, vec2)
		diffvec1 = diff(vec1);
		a = diff(vec0)./diffvec1;
		meana = sum(a)/length(a);
		@test all(diffvec1 .>= 0) # increaments have to be all positive
		@test isapprox.(a .- meana, 0; atol=meana*1e-6) |> all
	end

end


# Copyright (c) 2021, Joseph M. Rutherford
# All rights reserved.
#
# Code provided under the license contained in the LICENSE file.

include("../src/sudoku.jl") # Defines Sudoku module

# Puzzle example from
# https://en.wikipedia.org/wiki/Sudoku accessed 1/1/2022
# Puzzle image license is CC0 (Public Domain Dedication) per
# https://en.wikipedia.org/wiki/Sudoku#/media/File:Sudoku_Puzzle_by_L2G-20050714_standardized_layout.svg
puzzle_values = [
    5 3 0 0 7 0 0 0 0;
    6 0 0 1 9 5 0 0 0;
    0 9 8 0 0 0 0 6 0;
    8 0 0 0 6 0 0 0 3;
    4 0 0 8 0 3 0 0 1;
    7 0 0 0 2 0 0 0 6;
    0 6 0 0 0 0 2 8 0;
    0 0 0 4 1 9 0 0 5;
    0 0 0 0 8 0 0 7 9]
puzzle = Sudoku.SolvablePuzzle(3)
Sudoku.assign_values!(puzzle,puzzle_values)
print("\nUnsolved puzzle\n")
print(Sudoku.as_text_grid(puzzle))
iterations, uncertainty = Sudoku.iterative_solve!(puzzle,1)
print("\n\nSolution after $iterations steps\n")
print(Sudoku.as_text_grid(puzzle))

# Solution image is CC BY-SA 3.0 (Attribution-ShareAlike) per
# https://en.wikipedia.org/wiki/Sudoku#/media/File:Sudoku_Puzzle_by_L2G-20050714_solution_standardized_layout.svg
#
# It is an exercise for reader to compare the calculated solution to the reference
if uncertainty != 0
    println("Failure: puzzle solution is incomplete")
else
    result_values = Sudoku.as_values(puzzle)
    if Sudoku.satisfies(puzzle_values,result_values)
        println("Verified that solution satisfies initial puzzle")
    else
        println("Solution does not match original puzzle inputs")
    end
    if Sudoku.valid_puzzle(result_values)
        println("Verified that solution satisfies Suduko rules")
    else
        println("Solution is not a valid Sudoku")
    end
end


module ProductDistributions

import Distributions: ContinuousDistribution, ContinuousMultivariateDistribution,
## Univariate Distributions
logpdf, pdf, rand!
using Reexport
@reexport using Distributions
@reexport using BayesianTools

import Base: first, start, next, done

struct ProductDistribution{T} <: ContinuousMultivariateDistribution
    marginals::T
end

Base.length(d::ProductDistribution) = length(d.marginals)

function ProductDistribution(args...)
    ProductDistribution((args...,))
end

ProductDistribution(x::Distribution) = ProductDistribution([x])

function Distributions.insupport(d::ProductDistribution, x::AbstractVector{T}) where T<:Real
    insup = true
    for i in eachindex(d.marginals)
        insupport(d.marginals[i], x[i]) || (insup = false; break)
    end
    insup
end

function Distributions.logpdf(d::ProductDistribution, x::AbstractVector{T}) where T<:Real
    if Distributions.insupport(d, x)
        l = zero(T)
        @inbounds for i in 1:length(d.marginals)
            l += logpdf(d.marginals[i], x[i])
        end
    else
        l = convert(T, -Inf)
    end
    l
end

function Distributions.pdf(d::ProductDistribution, x::AbstractVector{T}) where T<:Real
    exp(logpdf(d, x))
end

function Distributions.rand!(d::ProductDistribution, x::DenseMatrix{T}) where T<:Real
    P, N = size(x)
    @assert P == length(d)
    @inbounds for i = 1:N
        for p = 1:P
            x[p, i] = rand(d.marginals[p])
        end
    end
    x
end

function Distributions.rand!(d::ProductDistribution, x::AbstractVector{T}) where T<:Real
    P = length(x)
    @assert P == length(d)
    @inbounds for p = 1:P
        x[p] = rand(d.marginals[p])
    end
    x
end

Distributions._rand!( d::ProductDistribution, a::AbstractArray ) = rand!( d, a )

# found this tidbit here: https://discourse.julialang.org/t/efficient-tuple-concatenation/5398/3
# this should probably go somewhere more generic...
@inline tuplejoin(x) = x
@inline tuplejoin(x, y) = (x..., y...)
@inline tuplejoin(x, y, z...) = tuplejoin(tuplejoin(x, y), z...)

Distributions.params( d::ProductDistribution ) = tuplejoin(params.(d.marginals)...)

# Base.start(d::ProductDistribution) = 1
# Base.first(d::ProductDistribution) = d.marginals[1]
# Base.done(d::ProductDistribution, state) = state == length(d) + 1
# Base.next(d::ProductDistribution, state) = (d.marginals[state], state + 1)

function Base.iterate(d::ProductDistribution{T}, state=1) where T

    if state > length(d)
        return nothing
    end

    return (d.marginals[state], state + 1)
end

export ProductDistribution

end


#=

  https://www.allendowney.com/blog/2018/10/18/how-tall-is-a/
  """
  Here are a series of problems I posed in my Bayesian statistics class:

  1) Suppose you meet an adult resident of the U.S. who is 170 cm tall.
     What is the probability that they are male?

  2) Suppose I choose two U.S. residents at random and A is taller than B.
     How tall is A?

  3) In a room of 10 randomly chosen U.S. residents, A is the second tallest.
     How tall is A?
     And what is the probability that A is male?

  As background: For adult male residents of the US, the mean and standard deviation of
  height are 178 cm and 7.7 cm. For adult female residents the corresponding stats
  are 163 cm and 7.3 cm.  And 51% of the adult population is female.

  """

  Cf ~/blog/how_tall_is_a.blog
     ~/webppl/how_tall_is_a.wppl

=#

using Turing, StatsPlots, DataFrames
using ReverseDiff, Zygote, Tracker
Turing.setadbackend(:forwarddiff)
# Turing.setadbackend(:reversediff)
# Turing.setadbackend(:tracker)
# Turing.setadbackend(:zygote)

include("jl_utils.jl")

#=
  Model 1:

  1) Suppose you meet an adult resident of the U.S. who is 170 cm tall.
  What is the probability that they are male?

  According to https://www.allendowney.com/blog/2018/10/18/how-tall-is-a/
    female 0.5432225483131837
    male 0.45677745168681627

  This model:
  Distributions of variable gender (num:0)
    2.00000 =>    5406  (0.540600)
    1.00000 =>    4594  (0.459400)
=#
@model function how_tall_is_1(height=170)
    male = 1
    female = 2
    gender ~ Categorical(simplex([0.49,0.51]))
    height ~ gender == male ? Normal(178,7.7) : Normal(163,7.3)

    # true ~ Dirac(height == 170.0)

end

#=
println("Model 1")
model = how_tall_is_1(170)
# chns = sample(model, Prior(), 10_000)
chns = sample(model, MH(), 100_000)
# chns = sample(model, PG(15), 1_000)
# chns = sample(model, SMC(1000), 10_000)
# chns = sample(model, IS(), 10_000)

#chns = sample(model, NUTS(), 1000)
# chns = sample(model, HMC(0.1, 5), 1000)
# chns = sample(model, HMCDA(0.15, 0.65), 1000)
# chns = sample(model, Gibbs(MH(:gender),NUTS(1000,0.65,:height)), 1_000)

display(chns)
println("male: 1 female: 2")
show_var_dist_pct(chns,:gender)
# show_var_dist_pct(chns,:height,20)
=#

#=

  2) Suppose I choose two U.S. residents at random and A is taller than B.
  How tall is A?

  Solution from https://www.allendowney.com/blog/2018/10/18/how-tall-is-a/
  A: 176.67506663725212
  B: 164.05298129722925

  This model:
    ...
    height[1]   176.4384    9.0135     0.2850    0.1777   824.6929    1.0007      179.5152
    height[2]   164.5469    8.4415     0.2669    0.2017   828.7247    0.9998      180.3928

=#

@model function how_tall_is_2()
    male = 1
    female = 2
    gender = tzeros(2)
    height = tzeros(2)
    for p in 1:2
        gender[p] ~ Categorical(simplex([0.49,0.51]))
        height[p] ~ gender[p] == male ? Normal(178,7.7) : Normal(163,7.3)
    end
    true ~ Dirac(height[1] > height[2])

end

#=
println("\n\nModel 2")
model = how_tall_is_2()
# chns = sample(model, Prior(), 10_000)
# chns = sample(model, MH(), 100_000)
chns = sample(model, PG(15), 1_000)
# chns = sample(model, SMC(1000), 10_000)
# chns = sample(model, IS(), 10_000)

display(chns)
println("male: 1 female: 2")
show_var_dist_pct(chns,Symbol("gender[1]"))
show_var_dist_pct(chns,Symbol("gender[2]"))
=#

#=

  3) In a room of 10 randomly chosen U.S. residents, A is the second tallest.
  How tall is A?
  And what is the probability that A is male?

    Solution from https://www.allendowney.com/blog/2018/10/18/how-tall-is-a/
    A: 181.60660153115973

=#

# Note:
# This model does NOT work as expected. It's fine for n=2 (see how_tall_is_2), but not the
# generalized version of n=10. See how_tall_is_3_2() for a working model (thanks to Seth Axen).
# See the discussion here:
#  https://discourse.julialang.org/t/help-with-turing-jl-model-using-categorical-normal-and-ordered-constraint/64765
#
@model function how_tall_is_3()
    male = 1
    female = 2
    n = 10
    gender = tzeros(n)
    height = tzeros(n)
    for p in 1:n
        gender[p] ~ Categorical(simplex([0.49,0.51]))
        height[p] ~ gender[p] == male ? Normal(178,7.7) : Normal(163,7.3)
    end

    for p in 1:n-1
        # true ~ Dirac(height[p] > height[p+1])
        height[p] > height[p+1] || begin Turing.@addlogprob! -Inf; end
        # height[p] > height[p+1] || begin Turing.@addlogprob! -2; end
    end

    a_is_male ~ Dirac(gender[2] == male)
end

#=
println("\n\nModel 3")
model = how_tall_is_3()
# chns = sample(model, Prior(), 10_000)
chns = sample(model, MH(), 10_000)
# chns = sample(model, PG(15), 10_000)
# chns = sample(model, SMC(), 10_000)
# chns = sample(model, IS(), 10_000)
display(chns)
println("male: 1 female: 2")
show_var_dist_pct(chns,:a_is_male)
=#

#
# Thanks to Seth Axen for this approach.
# From https://discourse.julialang.org/t/help-with-turing-jl-model-using-categorical-normal-and-ordered-constraint/64765/2
# Note that the heights[i] are all about 170cm; it's the height_ordered that we are
# interested in:
#=
   height_ordered[1]   186.706
   height_ordered[2]   181.5520
   height_ordered[3]   177.8786
   height_ordered[4]   174.6669
   height_ordered[5]   171.6923
   height_ordered[6]   168.7933
   height_ordered[7]   165.8899
   height_ordered[8]   162.8638
   height_ordered[9]   159.4338
  height_ordered[10]   154.5771
            height_A   181.5520
          isfemale_A     0.0893
=#
@model function how_tall_is_3_2()
    male = 1
    female = 2
    n = 10
    # This is from Seth Axen:
    isfemale ~ filldist(Bernoulli(0.51), n)
    height = tzeros(n) # Vector{T}(undef, n)
    for i in 1:n
        if isfemale[i]
            height[i] ~ Normal(163, 7.3)
        else
            height[i] ~ Normal(178, 7.7)
        end
    end
    # hakank: I added this since I actually wanted the
    #         weights of all people.
    height_ordered = tzeros(n) # Vector{T}(undef, n)
    perm = sortperm(height)
    for i in 1:n
        height_ordered[i] ~ Dirac(height[perm[n-i+1]])
    end

    # From Seth:
    id_A = perm[n-1]
    height_A ~ Dirac(height[id_A])
    isfemale_A ~ Dirac(isfemale[id_A])

end

model = how_tall_is_3_2()
# Comment from Seth:
# "no inference necessary! We just sample directly from the prior-predictive distribution"
chns = sample(model, Prior(), 10_000)
display(chns)
show_var_dist_pct(chns,:isfemale_A)
# Testing with IS(), same result
# chns = sample(model, IS(), 10_000)
# display(chns)


#
# This file is a part of MolecularGraph.jl
# Licensed under the MIT License http://opensource.org/licenses/MIT
#

@testset "graph.bipartite" begin

@testset "maxcardmatch" begin
    relation1 = Dict(
        1 => Set([1]), 2 => Set([1, 2, 3]),
        3 => Set([1, 2]), 4 => Set([2, 3, 4])
    )
    @test maxcard(Set(1:4), Set(1:4), relation1) == 4
    relation2 = Dict(
        1 => Set([2]), 2 => Set([2, 3]),
        3 => Set(1:5), 4 => Set([3]), 5 => Set([4])
    )
    @test maxcard(Set(1:5), Set(1:5), relation2) == 4
    relation3 = Dict(
        1 => Set([1, 2]), 2 => Set([3]),
        3 => Set([4, 5]), 4 => Set([3, 4]), 5 => Set([2, 5])
    )
    @test maxcard(Set(1:5), Set(1:5), relation3) == 5
    relation4 = Dict(
        1 => Set(1:6), 2 => Set([4]), 3 => Set([4]), 4 => Set([4])
    )
    @test maxcard(Set(1:4), Set(1:6), relation4) == 2
end

@testset "twocoloring" begin
    noedge = SubgraphView(pathgraph(5), Set(1:5), Set{Int}())
    a, b = twocoloring(noedge)
    @test length(a) == 5
    @test isempty(b)

    K5_5 = completebipartite(5, 5)
    a, b = twocoloring(K5_5)
    @test length(a) == 5
    @test length(b) == 5

    CL5 = circularladder(5)
    n = twocoloring(CL5)
    @test n === nothing

    CL6 = circularladder(6)
    a, b = twocoloring(CL6)
    @test length(a) == 6
    @test length(b) == 6
end

end # graph.bipartite


#--- load modules
using .LAP_julia
using Test, Images, Colors
using FileIO
using CSV
using Plots, ImageView: imshow

#--- load random gray ANHIR image
base_path = "/Users/MrTrololord/Google_Drive/cvut/bakalarka/anhir/"
dataset_path = "/Users/MrTrololord/Google_Drive/cvut/bakalarka/anhir/dataset/"
landmark_path = "/Users/MrTrololord/Google_Drive/cvut/bakalarka/anhir/landmarks/"
loc_table = CSV.read(base_path * "location_table.csv")

train_rows = loc_table[loc_table[:status] .== "training", :]

const TEST_SIZE = 1

random_rows = train_rows[rand(1:size(train_rows, 1), TEST_SIZE), :]

row_used = (random_rows[1, :])[Symbol("Column1")]

source_loc = (random_rows[1, :])[Symbol("Source image")]
target_loc = (random_rows[1, :])[Symbol("Target image")]
source_in = load(dataset_path * source_loc)
target_in = load(dataset_path * target_loc)

gray_source = Gray.(source_in)
gray_target = Gray.(target_in)

#--- load chessboard image1

tile_size = 4
board_size = 8 # must be even
mini_board = [zeros(tile_size, tile_size) ones(tile_size, tile_size);
              ones(tile_size, tile_size) zeros(tile_size, tile_size)]

chessboard = repeat(mini_board, outer=(convert(Integer,(board_size/2)), convert(Integer,(board_size/2))))
chessboard = repeat(mini_board, outer=(2, 2))

imshow(chessboard)

# imshow(chessboard);

#---
small_gray_source = imresize(gray_source, ratio = 0.1)

# LAP filters:
filter_size = 21
filter_half_size = (filter_size - 1) / 2
use_one_dim_filters = false
# Calculate separable filters from basis:
sigma = (filter_half_size + 2) / 4
centered_inds = ImageFiltering.centered(-filter_half_size:filter_half_size)
gaus = ImageFiltering.KernelFactors.gaussian((sigma, sigma), (filter_size, filter_size))
gaus = ImageFiltering.KernelFactors.gaussian(sigma, filter_size)
gaus_der = gaus .* centered_inds .* (-1)/sigma^2
gaus_der_flip = reverse(gaus_der, dims=1)

using Plots
# kernel basis:
# 1
kernf1 = kernelfactors((gaus, gaus))
kernp1 = broadcast(*, kernf1...)
surface(-filter_half_size:filter_half_size, -filter_half_size:filter_half_size, kernp1)

# 2
# forward
kernf2 = kernelfactors((gaus, gaus_der_flip))
kernp2 = broadcast(*, kernf2...)
surface(-filter_half_size:filter_half_size, -filter_half_size:filter_half_size, kernp2)
# backward
kernf2 = kernelfactors((gaus, gaus_der))
kernp2 = broadcast(*, kernf2...)
surface(-filter_half_size:filter_half_size, -filter_half_size:filter_half_size, kernp2)

# 3
# forward
kernf3 = kernelfactors((gaus_der_flip, gaus))
kernp3 = broadcast(*, kernf3...)
surface(-filter_half_size:filter_half_size, -filter_half_size:filter_half_size, kernp3)
# backward
kernf3 = kernelfactors((gaus_der, gaus))
kernp3 = broadcast(*, kernf3...)
surface(-filter_half_size:filter_half_size, -filter_half_size:filter_half_size, kernp3)

surface(-filter_half_size:filter_half_size, -filter_half_size:filter_half_size, kernp)

#--- application to images

# kernels:
# 1
@time imgfilt1 = ImageFiltering.imfilter(chessboard, kernf1, "symmetric")
imshow(imgfilt1)
surface(imgfilt1)

# 2
imgfilt2 = ImageFiltering.imfilter(chessboard, kernf2, "symmetric")
imshow(imgfilt2)
surface(imgfilt2)

# 3
imgfilt3 = ImageFiltering.imfilter(chessboard, kernf3, "symmetric")
imshow(imgfilt3)
surface(imgfilt3)

#--- imfilter with prealloc

imgfilt1_prealloc = similar(chessboard)

@time imgfilt1_tmp = ImageFiltering.imfilter!(imgfilt1_prealloc, chessboard, kernf1, "symmetric")
imshow(imgfilt1_prealloc)
surface(imgfilt1_prealloc)

#---

dif = zeros(4, 4, 3)

dif[:, :, 2] = 1:16
dif[:, :, 3] = ones(4, 4).+1

dif

dif = reshape(dif, (:, 3))

dif

#---

A = zeros(2, 2, 3)
b = zeros(2, 3)

A = rand(100, 100, 100)
b = rand(100, 100)

using BenchmarkTools
@benchmark outpin = pinv(A[:, :, 1])*b[:, 1]
@benchmark outqr = A[:, :, 1] \ b[:, 1]
@benchmark outqrpiv = qr(A[:, :, 1], Val(true)) \ b[:, 1]

outpin ≈ outqr

outqrpiv ≈ outpin

A[:, :, 1] = [1 0; 0 1]
A[:, :, 2] = [1 0; 0 1]
A[:, :, 3] = [1 0; 0 1]

b[:, 1] = [1; 1]
b[:, 2] = [0; 2]
b[:, 3] = [0; 4]

A[:, :, 1]\b[:, 1]

save = zeros(100, 100)

@time multi_mat_div!(save, A, b)

save

function multi_mat_div!(save, A, b)
    for k in axes(A)[3]
        @views save[:, k] = A[:, :, k] \ b[:, k]
    end
end

function multi_mat_div(A, b)
    res = zeros(size(A)[1], size(A)[3])
    # Threads.@threads for j in axes(A, 4) check: (https://stackoverflow.com/questions/57678890/batch-matrix-multiplication-in-julia)
    for k in axes(A)[3]
        @views res[:, k] = qr(A[:, :, k], Val(true)) \ b[:, k]
    end
    return transpose(res)
end

@code_warntype multi_mat_div(A, b)

#---

one = ones(12);
two = reshape(one, (3, 4))

one[:] = two

#--- pad with zeros

using ImageFiltering

image_size = size(chessboard)
filter_half_size = 5;

ones(image_size.-(2*filter_half_size))

border_mask = parent(padarray(ones(image_size.-(2*filter_half_size)), Fill(NaN, (filter_half_size, filter_half_size),(filter_half_size, filter_half_size))))

new_board = chessboard .* border_mask

imshow(new_board)

all_coeffs = ones(broadcast((*), image_size...), 3)

new_coeffs = all_coeffs .* reshape(border_mask, (:, 1))

imshow(reshape(new_coeffs[:, 1], image_size))

imshow(reshape(chessboard[border .== 1.0] .= NaN, (170, 170)))

#---

using BenchmarkTools
using LinearAlgebra

A = rand(11, 11);

A[1, :] = ones(11)

k = (-filter_half_size:filter_half_size)

k1 = repeat(k, outer=11)

sum(transpose(A) * k)

dot(A[:], repeat(k, outer=11))

B = A[:]

#---

u_est, coeffs = single_lap(chessboard, chessboard, 3, 15, [15, 15])

#---

ones_arr_1 = ones(window_size[1])
ones_arr_2 = ones(window_size[2])
ones_kernel = kernelfactors((ones_arr_1, ones_arr_2))
# filtering gets a sum of pixels of window size in each coord

pixels = reshape(chessboard, (:, 1))
image_size = size(chessboard)
window_size = [15, 15]

filter_result = similar(chessboard)

ImageFiltering.imfilter!(filter_result, reshape(pixels, image_size), shit, "symmetric")

imshow(filter_result)

#---

include("../src/lap.jl")

lap.window_sum!(filter_result, pixels, image_size, window_size)

imshow(filter_result)

#---

uv = zeros(size(chessboard)..., 2)
uv[:, :, 1] = real(u_est)
uv[:, :, 2] = imag(u_est)

Plots.quiver([1,2,3], [1,2,3], quiver=([NaN,1,1],[NaN,2,3]))

u_est, coeffs = single_lap(chessboard, chessboard, 3, 15, [15, 15])

p1 = figure()
uv_flow = zeros(size(u_est)..., 2)
uv_flow[:, :, 1] = real(u_est)
uv_flow[:, :, 2] = imag(u_est)

skip_count = 20
n = skip_count

trimmed_real = zeros(size(uv_flow[:,:,1]))
trimmed_real[1:n:end, 1:n:end] = uv_flow[1:n:end, 1:n:end, 1]
trimmed_imag = zeros(size(uv_flow[:,:,2]))
trimmed_imag[1:n:end, 1:n:end] = uv_flow[1:n:end, 1:n:end, 2]

PyPlot.quiver(trimmed_real, trimmed_imag)
gcf()

#---
using FileIO: load

board_path = "/Users/MrTrololord/Google_Drive/cvut/bakalarka/PolyFilter_LAP/test_boards/"

chess_norm = load(board_path * "chessboard.png")
chess_warp = load(board_path * "chess_warp.png")

chess_norm_f = Float32.(chess_norm)
chess_warp_f = Float32.(chess_warp)

u_est, coeffs = single_lap(chess_norm_f, chess_warp_f, 3, 15, [15, 15])

using LinearAlgebra: qr

qr([1 1; 1 1], Val(true)) \ [1; 1]

#--- TEST single_lap
using PyPlot: quiver, gcf, figure

# *************************
# 1) MATLAB RANDOM FLOW
# *************************
board_path = "/Users/MrTrololord/Google_Drive/cvut/bakalarka/PolyFilter_LAP/test_boards/"

# load chessboards from matlab generated flow
chess_norm = load(board_path * "chessboard.png")
chess_warp = load(board_path * "chess_warp.png")

chess_norm_f = Float32.(chess_norm)
chess_warp_f = Float32.(chess_warp)

u_est, coeffs = LAP_julia.registration.single_lap(chess_norm_f, chess_warp_f, 3, 15, [15, 15])

# *************************
# 2) ONE PIXEL IMAGE SHIFT
# *************************

chess_norm = chessboard
chess_norm_f = Float32.(chess_norm)

chess_warp_f = copy(chess_norm_f)
for k in 2:size(chess_warp_f)[1]
    for l in 2:size(chess_warp_f)[2]
        chess_warp_f[k-1, l-1] = chess_warp_f[k, l]
        if l == size(chess_warp_f)[2]
            chess_warp_f[k-1, l] = chess_warp_f[k, l]
        end
        if k == size(chess_warp_f)[1]
            chess_warp_f[k, l-1] = chess_warp_f[k, l]
        end
    end
end

filter_half_size = 2
filter_size = filter_half_size * 2 + 1

u_est, coeffs = LAP_julia.registration.single_lap(chess_norm_f, chess_warp_f, 3, filter_half_size, [filter_size, filter_size])

# using ImageFiltering: Pad
#
# @views u_est = u_est[]
# u_est = parent(padarray(u_est, Pad(:replicate, filter_half_size, filter_half_size)))

p1 = figure()
PyPlot.quiver(real(u_est), imag(u_est))
gfc()

#---
using Statistics: mean

average = mean(filter(!isnan, u_est))
sqrt(real(average)^2 + imag(average)^2)

#---
image_size = size(chessboard)

window_size = ones(2) .* filter_size
window_half_size = Int64.((window_size .- 1) ./ 2)

Δ_u = ones(10,10)

middle_vals = Δ_u[window_half_size[1]+1:end-window_half_size[1],
                 window_half_size[2]+1:end-window_half_size[2]]
Δ_u = parent(padarray(middle_vals, Pad(:replicate, window_half_size...)))

Δ_u[3, 3] = NaN

Δ_u

any(isnan.(Δ_u))

#---

u = ones(10,10) .+ 2im .* ones(10,10)
u[3:5, 3:5] .= NaN + 1im .* NaN

u

out = LAP_julia.interpolate.interpolate_nans(u)

mask = .!isnan.(u)

u0 = u
u0[.!mask] .= 0

div_coef = similar(u, Int16)

#--- TEST INPAINT

u = ones(10,10) .+ 2im .* ones(10,10);
u[3:8, 3:8] .= NaN .+ NaN .* 1im;

#mask = .!isnan.(real(u))

out = LAP_julia.inpaint.inpaint_nans!(u)

#---

u = ones(10,10) .+ 2im .* ones(10,10);
u[3:8, 3:8] .= NaN .+ NaN .* 1im;

mask = .!isnan.(real(u))

div_coef = ones(size(u))
using ImageFiltering: kernelfactors, imfilter!, Pad, centered

using ImageFiltering: kernelfactors, imfilter, Pad, centered
mask = ones(10, 10)
mask[3:8, 3:8] .= 0

ker = centered([0.5, 2, 0.5])
kernf = kernelfactors((ker, ker));

out = imfilter(mask, kernf, Pad(:symmetric))

#--- Test Makie

using Makie
using TestImages

img = chessboard

scene = Scene()

img1 = reverse(img, dims=1)
img1 = transpose(img1)

scene = image(img, show_axis = false, interpolate = false)

scene = heatmap(img, colormap = :grayscale)

display(scene)

#--- test PyPlot

tile_size = 4
board_size = 8 # must be even
mini_board = [zeros(tile_size, tile_size) ones(tile_size, tile_size);
              ones(tile_size, tile_size) zeros(tile_size, tile_size)]

chessboard = repeat(mini_board, outer=(convert(Integer,(board_size/2)), convert(Integer,(board_size/2))))
chessboard = repeat(mini_board, outer=(2, 2))

img  = chessboard

using Plots

plotly()

Plots.heatmap(img; hover=img)
gui()

###
using PlotlyJS

PlotlyJS.plot(heatmap(z = img))

using Plots
plotly()
data = 1:5
plt = plot(data; hover=data )
display(plt)


module FESpacesTests

using Test
using Gridap

model = CartesianDiscreteModel(partition=(2,2))

order = 1
tags = [1,2,3,4]
fespace = H1ConformingFESpace(Float64,model,order,tags)

@test string(fespace) == "ConformingFESpace object"
s = "ConformingFESpace object:\n physdim: 2\n refdim: 2\n valuetype: Float64\n nfree: 5\n ndiri: 4"
@test sprint(show,"text/plain",fespace) == s

trian = Triangulation(model)
quad = CellQuadrature(trian,degree=2)

cellbasis = CellBasis(fespace)
a(v,u) = varinner(v,u)
bfun(x) = x[2]
b(v) = varinner(v,CellField(trian,bfun))
mmat = integrate(a(cellbasis,cellbasis),trian,quad)
bvec = integrate(b(cellbasis),trian,quad)

fun(x) = x[1] + x[2]

test_fe_space(fespace,5,4,mmat,bvec,fun)

V = TestFESpace(fespace)
test_fe_space(V,5,4,mmat,bvec,fun)

U = TrialFESpace(fespace,fun)
test_fe_space(U,5,4,mmat,bvec,fun)

fun(x) = x[1] + x[2]
z = 2.0
U = TrialFESpace(fespace,z)
test_fe_space(U,5,4,mmat,bvec,fun)

U = TrialFESpace(fespace,[z,z,z,z])
test_fe_space(U,5,4,mmat,bvec,fun)

end # module


__precompile__()

module NearestNeighbors

using Distances
import Distances: Metric, result_type, eval_reduce, eval_end, eval_op, eval_start, evaluate

using StaticArrays
import Base.show

export NNTree, BruteTree, KDTree, BallTree, DataFreeTree
export knn, inrange # TODOs? , allpairs, distmat, npairs
export injectdata

export Euclidean,
       Cityblock,
       Minkowski,
       Chebyshev,
       Hamming

# Change this to enable debugging
const DEBUG = false

abstract type NNTree{V <: AbstractVector,P <: Metric} end

const MinkowskiMetric = Union{Euclidean,Chebyshev,Cityblock,Minkowski}

function check_input(::NNTree{V1}, ::AbstractVector{V2}) where {V1, V2 <: AbstractVector}
    if length(V1) != length(V2)
        throw(ArgumentError(
            "dimension of input points:$(length(V2)) and tree data:$(length(V1)) must agree"))
    end
end

function check_input(::NNTree{V1}, point::AbstractVector{T}) where {V1, T <: Number}
    if length(V1) != length(point)
        throw(ArgumentError(
            "dimension of input points:$(length(point)) and tree data:$(length(V1)) must agree"))
    end
end

get_T(::Type{T}) where {T <: AbstractFloat} = T
get_T(::T) where {T} = Float64

include("debugging.jl")
include("evaluation.jl")
include("tree_data.jl")
include("datafreetree.jl")
include("knn.jl")
include("inrange.jl")
include("hyperspheres.jl")
include("hyperrectangles.jl")
include("utilities.jl")
include("brute_tree.jl")
include("kd_tree.jl")
include("ball_tree.jl")
include("tree_ops.jl")

end # module


## Fun.jl docs

doc"""
    domain(::Fun)

returns the domain that a `Fun` is defined on
"""
domain(fun::Fun)

doc"""
    setdomain(fun,domain)

returns `fun` projected onto `domain`
"""
domain(fun::Fun,domain::Domain)

doc"""
    space(fun)

returns the space of `fun`
"""
space(fun::Fun)

doc"""
    values(fun)

returns `fun` evaluated at `points(fun)`
"""
Base.values(fun::Fun)

doc"""
    points(fun)

returns a grid of points that the fun can be transformed into values
and back
"""
points(fun::Fun)

doc"""
    ncoefficients(fun) -> Integer

returns the number of coefficients of a fun
"""
ncoefficients(fun::Fun)

doc"""
    stride(fun)

returns the stride of the coefficients, checked
numerically
"""
Base.stride(fun::Fun)

doc"""
    reverseorientation(fun)

return `fun` on a reversed orientated contour
"""
reverseorientation(f::Fun)

## Operator.jl docs

doc"""
    `Operator{T}` represents a general operator
"""
Operator


export zoom
function zoom(h :: LineModel,
              t₀ :: Float64,
              t₁ :: Float64;
              c₁ :: Float64=0.01,
              tol :: Float64=1e-7,
              ϵ :: Float64=1e-10,
              maxiter :: Int=50,
              verbose :: Bool=false)

        if obj(h,t₀)<obj(h,t₁)
          tl=t₀
          th=t₁
        else
          tl=t₁
          th=t₀
        end

        h₀=obj(h,t₀)
        dh₀=grad(h,t₀)

        hl=obj(h,tl)
        dhl=grad(h,tl)
        hh=obj(h,th)
        dhh=grad(h,th)

        i=0

        ti= (tl+th)/2
        hi=obj(h,ti)
        dhi=grad(h,ti)

        verbose && println("on a les paramètre de zoom")

        verbose && @printf(" iter        tₗ        tₕ         tᵢ        hₗ        hₕ         hᵢ         dhᵢ\n")
        verbose && @printf(" %7.2e %7.2e  %7.2e  %7.2e  %7.2e %7.2e %7.2e %7.2e\n", i,tl,th,ti,hl,hh,hi,dhi)

        while ((i<maxiter) & (abs(dhi)>tol)) || i==0
          # if i==0
          #   verbose && println("on est dans le while de zoom")
          # end
          hi=obj(h,ti)
          if (hi>h₀+c₁*(ti-t₀)*dh₀) || (hl<=hi)
            th=ti
            hh=hi
            dhh=hi
          else
            dhi=grad(h,ti)
            if (abs(dhi)<ϵ)
              topt=ti
              iter=i
              return (topt,iter)
            elseif dhi*(th-tl)>=0
              th=tl
              hh=hl
              dhh=dhl
            end
            tl=ti
            hl=hi
            dhl=dhi
          end

          ti= (tl+th)/2

          i=i+1
          verbose && @printf(" %7.2e %7.2e  %7.2e  %7.2e %7.2e %7.2e %7.2e %7.2e\n", i,tl,th,ti,hl,hh,hi,dhi)

        end
        topt=ti
        iter=i

        return (topt,iter)

end


# # BESO example: HalfMBB Beam
#-
#md # !!! tip
#md #     This example is also available as a Jupyter notebook:
#md #     [`beso.ipynb`](@__NBVIEWER_ROOT_URL__/examples/beso.ipynb)
#-
# ## Commented Program
#
# What follows is a program spliced with comments.
#md # The full program, without comments, can be found in the next [section](@ref beso-plain-program).

using TopOpt

# ### Define the problem
E = 1.0 # Young’s modulus
v = 0.3 # Poisson’s ratio
f = 1.0; # downward force

nels = (160, 40)
problem = HalfMBB(Val{:Linear}, nels, (1.0, 1.0), E, v, f)

# ### Define the FEA Solver and penalty functions
solver = FEASolver(Displacement, Direct, problem, xmin = 0.01, penalty = TopOpt.PowerPenalty(3.0))

# ### Define the compliance objective function and volume fraction constraint
comp = Compliance(problem, solver)
volfrac = Volume(problem, solver)
sensfilter = SensFilter(solver, rmin = 4.0)
beso = BESO(comp, volfrac, 0.5, sensfilter)

# ### Run optimization
x0 = ones(length(solver.vars))
result = beso(x0)

# ### (Optional) Visualize the result using Makie.jl
# Need to run `using Pkg; Pkg.add(Makie)` first
# ```julia
# using TopOpt.TopOptProblems.Visualization: visualize
# fig = visualize(problem; topology = result.topology)
# Makie.display(fig)
# ```

#md # ## [Plain Program](@id beso-plain-program)
#md #
#md # Below follows a version of the program without any comments.
#md # The file is also available here: [beso.jl](beso.jl)
#md #
#md # ```julia
#md # @__CODE__
#md # ```


using Test, Interpolations, DualNumbers, Unitful

@testset "ExtrapTypeStability" begin

    # Test type-stability of 1-dimensional extrapolation
    f(x) = sin((x-3)*2pi/9 - 1)
    xmax = 10
    A = Float64[f(x) for x in 1:xmax]
    itpg = interpolate(A, BSpline(Linear()))

    schemes = (
        Flat,
        Line,
        Reflect,
        Periodic
    )

    for etp in map(E -> @inferred(extrapolate(itpg, E())), schemes),
        x in (
            # In-bounds evaluation
            3.4, 3, dual(3.1),
            # Out-of-bounds evaluation
            -3.4, -3, dual(-3,1),
            13.4, 13, dual(13,1)
        )
        @inferred(etp(x))
    end

    # Test type-stability of 2-dimensional extrapolation with homogeneous scheme
    g(y) = (y/100)^3
    ymax = 4
    A = Float64[f(x)*g(y) for x in 1:xmax, y in 1:ymax]
    itp2 = interpolate(A, BSpline(Linear()))

    for (etp2,E) in map(E -> (extrapolate(itp2, E()), E), schemes),
        x in (
            # In-bounds evaluation
            3.4, 3, dual(3.1),
            # Out-of-bounds evaluation
            -3.4, -3, dual(-3,1),
            13.4, 13, dual(13,1)
        ),
        y in (
            # In-bounds evaluation
            2.1, 2, dual(2.3, 1),
            # Out-of-bounds evaluation
            -2.1, -2, dual(-2.3, 1),
            12.1, 12, dual(12.1, 1)
        )
        @inferred(etp2(x, y))
    end

    A = [1 2; 3 4]
    Af = Float64.(A)
    for B in (A, Af)
        itpg2 = interpolate(B, BSpline(Linear()))
        etp = extrapolate(itpg2, NaN)
        @test typeof(@inferred(etp(dual(1.5,1), dual(1.5,1)))) ==
              typeof(@inferred(etp(dual(6.5,1), dual(3.5,1))))
    end

    # issue #258
    t = [2017, 2018]
    u = u"V"
    y = [1.f0, 2.f0] .* u
    how = Interpolations.Gridded(Interpolations.Linear())
    itp = interpolate((t,), y, Gridded(Linear()))
    etp = extrapolate(itp, 1.f0 * u)
    @test @inferred(itp(2018)) === 2.0f0 * u
    @test @inferred(etp(2018)) === 2.0f0 * u
    @test @inferred(etp(2019)) === 1.0f0 * u

    # Unitful vector knots
    t = [0.0, 1.0] * u"m"
    itp = LinearInterpolation(t, y)
    @test itp(0.5u"m") == 1.5f0 * u
    @test itp([0.5, 0.75]u"m") == [1.5f0, 1.75f0] * u

    # issues #370, #424
    itp = LinearInterpolation([1.0, 2.0, 3.0], [1.0, 2.0, 3.0], extrapolation_bc=missing)
    @test ismissing(itp(4.0))
    @test @inferred(eltype(itp),itp(3.0)) === 3.0
    @test @inferred(eltype(itp),itp(4.0)) === missing
end


# Point in 2D space

export Point2D

mutable struct Point2D <: Point
    x :: Float64
    y :: Float64
end

point(x,y) = Point2D(x,y)
zero(::Type{Point2D}) = Point2D(0.0, 0.0)
function rand(r::AbstractRNG, ::Random.SamplerType{Point2D})
    t = 2pi*rand(r)
    Point2D(cos(t),sin(t))
end

+(lhs :: Point2D, rhs :: Point2D) = Point2D(lhs.x + rhs.x, lhs.y + rhs.y)
-(lhs :: Point2D, rhs :: Point2D) = Point2D(lhs.x - rhs.x, lhs.y - rhs.y)
-(p :: Point2D) = Point2D(-p.x, -p.y)
*(a :: Real, p :: Point2D) = Point2D(a*p.x, a*p.y)
/(p :: Point2D, a :: Real) = Point2D(p.x/a, p.y/a)

abs(p :: Point2D) = hypot(p.x, p.y)
abs2(p :: Point2D) = p.x*p.x + p.y*p.y

plot_str(p :: Point2D) = string(p.x, " ", p.y)


function compress(X, i_c, j_c, ij_d)
# Compresses a sparse, square matrix X into a rectangular matrix where each column is a full layer
    X_comp = sparse(i_c, j_c, X[ij_d])
    return full(X_comp)
end

function decompress(X_comp, i_d, j_d, ij_c)
# Decompresses a rectangular compressed matrix X into
# a square sparse matrix.
    X = sparse(i_d, j_d, X_comp[ij_c])
    return X
end

# TODO Make sure I can use these tricks to compute a quick analytical Jacobian
# for use in ∂xregen!
function waterColumnStructure(wet3d)
    Iabove = buildIabove(wet3d)
  # loop to make all powers of A
    B = spones(Iabove + I)
    foo = true
    while foo
        last_B = copy(B)
        B = spones(B^2)
        foo = B != last_B
    end
    return B
end

function buildCompressionIndices(wet3d)
  # TODO describe what this does
    layer_idx_3d = zeros(wet3d, Int64)
    ndepth = size(wet3d, 3)
    iwet = find(wet3d)
    for idepth = 1:ndepth
        layer_idx_3d[:,:,idepth] .= idepth
    end
    layer_idx = layer_idx_3d[iwet]

    B = waterColumnStructure(wet3d)
    ij_d = find(B)
    i_d, j_d = ind2sub(size(B), ij_d)

    i_c = copy(i_d)
    j_c = layer_idx[j_d] # This is because I know that this is the right scheme
    B_c = compress(B, i_c, j_c, ij_d)
    ij_c = sub2ind(size(B_c), i_c, j_c)
    return i_c, j_c, ij_c, i_d, j_d, ij_d
end

export buildCompressionIndices, compress, decompress


function solve()
    n, m  = [parse(Int, x) for x in split(readline())]
    a = [parse(Int, x) for x in split(readline())]
    a = sort(a)
    t = 0
    b = zeros(Int, 0)
    for x in a
        push!(b, x-t-1)
        t = x
    end
    push!(b, n-t)
    b = [x for x in b if x≠0]
    if length(b) == 0
        0
    else
        k = minimum(b)
        sum([(x-1)÷k+1 for x in b])
    end
end

println(solve())


is_valid(sides) = sort(sides)[1]+sort(sides)[2] > sort(sides)[3]
n_unique(sides) = length(Set(sides))
is_equilateral(sides) = n_unique(sides) == 1 && is_valid(sides)
is_isosceles(sides) = n_unique(sides) ≤ 2 && is_valid(sides)
is_scalene(sides) = n_unique(sides) == 3 && is_valid(sides)


notes_pattern = r"((?:[^:]+: \d+-\d+ or \d+-\d+\n)+)\nyour ticket:\n((?:\d+,?)+)\n\nnearby tickets:\n((?:(?:\d+,?)+\n?)+)"
fields, my_ticket, tickets = open("ticket_notes.txt") do f
    fields, my_ticket, tickets = match(notes_pattern, read(f, String)).captures
    my_ticket = parse.(Int, split(my_ticket, ","))
    tickets = map(ticket -> parse.(Int, split(ticket, ",")), split(tickets, "\n"))
    fields = map(r -> (r.captures[1], [parse(Int, r.captures[2]):parse(Int, r.captures[3]); parse(Int, r.captures[4]):parse(Int, r.captures[5])]), match.(r"([^:]+): (\d+)-(\d+) or (\d+)-(\d+)", split(strip(fields), "\n")))
    fields, my_ticket, tickets
end

function valid_ticket(ticket, ranges)::Bool
    for val in ticket
        if !any(r -> val in r, ranges)
            return false
        end
    end
    return true
end

ranges = Iterators.flatten(map(x -> (x[2]), fields))
invalid = Vector{Int64}()
for ticket in tickets
    for val in ticket
        if !any(r -> val in r, ranges)
            # println("$(ticket) is invalid ticket ($(val))")
            push!(invalid, val)
        end
    end
end

println(sum(invalid))

candidates = map(r -> map(f -> f[1], r), map(r -> filter(f -> all(x -> x in f[2], r), fields), eachrow(hcat(filter(t -> valid_ticket(t, ranges), tickets)...))))

while any(x -> length(x) > 1, candidates)
    solved = Iterators.flatten(filter(x -> length(x) == 1, candidates))
    for i in 1:length(candidates)
        if length(candidates[i]) == 1
            continue
        end
        candidates[i] = filter(c -> c ∉ solved, candidates[i])
    end
end

println(prod([my_ticket[i] for i in findall(x -> startswith(x, "departure"), collect(Iterators.flatten(candidates)))]))


using AmbientForcing
using NetworkDynamics, Distributions
using LightGraphs
using OrdinaryDiffEq

const N_plastic = 10 # number of nodes
k = 4  # average degree
g = barabasi_albert(N_plastic, k)

function kuramoto_plastic_edge!(de, e, v_s, v_d, p, t)
    de[1] =  e[2] * sin(v_s[1] - v_d[1] + α) / N_plastic - e[1]
    de[2] = - ϵ * (sin(v_s[1] - v_d[1] + β) + e[2])

    nothing
end

function kuramoto_plastic_vertex!(dv, v, edges, p, t)
    dv .= 0
    for e in edges
        dv .-= e[1]
    end
end

const ϵ = 0.1
const α = .2π
const β = -.95π

plasticvertex = ODEVertex(f! = kuramoto_plastic_vertex!, dim =1)
mass_matrix_plasticedge = zeros(2,2)
mass_matrix_plasticedge[2,2] = 1. # First variables is set to 0

plasticedge = ODEEdge(f! = kuramoto_plastic_edge!, dim=2, sym=[:e, :de], coupling=:undirected,mass_matrix = mass_matrix_plasticedge);
kuramoto_plastic! = network_dynamics(plasticvertex, plasticedge, g)

x0_plastic = rand(106)
g_nd = constraint_equations(kuramoto_plastic!)
sum(g_nd(x0_plastic))

x0_plastic = zeros(106)
g_nd = constraint_equations(kuramoto_plastic!)
g_nd(x0_plastic)

Frand = random_force(kuramoto_plastic!, [0.0, 1.0], Uniform)
z_new = ambient_forcing(kuramoto_plastic!, x0_plastic, 2.0, Frand)

sum(g_nd(z_new))

idx = idx_exclusive(kuramoto_plastic!, ["e_22", "de_22"])
Frand = random_force(kuramoto_plastic!, [0.0, 1.0], Uniform, idx)
z_new = ambient_forcing(kuramoto_plastic!, x0_plastic, 2.0, Frand)

sum(g_nd(z_new))

# This file was generated using Literate.jl, https://github.com/fredrikekre/Literate.jl


julia> using Statistics; mean([1,2,3])
2.0
julia> mean(1:10)
5.5
julia> mean([])
ERROR: mean of empty collection undefined: []


#=
Serial direct estimation of the probability for the Schlögl chemical reaction network
with X(0)= x₀ to reach X(T) ∈ A, the target set.

This reaction network corresponds to:
A + 2S → 3S
3S → A + 2S
B → S
S → B

Parameters were chosen so that the system is bistable, as in:
https://doi.org/10.1063/1.5017955
=#

using Statistics
using HypothesisTests
using Printf

include("schlogl_setup.jl");

Random.seed!(100);
n_trials = 10^6;
n_in_set = 0;
for j in 1:n_trials
    X₀ = copy(x₀);
    global n_in_set+= f(integrate_schlogl(X₀, Tmax));
end
CI = confint(BinomialTest(n_in_set, n_trials));
@printf("P Estimate = %g\n",n_in_set/n_trials);
@printf("95%% CI = (%g,%g)\n", CI[1], CI[2]);


adv_eq_dir  = joinpath(output_root,"AdvectionEquation")

print_norms = false

@testset "Linear advection, ∂_t ϕ + ∇•(cϕ) = 0 ∈ ∂Ω, ϕ(t=0) = Gaussian(σ, μ), ConservativeForm, explicit Euler" begin
  dd = DomainDecomp(gm=1)
  dss = DomainSubSet(gm=true)
  σ, μ = .1, 0.5
  δz = 0.2
  Triangle(z) = μ-δz < z < μ+δz ? ( μ > z ? (z-(μ-δz))/δz : ((μ+δz)-z)/δz) : 0.0
  Square(z) = μ-δz < z < μ+δz ? 1 : 0.0
  Gaussian(z) = exp(-1/2*((z-μ)/σ)^2)
  for n_elems_real in (64, 128)
    grid = Grid(0.0, 1.0, n_elems_real)
    domain_range = over_elems_real(grid)
    x = [grid.zc[k] for k in domain_range]
    unknowns = ( (:ϕ, dss), )
    vars = ( (:ϕ_initial, dss), (:ϕ_error, dss), (:ϕ_analytic, dss), )
    q = StateVec(unknowns, grid, dd)
    tmp = StateVec(vars, grid, dd)
    rhs = deepcopy(q)
    CFL = 0.1
    Δt = CFL*grid.Δz
    T = 0.25
    maxiter = Int(T/Δt)
    for scheme in (
                   UpwindAdvective(),
                   UpwindCollocated(),
                   CenteredUnstable(),
                   )
      for distribution in (Triangle,
                           Square,
                           Gaussian)
        for wave_speed in (-1, 1)

          scheme_name = replace(string(scheme), "()"=>"")
          distribution_name = joinpath("LinearAdvection",string(distribution))
          directory = joinpath(adv_eq_dir, distribution_name, scheme_name)
          export_plots && mkpath(directory)
          print_norms && print("\n", directory, ", ", n_elems_real, ", ", wave_speed, ", ")

          for k in over_elems_real(grid)
            tmp[:ϕ_initial, k] = distribution(grid.zc[k])
            q[:ϕ, k] = tmp[:ϕ_initial, k]
          end
          amax_w = max([max(tmp[:ϕ_initial, k]) for k in over_elems_real(grid)]...)
          for i in 1:maxiter
            for k in over_elems_real(grid)
              ϕ = q[:ϕ, Cut(k)]
              ϕ_dual = q[:ϕ, Dual(k)]
              w = [wave_speed, wave_speed, wave_speed]
              w_dual = [wave_speed, wave_speed]
              rhs[:ϕ, k] = - advect_old(ϕ, ϕ_dual, w, w_dual, grid, scheme, Δt)
            end
            for k in over_elems(grid)
              q[:ϕ, k] += Δt*rhs[:ϕ, k]
            end
            apply_Dirichlet!(q, :ϕ, grid, 0.0, Zmax())
            apply_Dirichlet!(q, :ϕ, grid, 0.0, Zmin())
          end
          sol_analtyic = [distribution(grid.zc[k] - wave_speed*T) for k in over_elems(grid)]
          sol_error = [sol_analtyic[k] - q[:ϕ, k] for k in over_elems(grid)]
          L2_norm = sum(sol_error.^2)/length(sol_error)

          if !(scheme==CenteredUnstable())
            @test all(abs.(sol_error) .< 100*grid.Δz)
            @test all(L2_norm < 100*grid.Δz)
          end
          for k in over_elems(grid)
            tmp[:ϕ_error, k] = sol_error[k]
            tmp[:ϕ_analytic, k] = sol_analtyic[k]
          end

          print_norms && print("L2_norm(err) = ", L2_norm)

          name = string(n_elems_real)
          markershape = wave_speed==-1 ? :dtriangle : :utriangle
          if export_plots
            if wave_speed==-1
              y = [tmp[:ϕ_initial , k] for k in domain_range]; plot(y, x, label="initial",
                markercolor="blue", linecolor="blue", markershapes=markershape, markersize=2, legend=:topleft)
              y = [  q[:ϕ         , k] for k in domain_range]; plot!(y, x, label="numerical",
                markercolor="black", linecolor="black", markershapes=markershape, markersize=2, legend=:topleft)
              y = [tmp[:ϕ_error   , k] for k in domain_range]; plot!(y, x, label="error",
                markercolor="red", linecolor="red", markershapes=markershape, markersize=2, legend=:topleft)
              y = [tmp[:ϕ_analytic, k] for k in domain_range]; plot!(y, x, label="analytic",
                markercolor="green", linecolor="green", markershapes=markershape, markersize=2, legend=:topleft)
            else
              y = [  q[:ϕ         , k] for k in domain_range]; plot!(y, x, label="",
                markercolor="black", linecolor="black", markershapes=markershape, markersize=2, legend=:topleft)
              y = [tmp[:ϕ_error   , k] for k in domain_range]; plot!(y, x, label="",
                markercolor="red", linecolor="red", markershapes=markershape, markersize=2, legend=:topleft)
              y = [tmp[:ϕ_analytic, k] for k in domain_range]; plot!(y, x, label="",
                markercolor="green", linecolor="green", markershapes=markershape, markersize=2, legend=:topleft)
            end
            wave_speed==1 && savefig(joinpath(directory, name))
          end
        end
      end
    end
  end
end

@testset "Non-linear Bergers, ∂_t w + ∇•(ww) = 0 ∈ ∂Ω, u(t=0) = Gaussian(σ, μ), ConservativeForm, explicit Euler" begin
  dd = DomainDecomp(gm=1)
  dss = DomainSubSet(gm=true)
  σ, μ = .1, 0.5
  δz = 0.2
  Triangle(z, velocity_sign)   = μ-δz < z < μ+δz ? ( μ > z ? velocity_sign*(z-(μ-δz))/δz : velocity_sign*((μ+δz)-z)/δz  ) : 0.0
  Square(z, velocity_sign) = μ-δz < z < μ+δz ? velocity_sign : 0.0
  Gaussian(z, velocity_sign) = velocity_sign*exp(-1/2*((z-μ)/σ)^2)
  for n_elems_real in (64, 128)
    grid = Grid(0.0, 1.0, n_elems_real)
    domain_range = over_elems_real(grid)
    x = [grid.zc[k] for k in domain_range]
    unknowns = ( (:w, dss), )
    vars = ( (:w_initial, dss), (:w_error, dss), (:w_analytic, dss), )
    q = StateVec(unknowns, grid, dd)
    tmp = StateVec(vars, grid, dd)
    rhs = deepcopy(q)
    CFL = 0.1
    Δt = CFL*grid.Δz
    T = 0.25
    maxiter = Int(T/Δt)
    for scheme in (
                   UpwindAdvective(),
                   UpwindCollocated(),
                   CenteredUnstable(),
                   )
      for distribution in (Triangle,
                           Square,
                           Gaussian)
        for velocity_sign in (-1, 1)
          for k in over_elems_real(grid)
            tmp[:w_initial, k] = distribution(grid.zc[k], velocity_sign)
            q[:w, k] = tmp[:w_initial, k]
          end
          amax_w = max([max(tmp[:w_initial, k]) for k in over_elems_real(grid)]...)
          for i in 1:maxiter
            for k in over_elems_real(grid)
              w = q[:w, Cut(k)]
              w_dual = (w[1:2]+w[2:3])/2
              rhs[:w, k] = - advect_old(w, w_dual, w, w_dual, grid, scheme, Δt)
            end
            for k in over_elems(grid)
              q[:w, k] += Δt*rhs[:w, k]
            end
            apply_Dirichlet!(q, :w, grid, 0.0, Zmax())
            apply_Dirichlet!(q, :w, grid, 0.0, Zmin())
          end

          scheme_name = replace(string(scheme), "()"=>"")
          distribution_name = joinpath("BurgersEquation",string(distribution))
          directory = joinpath(adv_eq_dir, distribution_name, scheme_name)
          export_plots && mkpath(directory)
          name = string(n_elems_real)
          markershape = velocity_sign==-1 ? :dtriangle : :utriangle
          if export_plots
            if velocity_sign==-1
              y = [tmp[:w_initial , k] for k in domain_range]; plot(y, x, label="initial",
                markercolor="blue", linecolor="blue", markershapes=markershape, markersize=2, legend=:topleft)
              y = [  q[:w         , k] for k in domain_range]; plot!(y, x, label="numerical",
                markercolor="black", linecolor="black", markershapes=markershape, markersize=2, legend=:topleft)
            else
              y = [tmp[:w_initial , k] for k in domain_range]; plot!(y, x, label="",
                markercolor="blue", linecolor="blue", markershapes=markershape, markersize=2, legend=:topleft)
              y = [  q[:w         , k] for k in domain_range]; plot!(y, x, label="",
                markercolor="black", linecolor="black", markershapes=markershape, markersize=2, legend=:topleft)
            end
            velocity_sign==1 && savefig(joinpath(directory, name))
          end
        end
      end
    end
  end

  end


using Reexport

@reexport module TriangulationPartitions
    # Necessary modules
    include("Simplices/Simplices.jl")
    include("DelaunayTriangulations/DelaunayTriangulations.jl")

	dimension(DT::DelaunayTriangulation) = DelaunayTriangulations.dimension(DT)

    # Abstract types and functions that work on all subtypes
    include("AbstractTriangulationPartition.jl")

    # Composite types for different types of data
    include("composite_types/DatasetTriangulationPartition.jl")
    include("composite_types/EmbeddingTriangulationPartition.jl")
    include("composite_types/PointTriangulationPartition.jl")
	include("composite_types/MatrixTriangulationPartition.jl")

    # Functions to create triangulations from data
    include("triangulate.jl")

    # Plot triangulations
    include("plot_recipes/plot_recipes.jl")

end

"""
    TriangulationPartitions

A module handling partitioned formed by triangulations.
"""
TriangulationPartitions


# Advent of code 2021, day 4

function parseinput()
    paragraphs = split(strip(read("AoCdata/AoC_2021_day4.txt", String)), "\n\n")
    draworder = [parse(Int, num) for num in split(strip(first(paragraphs)), ",")]
    matrices, picks = Matrix{Int}[], Matrix{Int}[]
    for mattxt in paragraphs[(begin + 1):end]
        push!(matrices, reshape([parse(Int, s) for s in split(strip(mattxt))]', (5, 5)))
        push!(picks, zeros(Int, 5, 5))
    end
    return draworder, matrices, picks
end

function part1and2()
    win(mat) = any(sum(mat; dims=1) .== 5) || any(sum(mat; dims=2) .== 5)
    draw, mats, picks = parseinput()
    totalwinners, part1, winners, pastwinners = 0, 0, Int[], Int[]
    for d in draw
        for n in eachindex(mats)
            if (idx = findfirst(j -> j == d, mats[n])) != nothing
                picks[n][idx] = 1
            end
        end
        winners = filter(
            x -> !(x in pastwinners), findall(i -> win(picks[i]), eachindex(mats))
        )
        isempty(winners) && continue
        push!(pastwinners, winners...)
        winmat, winpick = mats[winners[end]], picks[winners[end]]
        unpicksum = sum([winmat[n] for n in eachindex(winmat) if winpick[n] == 0])
        if part1 == 0
            part1 = winners[1]
            println("Part 1: ", unpicksum * d)
        elseif length(pastwinners) == length(mats)  # all have won
            lastwinner = last(pastwinners)
            println("Last winner is $lastwinner: ", unpicksum * d)
            break
        end
    end
end

part1and2()


"""
Code for estimating the radial velocity based on projecting the flux onto the derivative of a template spectrum.

Author: Eric Ford
"""

""" Return mean flux (averaging over observations at different times, variance weighted) based on a common set of wavelengths.
   Inputs: flux & var (2d: pixel, time)
"""
function calc_mean_spectrum(flux::AbstractArray{T1,2}, var::AbstractArray{T2,2} ) where { T1<:Real, T2<:Real }
    flux_mean = vec(sum(flux./var,dims=2)./sum(1.0./var,dims=2))
end

""" Estimate numerical derivative of fluxes given wavelengths. """
function calc_dfluxdlnlambda(flux::AbstractArray{T1,1}, λ::AbstractArray{T2,1}) where { T1<:Real, T2<:Real }
    @assert size(flux) == size(λ)
    @assert length(flux) >= 3
    dfdlogλ = Array{T1,1}(undef,length(flux))
    calc_dfluxdlnlambda!(dfdlogλ,flux,λ)
    #dfdlogλ[1] = 0.5 * (flux[2]-flux[1])/(λ[2]-λ[1])*(λ[2]+λ[1])
    #dfdlogλ[2:end-1] .= 0.5 .* (flux[3:end].-flux[1:end-2])./(λ[3:end].-λ[1:end-2]).*(λ[3:end].+λ[1:end-2])
    #dfdlogλ[end] = 0.5 * (flux[end]-flux[end-1])/(λ[end]-λ[end-1])*(λ[end]+λ[end-1])
    return dfdlogλ
end

""" Estimate numerical derivative of fluxes given wavelengths. """
function calc_dfluxdlnlambda!(dfdlogλ::AbstractArray{T1,1}, flux::AbstractArray{T2,1}, λ::AbstractArray{T3,1}) where { T1<:Real, T2<:Real, T3<:Real }
    @assert size(flux) == size(λ)
    @assert size(dfdlogλ) == size(flux)
    @assert length(flux) >= 3
    #dfdlogλ = Array{T1,1}(undef,length(flux))
    dfdlogλ[1] = 0.5 * (flux[2]-flux[1])/(λ[2]-λ[1])*(λ[2]+λ[1])
    dfdlogλ[2:end-1] .= 0.5 .* (flux[3:end].-flux[1:end-2])./(λ[3:end].-λ[1:end-2]).*(λ[3:end].+λ[1:end-2])
    dfdlogλ[end] = 0.5 * (flux[end]-flux[end-1])/(λ[end]-λ[end-1])*(λ[end]+λ[end-1])
    return dfdlogλ
end

""" Estimate numerical second derivative of fluxes given wavelengths. """
function calc_d2fluxdlnlambda2(flux::AbstractArray{T1,1}, λ::AbstractArray{T2,1}) where { T1<:Real, T2<:Real }
    @assert size(flux) == size(λ)
    @assert length(flux) >= 3
    logλ = log.(λ)
    d2fdlogλ2 = Array{T1,1}(undef,length(flux))
    #d2fdlogλ2[2:end-1] .= 0.25*(flux[3:end].+flux[1:end-2].-2.0.*flux[2:end-1])./(λ[3:end].+λ[1:end-2].-2.0*λ[2:end-1]).*(λ[3:end].+λ[end-2]).^2
    #d2fdlogλ2[2:end-1] .= 0.5 * (flux[3:end].+flux[1:end-2].-2.0.*flux[2:end-1]).* ((λ[3:end].+λ[1:end-2])./(logλ[3:end].-logλ[1:end-2])).^2
    #d2fdlogλ2[end] = d2fdlogλ2[end-1]
    #d2fdlogλ2[1] = d2fdlogλ2[2]
    calc_d2fluxdlnlambda2!(d2fdlogλ2,flux,λ)
    return d2fdlogλ2
end

""" Estimate numerical second derivative of fluxes given wavelengths. """
function calc_d2fluxdlnlambda2!(d2fdlogλ2::AbstractArray{T1,1}, flux::AbstractArray{T1,1}, λ::AbstractArray{T2,1}) where { T1<:Real, T2<:Real, T3<:Real }
    @assert size(flux) == size(λ)
    @assert size(d2fdlogλ2) == size(flux)
    @assert length(flux) >= 3
    #logλ = log.(λ)
    #d2fdlogλ2 = Array{T1,1}(undef,length(flux))
    #d2fdlogλ2[2:end-1] .= 0.25*(flux[3:end].+flux[1:end-2].-2.0.*flux[2:end-1])./(λ[3:end].+λ[1:end-2].-2.0*λ[2:end-1]).*(λ[3:end].+λ[end-2]).^2
    d2fdlogλ2[2:end-1] .= 0.5 * (flux[3:end].+flux[1:end-2].-2.0.*flux[2:end-1]).* ((λ[3:end].+λ[1:end-2]).^2 ./ (2.0.*log.(λ[3:end]./λ[1:end-2])))
    d2fdlogλ2[end] = d2fdlogλ2[end-1]
    d2fdlogλ2[1] = d2fdlogλ2[2]
    return d2fdlogλ2
end

""" Return mean numerical derivative of fluxes based on a common set of wavelengths.
    Inputs: flux & var (2d) and λ (1d)
 """
function calc_mean_dfluxdlnlambda(flux::AbstractArray{T1,2}, var::AbstractArray{T1,2}, λ::AbstractArray{T3,1},
        chunk_map::AbstractArray{URT,1}) where
    { T1<:Real, T2<:Real, T3<:Real, URT<:AbstractUnitRange} #, RT<:AbstractRange }
    flux_mean = calc_mean_spectrum(flux,var)
    deriv = Array{T1,1}(undef,length(flux_mean))
    #map(c->deriv[c] .= calc_dfluxdlnlambda(flux_mean[c],λ[c]),chunk_map )
    for c in chunk_map
         calc_dfluxdlnlambda!(deriv[c],flux_mean[c],λ[c])
    end
    return deriv
end

function calc_mean_d2fluxdlnlambda2(flux::AbstractArray{T1,2}, var::AbstractArray{T1,2}, λ::AbstractArray{T3,1},
        chunk_map::AbstractArray{URT,1}) where
    { T1<:Real, T2<:Real, T3<:Real, URT<:AbstractUnitRange} #, RT<:AbstractRange }
    flux_mean = calc_mean_spectrum(flux,var)
    deriv2 = Array{T1,1}(undef,length(flux_mean))
    #map(c->deriv2[c] .= calc_d2fluxdlnlambda2(flux_mean[c],λ[c]),chunk_map )
    for c in chunk_map
         calc_d2fluxdlnlambda2!(deriv2[c],flux_mean[c],λ[c])
    end
    return deriv2
end

function calc_rvs_from_taylor_expansion(spectra::STS; mean::MT = calc_mean_spectrum(spectra),
                deriv::DT = calc_mean_dfluxdlnlambda(spectra), idx::RT = 1:length(mean),
                equal_weight::Bool = true ) where {
                    STS<:AbstractSpectralTimeSeriesCommonWavelengths, T1<:Real, MT<:AbstractVector{T1},
                    T2<:Real, DT<:AbstractVector{T2}, RT<:AbstractUnitRange }
   @assert length(mean) == length(deriv)
   @assert size(spectra.flux,1) == length(mean)
   @assert minimum(idx) >=1
   @assert maximum(idx) <= length(mean)

   if equal_weight # Pixels equally-weighted
      norm = sum(abs2.(deriv[idx]))
      rv = sum((spectra.flux[idx,:].-mean[idx]).*deriv[idx],dims=1).*(speed_of_light_mps/norm)
      # TODO: WARN: Uncertinaties ignores correlations between pixels, particularly problematic when oversample pixels
      σ_rv = sqrt.(sum(spectra.var[idx,:].*abs2.(deriv[idx]),dims=1)) .*(speed_of_light_mps/norm)
  else # Pixels inverse variance weighted
      # TODO:  CHECK/FIX? math for inverse variance weighting.
      # Or maybe the math is right, but it's just a bad idea to have different weighting within one line/chunk
      @info "Warning: I think this is either wrong or a bad idea.  Need to check."
      norm = sum(abs2.(deriv[idx])./spectra.var[idx,:])
      rv = sum((spectra.flux[idx,:].-mean[idx]).*deriv[idx]./spectra.var[idx,:],dims=1).*(speed_of_light_mps/norm)
      σ_rv = sqrt.(sum(abs2.(deriv)./spectra.var[idx,:],dims=1)).*(speed_of_light_mps/norm)
   end
   return (rv=-vec(rv), σ_rv=vec(σ_rv))
end

function calc_rvs_from_taylor_expansion_alt(spectra::STS; mean::MT = calc_mean_spectrum(spectra),
                deriv::DT = calc_mean_dfluxdlnlambda(spectra),
                deriv2::DT = calc_d2fluxdlnlambda2(spectra), idx::RT = 1:length(mean),
                equal_weight::Bool = true ) where {
                    STS<:AbstractSpectralTimeSeriesCommonWavelengths, T1<:Real, MT<:AbstractVector{T1},
                    T2<:Real, DT<:AbstractVector{T2}, RT<:AbstractUnitRange }
   @assert length(mean) == length(deriv)
   @assert size(spectra.flux,1) == length(mean)
   @assert minimum(idx) >=1
   @assert maximum(idx) <= length(mean)

   if equal_weight # Pixels equally-weighted
      #norm = sum(abs2.(deriv[idx]))
      norm = sum(deriv[idx].^2 .- (spectra.flux[idx,:].-mean[idx]).*deriv2[idx],dims=1)
      rv = sum(((spectra.flux[idx,:].-mean[idx]).*deriv[idx]),dims=1)./norm
      rv *= speed_of_light_mps
      # TODO: WARN: Uncertinaties ignores correlations between pixels, particularly problematic when oversample pixels
      σ_rv = sqrt.(sum(spectra.var[idx,:].*abs2.(deriv[idx]),dims=1)./norm)
      σ_rv *= speed_of_light_mps
  else # Pixels inverse variance weighted
      # TODO:  CHECK/FIX? math for inverse variance weighting.
      # Or maybe the math is right, but it's just a bad idea to have different weighting within one line/chunk
      @info "Warning: I think this is either wrong or a bad idea.  Need to check."
      norm = sum(abs2.(deriv[idx]./spectra.var[idx,:]))  # TODO comparable to norm_old, if new norm works, then check updating this also helps
      rv = sum((spectra.flux[idx,:].-mean[idx]).*deriv[idx]./spectra.var[idx,:],dims=1).*(speed_of_light_mps/norm)
      σ_rv = sqrt.(sum(abs2.(deriv)./spectra.var[idx,:],dims=1)).*(speed_of_light_mps/norm)
   end
   return (rv=-vec(rv), σ_rv=vec(σ_rv))
end

function calc_chunk_rvs_from_taylor_expansion(spectra::STS; mean::MT = calc_mean_spectrum(spectra),
                deriv::DT = calc_mean_dfluxdlnlambda(spectra),
                equal_weight::Bool = false ) where {
                    STS<:AbstractSpectralTimeSeriesCommonWavelengths, T1<:Real, MT<:AbstractVector{T1},
                    T2<:Real, DT<:AbstractVector{T2}, RT<:AbstractUnitRange }
   @assert length(mean) == length(deriv)
   @assert size(spectra.flux,1) == length(mean)

   map(idx->calc_rvs_from_taylor_expansion(spectra,mean=mean,deriv=deriv,idx=idx),spectra.chunk_map)
end


using MPI
using CLIMA
using CLIMA.Mesh.Topologies
using CLIMA.Mesh.Grids
using CLIMA.DGBalanceLawDiscretizations
using Printf
using LinearAlgebra
using Logging
using CLIMA.MPIStateArrays

if !@isdefined integration_testing
  const integration_testing =
    parse(Bool, lowercase(get(ENV,"JULIA_CLIMA_INTEGRATION_TESTING","false")))
end

const ArrayTypes = (CLIMA.array_type(),)

const _nauxstate = 5
const _a_r, _a_θ, _a_ϕ, _a_f, _a_fdown = 1:_nauxstate
@inline function auxiliary_state_initialization!(aux, x, y, z, Rinner, Router)
  @inbounds begin
    r = hypot(x, y, z)
    θ = atan(y , x)
    ϕ = asin(z / r)
    aux[_a_r] = r
    aux[_a_θ] = θ
    aux[_a_ϕ] = ϕ

    # Exact integral
    a = 1 + sin(θ)^2 + sin(ϕ)^2
    aux[_a_f] = exp(-a * r^2) - exp(-a * Rinner^2)
    aux[_a_fdown] = exp(-a * Router^2) - exp(-a * r^2)
  end
end

@inline function integral_knl(val, Q, aux)
  @inbounds begin
    r, θ, ϕ = aux[_a_r], aux[_a_θ], aux[_a_ϕ]
    a = 1 + sin(θ)^2 + sin(ϕ)^2
    val[1] = -2r * a * exp(-a * r^2)
  end
end

using Test
function run(mpicomm, topl, ArrayType, N, FT, Rinner, Router)
  grid = DiscontinuousSpectralElementGrid(topl,
                                          FloatType = FT,
                                          DeviceArray = ArrayType,
                                          polynomialorder = N,
                                          meshwarp = Topologies.cubedshellwarp,
                                         )

  spacedisc = DGBalanceLaw(grid = grid,
                           length_state_vector = 0,
                           flux! = (x...) -> (),
                           numerical_flux! = (x...) -> (),
                           auxiliary_state_length = _nauxstate,
                           auxiliary_state_initialization! = (x...) ->
                           auxiliary_state_initialization!(x..., Rinner,
                                                           Router),
                           numerical_boundary_flux! = (x...) -> ())

  Q = MPIStateArray(spacedisc)
  exact_aux = copy(spacedisc.auxstate)

  DGBalanceLawDiscretizations.indefinite_stack_integral!(spacedisc,
                                                         integral_knl, Q,
                                                         _a_f)
  DGBalanceLawDiscretizations.reverse_indefinite_stack_integral!(spacedisc,
                                                                 _a_fdown, _a_f)

  euclidean_distance(exact_aux, spacedisc.auxstate)
end

let
  CLIMA.init()
  mpicomm = MPI.COMM_WORLD
  ll = uppercase(get(ENV, "JULIA_LOG_LEVEL", "INFO"))
  loglevel = ll == "DEBUG" ? Logging.Debug :
  ll == "WARN"  ? Logging.Warn  :
  ll == "ERROR" ? Logging.Error : Logging.Info
  logger_stream = MPI.Comm_rank(mpicomm) == 0 ? stderr : devnull
  global_logger(ConsoleLogger(logger_stream, loglevel))

  polynomialorder = 4

  base_Nhorz = 4
  base_Nvert = 2
  Rinner = 1//2
  Router = 1

  polynomialorder = 4

  expected_result = [6.228615762850257e-7
                     9.671308320438864e-9
                     1.5102832678375277e-10
                     2.359860999112363e-12]

  lvls = integration_testing ? length(expected_result) : 1

  @testset "$(@__FILE__)" for ArrayType in ArrayTypes
    for FT in (Float64,) #Float32)
      err = zeros(FT, lvls)
      for l = 1:lvls
        @info (ArrayType, FT, "sphere", l)
        Nhorz = 2^(l-1) * base_Nhorz
        Nvert = 2^(l-1) * base_Nvert
        Rrange = range(FT(Rinner); length=Nvert+1, stop=Router)
        topl = StackedCubedSphereTopology(mpicomm, Nhorz, Rrange)
        err[l] = run(mpicomm, topl, ArrayType, polynomialorder, FT,
                     FT(Rinner), FT(Router))
        @test expected_result[l] ≈ err[l]
      end
      if integration_testing
        @info begin
          msg = ""
          for l = 1:lvls-1
            rate = log2(err[l]) - log2(err[l+1])
            msg *= @sprintf("\n  rate for level %d = %e\n", l, rate)
          end
          msg
        end
      end
    end
  end
end

nothing


include("test_schurparlett.jl")
include("test_aaa.jl")
include("test_ratkrylov.jl")