struct ClusterCovariance{T} <: CovarianceEstimator
    clusters::T
end

cluster(x::Symbol) = ClusterCovariance((x,))
cluster(args...) = ClusterCovariance(args)

Vcov.completecases(df::AbstractDataFrame, v::ClusterCovariance) = DataFrames.completecases(df, collect(v.clusters))

function materialize(df::AbstractDataFrame, v::ClusterCovariance)
    ClusterCovariance(NamedTuple{v.clusters}(ntuple(i -> group(df[!, v.clusters[i]]), length(v.clusters))))
end

function df_FStat(::RegressionModel, v::ClusterCovariance, ::Bool)
    minimum((length(c.pool) for c in values(v.clusters))) - 1
end

function S_hat(x::RegressionModel, v::ClusterCovariance)
    # Cameron, Gelbach, & Miller (2011): section 2.3
    dim = size(modelmatrix(x), 2) * size(residuals(x), 2)
    S = zeros(dim, dim)
    G = typemax(Int)
    for c in combinations(keys(v.clusters))
        # no need for group in case of one fixed effect, since was already done in VcovMethod
        f = (length(c) == 1) ? v.clusters[c[1]] : group((v.clusters[var] for var in c)...)
        # capture length of smallest dimension of multiway clustering in G
        G = min(G, length(f.pool))
        S += (-1)^(length(c) - 1) * helper_cluster(modelmatrix(x), residuals(x), f)
    end
    # scale total vcov estimate by ((N-1)/(N-K)) * (G/(G-1))
    # another option would be to adjust each matrix given by helper_cluster by number of its categories
    # both methods are presented in Cameron, Gelbach and Miller (2011), section 2.3
    # I choose the first option following reghdfe
    rmul!(S, (size(modelmatrix(x), 1) - 1) / dof_residual(x) * G / (G - 1))
end

# res is a Vector in OLS, Matrix in IV
function helper_cluster(X::Matrix, res::Union{Vector, Matrix}, f::CategoricalVector)
    X2 = zeros(eltype(X), length(f.pool), size(X, 2) * size(res, 2))
    index = 0
    for k in 1:size(res, 2)
        for j in 1:size(X, 2)
            index += 1
            @inbounds @simd for i in 1:size(X, 1)
                X2[f.refs[i], index] += X[i, j] * res[i, k]
            end
        end
    end
    return Symmetric(X2' * X2)
end

function StatsBase.vcov(x::RegressionModel, v::ClusterCovariance)
    xtx = inv(crossmodelmatrix(x))
    pinvertible(Symmetric(xtx * S_hat(x, v) * xtx))
end


stats_names2D = ["mode","mean","var"]
run_stats_names2D = ["nsteps","nchains","Times"]
##########################################multi normal###########################################
sig = Matrix{Float64}([1.5^2 1.5*2.5*0.4 ; 1.5*2.5*0.4 2.5^2])
analytical_stats_gauss2D = Vector{Any}(undef,length(stats_names2D))
analytical_stats_gauss2D[1] = [15,10]       #analytical_mode_gauss2D
analytical_stats_gauss2D[2] = [15,10]       #analytical_mean_gauss2D
analytical_stats_gauss2D[3] = [1.5^2,2.5^2] #analytical_var_gauss2D
gauss2D = (posterior=MvNormal([15,10],sig),mode=analytical_stats_gauss2D[1],mean=analytical_stats_gauss2D[2],var=analytical_stats_gauss2D[3],chi2=[9999],ks=[9999.,9999.],ahmi=[9999.,9999.])

########################################multi cauchy###########################################
analytical_stats_multi_cauchy2D = Vector{Any}(undef,length(stats_names2D))
analytical_stats_multi_cauchy2D[1] = [[-5,-5],[-5,5],[5,-5],[5,5]]      #analytical_mode_multi_cauchy2D
analytical_stats_multi_cauchy2D[2] = [0,0]                              #analytical_mean_multi_cauchy2D
analytical_stats_multi_cauchy2D[3] = [Inf,Inf]                #analytical_var_multi_cauchy2D
multi_cauchy2D = (posterior=BAT.MultimodalCauchy(µ=5.,σ=4.,n=2),mode=analytical_stats_multi_cauchy2D[1],mean=analytical_stats_multi_cauchy2D[2],var=analytical_stats_multi_cauchy2D[3],chi2=[9999],ks=[9999.,9999.],ahmi=[9999.,9999.])

################function funnel 2D##############################################################
analytical_stats_funnel2D = Vector{Any}(undef,length(stats_names2D))
analytical_stats_funnel2D[1] = [-1.0,0.0]
analytical_stats_funnel2D[2] = [0.0,0.0]
analytical_stats_funnel2D[3] = [1.0,7.406718]
funnel2D = (posterior=BAT.FunnelDistribution(1.,0.5,2),mode=analytical_stats_funnel2D[1],mean=analytical_stats_funnel2D[2],var=analytical_stats_funnel2D[3],chi2=[9999],ks=[9999.,9999.],ahmi=[9999.,9999.])

################################################################################################
testfunctions_2D = Dict(
	"normal" 		=> gauss2D,
	"multi cauchy"	=> multi_cauchy2D,
	"funnel"		=> funnel2D
)

sample_stats2D=[Vector{Any}(undef,length(stats_names2D)) for i in 1:length(testfunctions_2D)]
run_stats2D=[Vector{Any}(undef,length(run_stats_names2D)) for i in 1:length(testfunctions_2D)]


module FluxUtils

using Flux
using Flux.Tracker
using Random
# import Reproduce: ArgParseSettings, @add_arg_table

glorot_uniform(rng::Random.AbstractRNG, T::Type, dims...) = (rand(rng, T, dims...) .- 0.5f0) .* sqrt(24.0f0/sum(dims))
glorot_normal(rng::Random.AbstractRNG, T::Type, dims...) = randn(rng, T, dims...) .* sqrt(2.0f0/sum(dims))

glorot_uniform(rng::Random.AbstractRNG, dims::Vararg{Int64}) = (rand(rng, Float32, dims...) .- 0.5f0) .* sqrt(24.0f0/sum(dims))
glorot_normal(rng::Random.AbstractRNG, dims...) = randn(rng, Float32, dims...) .* sqrt(2.0f0/sum(dims))

function get_activations(model::M, data) where {M<:Flux.Chain}
    # Assume it is a chain of things
    activations = [data]
    for (idx, layer) in enumerate(model)
        push!(activations, layer(activations[idx]))
    end
    return activations
end

struct ConcatStreams{M1, M2}
    m1::M1
    m2::M2
end

Flux.@treelike ConcatStreams
(l::ConcatStreams)(x) = vcat(l.m1(x), l.m2(x))

function Base.show(io::IO, l::ConcatStreams)
  print(io, "ConcatStreams(", string(l.m1), ", ", string(l.m2), ")")
end

end


module HungarianAlg

    mutable struct Result
        #= A simple result object. =#
        match
        revenues
        row_weights
        col_weights
        revenue_sum
    end

    function hungarian(matrx)
        #= Runs the Hungarian Algorithm on a given matrix and returns the optimal matching with potentials. =#

        # Step 1: Prep matrix, get size
        matchsize = size(matrx)[1]

        # Step 2+3: Generate trivial potentials and alternating tree
        rpotentials = []
        cpotentials = []
        matching = -1 * ones(Int64, 1, matchsize)
        x_parents =  -1 * ones(Int64, 1, matchsize)
        y_parents =  -1 * ones(Int64, 1, matchsize)
        for i = 1:matchsize
            # the row weight is the maximum revenue in a row
            push!(rpotentials, maximum(matrx[i,:]))
            push!(cpotentials, 0)
        end

        treed_x = BitSet([1])
        untreed_y = BitSet(1:matchsize)
        treed_y = BitSet()

        # Step 4: Loop while our matching is too small
        while -1 in matching
            # Step A: Find any neighbour in equality graph
            # where a row is in the tree and a col is not in the tree
            pair = nothing
            for x in treed_x, y in untreed_y
                if matrx[x, y] == rpotentials[x] + cpotentials[y]
                    pair = [x, y]
                    break
                end
            end

            if isnothing(pair)
                # Step B: If all firms are in the tree, update potentials to get a new one
                big = Inf
                # iterate over relevant pairs
                for dx in treed_x, dy in untreed_y
                    # find the difference and check if its smaller than any we found before
                    alpha = rpotentials[dx] + cpotentials[dy] - matrx[dx, dy]
                    if alpha < big
                      big = alpha
                      pair = [dx, dy]
                    end
                end

                # apply difference to potentials as needed
                for dx in treed_x
                    rpotentials[dx] -= big
                end

                for dy in treed_y
                    cpotentials[dy] += big
                end
            end

            if !(pair[2] in matching)
                # Step D: Firm is not matched so add it to matching
                matching[pair[1]] = pair[2]
                # Step E: Swap the alternating path in our alternating tree attached to the worker we matched
                source = pair[1]
                matched = 1
                while true
                    if matched == 1
                        if x_parents[source] == -1
                            break
                        end
                        above = x_parents[source]
                    else
                        above = y_parents[source]
                        matching[above] = source
                    end
                    matched = 1 - matched
                    source = above
                end

                # Step F: Destroy the tree, go to Step 4 to check completion, and possibly go to Step A
                if -1 in matching
                    x_parents =  -1 * ones(Int64, 1, matchsize)
                    y_parents =  -1 * ones(Int64, 1, matchsize)
                    free = findfirst(isequal(-1), matching)[2]
                    treed_x = BitSet([free])
                    untreed_y = BitSet(1:matchsize)
                    treed_y = BitSet()
                end

            else
                # Step C: Firm is matched so add it to the tree and go back to Step A
                wasMatchedTo = findfirst(isequal(pair[2]), matching)[2]
                push!(treed_x, wasMatchedTo)
                push!(treed_y, pair[2])
                delete!(untreed_y, pair[2])
                y_parents[pair[2]] = pair[1]
                x_parents[wasMatchedTo] = pair[2]
            end

        end
        revenues = [matrx[i, matching[i]] for i = 1:matchsize]
        return Result([[i, matching[i]] for i = 1:matchsize], revenues, rpotentials, cpotentials, sum(revenues))
    end

    function display(self)
        #= A function for pretty-printing a Hungarian Algorithm Result =#
        matchsize = length(self.match)
        maxlen = max(length(string(maximum(self.revenues))), length(string(minimum(self.revenues))))
        baselist = [[" "^maxlen for i = 1:matchsize] for j = 1:matchsize]
        for i = 1:matchsize
            entry = self.match[i]
            newstring = string(self.revenues[i])
            if length(newstring) > maxlen
                appender = ""
            else
                appender = " "^(maxlen - length(newstring))
            end
            baselist[entry[1]][entry[2]] = appender * newstring
        end
        formatted_list = join([string(row) for row in baselist], "\n")
        return "Matching:\n" * formatted_list*"\n\nRow Potentials: " * string(self.row_weights) * "\nColumn Potentials: " * string(self.col_weights)
    end

end


@testset "WeightedCovMatrix fit!" begin
    s, m, c = (map(x -> sum(x .* w), eachcol(x2)),
               map(x -> mean(x, weights(w)), eachcol(x2)),
               cov(x2, aweights(w), corrected = true))
    ma, ca = map(x -> mean(x, weights(w)), eachcol(x2)), cov(x2, aweights(w), corrected = true)
    mf, cf = map(x -> mean(x, weights(w)), eachcol(x2)), cov(x2, fweights(w), corrected = true)
    mp, cp = map(x -> mean(x, weights(w)), eachcol(x2)), cov(x2, pweights(w), corrected = true)

    o = WeightedCovMatrix()
    for i in 1:l
        fit!(o, x2[i,:], w[i])
    end
    o
    sfor, mfor, cfor = sum(o), mean(o), cov(o)

    o_32 = WeightedCovMatrix(Float32)
    for i in 1:l
        fit!(o_32, x2[i,:], w[i])
    end
    o_32
    sfor_32, mfor_32, cfor_32 = sum(o_32), mean(o_32), cov(o_32)

    sval, mval, cval = fit!(WeightedCovMatrix(), x2, w) |> x -> (sum(x), mean(x), cov(x))
    szip, mzip, czip = fit!(WeightedCovMatrix(), zip(eachrow(x2), w)) |> x -> (sum(x), mean(x), cov(x))

    mvala, cvala = fit!(WeightedCovMatrix(), x2, w) |>
        x -> (mean(x), cov(x, corrected = true, weight_type = :analytic))
    mvalf, cvalf = fit!(WeightedCovMatrix(), x2, w) |>
        x -> (mean(x), cov(x, corrected = true, weight_type = :frequency))
    @test_throws ArgumentError fit!(WeightedCovMatrix(), x2, w) |>
        x -> (mean(x), cov(x, corrected = true, weight_type = :something))

    @test c ≈ czip
    @test c ≈ cfor
    @test c ≈ cfor_32
    @test c ≈ cval
    @test m ≈ mzip
    @test m ≈ mfor
    @test m ≈ mfor_32
    @test m ≈ mval
    @test ca ≈ cvala
    @test cf ≈ cvalf
    @test ma ≈ mvala
    @test mf ≈ mvalf

    @test s ≈ sfor
    @test s ≈ sfor_32
    @test s ≈ sval
    @test s ≈ szip

    # After implementing :probability, these should pass/not throw any more:
    @test_throws ErrorException mvalp, vvalp = fit!(WeightedCovMatrix(), x2, w) |>
        x -> (mean(x), cov(x, corrected = true, weight_type = :probability))
    @test_broken vp ≈ vvalp
    @test_broken mp ≈ mvalp

    @test eltype(o) == Float64
    @test eltype(o_32) == Float32

    oold = copy(o)
    fit!(o, [missing, 1, 2, 3, 4], 1)
    @test o == oold
    fit!(o, [1, 2, 3, 4, 5], missing)
    @test o == oold

    # issue #29
    fit!(o, eachrow([missing 1 2 3 4;
                     missing 2 3 4 5]),
         [1.0, 1.0])
    @test o == oold
end

@testset "WeightedCovMatrix merge!" begin
    c = cov(x2, aweights(w), corrected = true)

    wc = fit!(WeightedCovMatrix(), x2, w)
    oc = map(eachrow(x2), w) do xi, wi
        fit!(WeightedCovMatrix(), xi, wi)
    end;
    rc = reduce(merge!, deepcopy(oc))
    rc2 = merge!(
        fit!(WeightedCovMatrix(), x2[1:end ÷ 2, :],       w[1:end ÷ 2]),
        fit!(WeightedCovMatrix(), x2[end ÷ 2 + 1:end, :], w[end ÷ 2 + 1:end]))

    rc_32 = reduce(merge!, deepcopy(oc), init = WeightedCovMatrix(Float32))
    rc2_32 = merge!(
        fit!(WeightedCovMatrix(Float32), x2[1:end ÷ 2, :],       w[1:end ÷ 2]),
        fit!(WeightedCovMatrix(),        x2[end ÷ 2 + 1:end, :], w[end ÷ 2 + 1:end]))

    @test rc.b ≈ wc.b
    @test rc.C ≈ wc.C
    @test rc.W ≈ wc.W
    @test rc.W2 ≈ wc.W2

    @test rc2.b ≈ wc.b
    @test rc2.C ≈ wc.C
    @test rc2.W ≈ wc.W
    @test rc2.W2 ≈ wc.W2

    @test cov(rc) ≈ c
    @test cov(rc2) ≈ c

    @test rc_32.b ≈ wc.b
    @test rc_32.C ≈ wc.C
    @test rc_32.W ≈ wc.W
    @test rc_32.W2 ≈ wc.W2

    @test rc2_32.b ≈ wc.b
    @test rc2_32.C ≈ wc.C
    @test rc2_32.W ≈ wc.W
    @test rc2_32.W2 ≈ wc.W2

    @test cov(rc_32) ≈ c
    @test cov(rc2_32) ≈ c

    @test eltype(rc) == Float64
    @test eltype(rc2) == Float64
    @test eltype(rc_32) == Float32
    @test eltype(rc2_32) == Float32
end

@testset "WeightedCovMatrix copy" begin
    o = fit!(WeightedCovMatrix(), x2, w)

    o2 = copy(o)

    @test Ref(o.C) != Ref(o2.C)
    @test Ref(o.A) != Ref(o2.A)
    @test Ref(o.b) != Ref(o2.b)
    @test Ref(o.W) != Ref(o2.W)
    @test Ref(o.W2) != Ref(o2.W2)
    @test Ref(o.n) != Ref(o2.n)

    @test o.C == o2.C
    @test o.A == o2.A
    @test o.b == o2.b
    @test o.W == o2.W
    @test o.W2 == o2.W2
    @test o.n == o2.n
end

@testset "WeightedCovMatrix constructor" begin
    @test WeightedCovMatrix{Float64}() == WeightedCovMatrix()
    @test WeightedCovMatrix{Float32}() == WeightedCovMatrix(Float32)
    @test WeightedCovMatrix() == WeightedCovMatrix(zeros(Float64, 0, 0),
                                                   zeros(Float64, 0, 0),
                                                   zeros(Float64, 0),
                                                   0.0, 0.0, 0)
    @test convert(WeightedCovMatrix{Float32}, WeightedCovMatrix()) ==
        WeightedCovMatrix{Float32}()
end


@testset "Estimators" begin
  # floating point tolerance
  tol = 100eps()

  # create some data
  dim = 3; nobs = 10
  pset = PointSet(10*rand(dim, nobs))
  data = georef((z=rand(nobs),), pset)

  γ = GaussianVariogram(sill=1., range=1., nugget=0.)
  simkrig = SimpleKriging(data, γ, mean(data[:z]))
  ordkrig = OrdinaryKriging(data, γ)
  unikrig = UniversalKriging(data, γ, 1)
  drikrig = ExternalDriftKriging(data, γ, [x->1.])

  # Kriging is an interpolator
  for j in 1:nobs
    SKestimate, SKvar = predict(simkrig, :z, pset[j])
    OKestimate, OKvar = predict(ordkrig, :z, pset[j])
    UKestimate, UKvar = predict(unikrig, :z, pset[j])
    DKestimate, DKvar = predict(drikrig, :z, pset[j])

    # estimate checks
    @test SKestimate ≈ data[:z][j]
    @test OKestimate ≈ data[:z][j]
    @test UKestimate ≈ data[:z][j]
    @test DKestimate ≈ data[:z][j]

    # variance checks
    @test SKvar + tol ≥ 0
    @test OKvar + tol ≥ 0
    @test UKvar + tol ≥ 0
    @test DKvar + tol ≥ 0
    @test SKvar ≤ OKvar + tol
  end

  # save results on a particular location pₒ
  pₒ = rand(Point3)
  SKestimate, SKvar = predict(simkrig, :z, pₒ)
  OKestimate, OKvar = predict(ordkrig, :z, pₒ)
  UKestimate, UKvar = predict(unikrig, :z, pₒ)
  DKestimate, DKvar = predict(drikrig, :z, pₒ)

  # Kriging is translation-invariant
  h = rand(Vec3)
  pset_h = PointSet([pset[i] + h for i in 1:nelements(pset)])
  data_h = georef((z=data[:z],), pset_h)
  simkrig_h = SimpleKriging(data_h, γ, mean(data_h[:z]))
  ordkrig_h = OrdinaryKriging(data_h, γ)
  unikrig_h = UniversalKriging(data_h, γ, 1)
  drikrig_h = ExternalDriftKriging(data_h, γ, [x->1.])
  SKestimate_h, SKvar_h = predict(simkrig_h, :z, pₒ + h)
  OKestimate_h, OKvar_h = predict(ordkrig_h, :z, pₒ + h)
  UKestimate_h, UKvar_h = predict(unikrig_h, :z, pₒ + h)
  DKestimate_h, DKvar_h = predict(drikrig_h, :z, pₒ + h)
  @test SKestimate_h ≈ SKestimate
  @test SKvar_h ≈ SKvar
  @test OKestimate_h ≈ OKestimate
  @test OKvar_h ≈ OKvar
  @test UKestimate_h ≈ UKestimate
  @test UKvar_h ≈ UKvar
  @test DKestimate_h ≈ DKestimate
  @test DKvar_h ≈ DKvar

  # Kriging estimate is invariant under covariance scaling
  # Kriging variance is multiplied by the same factor
  α = rand()
  γ_α = GaussianVariogram(sill=α, range=1., nugget=0.)
  simkrig_α = SimpleKriging(data, γ_α, mean(data[:z]))
  ordkrig_α = OrdinaryKriging(data, γ_α)
  unikrig_α = UniversalKriging(data, γ_α, 1)
  drikrig_α = ExternalDriftKriging(data, γ_α, [x->1.])
  SKestimate_α, SKvar_α = predict(simkrig_α, :z, pₒ)
  OKestimate_α, OKvar_α = predict(ordkrig_α, :z, pₒ)
  UKestimate_α, UKvar_α = predict(unikrig_α, :z, pₒ)
  DKestimate_α, DKvar_α = predict(drikrig_α, :z, pₒ)
  @test SKestimate_α ≈ SKestimate
  @test SKvar_α ≈ α*SKvar
  @test OKestimate_α ≈ OKestimate
  @test OKvar_α ≈ α*OKvar
  @test UKestimate_α ≈ UKestimate
  @test UKvar_α ≈ α*UKvar
  @test DKestimate_α ≈ DKestimate
  @test DKvar_α ≈ α*DKvar

  # Kriging variance is a function of data configuration, not data values
  δ = rand(nobs)
  data_δ = georef((z=data[:z].+δ,), pset)
  simkrig_δ = SimpleKriging(data_δ, γ, mean(data_δ[:z]))
  ordkrig_δ = OrdinaryKriging(data_δ, γ)
  unikrig_δ = UniversalKriging(data_δ, γ, 1)
  drikrig_δ = ExternalDriftKriging(data_δ, γ, [x->1.])
  SKestimate_δ, SKvar_δ = predict(simkrig_δ, :z, pₒ)
  OKestimate_δ, OKvar_δ = predict(ordkrig_δ, :z, pₒ)
  UKestimate_δ, UKvar_δ = predict(unikrig_δ, :z, pₒ)
  DKestimate_δ, DKvar_δ = predict(drikrig_δ, :z, pₒ)
  @test SKvar_δ ≈ SKvar
  @test OKvar_δ ≈ OKvar
  @test UKvar_δ ≈ UKvar
  @test DKvar_δ ≈ DKvar

  # Ordinary Kriging ≡ Universal Kriging with 0th degree drift
  unikrig_0th = UniversalKriging(data, γ, 0)
  OKestimate, OKvar = predict(ordkrig, :z, pₒ)
  UKestimate, UKvar = predict(unikrig_0th, :z, pₒ)
  @test OKestimate ≈ UKestimate
  @test OKvar ≈ UKvar

  # Ordinary Kriging ≡ Kriging with constant external drift
  driftkrig_const = ExternalDriftKriging(data, γ, [x->1.])
  OKestimate, OKvar = predict(ordkrig, :z, pₒ)
  DKestimate, DKvar = predict(driftkrig_const, :z, pₒ)
  @test OKestimate ≈ DKestimate
  @test OKvar ≈ DKvar

  # Non-stationary variograms are allowed
  γ_ns = PowerVariogram()
  ordkrig_ns = OrdinaryKriging(data, γ_ns)
  unikrig_ns = UniversalKriging(data, γ_ns, 1)
  drikrig_ns = ExternalDriftKriging(data, γ_ns, [x->1.])
  for j in 1:nobs
    OKestimate, OKvar = predict(ordkrig_ns, :z, pset[j])
    UKestimate, UKvar = predict(unikrig_ns, :z, pset[j])
    DKestimate, DKvar = predict(drikrig_ns, :z, pset[j])

    # estimate checks
    @test OKestimate ≈ data[:z][j]
    @test UKestimate ≈ data[:z][j]
    @test DKestimate ≈ data[:z][j]

    # variance checks
    @test OKvar + tol ≥ 0
    @test UKvar + tol ≥ 0
    @test DKvar + tol ≥ 0
  end

  # Floating point precision checks
  X_f         = rand(Float32, dim, nobs)
  z_f         = rand(Float32, nobs)
  X_d         = Float64.(X_f)
  z_d         = Float64.(z_f)
  pset_f      = PointSet(X_f)
  data_f      = georef((z=z_f,), pset_f)
  pset_d      = PointSet(X_d)
  data_d      = georef((z=z_d,), pset_d)
  pₒ_f        = rand(Point{dim,Float32})
  pₒ_d        = convert(Point{dim,Float64}, pₒ_f)
  γ_f         = GaussianVariogram(sill=1f0, range=1f0, nugget=0f0)
  simkrig_f   = SimpleKriging(data_f, γ_f, mean(data_f[:z]))
  ordkrig_f   = OrdinaryKriging(data_f, γ_f)
  unikrig_f   = UniversalKriging(data_f, γ_f, 1)
  drikrig_f   = ExternalDriftKriging(data_f, γ_f, [x->1f0])
  γ_d         = GaussianVariogram(sill=1., range=1., nugget=0.)
  simkrig_d   = SimpleKriging(data_d, γ_d, mean(data_d[:z]))
  ordkrig_d   = OrdinaryKriging(data_d, γ_d)
  unikrig_d   = UniversalKriging(data_d, γ_d, 1)
  drikrig_d = ExternalDriftKriging(data_d, γ_d, [x->1.])
  SKestimate_f, SKvar_f = predict(simkrig_f, :z, pₒ_f)
  OKestimate_f, OKvar_f = predict(ordkrig_f, :z, pₒ_f)
  UKestimate_f, UKvar_f = predict(unikrig_f, :z, pₒ_f)
  DKestimate_f, DKvar_f = predict(drikrig_f, :z, pₒ_f)
  SKestimate_d, SKvar_d = predict(simkrig_d, :z, pₒ_d)
  OKestimate_d, OKvar_d = predict(ordkrig_d, :z, pₒ_d)
  UKestimate_d, UKvar_d = predict(unikrig_d, :z, pₒ_d)
  DKestimate_d, DKvar_d = predict(drikrig_d, :z, pₒ_d)
  @test isapprox(SKestimate_f, SKestimate_d, atol=1e-4)
  @test isapprox(SKvar_f, SKvar_d, atol=1e-4)
  @test isapprox(OKestimate_f, OKestimate_d, atol=1e-4)
  @test isapprox(OKvar_f, OKvar_d, atol=1e-4)
  @test isapprox(UKestimate_f, UKestimate_d, atol=1e-4)
  @test isapprox(UKvar_f, UKvar_d, atol=1e-4)
  @test isapprox(DKestimate_f, DKestimate_d, atol=1e-4)
  @test isapprox(DKvar_f, DKvar_d, atol=1e-4)

  # ------------------
  # CHANGE OF SUPPORT
  # ------------------

  # create some data
  dim = 2; nobs = 10
  pset = PointSet(10*rand(dim, nobs))
  data = georef((z=rand(nobs),), pset)

  # basic estimators
  γ = GaussianVariogram(sill=1., range=1., nugget=0.)
  simkrig = SimpleKriging(data, γ, mean(data[:z]))
  ordkrig = OrdinaryKriging(data, γ)
  unikrig = UniversalKriging(data, γ, 1)
  drikrig = ExternalDriftKriging(data, γ, [x->1.])

  # prediction on a quadrangle
  uₒ = Quadrangle((0.,0.), (1.,0.), (1.,1.), (0.,1.))
  _, SKvar = predict(simkrig, :z, uₒ)
  _, OKvar = predict(ordkrig, :z, uₒ)
  _, UKvar = predict(unikrig, :z, uₒ)
  _, DKvar = predict(drikrig, :z, uₒ)

  # variance checks
  @test SKvar + tol ≥ 0
  @test OKvar + tol ≥ 0
  @test UKvar + tol ≥ 0
  @test DKvar + tol ≥ 0
  @test SKvar ≤ OKvar + tol
end


using CommutativeRings
using Test
using Polynomials

if !isdefined(Polynomials, :Poly)
    Poly = Polynomials.Polynomial
end

ismonic = CommutativeRings.ismonic

import Base: ==
function ==(a::UnivariatePolynomial, b::Poly)
    va = getproperty.(a.coeff, :val)
    vb = coeffs(b)
    n = length(vb)
    while n > 0 && vb[n] == 0
        n -= 1
    end
    va == vb[1:n]
end

let S = ZZ{Int}, P = UnivariatePolynomial{S,:x}, PQ = UnivariatePolynomial{QQ{Int},:x}

x = P([0, 1])
CP = (Int[], [1], [0, 0, 4], [2, 1], [1,0,30])

@testset "constructors" begin
    @test basetype(P) == S
    @test depth(P) == 2
    @test P([S(0), S(1)]).coeff[2] == S(1)
    @test P([1]).coeff[1] == S(1)
    @test length(P(Int[]).coeff) == 0
    @test length(P(Int[0]).coeff) == 0
    @test eltype(UnivariatePolynomial(:x, S[]).coeff) == S
    @test iscoprime(x^2 + 1, x + 1)
    @test !iscoprime(4x^2 + 4x + 1, 2x + 1)

    @test typeof(UnivariatePolynomial(:x, ZZ(1))) == P
    co = [1,2,3]
    @test copy(P(co)) == P(co)
    @test copy(P(co)).coeff !== P(co).coeff

    @test isunit(P([-1]))
    @test isunit(P([1]))
    @test !isunit(P([0]))
    @test length(zero(P).coeff) == 0
    z = zero(P)
    o = one(P)
    @test length(o.coeff) == 1 && o.coeff[1] == S(1)
    @test iszero(z)
    @test !iszero(o)
    @test isone(o)
    @test !isone(z)

    @test P([1,2,3]) == P([1,2,3])
    @test P([1,2]) != P([1])

    @test P(1) == 1
    @test P(ZZ(Int8(1))) == 1
    @test 1 + P(1) == P(2)
    @test ZZ(1) + P(1) == P(2)
    @test ZZ(1) - P(1) == z
    @test P(1) == UnivariatePolynomial{ZZ{Int8},:y}(1)
    @test P([0,1]) != UnivariatePolynomial{ZZ{Int}, :y}([0,1])
    @test P(1) == UnivariatePolynomial{ZZ{Int},:x}(1)
    @test UnivariatePolynomial{ZZ{Int8},:x}([1]) == P([1])
    @test hash(UnivariatePolynomial{ZZ{Int8},:x}([1])) == hash(P([1]))
end

@testset "varnames and generators" for P in (S[:x],)
    x = monom(P)
    @test [x] == generators(P)
    @test varnames(P) == [:x]
    @test varnames(P(0)) == [:x]

    p = P(S[1,2,3])
    @test deg(p) == 2
    @test multideg(p) == [2]
    @test LC(p) == S(3)
    @test LM(p) == x^2
    @test LT(p) == 3*x^2
end

@testset "promotion of types" begin
    @test promote_type(P, PQ) == PQ
    @test promote_type(P, ZZ{Int}) == P
    @test promote_type(P, Int) == P
    @test promote_type(P, Rational{Int8}) == UnivariatePolynomial{QQ{Int},:x}
    @test promote_rule(P, UnivariatePolynomial{S,:y}) == Base.Bottom
end

@testset "operation $cp $(string(op)) $cq" for op in (+, -, *), cp in CP, cq in CP
    p = P(S.(cp))
    q = P(S.(cq))
    @test op(p, q) == op(Poly(cp), Poly(cq))
end

@testset "operation unary - and ^ $cp" for cp in CP
    p = P(S.(cp))
    @test -p == -Poly(cp)
    @test p^10 == Poly(cp)^10
    @test p^0 == one(P)
    !isunit(p) && @test_throws DomainError p^-1
    @test zero(p)^0 == one(P)
    @test one(p)^-1 == one(P)
    @test zero(p)^0 == one(P)
end

hasunitlead(p::UnivariatePolynomial) = !iszero(p) && isunit(p.coeff[end])

@testset "operation divrem($cp,$cq)" for cp in CP, cq in CP
    p = P(S.(cp))
    q = P(S.(cq))
    if hasunitlead(q) || deg(p) < deg(q) || ( p == q && !iszero(q))
        @test divrem(p, q) == divrem(Poly(cp), Poly(cq))
        @test deg(rem(p, q)) < deg(q)
    elseif !iszero(q)
        @test deg(rem(p, q)) >= deg(q)
    else
        @test_throws DomainError divrem(p, q)
    end
end

@testset "gcd calculations" begin

    P = UnivariatePolynomial{ZZ{Int},:x}
    @test P([0, 1]) == monom(P)
    x = P([0,1])
    @test deg(x + 2x^2 + 1 +0*x^3) == 2
    p = x^8 + x^6 -3x^4 - 3x^3 + 8x^2 +2x - 5
    q = 21 + 3x^6 + 5x^4 - 4*x^2 -9x
    @test deg(pgcd(p, q)) == 0
    r1 = pgcd(p, q)
    s = (ZZ(1) + 3x + x^2)^5
    @test pgcd(p * s, q * s) / r1 == s

    p = P([1, 2, 1])
    q = P([1, 1])
    @test gcd(p, q) == q
    @test gcd([p, p, q]) == q
    @test gcdx(p, q) == (q, P(Int[]), P([1]))
    @test lcm(p, q) == p
    @test lcm([p, q, p]) == p

    @test UnivariatePolynomial{ZZ{Int},:x}(p) === p
    @test UnivariatePolynomial{ZZ{Int32},:x}(p) == p
    @test div(p, q) == q
    @test content(p) == ZZ(1)
    @test primpart(p) == p
    @test content(-p) == ZZ(-1)
    @test primpart(-p) == p

    Q = UnivariatePolynomial{QQ{Int},:x}
    qp = Q(-p) / 12
    @test primpart(qp) == Q(p)
    @test content(qp) == -1//12

    @test iszero(LC(zero(P)))
    @test isone(LC(one(P)))
    @test LC(p) == 1
    @test multideg(one(P)) == [0]
    @test multideg(p) == [2]

    PP = UnivariatePolynomial{P,:z}
    pp = PP([p, q, s])

    # call show
    @test sprint(show, p) !== nothing
    @test sprint(show, zero(P)) !== nothing
    @test sprint(show, one(P)) !== nothing
    @test sprint(show, x^3 + x) !== nothing
end

@testset "pseudo gcd" begin
    S = ZZ{BigInt}
    P = UnivariatePolynomial{S,:x}
    p = P([1, 4, 5, 1, 6, 0, 3])
    q = P([5, 1, 1, 3])
    x = P([0, 1])
    @test content(2p) == S(2)
    @test primpart(12p) == p
    @test LC(p) == S(3)
    a, b, f = pdivrem(p, q)
    @test f * p == a * q + b
    a, b, f = pdivrem(p, P([2]))
    @test f * p == a * P([2]) + b

    @test gcd(p, q) == one(P)
    @test pgcd(p, q) == P([1])
    @test_throws DomainError gcdx(p, q)
    g, u, v, f = pgcdx(p, q)
    @test iszero(p * u + q * v - g * f)
    @test deg(g) == 0
    p = 3x^10 - 3x^9 - 3x^8 + 3x^7 - 2x^5 - 3x^4 - 2x^3 - 1x - 3
    q = 2x^4 + 3x^3 - 1x^2 + x - 2
    g, u, v, f = pgcdx(q, p)
    @test iszero(q * u + p * v - g * f)
    @test deg(g) == 0
    s = 3x^3 - 2x^2 - ZZ(1)
    p *= s
    q *= s
    g, u, v, f = pgcdx(p, q)
    @test iszero(p * u + q * v - g * f)
    @test isdiv(g, s)

    S = ZZmod{31,Int}
    P = UnivariatePolynomial{S,:x}
    x = P([0, 1])
    p = x^3 + 3x^2 - 4
    q = 4 - (-x)
    g, u, v, f = pgcdx(p, q)
    @test iszero(p * u + q * v - g * f)
    p2 = 2p
    @test p2 / 2 == p

    p = 12x^2 + 45x + 60
    q = 5x^2 + 1
    g, u, v, f = pgcdx(p, q)
    @test iszero(p * u + q * v - g * f)

    qq = [q, P([2,1])]
    a, r = divrem(p, qq)
    @test sum(a .* qq) + r == p

    x = UnivariatePolynomial{ZZ{Int},:x}([0, 1])
    @test CommutativeRings.invert(x, x^2 + 1) == -x
    @test_throws DomainError CommutativeRings.invert(x + 1, x^2 + 1)

    @test CommutativeRings.issimple(1.0)
    @test CommutativeRings.issimple(x^3)
    @test !CommutativeRings.issimple(x+1)
end

@testset "polynomials of other structures" begin

    Zp = ZZmod{17, Int8}
    P = Zp[:x]
    @test P([1,2]) != Zp[:y]([1,2])
    x = P([0, 1])
    p = (3x^2 + 2x + 1) * (x - 1)
    q = x^2 - x
    @test p / (x-1) == (3x^2 + 2x + 1)
    @test gcd(p, q) == x - 1

    @test ismonom(P([0,2]))
    @test !ismonom(P([1,2]))
    @test ismonic(P([2,1]))
    @test !ismonic(P([2,-1]))

end

@testset "apply polynomial" begin
    x = monom(ZZ{Int}[:x], 1)
    y = monom(ZZ{Int}[:y], 1)
    @test x(1) == 1
    @test x(0) == 0
    @test (y^2)(x+1) == x^2 + 2x + 1
    @test (0*y)(x^2) == 0
end

@testset "abstract types" begin
    @test QQ[:x]([1//2, 2//3]) == QQ{Int}[:x]([1//2, 2//3])
end

@testset "GF[:x]" begin
    G = GF(7, 3)
    GX = G[:x]
    x = monom(GX)
    p = x^3 + G[(7*2+1)*7+5]x + 1
    q = x^3 + 1 + G[200]*x
    @test isirreducible(p)
    @test factor(q) |> length == 3
    @test deg(p^2) == 6
    @test p * q == q * p
    @test p + q - p == q
    @test q / (x + G[28]) |> deg == 2
end

x = monom(ZZ{BigInt}[:x])
@testset "resultant($a,$b)" for (a, b, c) in ((0,0,0), (0,1,0), (0, x,0), (3,0,0),
                                       (3,2,1), (3,x,3), (3,5x^2+1,3^2), (7,5x^7-1,7^7),
                                       (2x^32+1, 3*x^11+13, 906811932214969750127235270540901107329),
                                       (2x^31+1, 3x^11+1,-617673396281899))

    @test resultant(a, b) == c
    @test resultant(b, a) == c * (-1)^(deg(a)*deg(b)+2)
    @test resultant(11a, 13b) == (a*b == 0 ? 0 : c * big(11)^deg(b) * big(13)^deg(a))

    ref = CommutativeRings.resultant_naive
    @test a*b == 0 || ref(a*x^0, b*x^0) == c
    @test a*b == 0 || ref(b*x^0, a*x^0) == c * (-1)^(deg(a)*deg(b)+2)
end

end


"""
    MatrixWorkspace

This is a data structure for the efficient repeated solution of a square or
overdetermined linear system `Ax=b`.
"""
struct MatrixWorkspace{M<:AbstractMatrix{ComplexF64}} <: AbstractMatrix{ComplexF64}
    A::M # Matrix
    d::Vector{Float64} # Inverse of scaling factors
    factorized::Base.RefValue{Bool}
    lu::LA.LU{ComplexF64,M} # LU Factorization of D * J
    qr::LA.QR{ComplexF64,Matrix{ComplexF64}}
    row_scaling::Vector{Float64}
    scaled::Base.RefValue{Bool}
    # mixed precision iterative refinement
    x̄::Vector{ComplexDF64}
    r::Vector{ComplexF64}
    r̄::Vector{ComplexDF64}
    δx::Vector{ComplexF64}
    inf_norm_est_work::Vector{ComplexF64}
    inf_norm_est_rwork::Vector{Float64}
end

MatrixWorkspace(m::Integer, n::Integer; kwargs...) =
    MatrixWorkspace(zeros(ComplexF64, m, n); kwargs...)
function MatrixWorkspace(Â::AbstractMatrix; optimize_data_structure = true)
    m, n = size(Â)
    m ≥ n || throw(ArgumentError("Expected system with more rows than columns."))

    A = Matrix{ComplexF64}(Â)
    d = ones(m)
    factorized = Ref(false)
    qr = LA.qrfactUnblocked!(copy(A))
    # experiments show that for m > 25 the data layout as a
    # struct array is beneficial
    if m > 25 && optimize_data_structure
        A = StructArrays.StructArray(A)
    end
    row_scaling = ones(m)
    scaled = Ref(false)

    lu = LinearAlgebra.LU{eltype(A),typeof(A)}(copy(A), zeros(Int, m), 0)
    r = zeros(ComplexF64, m)
    r̄ = zeros(ComplexDF64, m)
    x̄ = zeros(ComplexDF64, n)
    δx = zeros(ComplexF64, n)
    inf_norm_est_work = Vector{ComplexF64}(undef, n)
    inf_norm_est_rwork = Vector{Float64}(undef, n)

    MatrixWorkspace(
        A,
        d,
        factorized,
        lu,
        qr,
        row_scaling,
        scaled,
        x̄,
        r,
        r̄,
        δx,
        inf_norm_est_work,
        inf_norm_est_rwork,
    )
end

Base.size(MW::MatrixWorkspace) = size(MW.A)

import Base: @propagate_inbounds
@propagate_inbounds Base.getindex(MW::MatrixWorkspace, i::Integer) = getindex(MW.A, i)
@propagate_inbounds Base.getindex(MW::MatrixWorkspace, i::Integer, j::Integer) =
    getindex(MW.A, i, j)
@propagate_inbounds Base.setindex!(MW::MatrixWorkspace, x, i::Integer) =
    setindex!(MW.A, x, i)
@propagate_inbounds Base.setindex!(MW::MatrixWorkspace, x, i::Integer, j::Integer) =
    setindex!(MW.A, x, i, j)
@propagate_inbounds Base.copyto!(MW::MatrixWorkspace, A::AbstractArray) = copyto!(MW.A, A)

matrix(M::MatrixWorkspace) = M.A
matrix(M::AbstractMatrix) = M

"""
    updated!(MW::MatrixWorkspace)

Indicate that the matrix `MW` got updated.
"""
function updated!(MW::MatrixWorkspace)
    MW.factorized[] = false
    MW.scaled[] = false
    m, n = size(MW)
    if m == n
        @inbounds copyto!(MW.lu.factors, MW.A)
    else
        @inbounds copyto!(MW.qr.factors, MW.A)
    end
    MW
end
updated!(M::AbstractMatrix) = M

"""
    update!(MW::MatrixWorkspace, A::Matrix)

Update the matrix in `MW` with `A`.
"""
@inline function update!(MW::AbstractMatrix, A::Matrix)
    @boundscheck (size(MW.A) == size(A) || throw(ArgumentError("Matrix of invalid size.")))
    @inbounds copyto!(MW.A, A)
    updated!(MW)
end

##################
# Factorizations #
##################

###########################
# CUSTOM LU Factorization #
###########################
# This is an adoption of LA.generic_lufact!
# with 3 changes:
# 1) For choosing the pivot we use abs2 instead of abs
# 2) Instead of using the robust complex division we
#    just use the naive division. This shouldn't
#    lead to problems since numerical diffuculties only
#    arise for very large or small exponents
# 3) We fold lu! and ldiv! into one routine
#    this has the effect that we do not need to allocate
#    the pivot vector anymore and also avoid the allocations
#    coming from the LU wrapper
function lu!(
    A::AbstractMatrix{T},
    ipiv::Union{Vector{I},Nothing} = nothing,
) where {T,I<:Integer}
    m, n = size(A)
    minmn = min(m, n)
    # LU Factorization
    @inbounds begin
        for k = 1:minmn
            # find index max
            kp = k
            if T <: Complex
                amax = abs2(A[k, k])
            else
                amax = abs(A[k, k])
            end
            for i = k+1:m
                if T <: Complex
                    absi = abs2(A[i, k])
                else
                    absi = abs(A[i, k])
                end
                kp, amax = ifelse(absi > amax, (i, absi), (kp, amax))
            end
            if !(ipiv isa Nothing)
                ipiv[k] = kp
            end
            if !iszero(amax)
                if k != kp
                    # Interchange
                    for i = 1:n
                        tmp = A[k, i]
                        A[k, i] = A[kp, i]
                        A[kp, i] = tmp
                    end
                end
                # Scale first column
                Akk = A[k, k]
                for i = k+1:m
                    # we assume the matrix is decently scaled
                    # so that the naive division algorithm doesn't overflow
                    @fastmath A[i, k] = A[i, k] / Akk
                end
            end
            # Update the rest
            for j = k+1:n
                A_kj = A[k, j]
                for i = k+1:m
                    A[i, j] -= A[i, k] * A_kj
                end
            end
        end
    end
    A
end

####################
# QR Factorization #
####################
# This is the generic base implementation with the small change that we introduce
# a @fastmath for the division step since results in an approx 50% speedup

# Elementary reflection similar to LAPACK. The reflector is not Hermitian but
# ensures that tridiagonalization of Hermitian matrices become real. See lawn72
@inline function reflector!(x::AbstractVector)
    n = length(x)
    @inbounds begin
        ξ1 = x[1]
        normu = abs2(ξ1)
        for i = 2:n
            normu += abs2(x[i])
        end
        if iszero(normu)
            return zero(ξ1 / normu)
        end
        normu = sqrt(normu)
        ν = copysign(normu, real(ξ1))
        ξ1 += ν
        x[1] = -ν
        for i = 2:n
            @fastmath x[i] = x[i] / ξ1
        end
    end
    ξ1 / ν
end

# apply reflector from left
@inline function reflectorApply!(x::AbstractVector, τ::Number, A::StridedMatrix)
    m, n = size(A)
    @inbounds begin
        for j = 1:n
            # dot
            vAj = A[1, j]
            for i = 2:m
                vAj += x[i]' * A[i, j]
            end

            vAj = conj(τ) * vAj

            # ger
            A[1, j] -= vAj
            for i = 2:m
                A[i, j] -= x[i] * vAj
            end
        end
    end
    return A
end

function qr!(qr::LA.QR{T}) where {T}
    A = qr.factors
    τ = qr.τ
    m, n = size(A)
    for k = 1:min(m - 1 + !(T <: Real), n)
        x = view(A, k:m, k)
        τk = reflector!(x)
        τ[k] = τk
        reflectorApply!(x, τk, view(A, k:m, k+1:n))
    end
    qr
end

function factorize!(WS::MatrixWorkspace)
    m, n = size(WS)
    if m == n
        lu!(WS.lu.factors, WS.lu.ipiv)
    else
        qr!(WS.qr)
    end
    WS.factorized[] = true
    WS
end

##########
## ldiv ##
##########
@inline _ipiv!(A::LA.LU, b::AbstractVector) = apply_ipiv!(b, 1:length(A.ipiv), A.ipiv)
@inline _inverse_ipiv!(A::LA.LU, b::StridedVecOrMat) =
    apply_ipiv!(b, length(A.ipiv):-1:1, A.ipiv)
@inline function apply_ipiv!(b::AbstractVector, range::OrdinalRange, ipiv)
    @inbounds for i in range
        if i != ipiv[i]
            _swap_rows!(b, i, ipiv[i])
        end
    end
    b
end
@propagate_inbounds function _swap_rows!(B::StridedVector, i::Integer, j::Integer)
    B[i], B[j] = B[j], B[i]
    B
end

@inline function ldiv_upper!(A::AbstractMatrix, b::AbstractVector, x::AbstractVector = b)
    n = size(A, 2)
    @inbounds for j = n:-1:1
        # singular_exception && iszero(A[j, j]) && throw(LA.SingularException(j))
        xj = x[j] = (@fastmath A[j, j] \ b[j])
        for i = 1:(j-1)
            b[i] -= A[i, j] * xj
        end
    end
    b
end
@inline function ldiv_unit_lower!(
    A::AbstractMatrix,
    b::AbstractVector,
    x::AbstractVector = b,
)
    n = size(A, 2)
    @inbounds for j = 1:n
        xj = x[j] = b[j]
        for i = j+1:n
            b[i] -= A[i, j] * xj
        end
    end
    x
end

function lu_ldiv!(x, LU::LA.LU, b::AbstractVector)
    x === b || copyto!(x, b)
    _ipiv!(LU, x)
    ldiv_unit_lower!(LU.factors, x)
    ldiv_upper!(LU.factors, x)
    x
end

function ldiv_adj_unit_lower!(A::AbstractMatrix, b::AbstractVector, x::AbstractVector = b)
    n = size(A, 1)
    @inbounds for j = n:-1:1
        z = b[j]
        for i = n:-1:j+1
            z -= conj(A[i, j]) * x[i]
        end
        x[j] = z
    end
    x
end

function ldiv_adj_upper!(
    A::AbstractMatrix,
    b::AbstractVector,
    x::AbstractVector = b;
    singular_exception::Bool = false,
)
    n = size(A, 1)
    @inbounds for j = 1:n
        z = b[j]
        for i = 1:j-1
            z -= conj(A[i, j]) * x[i]
        end
        iszero(A[j, j]) && singular_exception && throw(SingularException(j))
        x[j] = @fastmath conj(A[j, j]) \ z
    end
    x
end

function lu_ldiv_adj!(x, LU::LA.LU, b::AbstractVector; check::Bool = false)
    x === b || copyto!(x, b)
    ldiv_adj_upper!(LU.factors, x; singular_exception = check)
    ldiv_adj_unit_lower!(LU.factors, x)
    _inverse_ipiv!(LU, x)
    x
end

# QR
# adaption of base implementation
function lmul_Q_adj!(A::LA.QR, b::AbstractVector)
    mA, nA = size(A.factors)
    mB = length(b)
    Afactors = A.factors
    @inbounds begin
        for k = 1:min(mA, nA)
            vBj = b[k]
            for i = k+1:mB
                vBj += conj(Afactors[i, k]) * b[i]
            end
            vBj = conj(A.τ[k]) * vBj
            b[k] -= vBj
            for i = k+1:mB
                b[i] -= Afactors[i, k] * vBj
            end
        end
    end
    b
end
function qr_ldiv!(x, QR::LA.QR, b::AbstractVector) where {T}
    # overwrites b
    # assumes QR is a tall matrix
    lmul_Q_adj!(QR, b)
    @inbounds for i = 1:length(x)
        x[i] = b[i]
    end
    ldiv_upper!(QR.factors, x)
    return x
end

function LA.ldiv!(x::AbstractVector, WS::MatrixWorkspace, b::AbstractVector)
    m, n = size(WS)
    if (m, n) == (1, 1)
        x[1] = b[1] / WS[1, 1]
        return x
    end
    WS.factorized[] || factorize!(WS)
    if m == n
        if WS.scaled[]
            x .= WS.row_scaling .* b
            lu_ldiv!(x, WS.lu, x)
        else
            lu_ldiv!(x, WS.lu, b)
        end
    else
        WS.r .= b
        qr_ldiv!(x, WS.qr, WS.r)
    end
    x
end

function LA.inv!(B::AbstractMatrix, M::MatrixWorkspace)
    M.factorized[] || factorize!(M)
    B .= M.lu.factors
    LA.inv!(M.lu)
    @inbounds for i in eachindex(B)
        M.lu.factors[i], B[i] = B[i], M.lu.factors[i]
    end
    return B
end

"""
    skeel_row_scaling!(W::MatrixWorkspace, c)
    skeel_row_scaling!(d, A, c)

Compute optimal scaling factors `d` for the matrix `A` following Skeel [^S79]
if `c` is approximately of the order of the solution of the linear system
of interest.
The scaling factors are rounded to powers of the base radix 2.

[^S79]: Skeel, Robert D. "Scaling for numerical stability in Gaussian elimination." Journal of the ACM (JACM) 26.3 (1979): 494-526.
"""
function skeel_row_scaling!(
    d::AbstractVector{<:Real},
    A::AbstractMatrix{<:Complex},
    c::AbstractVector{<:Real},
)
    n = length(c)
    @inbounds d .= zero(eltype(d))
    @inbounds for j = 1:n
        cj = c[j]
        for i = 1:n
            d[i] += fast_abs(A[i, j]) * cj
        end
    end
    @inbounds for i = 1:n
        d[i] = exp2(-last(frexp(d[i])))
    end

    d
end

function skeel_row_scaling!(W::MatrixWorkspace, c::AbstractVector{<:Real})
    skeel_row_scaling!(W.row_scaling, W.A, c)
    W
end

function row_scaling!(
    d::AbstractVector{<:Real},
    WS::MatrixWorkspace,
    c::AbstractVector{<:Real},
)
    m, n = size(WS)
    if m == n
        skeel_row_scaling!(d, WS.A, c)
    else
        d .= 1.0
    end
    d
end

"""
    apply_row_scaling!(W::MatrixWorkspace)

Apply the computed row scaling.
"""
function apply_row_scaling!(W::MatrixWorkspace)
    A, d = W.lu.factors, W.row_scaling
    m, n = size(A)
    @inbounds for j = 1:n, i = 1:m
        A[i, j] = A[i, j] * d[i]
    end
    W.scaled[] = true
    W
end

## Iterative Refinement
"""
    residual!(r, A, x, b)

Compute the residual `Ax-b` and store it in `r`
"""
function residual!(r::AbstractVector, A::AbstractMatrix, x::AbstractVector, b)
    m, n = size(A)
    @boundscheck m == length(b) == length(r) && n == length(x)
    r .= 0
    @inbounds for j = 1:n
        x_j = x[j]
        for i = 1:m
            r[i] += A[i, j] * x_j
        end
    end
    @inbounds for i = 1:m
        r[i] -= b[i]
    end
    r
end

"""
    mixed_precision_iterative_refinement!(x, A, b, norm = nothing)

Perform one step of mixed precision iterative refinement.
If `norm` is an `AbstractNorm` then the normwise relative error before
the refinement step is returned otherwise `x` is returned.
"""
function mixed_precision_iterative_refinement!(
    x::AbstractVector,
    M::MatrixWorkspace,
    b::AbstractVector,
    norm::Union{AbstractNorm,Nothing} = nothing,
)
    M.x̄ .= x
    residual!(M.r̄, M.A, M.x̄, b)
    M.r .= M.r̄
    LA.ldiv!(M.δx, M, M.r)
    x .-= M.δx
    if norm isa Nothing
        x
    else
        norm(M.δx) / norm(x)
    end
end

"""
    fixed_precision_iterative_refinement!(x, A, b, norm = nothing)

Perform one step of mixed precision iterative refinement.
If `norm` is an `AbstractNorm` then the normwise relative error before
the refinement step is returned otherwise `x` is returned.
"""
function fixed_precision_iterative_refinement!(
    x::AbstractVector,
    M::MatrixWorkspace,
    b::AbstractVector,
    norm::Union{AbstractNorm,Nothing} = nothing,
)
    residual!(M.r, M.A, x, b)
    LA.ldiv!(M.δx, M, M.r)
    x .-= M.δx
    if norm isa Nothing
        x
    else
        norm(M.δx) / norm(x)
    end
end

# CONDITION NUMBERS

"""
    inverse_inf_norm_est(WS::MatrixWorkspace,
                         d_l::Union{Nothing,Vector{<:Real}}=nothing
                         d_r::Union{Nothing,Vector{<:Real}}=nothing)

Estimate of the infinity norm of `diag(d_r)⁻¹A⁻¹diag(d_l)⁻¹` where `d_l` and `d_r` are
optional positive vectors.
If `d_l` or `d_r` is `nothing` the all one vector is used.
This uses the 1-norm lapack condition estimator described by Highahm in [^H88].

[^H88]: Higham, Nicholas J. "FORTRAN codes for estimating the one-norm of a real or complex matrix, with applications to condition estimation." ACM Transactions on Mathematical Software (TOMS) 14.4 (1988): 381-396.
"""
function inverse_inf_norm_est(
    WS::MatrixWorkspace,
    d_l::Union{Nothing,Vector{<:Real}} = nothing,
    d_r::Union{Nothing,Vector{<:Real}} = nothing,
)
    WS.factorized[] || factorize!(WS)
    work, rwork = WS.inf_norm_est_work, WS.inf_norm_est_rwork
    if !WS.scaled[]
        inverse_inf_norm_est(WS.lu, d_l, d_r, nothing, work, rwork)
    else
        inverse_inf_norm_est(WS.lu, d_l, d_r, WS.row_scaling, work, rwork)
    end
end
function inverse_inf_norm_est(
    lu::LA.LU,
    d_l::Union{Nothing,Vector{<:Real}},
    d_r::Union{Nothing,Vector{<:Real}},
    row_scaling::Union{Nothing,Vector{<:Real}},
    work::Vector{<:Complex},
    rwork::Vector{<:Real},
)
    z = ξ = y = work
    x = rwork

    n = size(lu.factors, 1)
    x .= inv(n)
    if d_r !== nothing
        x ./= d_r
    end
    lu_ldiv_adj!(y, lu, x)
    if d_l !== nothing
        y ./= d_l
    end
    if row_scaling !== nothing
        y .*= row_scaling
    end
    γ = sum(fast_abs, y)
    ξ ./= fast_abs.(y)
    if d_l !== nothing
        ξ ./= d_l
    end
    if row_scaling !== nothing
        ξ ./= row_scaling
    end
    lu_ldiv!(z, lu, ξ)

    if d_r !== nothing
        x .= real.(z) ./ d_r
    else
        x .= real.(z)
    end
    k = 2
    while true
        j, max_xᵢ = 1, abs(x[1])
        for i = 2:n
            abs_xᵢ = abs(x[i])
            if abs_xᵢ > max_xᵢ
                j, max_xᵢ = i, abs_xᵢ
            end
        end
        x .= zero(eltype(x))
        x[j] = one(eltype(x))
        if d_r !== nothing
            x ./= d_r
        end
        lu_ldiv_adj!(y, lu, x)
        if d_l !== nothing
            y ./= d_l
        end
        if row_scaling !== nothing
            y .*= row_scaling
        end
        γ̄ = γ
        γ = sum(fast_abs, y)
        if γ ≤ γ̄
            γ = γ̄
            break
        end
        ξ .= y ./ fast_abs.(y)
        if d_l !== nothing
            ξ ./= d_l
        end
        if row_scaling !== nothing
            ξ ./= row_scaling
        end
        lu_ldiv!(z, lu, ξ)
        if d_r !== nothing
            x .= real.(z) ./ d_r
        else
            x .= real.(z)
        end
        k += 1
        if x[j] == LA.norm(x, Inf) || k > 2
            break
        end
    end
    nanmin(γ, Inf)
end

function inf_norm(
    WS::MatrixWorkspace,
    d_l::Union{Nothing,Vector{<:Real}} = nothing,
    d_r::Union{Nothing,Vector{<:Real}} = nothing,
)
    norm = -Inf
    A = WS.A
    n, m = size(A)
    @inbounds for i = 1:n
        normᵢ = 0.0
        for j = 1:m
            if d_r isa Nothing
                normᵢ += fast_abs(A[i, j])
            else
                normᵢ += fast_abs(A[i, j]) * d_r[j]
            end
        end
        if !(d_l isa Nothing)
            normᵢ *= d_l[i]
        end
        norm = @fastmath max(norm, normᵢ)
    end
    norm
end

function max_min_row(
    WS::MatrixWorkspace,
    d_l::Union{Nothing,Vector{<:Real}} = nothing,
    d_r::Union{Nothing,Vector{<:Real}} = nothing,
)
    max_row = -Inf
    min_row = Inf
    A = WS.A
    n, m = size(A)
    @inbounds for i = 1:n
        normᵢ = 0.0
        for j = 1:m
            if d_r isa Nothing
                normᵢ += fast_abs(A[i, j])
            else
                normᵢ += fast_abs(A[i, j]) * d_r[j]
            end
        end
        if !(d_l isa Nothing)
            normᵢ *= d_l[i]
        end
        max_row = @fastmath max(max_row, normᵢ)
        min_row = @fastmath min(min_row, normᵢ)
    end
    max_row, min_row
end

"""
    cond(A::MatrixWorkspace,
         d_l::Union{Nothing,Vector{<:Real}} = nothing,
         d_r::Union{Nothing,Vector{<:Real}} = nothing)

Compute the condition number w.r.t. the infinity-norm of `diag(d_l) * A * diag(d_r)`.
If `d_l` or `d_r` is `nothing` the all one vector is used.
If `size(A) == (1,1)` then just the norm of the inverse is returned.
"""
function LA.cond(
    WS::MatrixWorkspace,
    d_l::Union{Nothing,Vector{<:Real}} = nothing,
    d_r::Union{Nothing,Vector{<:Real}} = nothing,
)
    m, n = size(WS)
    if m == n == 1
        if isa(d_l, Nothing) && isa(d_r, Nothing)
            inv(abs(WS.A[1, 1]))
        elseif isa(d_l, Nothing)
            inv(abs(WS.A[1, 1]) * d_r[1])
        elseif isa(d_r, Nothing)
            inv(d_l[1] * abs(WS.A[1, 1]))
        else
            inv(d_l[1] * abs(WS.A[1, 1]) * d_r[1])
        end
    elseif m > n
        WS.factorized[] || factorize!(WS)
        rmax, rmin = -Inf, Inf
        for i = 1:n
            rᵢ = fast_abs(WS.qr.factors[i, i]) * d_r[i]
            rmax = max(rmax, rᵢ)
            rmin = min(rmin, rᵢ)
        end
        rmax / rmin
    else
        inverse_inf_norm_est(WS, d_l, d_r) * inf_norm(WS, d_l, d_r)
    end
end

function egcond(
    WS::MatrixWorkspace,
    d_l::Union{Nothing,Vector{<:Real}} = nothing,
    d_r::Union{Nothing,Vector{<:Real}} = nothing,
)
    m, n = size(WS)
    if m == n == 1
        a = abs(WS.A[1, 1])
        if d_l !== nothing
            a *= d_l[1]
        end
        if d_r !== nothing
            a *= d_r[1]
        end
        max(a, 1.0) / a
    elseif m > n
        WS.factorized[] || factorize!(WS)
        rmax, rmin = -Inf, Inf
        for i = 1:n
            rᵢ = abs(WS.qr.factors[i, i]) * d_r[i]
            rmax = max(rmax, rᵢ)
            rmin = min(rmin, rᵢ)
        end
        rmax / rmin
    else
        max_row, min_row = max_min_row(WS, d_l, d_r)
        inverse_inf_norm_est(WS, d_l, d_r) * max(max_row, 1 / min_row)
    end
end

#####################
## Jacobian ##
#####################
struct Jacobian{M}
    workspace::MatrixWorkspace{M}
    # stats
    factorizations::Base.RefValue{Int}
    ldivs::Base.RefValue{Int}
end
Jacobian(A::AbstractMatrix) = Jacobian(MatrixWorkspace(A), Ref(0), Ref(0))

updated!(J::Jacobian) = (updated!(J.workspace); J)

workspace(J::Jacobian) = J.workspace
matrix(J::Jacobian) = matrix(workspace(J))

function Base.show(io::IO, J::Jacobian{T}) where {T}
    println(io, "Jacobian{$T}:")
    println(io, " • # factorizations → ", J.factorizations[])
    println(io, " • # ldivs → ", J.ldivs[])
end

"""
    ldiv!(x̂::AbstractVector, J::Jacobian, b::AbstractVector)

solve the linear system `matrix(J)x̂=b`.
"""
function LA.ldiv!(
    x̂::AbstractVector,
    J::Jacobian,
    b::AbstractVector,
    norm = nothing;
    row_scaling::Bool = true,
)
    # stats update
    J.factorizations[] += !workspace(J).factorized[]
    J.ldivs[] += 1
    LA.ldiv!(x̂, workspace(J), b)
    x̂
end
function LA.ldiv!(
    x̂::AbstractVector,
    J::Jacobian,
    b::AbstractVector,
    norm::WeightedNorm;
    row_scaling::Bool = true,
)
    if !workspace(J).factorized[] && row_scaling
        skeel_row_scaling!(workspace(J), weights(norm))
        apply_row_scaling!(workspace(J))
    end
    # stats update
    J.factorizations[] += !workspace(J).factorized[]
    J.ldivs[] += 1
    LA.ldiv!(x̂, workspace(J), b)
    x̂
end

function iterative_refinement!(
    x::AbstractVector,
    J::Jacobian,
    b::AbstractVector,
    norm::AbstractNorm;
    max_iters::Int = 3,
    tol::Float64 = sqrt(eps()),
)
    J.ldivs[] += 1
    δ = mixed_precision_iterative_refinement!(x, workspace(J), b, norm)
    for i = 2:max_iters
        J.ldivs[] += 1
        δ′ = mixed_precision_iterative_refinement!(x, workspace(J), b, norm)
        if δ′ < tol
            return (accuracy = δ′, diverged = false)
        elseif δ′ > 0.5 * δ
            return (accuracy = δ′, diverged = true)
        end
        δ = δ′
    end
    (accuracy = δ, diverged = false)
end

"""
    init!(J::Jacobian)

(Re-)initialize the `Jacobian`.
"""
function init!(J::Jacobian; keep_stats::Bool = false)
    J.factorizations[] = 0
    J.ldivs[] = 0
    J
end

function LA.cond(
    J::Jacobian,
    d_l::Union{Nothing,Vector{<:Real}} = nothing,
    d_r::Union{Nothing,Vector{<:Real}} = nothing,
)
    LA.cond(workspace(J), d_l, d_r)
end

function egcond(
    J::Jacobian,
    d_l::Union{Nothing,Vector{<:Real}} = nothing,
    d_r::Union{Nothing,Vector{<:Real}} = nothing,
)
    egcond(workspace(J), d_l, d_r)
end


using LinearAlgebra
using Plots
using Revise

includet("src/IterativeMethods.jl")
includet("src/InternalArray.jl")

IArray = InternalArray

t0 = 0
x0 = [0, 100]
v0 = [300, 0]
m = 1
c = 0.1
g = [0, -9.8]

# Solving without friction
x_nofriction = eulermethod(t0, x0) do t, x
	v0 + t*g
end
plot_nofriction = plot(IArray(x_nofriction, 1), IArray(x_nofriction, 2), label = "without friction")

# Solving with friction
x_friction = eulermethod(t0, [x0; v0]) do t, x
	v = [x[3], x[4]]
	return [v; g-c/m*norm(v)*v]
end
plot_friction = plot(IArray(x_friction, 1), IArray(x_friction, 2), label = "with friction")

plot(plot_nofriction, plot_friction, plot_title = "comparing with and without friction")

# Solving for collision with floor
t, x = targetting_eulermethod(x -> x[2], t0, [x0; v0]; h = 5e-3) do t, x
	v = [x[3], x[4]]
	return [v; g-c/m*norm(v)*v]
end
p = plot(IArray(x, 1), IArray(x, 2), label = "landing")

# Needed to stop showing due to error
using IJulia

x1 = IArray(x, 1)
x2 = IArray(x, 2)
v1 = IArray(x, 3)
v2 = IArray(x, 4)

ppp = 35 # points per frame
lims(v) = (minimum(v), maximum(v))
p1 = plot(1, xlim = lims(x1), ylim = lims(x2), label = "position")
p2 = plot(1, xlim = lims(v1), ylim = lims(v2), label = "speed")
anim = Animation()
for i = ppp:ppp:length(x)
	for j = i-ppp+1:i
    	push!(p1, x1[j], x2[j])
    	push!(p2, v1[j], v2[j])
    end
	plot(p1, p2)
	frame(anim)
end
gif(anim, "media/dropped-object-trayectory-to-collision.gif", fps = 24)


"""
    max_min_driving_force(
        model::MetabolicModel,
        reaction_standard_gibbs_free_energies::Dict{String,Float64},
        optimizer;
        flux_solution::Dict{String,Float64} = Dict{String,Float64}(),
        proton_ids::Vector{String} = ["h_c", "h_e"],
        water_ids::Vector{String} = ["h2o_c", "h2o_e"],
        constant_concentrations::Dict{String,Float64} = Dict{String,Float64}(),
        concentration_ratios::Dict{Tuple{String,String},Float64} = Dict{
            Tuple{String,String},
            Float64,
        }(),
        concentration_lb = 1e-9,
        concentration_ub = 100e-3,
        T = _constants.T,
        R = _constants.R,
        small_flux_tol = 1e-6,
        modifications = [],
        ignore_reaction_ids = [],
    )

Perform a max-min driving force analysis on the `model`, as defined by Noor, et al.,
"Pathway thermodynamics highlights kinetic obstacles in central metabolism.", PLoS
computational biology, 2014.

The function uses the supplied `optimizer` and `reaction_standard_gibbs_free_energies`.
Optionally, `flux_solution` can be used to set the directions of each reaction in `model`
(all reactions are assumed to proceed forward and are active by default). The supplied
`flux_solution` should be free of internal cycles i.e. thermodynamically consistent. This
optional input is important if a reaction in `model` normally runs in reverse (negative
flux). Note, reactions in `flux_solution` that are smaller than `small_flux_tol` are also
ignored in the analysis function (for numerical stability).

The max-min driving force algorithm returns the Gibbs free energy of the reactions, the
concentrations of metabolites and the actual maximum minimum driving force. The optimization
problem solved is:
```
max min -ΔᵣG
s.t. ΔᵣG = ΔrG⁰ + R T S' ln(C)
     ΔᵣG ≤ 0
     ln(Cₗ) ≤ ln(C) ≤ ln(Cᵤ)
```
where `ΔrG` are the Gibbs energies dissipated by the reactions, R is the gas constant, T is
the temperature, S is the stoichiometry of the model, and C is the vector of metabolite
concentrations (and their respective lower and upper bounds).

In case no feasible solution exists, `nothing` is returned.

Reactions specified in `ignore_reaction_ids` are internally ignored when calculating the
max-min driving force. This should include water and proton importers.

Since biochemical thermodynamics are assumed, the `proton_ids` and `water_ids` need to be
specified so that they can be ignored in the calculations. Effectively this assumes an
aqueous environment at constant pH is used.

`constant_concentrations` is used to fix the concentrations of certain metabolites (such as
CO₂). `concentration_ratios` is used to specify additional constraints on metabolite pair
concentrations (typically, this is done with various cofactors such as the ATP/ADP ratio.
For example, you can fix the concentration of ATP to be always 5× higher than of ADP by
specifying `Dict(("ATP","ADP") => 5.0)`

`concentration_lb` and `concentration_ub` set the `Cₗ` and `Cᵤ` in the
optimization problems.

`T` and `R` can be specified in the corresponding units; defaults are K and kJ/K/mol.
"""
function max_min_driving_force(
    model::MetabolicModel,
    reaction_standard_gibbs_free_energies::Dict{String,Float64},
    optimizer;
    flux_solution::Dict{String,Float64} = Dict{String,Float64}(),
    proton_ids::Vector{String} = ["h_c", "h_e"],
    water_ids::Vector{String} = ["h2o_c", "h2o_e"],
    constant_concentrations::Dict{String,Float64} = Dict{String,Float64}(),
    concentration_ratios::Dict{Tuple{String,String},Float64} = Dict{
        Tuple{String,String},
        Float64,
    }(),
    concentration_lb = 1e-9,
    concentration_ub = 100e-3,
    T = _constants.T,
    R = _constants.R,
    small_flux_tol = 1e-6,
    modifications = [],
    ignore_reaction_ids = [],
)
    opt_model = Model(optimizer)

    @variables opt_model begin
        mmdf
        logcs[1:n_metabolites(model)]
        dgrs[1:n_reactions(model)]
    end

    # set proton log concentration to zero so that it won't impact any calculations (biothermodynamics assumption)
    proton_idxs = Int.(indexin(proton_ids, metabolites(model)))
    for idx in proton_idxs
        JuMP.fix(logcs[idx], 0.0)
    end

    # set water concentrations to zero (aqueous condition assumptions)
    water_idxs = Int.(indexin(water_ids, metabolites(model)))
    for idx in water_idxs
        JuMP.fix(logcs[idx], 0.0)
    end

    # only consider reactions with supplied thermodynamic data AND have a flux bigger than
    # small_flux_tol => finds a thermodynamic profile that explains flux_solution
    active_rids = filter(
        rid ->
            haskey(reaction_standard_gibbs_free_energies, rid) &&
                abs(get(flux_solution, rid, small_flux_tol / 2)) > small_flux_tol &&
                !(rid in ignore_reaction_ids),
        reactions(model),
    )
    active_ridxs = Int.(indexin(active_rids, reactions(model)))

    # give dummy dG0 for reactions that don't have data
    dg0s =
        [get(reaction_standard_gibbs_free_energies, rid, 0.0) for rid in reactions(model)]

    S = stoichiometry(model)

    @constraint(opt_model, dgrs .== dg0s .+ (R * T) * S' * logcs)

    # thermodynamics should correspond to the fluxes
    flux_signs = [sign(get(flux_solution, rid, 1.0)) for rid in reactions(model)]

    # only constrain reactions that have thermo data
    @constraint(opt_model, dgrs[active_ridxs] .* flux_signs[active_ridxs] .<= 0)

    # add the absolute bounds
    missing_mets =
        [mid for mid in keys(constant_concentrations) if !(mid in metabolites(model))]
    !isempty(missing_mets) &&
        throw(DomainError(missing_mets, "metabolite(s) not found in model."))
    for (midx, mid) in enumerate(metabolites(model)) # idx in opt_model (missing ignore_metabolites)
        midx in water_idxs && continue
        midx in proton_idxs && continue
        if haskey(constant_concentrations, mid)
            JuMP.fix(logcs[midx], log(constant_concentrations[mid]))
        else
            # this metabolite needs bounds
            @constraint(
                opt_model,
                log(concentration_lb) <= logcs[midx] <= log(concentration_ub)
            )
        end
    end

    # add the relative bounds
    for ((mid1, mid2), val) in concentration_ratios
        idxs = indexin([mid1, mid2], metabolites(model)) # TODO: this is not performant
        any(isnothing.(idxs)) &&
            throw(DomainError((mid1, mid2), "metabolite pair not found in model."))
        @constraint(opt_model, logcs[idxs[1]] == log(val) + logcs[idxs[2]])
    end

    @constraint(opt_model, mmdf .<= -dgrs[active_ridxs] .* flux_signs[active_ridxs])

    @objective(opt_model, Max, mmdf)

    # apply the modifications, if any
    for mod in modifications
        mod(model, opt_model)
    end

    optimize!(opt_model)

    is_solved(opt_model) || return nothing

    return (
        mmdf = value(opt_model[:mmdf]),
        dg_reactions = Dict(
            rid => value(opt_model[:dgrs][i]) for (i, rid) in enumerate(reactions(model))
        ),
        concentrations = Dict(
            mid => exp(value(opt_model[:logcs][i])) for
            (i, mid) in enumerate(metabolites(model))
        ),
    )
end

"""
    max_min_driving_force_variability(
        model::MetabolicModel,
        reaction_standard_gibbs_free_energies::Dict{String,Float64},
        optimizer;
        workers =[myid()],
        optimal_objective_value = nothing,
        bounds = z -> (z, Inf),
        flux_solution::Dict{String,Float64} = Dict{String,Float64}(),
        proton_ids::Vector{String} = ["h_c", "h_e"],
        water_ids::Vector{String} = ["h2o_c", "h2o_e"],
        constant_concentrations::Dict{String,Float64} = Dict{String,Float64}(),
        concentration_ratios::Dict{Tuple{String,String},Float64} = Dict{
            Tuple{String,String},
            Float64,
        }(),
        concentration_lb = 1e-9,
        concentration_ub = 100e-3,
        T = _constants.T,
        R = _constants.R,
        small_flux_tol = 1e-6,
        modifications = [],
        ignore_reaction_ids = [],
    )

Perform a variant of flux variability analysis on a max min driving force type problem.
Arguments are forwarded to [`max_min_driving_force`](@ref). Calls [`screen`](@ref)
internally and possibly distributes computation across `workers`. If
`optimal_objective_value = nothing`, the function first performs regular max min driving
force analysis to find the max min driving force of the model and sets this to
`optimal_objective_value`. Then iteratively maximizes and minimizes the driving force across
each reaction, and then the concentrations while staying close to the original max min
driving force as specified in `bounds`.

The `bounds` is a user-supplied function that specifies the max min driving force bounds for
the variability optimizations, by default it restricts the flux objective value to the
precise optimum reached in the normal max min driving force analysis. It can return `-Inf`
and `Inf` in first and second pair to remove the limit. Use [`gamma_bounds`](@ref) and
[`objective_bounds`](@ref) for simple bounds.

Returns a matrix of solutions to [`max_min_driving_force`](@ref) additionally constrained as
described above, where the rows are in the order of the reactions and then the metabolites
of the `model`. For the reaction rows the first column is the maximum dG of that reaction,
and the second column is the minimum dG of that reaction subject to the above constraints.
For the metabolite rows, the first column is the maximum concentration, and the second column
is the minimum concentration subject to the constraints above.
"""
function max_min_driving_force_variability(
    model::MetabolicModel,
    reaction_standard_gibbs_free_energies::Dict{String,Float64},
    optimizer;
    workers = [myid()],
    optimal_objective_value = nothing,
    bounds = z -> (z, Inf),
    modifications = [],
    kwargs...,
)
    if isnothing(optimal_objective_value)
        initsol = max_min_driving_force(
            model,
            reaction_standard_gibbs_free_energies,
            optimizer;
            modifications,
            kwargs...,
        )
        mmdf = initsol.mmdf
    else
        mmdf = optimal_objective_value
    end

    lb, ub = bounds(mmdf)

    dgr_variants = [
        [[_mmdf_add_df_bound(lb, ub), _mmdf_dgr_objective(ridx, sense)]] for
        ridx = 1:n_reactions(model), sense in [MOI.MAX_SENSE, MOI.MIN_SENSE]
    ]
    concen_variants = [
        [[_mmdf_add_df_bound(lb, ub), _mmdf_concen_objective(midx, sense)]] for
        midx = 1:n_metabolites(model), sense in [MOI.MAX_SENSE, MOI.MIN_SENSE]
    ]

    return screen(
        model;
        args = [dgr_variants; concen_variants],
        analysis = (m, args) -> max_min_driving_force(
            m,
            reaction_standard_gibbs_free_energies,
            optimizer;
            modifications = [args; modifications],
            kwargs...,
        ),
        workers,
    )
end

"""
Helper function to change the objective to optimizing some dG.
"""
function _mmdf_dgr_objective(ridx, sense)
    (model, opt_model) -> begin
        @objective(opt_model, sense, opt_model[:dgrs][ridx])
    end
end

"""
Helper function to change the objective to optimizing some concentration.
"""
function _mmdf_concen_objective(midx, sense)
    (model, opt_model) -> begin
        @objective(opt_model, sense, opt_model[:logcs][midx])
    end
end

"""
Helper function to add a new constraint on the driving force.
"""
function _mmdf_add_df_bound(lb, ub)
    (model, opt_model) -> begin
        if lb == ub
            fix(opt_model[:mmdf], lb; force = true)
        else
            @constraint(opt_model, lb <= opt_model[:mmdf] <= ub)
        end
    end
end


using Test
using TopOpt
import Ferrite
import GeometryBasics
import Makie

# Define problem, can also be imported from .inp files
# nels = (30,10)
# sizes = (1.0,1.0)
nels = (20, 10, 10)
sizes = (1.0,1.0,1.0)
E = 1.0;
ν = 0.3;
force = -1.0;
# problem = PointLoadCantilever(nels, sizes, E, ν, force)
problem = PointLoadCantilever(Val{:Linear}, nels, sizes, E, ν, force);
# @show TopOpt.TopOptProblems.getdim(problem)

# Build element stiffness matrices and force vectors
# einfo = ElementFEAInfo(problem);

# Assemble global stiffness matrix and force vector
# ginfo = assemble(problem, einfo);

# Solve for node displacements
# u = ginfo.K \ ginfo.f
# @show u

solver = FEASolver(Displacement, Direct, problem)
solver()
u = solver.u

# ncells = Ferrite.getncells(problem)
# result_mesh = GeometryBasics.Mesh(problem, ones(ncells));
# GLMakie.mesh(result_mesh);

if get(ENV, "CI", nothing) != "true"
    @eval begin
        using TopOpt.TopOptProblems.Visualization: visualize
        # fig = visualize(problem)
        fig = visualize(problem, u; default_exagg_scale=0.07, scale_range=10.0, vector_linewidth=3, vector_arrowsize=0.5)
        # Makie.display(fig)
    end
end

# using GeometryBasics, Makie
# using Ferrite
# full_topology = ones(Float64, Ferrite.getncells(problem))
# result_mesh = GeometryBasics.Mesh(problem, full_topology);
# mesh(result_mesh);

# # Get the stiffness matrices for buckling analysis
# K, Kσ = TopOpt.TopOptProblems.buckling(problem, ginfo, einfo)

# using IterativeSolvers

# # Find the maximum eigenvalue of system Kσ x = 1/λ K x
# r = lobpcg(Kσ.data, K.data, true, 1)

# # Minimum eigenvalue of the system K x = λ Kσ x
# λ = 1/r.λ[1]

# https://mohamed82008.github.io/ScienceLounge/text/2018/07/11/The-buckle-dance/


using Test
using LinearAlgebra
using GeoParams

@testset "MeltingParam.jl" begin

    #Make sure structure is isbits
    x = MeltingParam_Caricchi()
    @test isbits(x)

    # This tests the various melting parameterizations
    CharUnits_GEO = GEO_units(; viscosity=1e19, length=10km)

    T = collect(250:100:1250) * K .+ 273.15K
    T_nd = Float64.(T / CharUnits_GEO.Temperature)

    # Caricchi parameterization [in ND numbers, which is anyways the typical use case]
    p = MeltingParam_Caricchi()
    phi_dim = zeros(size(T))
    args=(;T=ustrip.(T))
    compute_meltfraction!(phi_dim, p, args)

    phi_dim1 = zeros(size(phi_dim))
    compute_meltfraction!(phi_dim1, p, args) # in-place routine

    p_nd = p
    p_nd = nondimensionalize(p_nd, CharUnits_GEO)
    phi_nd = zeros(size(T))
    args=(;T=T_nd)
    compute_meltfraction!(phi_nd, p_nd, args)

    # Do this computation manually, using the actual expression of Caricchi
    T_C = collect(250:100:1250)    # in celcius
    Phi_solid = 1.0 .- 1.0 ./ (1.0 .+ exp.((800.0 .- T_C) ./ 23.0))
    Phi_anal = 1.0 .- Phi_solid

    @test sum(phi_dim - Phi_anal) < 1e-12
    @test sum(phi_dim1 - Phi_anal) < 1e-12
    @test sum(phi_nd - Phi_anal) < 1e-12

    # test derivative vs T
    dϕdT_dim = zeros(size(T))
    args=(;T=ustrip.(T))
    compute_dϕdT!(dϕdT_dim, p, args)
    @test sum(dϕdT_dim) ≈ 0.008102237679214096

    #------------------------------
    # 5th order polynomial
    p = MeltingParam_5thOrder()
    compute_meltfraction!(phi_dim, p, args)
    @test sum(phi_dim) ≈ 4.708427909521561

    # experimental data to create the fit
    data = [
        1115 100
        1050 90
        969 85
        932 80
        907 70
        880 60
        850 54
        825 51.7
        800 52.9
        775 46.3
        750 44.9
        725 29.9
        700 14.9
        690 0
    ]

    data[:, 2] = data[:, 2] / 100
    Tdata = data[:, 1] .+ 273.15
    phi = zeros(size(Tdata))
    args = (;T=ustrip.(Tdata))
    compute_meltfraction!(phi, p, args)

    @test norm(data[:, 2] - phi) ≈ 0.07151515017819135

    # test derivative vs T
    dϕdT_dim = zeros(size(T))
    args = (;T=ustrip.(T))
    compute_dϕdT!(dϕdT_dim, p, args)
    @test sum(dϕdT_dim) ≈ 0.006484458453421382

    #------------------------------

    #------------------------------
    # 4th order polynomial
    p = MeltingParam_4thOrder()
    compute_meltfraction!(phi_dim, p, args)
    @test sum(phi_dim) ≈ 4.853749635538406

    # experimental data to create the fit
    data = [
        1000 100
        990 100
        975 93
        950 89.2
        925 76.3
        900 69.6
        875 59
        850 54
        825 51.7
        800 52.9
        775 46.3
        750 44.9
        725 29.9
        700 14.9
        690 0
    ]

    data[:, 2] = data[:, 2] / 100
    Tdata = data[:, 1] .+ 273.15
    phi = zeros(size(Tdata))
    args = (;T=ustrip.(Tdata))
    compute_meltfraction!(phi, p, args)

    @test norm(data[:, 2] - phi) ≈ 0.0678052542705406

    # test derivative vs T
    dϕdT_dim = zeros(size(T))
    args = (;T=ustrip.(T))
    compute_dϕdT!(dϕdT_dim, p, args)
    @test sum(dϕdT_dim) ≈ 0.00830985782591842

    #------------------------------

    #------------------------------
    # Quadratic parameterisation
    p = MeltingParam_Quadratic()
    compute_meltfraction!(phi_dim, p, args)
    @test sum(phi_dim) ≈ 5.0894901144641

    dϕdT_dim = zeros(size(T))
    compute_dϕdT!(dϕdT_dim, p, args)
    @test sum(dϕdT_dim) ≈ 0.009365244536940681
    #------------------------------

    #------------------------------
    # Assimilation parameterisation
    p        =  MeltingParam_Assimilation();
    compute_meltfraction!(phi_dim, p, args)
    @test sum(phi_dim) ≈ 4.995
    dϕdT_dim =  zeros(size(T))
    compute_dϕdT!(dϕdT_dim, p, args)
    @test sum(abs.(dϕdT_dim)) ≈0.004605170185988078
    #------------------------------

    # Test computation of melt parameterization for the whole computational domain, using arrays
    MatParam = Vector{MaterialParams}(undef, 4)
    MatParam[1] = SetMaterialParams(;
        Name="Mantle", Phase=1, Melting=PerpleX_LaMEM_Diagram("test_data/Peridotite.in")
    )

    MatParam[2] = SetMaterialParams(;
        Name="Crust", Phase=2, Melting=MeltingParam_Caricchi()
    )

    # No melting parameterization for this phase
    MatParam[3] = SetMaterialParams(;
        Name="UpperCrust", Phase=3, Melting=MeltingParam_5thOrder(), Density=PT_Density()
    )

    # No melting parameterization for this phase
    MatParam[4] = SetMaterialParams(; Name="LowerCrust", Phase=4, Density=PT_Density())

    Mat_tup = Tuple(MatParam)

    # test computing material properties
    n = 100
    Phases = ones(Int64, n, n, n)
    Phases[:, :, 20:end] .= 2
    Phases[:, :, 80:end] .= 3
    Phases[:, :, 90:end] .= 4

    ϕ = zeros(size(Phases))
    dϕdT = zeros(size(Phases))
    T = ones(size(Phases)) * 1500
    P = ones(size(Phases)) * 10
    args = (P=P, T=T)
    compute_meltfraction!(ϕ, Mat_tup, Phases, args) #allocations coming from computing meltfraction using PhaseDiagram_LookupTable
    @test sum(ϕ) / n^3 ≈ 0.7463001302812086

    compute_dϕdT!(dϕdT, Mat_tup, Phases, args)
    @test sum(dϕdT) / n^3 ≈ 0.000176112129245805

    # test computing material properties when we have PhaseRatios, instead of Phase numbers
    PhaseRatio = zeros(n, n, n, 4)
    for i in CartesianIndices(Phases)
        iz = Phases[i]
        I = CartesianIndex(i, iz)
        PhaseRatio[I] = 1.0
    end

    compute_meltfraction!(ϕ, Mat_tup, PhaseRatio, args)
    @test sum(ϕ) / n^3 ≈ 0.7463001302812086

    compute_dϕdT!(dϕdT, Mat_tup, PhaseRatio, args)
    @test sum(dϕdT) / n^3 ≈ 0.000176112129245805

    # Test smoothening of the melting curves:
    p = SmoothMelting(p=MeltingParam_5thOrder())
    @test isbits(p)
    T = collect(250:100:1250) * K .+ 273.15K
    phi_dim = zeros(size(T))
    args=(;T=ustrip.(T))
    compute_meltfraction!(phi_dim, p, args)
    @test sum(phi_dim) ≈ 4.7084279086574226

    dϕdT = zeros(size(T))
    compute_dϕdT!(dϕdT, p, args)
    @test sum(dϕdT) ≈ 0.006484456307540648

    # try non-dimensionalisation
    p_nd = nondimensionalize(p, CharUnits_GEO)
    @test isbits(p)
    @test isdimensional(p_nd.p.a)==false
    @test p_nd.p.a ≈  6968.639721576996

    p1= dimensionalize(p_nd,CharUnits_GEO)
    @test isdimensional(p1.p.a)==true
    @test Value(p1.p.a) ≈ Value(p.p.a)

    #Mat_tup = ( SetMaterialParams( Name="Mantle", Phase=1, Melting=SmoothMelting(MeltingParam_4thOrder())),
    #            SetMaterialParams( Name="Crust", Phase=2, Melting=MeltingParam_5thOrder()),
    #            SetMaterialParams( Name="UpperCrust", Phase=3, Melting=SmoothMelting(MeltingParam_5thOrder()), Density=PT_Density()),
    #            SetMaterialParams( Name="LowerCrust", Phase=4, Density=PT_Density())
    #                )
    Mat_tup = ( SetMaterialParams( Name="Mantle",       Phase=1, Melting=SmoothMelting(MeltingParam_4thOrder())),
                SetMaterialParams( Name="Crust",        Phase=2, Melting=MeltingParam_5thOrder()),
                SetMaterialParams( Name="UpperCrust",   Phase=3, Melting=SmoothMelting(MeltingParam_5thOrder()), Density=PT_Density()),
                SetMaterialParams( Name="LowerCrust",   Phase=4, Density=PT_Density())
                )

    ϕ = zeros(size(Phases));
    dϕdT = zeros(size(Phases));
    T = ones(size(Phases)) * (800+273.15);
    P = ones(size(Phases)) * 10;
    args = (P=P, T=T);
    compute_meltfraction!(ϕ, Mat_tup, Phases, args); #allocation free
    @test sum(ϕ) / n^3 ≈ 0.4409766063389861

    compute_dϕdT!(dϕdT, Mat_tup, Phases, args) #allocation free
    @test sum(dϕdT) / n^3 ≈ 0.0006838372430250584

end


#### Build full Gram matrix (ReproducingKernel_0)

function _gram!(
        A::AbstractMatrix,
        nodes::AbstractVecOfSVecs,
        kernel::ReproducingKernel_0,
    )
    n₁ = length(nodes)
    @inbounds for j in 1:n₁
        for i in 1:j
            A[i,j] = _rk(kernel, nodes[i], nodes[j])
        end
    end
    return Hermitian(A, :U)
end

function _gram(
        nodes::AbstractVecOfSVecs,
        kernel::ReproducingKernel_0,
    )
    n₁ = length(nodes)
    T  = eltype(eltype(nodes))
    _gram!(zeros(T, n₁, n₁), nodes, kernel)
end

#### Incrementally add column to Gram matrix (ReproducingKernel_0)

function _gram!(
        A::AbstractMatrix,
        new_node::SVector,
        curr_nodes::AbstractVecOfSVecs,
        kernel::ReproducingKernel_0,
    )
    n₁ = length(curr_nodes)
    @inbounds for i in 1:n₁
        A[i, n₁+1] = _rk(kernel, curr_nodes[i], new_node)
    end
    @inbounds A[n₁+1, n₁+1] = _rk(kernel, new_node, new_node)
    return Hermitian(A, :U)
end

#### Build full Gram matrix (ReproducingKernel_1)

function _gram!(
        A::AbstractMatrix,
        nodes::AbstractVecOfSVecs{n},
        d_nodes::AbstractVecOfSVecs{n},
        d_dirs::AbstractVecOfSVecs{n},
        kernel::ReproducingKernel_1,
    ) where {n}
    n₁  = length(nodes)
    n₂  = length(d_nodes)
    A11 = A
    A12 = uview(A, 1    : n₁,    n₁+1 : n₁+n₂)
    A22 = uview(A, n₁+1 : n₁+n₂, n₁+1 : n₁+n₂)

    @inbounds for j in 1:n₁
        # Top-left block (n₁ × n₁)
        for i in 1:j
            A11[i,j] = _rk(kernel, nodes[i], nodes[j])
        end
    end

    ε² = kernel.ε^2
    @inbounds for j in 1:n₂
        # Top-right block (n₁ × n₂)
        for i in 1:n₁
            A12[i,j] = _∂rk_∂e(kernel, nodes[i], d_nodes[j], d_dirs[j])
        end

        # Bottom-right block (n₂ × n₂)
        for i in 1:j-1
            A22[i,j] = _∂²rk_∂²e(kernel, d_nodes[j], d_nodes[i], d_dirs[j], d_dirs[i])
        end
        A22[j,j] = ε²
    end

    return Hermitian(A, :U)
end

function _gram(
        nodes::AbstractVecOfSVecs{n},
        d_nodes::AbstractVecOfSVecs{n},
        d_dirs::AbstractVecOfSVecs{n},
        kernel::ReproducingKernel_1,
    ) where {n}
    n₁ = length(nodes)
    n₂ = length(d_nodes)
    T  = promote_type(eltype(eltype(nodes)), eltype(eltype(d_nodes)), eltype(eltype(d_dirs)))
    _gram!(zeros(T, n₁+n₂, n₁+n₂), nodes, d_nodes, d_dirs, kernel)
end

#### Incrementally add column to Gram matrix (ReproducingKernel_1)

function _gram!(
        A::AbstractMatrix,
        new_node::SVector{n},
        curr_nodes::AbstractVecOfSVecs{n},
        curr_d_nodes::AbstractVecOfSVecs{n},
        curr_d_dirs::AbstractVecOfSVecs{n},
        kernel::ReproducingKernel_1,
    ) where {n}
    n₁ = length(curr_nodes)
    n₂ = length(curr_d_nodes)
    @assert size(A) == (n₁+1+n₂, n₁+1+n₂)

    # Top-left block (n₁+1 × n₁+1), right column (n₁+1 terms)
    @inbounds for i in 1:n₁
        A[i, n₁+1] = _rk(kernel, new_node, curr_nodes[i])
    end
    @inbounds A[n₁+1, n₁+1] = _rk(kernel, new_node, new_node)

    # Top-right block (n₁+1 × n₂), bottom row (n₂ terms)
    @inbounds for j in 1:n₂
        A[n₁+1, n₁+1+j] = _∂rk_∂e(kernel, new_node, curr_d_nodes[j], curr_d_dirs[j])
    end

    return Hermitian(A, :U)
end

function _gram!(
        A::AbstractMatrix,
        d_node::SVector{n},
        d_dir::SVector{n},
        curr_nodes::AbstractVecOfSVecs{n},
        curr_d_nodes::AbstractVecOfSVecs{n},
        curr_d_dirs::AbstractVecOfSVecs{n},
        kernel::ReproducingKernel_1,
    ) where {n}
    n₁ = length(curr_nodes)
    n₂ = length(curr_d_nodes)
    @assert size(A) == (n₁+n₂+1, n₁+n₂+1)

    # Top-right block, (n₁ × n₂+1), right column (n₁ terms)
    @inbounds for i in 1:n₁
        A[i, n₁+n₂+1] = _∂rk_∂e(kernel, curr_nodes[i], d_node, d_dir)
    end

    # Bottom-right block (n₂+1 × n₂+1), right column (n₂+1 terms)
    ε² = kernel.ε^2
    @inbounds for i in 1:n₂
        A[n₁+i, n₁+n₂+1] = _∂²rk_∂²e(kernel, d_node, curr_d_nodes[i], d_dir, curr_d_dirs[i])
    end
    @inbounds A[n₁+n₂+1, n₁+n₂+1] = ε²

    return Hermitian(A, :U)
end

#### Elastic Cholesky

Base.@kwdef struct ElasticCholesky{T, AType <: AbstractMatrix{T}} <: LinearAlgebra.Factorization{T}
    maxcols::Int
    ncols::Base.RefValue{Int} = Ref(0)
    colperms::Vector{Int}     = zeros(Int, maxcols)
    A::AType                  = zeros(T, maxcols, maxcols)
    U::Matrix{T}              = zeros(T, maxcols, maxcols)
    U⁻ᵀb::Vector{T}           = zeros(T, maxcols)
end
ElasticCholesky{T}(maxcols::Int) where {T} = ElasticCholesky{T,Matrix{T}}(; maxcols = maxcols)
ElasticCholesky(A::AbstractMatrix{T}) where {T} = ElasticCholesky{T,typeof(A)}(; maxcols = size(A,2), A = A)

Base.eltype(::ElasticCholesky{T}) where {T} = T
Base.size(C::ElasticCholesky) = (C.ncols[], C.ncols[])
Base.parent(C::ElasticCholesky) = C.A
Base.empty!(C::ElasticCholesky) = (C.ncols[] = 0; C)
Base.show(io::IO, mime::MIME"text/plain", C::ElasticCholesky{T}) where {T} = (print(io, "ElasticCholesky{T}\nU factor:\n"); show(io, mime, UpperTriangular(C.U[C.colperms[1:C.ncols[]], C.colperms[1:C.ncols[]]])))

function LinearAlgebra.ldiv!(x::AbstractVector{T}, C::ElasticCholesky{T}, b::AbstractVector{T}, ::Val{permview} = Val(false)) where {T, permview}
    @unpack U, U⁻ᵀb, colperms, ncols = C
    J = uview(colperms, 1:ncols[])
    U = UpperTriangular(uview(U, J, J))
    U⁻ᵀb = uview(U⁻ᵀb, 1:ncols[])
    if permview
        x = uview(x, J)
        b = uview(b, J)
    end
    ldiv!(U⁻ᵀb, U', b)
    ldiv!(x, U, U⁻ᵀb)
    return x
end

function Base.insert!(C::ElasticCholesky{T}, j::Int, B::AbstractMatrix{T}) where {T}
    @unpack A, colperms, ncols = C
    @inbounds colperms[ncols[] + 1] = j
    rows = uview(colperms, 1 : ncols[] + 1)
    @inbounds for i in rows
        A[i,j] = B[i,j]
    end
    return C
end

"""
    LinearAlgebra.cholesky!(C::ElasticCholesky, v::AbstractVector{T}) where {T}

Update the Cholesky factorization `C` as if the column `v` (and by symmetry, the corresponding row `vᵀ`)
were inserted into the underlying matrix `A`. Specifically, let `L` be the lower-triangular cholesky factor
of `A` such that `A = LLᵀ`, and let `v = [d; γ]` such that the new matrix `A⁺` is given by

```
A⁺ = [A  d]
     [dᵀ γ].
```

Then, the corresponding updated cholesky factor `L⁺` of `⁺` is:

```
L⁺ = [L  e]
     [eᵀ α]
```

where `e = L⁻¹d`, `α = √τ`, and `τ = γ - e⋅e > 0`. If `τ ≤ 0`, then `A⁺` is not positive definite.

See:
    https://igorkohan.github.io/NormalHermiteSplines.jl/dev/Normal-Splines-Method/#Algorithms-for-updating-Cholesky-factorization
"""
function LinearAlgebra.cholesky!(
        C::ElasticCholesky{T},
        j::Int,
        v::AbstractVector{T},
        ::Val{fill_parent},
    ) where {T, fill_parent}
    @unpack maxcols, A, U, colperms, ncols = C
    @assert length(v) == ncols[] + 1 <= maxcols

    @inbounds if ncols[] == 0
        # Initialize first entry of `A`
        colperms[1] = j
        if fill_parent
            A[j,j] = v[1]
        end
        U[j,j] = sqrt(v[1])
        ncols[] = 1
    else
        # Fill `A` with new column
        colperms[ncols[] + 1] = j
        if fill_parent
            rows = uview(colperms, 1:ncols[] + 1)
            copyto!(uview(A, rows, j), v)
        end

        # Update `U` with new column
        J = uview(colperms, 1:ncols[])
        d = uview(A, J, j)
        γ = A[j,j]
        e = uview(U, J, j)
        Uᵀ = UpperTriangular(uview(U, J, J))'
        ldiv!(e, Uᵀ, d)
        τ = γ - e⋅e
        α = √max(τ, 0) # `τ` should be positive by construction
        U[j,j] = max(α, eps(T)) # if `α < ϵ` you have bigger problems...

        # Increment column counter
        ncols[] += 1
    end

    return C
end

# Update the `j`th column of the factorization `C.U`, assuming the corresponding column `j` of `C.A` has been filled
function LinearAlgebra.cholesky!(C::ElasticCholesky{T}, j::Int) where {T}
    @unpack maxcols, A, colperms, ncols = C
    @assert ncols[] + 1 <= maxcols

    @inbounds colperms[ncols[] + 1] = j
    rows = uview(colperms, 1:ncols[] + 1)
    v = uview(A, rows, j)
    cholesky!(C, j, v, Val(false))

    return C
end

# Update columns `J` of the factorization `C.U`, assuming the corresponding columns `J` of `C.A` have been filled
function LinearAlgebra.cholesky!(C::ElasticCholesky, J = axes(C.A, 2))
    for j in J
        cholesky!(C, j)
    end
    return C
end


# This file was generated, do not modify it.

using MLJ, Statistics
X, y = @load_iris;

@load DecisionTreeClassifier

tree_model = DecisionTreeClassifier()
tree = machine(tree_model, X, y)

train, test = partition(eachindex(y), 0.7, shuffle=true)
test[1:3]

fit!(tree, rows=train)

tree.fitresult

ŷ = predict(tree, rows=test)
@show ŷ[1]

ȳ = predict_mode(tree, rows=test)
@show ȳ[1]
@show mode(ŷ[1])

mce = cross_entropy(ŷ, y[test]) |> mean
round(mce, digits=4)

v = [1, 2, 3, 4]
stand_model = UnivariateStandardizer()
stand = machine(stand_model, v)

fit!(stand)
w = transform(stand, v)
@show round.(w, digits=2)
@show mean(w)
@show std(w)

vv = inverse_transform(stand, w)
sum(abs.(vv .- v))


module UnitCells

using LinearAlgebra

export UnitCell, monkhorst_pack_mesh, calc_cell_pos!, calc_cell_pos, calc_site_pos!, calc_site_pos

"""
Represents a specfied lattice geometry.
"""
struct UnitCell{T<:AbstractFloat}

    "ndim: the number of dimensions that the geometry lives in."
    ndim::Int

    "norbits: the number of sites per unit cell."
    norbits::Int

    "lvecs:: (3 x 3) matrix where the columns give the lattice vectors."
    lvecs::Matrix{T}

    "rlvecs: (3 x 3) matrix where the columns give the reciprocal lattice vectors."
    rlvecs::Matrix{T}

    "bvecs: (3 x norbits) matrix where the columns give the basis vectors."
    bvecs::Matrix{T}

    #######################
    ## INNER CONSTRUCTOR ##
    #######################
    """
        UnitCell(ndim::Int,norbits::Int,lvecs::Matrix{T},bvecs::Matrix{T}) where {T<:AbstractFloat}

    Constructor for UnitCell type.
    """
    function UnitCell(ndim::Int,norbits::Int,lvecs::Matrix{T},bvecs::Matrix{T}) where {T<:AbstractFloat}

        # constructing matrix containing lattice vectors
        nrows = size(lvecs,1)
        ncols = size(lvecs,2)
        Lvecs = Matrix{T}(I,3,3) # intialized as identity matri
        Lvecs[1:nrows,1:ncols] = lvecs

        # calculating reciprocal lattice vectors
        RLvecs = 2*π*inv(Lvecs)

        # constucting matrix to contain basis vectors
        nrows = size(bvecs,1)
        ncols = size(bvecs,2)
        Bvecs = zeros(T,3,ncols)
        Bvecs[1:nrows,:] = bvecs

        new{T}(ndim,norbits,Lvecs,RLvecs,Bvecs)
    end
end

########################
## OUTER CONSTRUCTORS ##
########################

function UnitCell(ndim::Int, norbits::Int, lvecs::Matrix{T}, bvecs::Vector{Vector{T}}) where {T<:AbstractFloat}
    UnitCell(ndim, norbits, lvecs, hcat(bvecs...))
end

function UnitCell(ndim::Int, norbits::Int, lvecs::Vector{Vector{T}}, bvecs::Matrix{T}) where {T<:AbstractFloat}
    UnitCell(ndim, norbits, hcat(lvecs...), bvecs)
end

function UnitCell(ndim::Int, norbits::Int, lvecs::Vector{Vector{T}}, bvecs::Vector{Vector{T}}) where {T<:AbstractFloat}
    UnitCell(ndim, norbits, hcat(lvecs...), hcat(bvecs...))
end

#############################################
## DEFINING METHODS THAT USE GEOMETRY TYPE ##
#############################################

"""
Calculates the position of a unit cell in a lattice.
"""
function calc_cell_pos!(pos::AbstractVector{T},geom::UnitCell,l1::Int,l2::Int=0,l3::Int=0) where {T<:AbstractFloat}

    @assert length(pos)==3
    lv1 = @view geom.lvecs[:,1] # first lattice vector
    lv2 = @view geom.lvecs[:,2] # second lattice vector
    lv3 = @view geom.lvecs[:,3] # third lattice vector
    @. pos = l1*lv1 + l2*lv2 + l3*lv3 # calculating position
    return nothing
end

function calc_cell_pos(geom::UnitCell,l1::Int,l2::Int=0,l3::Int=0)::Vector{T} where {T<:AbstractFloat}

    pos = zeros(3)
    calc_cell_pos!(pos,geom,l1,l2,l3)
    return pos
end

"""
Calculates the position of a site in a lattice.
"""
function calc_site_pos!(pos::AbstractVector{T},geom::UnitCell,orbit::Int,l1::Int,l2::Int=0,l3::Int=0) where {T<:AbstractFloat}

    @assert orbit>0
    # calculating position of unit cell that site lives in
    calc_cell_pos!(pos,geom,l1,l2,l3)
    # adding basis vector for given orbital to position
    pos .+= @view geom.bvecs[:,orbit]
    return nothing
end

function calc_site_pos(geom::UnitCell{T},orbit::Int,l1::Int,l2::Int=0,l3::Int=0)::Vector{T}  where {T<:AbstractFloat}

    pos = zeros(3)
    calc_site_pos!(pos,geom,orbit,l1,l2,l3)
    return pos
end

"""
    monkhorst_pack_mesh(geom::UnitCell, L1::Int, L2::Int=1, L3::Int=1)::Matrix{AbstractFloat}

Returns a matrix where each column is a k-point vector in a Monkhort-Pack meshgrid over the full Brillouin Zone.
"""
function monkhorst_pack_mesh(geom::UnitCell{T}, L1::Int, L2::Int=1, L3::Int=1)::Matrix{T}  where {T<:AbstractFloat}

    kpoints = zeros(T,3,L1*L2*L3)
    v1 = @view geom.rlvecs[:,1]
    v2 = @view geom.rlvecs[:,2]
    v3 = @view geom.rlvecs[:,3]
    i = 1
    for l3=0:L3-1
        for l2=0:L2-1
            for l1=0:L1-1
                kpoints[:,i] = @. (l1/L1)*v1 + (l2/L2)*v2 + (l3/L3)*v3
                i += 1
            end
        end
    end
    return kpoints
end

end


export Half

struct Half{M} <: AbstractMeasure
    parent::M
end

function Base.show(io::IO, μ::Half)
    print(io, "Half")
    show(io, μ.parent)
end

unhalf(μ::Half) = μ.parent

function basemeasure(μ::Half)
    inbounds(x) = x > 0
    constℓ = logtwo
    varℓ() = 0.0
    base = basemeasure(unhalf(μ))
    FactoredBase(inbounds, constℓ, varℓ, base)
end

function Base.rand(rng::AbstractRNG, T::Type, μ::Half)
    return abs(rand(rng, T, unhalf(μ)))
end

logdensity(μ::Half, x) = logdensity(unhalf(μ), x)


struct And end

(a::And)(x::BOOL, y::BOOL) = x(y, FALSE())

struct Or end

(o::Or)(x::BOOL, y::BOOL) = x(TRUE(), y)

struct Not end

(n::Not)(x::BOOL) = x(FALSE(), TRUE())

struct Eq end

(e::Eq)(x::BOOL, y::BOOL) = x(y, Not()(y))

struct Xor end

(o::Xor)(x::BOOL, y::BOOL) = Not()(Eq()(x, y))


using LinearAlgebra
using DifferentialEquations
using Distributions
using Random
using Printf
using Plots
using RCall
using CPUTime
using DynamicalSystems
using StatsPlots
# using StatsPlots
include("utils.jl")
include("lorenz_ode.jl")
include("abc.jl")
include("abcMCMC.jl")
include("BayesianCalibration.jl")
include("abcSMC.jl")

# sigma, rho, beta, x0
# sigma_prior = Uniform(5,15)
# rho_prior = Uniform(20, 30)
# beta_prior = Uniform(8/3 - 1.5, 8/3 +1.5)
# x0_prior = Uniform(7, 13)

# sigma_prior = truncated(Normal(10., 5), 0, 20)
# rho_prior = truncated(Normal(26, 10), 0, 40)
# beta_prior = truncated(Normal(8/3., 2), 1, 4)

sigma_prior = Uniform(5,15)
rho_prior = Uniform(20, 30)
beta_prior = Uniform(8/3 - 1.5, 8/3 +1.5)

sd = 0.01
sigma_true = truncated(Normal(10., sd), 0, Inf)
rho_true = truncated(Normal(26., sd), 0, Inf)
beta_true = truncated(Normal(8/3., sd), 0, Inf)

function data_generator_lorenz(p)

    u0 = [10., 10., 10.]
    time_window = (0., 20.)

    y = solve_lorenz(u0,time_window,p)
end

function data_generator_lorenz10(p)

    u0 = [10., 10., 10.]
    time_window = (0., 10.)

    y = solve_lorenz(u0,time_window,p)
end

true_prior=(sigma_true, rho_true, beta_true)
true_p = map(x -> rand(x),true_prior)
y = data_generator_lorenz(true_p)

y1 = data_generator_lorenz((10., 26., 8/3))
y2 = data_generator_lorenz((10, 26, 8/3 - 1e-3))

sol1=hcat(y1[:]...)'
sol2=hcat(y2[:]...)'
l2_error = sqrt(sum(sum((sol1.-sol2).^2,dims=2)*0.01,dims=1))[1]

#
#
#
# ds_1 = Systems.lorenz([10., 10., 10.];σ= 10., ρ= 26., β= 8/3)
# ds_2 = Systems.lorenz([10., 10., 10.];σ= 10. , ρ= 26., β= 8/3 + 1)
# lyap_1 = ChaosTools.lyapunov(ds_1, 1000.0, dt=1.0, Ttr=10.0)
# lyap_2 = ChaosTools.lyapunov(ds_2, 1000.0, dt=1.0, Ttr=10.0)
# dist_pert = abs(lyap_1 - lyap_2)

#abc
algo_parameter_ABC = (prior = ( sigma_prior, rho_prior, beta_prior),epsilon = 45.,
eta= identity_mapping, d= compute_l2_norm_decay)

#  ABC MCMC
# algo_parameter_mcmc = (prior = ( sigma_prior, rho_prior, beta_prior) ,epsilon = 60., eta = identity_mapping,
#                         d= compute_l2_norm_decay, proposal = proposal_Normal,
#                         proposalRatio = proposalRatio_Normal,sd = (0.5,0.5,0.5),
#                         thinning = 100, burn_in = 100,verbose=true)

algo_parameter_mcmc = (prior = ( sigma_prior, rho_prior, beta_prior) ,epsilon = 60., eta = identity_mapping,
                        d= compute_l2_norm_decay, proposal = proposal_Normal,
                        proposalRatio = proposalRatio_Normal,sd = (0.5,0.5,0.5),
                        thinning = 20, burn_in = 100,verbose=true)

#smc
algo_parameter_smc = (prior = ( sigma_prior, rho_prior, beta_prior), time_final=8, eps_list = [80,70, 65,60,55,50,44,45],
eta= identity_mapping, d= compute_l2_norm_decay, kernel=proposal_Normal,
kernel_density=proposal_Normal_density,
sd=(0.3,0.3, 0.3), resample_method=systematic_resample, verbose=true)

output_abc = ABC(y, data_generator_lorenz, algo_parameter_ABC, 250)
output_smc = ABC_SMC(y, data_generator_lorenz, algo_parameter_smc, 250)
output_mcmc = ABC_MCMC(y, data_generator_lorenz, algo_parameter_mcmc, 250)

n=100

sigma_abc=output_abc[1][:,1]; rho_abc=output_abc[1][:,2]; beta_abc=output_abc[1][:,3]

ds_observed = Systems.lorenz([10., 10., 10.];σ= true_p[1], ρ= true_p[2], β= true_p[3])
lyap_observed = ChaosTools.lyapunov(ds_observed, 1000.0, dt=10.0, Ttr=100.0)
dist_lyapunov_abc=zeros(n)
lya_abc=zeros(n)
for j=1:n
    u0 = [10., 10., 10.]
    ds_fitted = Systems.lorenz(u0;σ= sigma_abc[j], ρ= rho_abc[j], β= beta_abc[j])
    lyap_fit = ChaosTools.lyapunov(ds_fitted, 1000.0, dt=10.0, Ttr=100.0)
    dist_lyapunov_abc[j] = abs(lyap_fit-lyap_observed)/lyap_observed
    lya_abc[j] = lyap_fit
end

sigma_mcmc=output_mcmc[1][:,1]; rho_mcmc=output_mcmc[1][:,2]; beta_mcmc=output_mcmc[1][:,3]
dist_lyapunov_mcmc=zeros(n)
lya_mcmc=zeros(n)
for j=1:n
    u0 = [10., 10., 10.]
    ds_fitted = Systems.lorenz(u0;σ= sigma_mcmc[j], ρ= rho_mcmc[j], β= beta_mcmc[j])
    lyap_fit = ChaosTools.lyapunov(ds_fitted, 1000.0, dt=1.0, Ttr=10.0)
    dist_lyapunov_mcmc[j] = abs(lyap_observed-lyap_fit)/lyap_observed
    lya_mcmc[j]=lyap_fit
end

sigma_smc=output_smc[1][:,1]; rho_smc=output_smc[1][:,2]; beta_smc=output_smc[1][:,3]
dist_lyapunov_smc=zeros(n)
lyapunov_smc=zeros(n)
for j=1:n
    u0 = [10., 10., 10.]
    ds_fitted = Systems.lorenz(u0;σ= sigma_smc[j], ρ= rho_smc[j], β= beta_smc[j])
    lyap_fit = ChaosTools.lyapunov(ds_fitted, 1000.0, dt=10.0, Ttr=10.0)
    dist_lyapunov_smc[j] = abs(lyap_fit-lyap_observed)/lyap_observed
    lyapunov_smc[j] = lyap_fit
end
using LaTeXStrings
lyap_dist_plot = plot(plot(1:n,dist_lyapunov_abc, seriestype=:scatter, markersize=3, alpha=0.5, label="ABC"),
plot(1:n,dist_lyapunov_mcmc,seriestype=:scatter, markersize=3, alpha=0.5, label="MCMC"),
plot(1:n,dist_lyapunov_smc,seriestype=:scatter,  markersize=3, alpha=0.5,label="SMC"),
ylim=(0,1.1), xlabel=L"n", ylabel=L"$|\lambda^* - \lambda^n|$",
legend=:outertopright,
 dpi=350, layout=(3,1),
 foreground_color_legend=nothing)
plot(lya_abc,zeros(n),seriestype=:scatter)
plot!([lyap_observed],[0], seriestype=:scatter, color="red")

lyap_density_plot = density(lya_abc); vline!([lyap_observed])

n= size(output_abc[1])[1]
p_abc=plot(y, vars=(1,3),w=1,legend=false,title=L"\textrm{ABC}", dpi=300)
abc_best_ind = sortperm(dist_lyapunov_abc)
abc_best_ind = randperm(n)
# abc_best_ind = 1:1:n
output_abc_best = output_abc[1][abc_best_ind,:]
for k in 1:20
    p_pred = (output_abc_best[k,1], output_abc_best[k,2], output_abc_best[k,3])
    y_fitted = data_generator_lorenz(p_pred)
    plot!(p_abc,y_fitted, vars=(1,3),alpha=0.3, color="#BBBBBB", legend=false, xlim=(-20, 20),xlabel=L"x", ylabel=L"z")
end
p_abc

# mcmc_best_ind = sortperm(dist_lyapunov_mcmc)
mcmc_best_ind =randperm(n)
output_mcmc_best = output_mcmc[1][mcmc_best_ind,:]
p_mcmc=plot(y, vars=(1,3),w=1,legend=false,title=L"\textrm{ABC-MCMC}", dpi=300)
for k in 1:20
    p_pred = (output_mcmc_best[k,1], output_mcmc_best[k,2], output_mcmc_best[k,3])
    y_fitted = data_generator_lorenz(p_pred)
    plot!(p_mcmc,y_fitted, vars=(1,3),alpha=0.3, color="#BBBBBB", legend=false, xlim=(-20, 20),xlabel=L"x", ylabel=L"z")
end
p_mcmc

smc_best_ind = sortperm(dist_lyapunov_smc)
mcmc_best_ind =randperm(n)
# smc_best_ind = 1:1:n
output_smc_best = output_smc[1][smc_best_ind,:]
p_smc=plot(y, vars=(1,3),w=1,legend=false,title=L"\textrm{ABC-SMC}",dpi=300)
for k in 1:20
    p_pred = (output_smc_best[k,1], output_smc_best[k,2], output_smc_best[k,3])
    y_fitted = data_generator_lorenz(p_pred)
    plot!(p_smc,y_fitted, vars=(1,3),alpha=0.3, color="#BBBBBB", legend=false, xlim=(-20, 20),xlabel=L"x", ylabel=L"z")
end
p_smc

sigma_density=density(output_abc[1][:,1], xlabel="σ", label="ABC")
density!(output_mcmc[1][:,1], label="MCMC")
density!(output_smc[1][:,1], label="SMC")

rho_density=density(output_abc[1][:,2], xlabel="ρ", label="ABC")
density!(output_mcmc[1][:,2], label="MCMC")
density!(output_smc[1][:,2], label="SMC")

b_density=density(output_abc[1][:,3], xlabel="b", label="ABC")
density!(output_mcmc[1][:,3], label="MCMC")
density!(output_smc[1][:,3], label="SMC")